From 4d5c7f32d6280c8c730ee838719b8fb5ef6f7760 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 12 Jun 2026 08:11:51 +0800 Subject: [PATCH] Keep CP shared-KV prefetch warnings actionable Expected no-prefetch paths were polluting production logs: no cache prefix, tiny/first-layer windows, and FP8 RAGGED top-k were being reported as fallback warnings. The prefetch contract now treats zero-prefix and first-layer misses as normal skips, while preserving warnings for non-zero misaligned prefixes and real consume misses after the first layer. The same change keeps RAGGED cache-hit prefetch eligible and records the CE/IPM prefetch contract in the plan doc. Constraint: FP8 sparse prefill uses RAGGED top-k, but CP shared-KV prefix materialization is still page-slot based Constraint: Layer 0 has no previous attention-window hook that can have prefetched the layer Rejected: Warn whenever a prefetcher is absent | no-cache and too-short requests are expected synchronous paths and make logs unusable Confidence: high Scope-risk: moderate Directive: Keep CP_SHARED_KV_FALLBACK warnings for unexpected contract failures only; use debug logs for expected skip paths Tested: Local py_compile for cp_shared_kv_prefetch.py, nsa_indexer.py, nsa_backend.py Tested: Remote cjy-glm5-new targeted regression: 3 passed, 21 warnings Tested: Remote cjy-glm5-new full test_cp_shared_kv_runtime.py: 156 passed, 21 warnings, 2 subtests passed Not-tested: New ETE run after prefill restart to confirm log volume reduction in production traffic (cherry picked from commit e08e321e5929fdbb30102ec0b19c6ff0ecac7e7e) --- ...d_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md | 138 ++- .../attention/nsa/cp_shared_kv_prefetch.py | 591 +++++------ .../attention/nsa/cp_shared_kv_runtime.py | 247 +++++ .../srt/layers/attention/nsa/nsa_indexer.py | 45 +- .../srt/layers/attention/nsa_backend.py | 104 +- .../mem_cache/test_cp_shared_kv_runtime.py | 966 +++++++++++++++--- 6 files changed, 1509 insertions(+), 582 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md index 236bf1dbb..e9abfa3e2 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md @@ -430,15 +430,37 @@ L2->L1 load finished on owner rank ### P0. 文档与现状锁定 -**文件：** +**状态（2026-06-12）：已完成。** -- 新增：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md` +完成内容： -**验收：** +1. L1 prefetch create 路径不再以 `get_attention_cp_group().pynccl_comm` 作为启用条件。 + prefetch 的新合同是 IPC/CE materialize，不再依赖 CP collective communicator。 +2. `cp_shared_kv_prefetch.py` 与 `nsa_backend.py` 中 prefetch hot path 已清理旧 + `launch_pending_reduce` / `prefix_reduce` / `reduce_enqueue` / `collective_disallowed` + 命名，统一改成 `finalize_pending_materialize` / `prefix_materialize` / + `materialize_enqueue` / `materialize_event_missing`。 +3. prefetch 文件中不再 import/use `_all_reduce_materialized_buffer_*` 或 + `get_attention_cp_group`。如果 IPC/CE materialize 不可用，CP>1 下应显式 + fail-fast，而不是 silent fallback 到 all-reduce。 +4. 测试已覆盖：`maybe_create()` 不依赖 pynccl，attention-window 不再触发 + prefix reduce 旧路径，tiny extend skip 测试也不再 mock pynccl。 -- 文档明确 L2->L1 `+2 layer`、L1 prefetch `+1 layer`。 -- 文档明确当前 prefetcher 的 bs=1 假设和风险。 -- 文档明确 0SM 与现有 SM IPC kernel 的区别。 +验收证据： + +```text +远端 cjy-glm5-new: +PYTHONPATH=python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +150 passed, 21 warnings, 2 subtests passed +``` + +静态检查： + +```text +rg "pynccl|launch_pending_reduce|prefix_reduce|collective_disallowed|reduce_enqueue|get_attention_cp_group|_all_reduce" \ + cp_shared_kv_prefetch.py nsa_backend.py +# no matches +``` ### P1. 补 bs>1 prefetch plan 单测 @@ -508,17 +530,17 @@ PYTHONPATH=python python -m pytest -q \ - `current_slot_spans` - `prefix_page_count` - `current_page_count` -3. `start_next_layer_prefix()` 只 materialize/reduce `prefix_slot_spans`， +3. `start_next_layer_prefix()` 只 materialize `prefix_slot_spans`， 不再把 batch flattened page table 当成 `[0:prefix_pages)`。 -4. `consume_prefix_with_current()` 只 reduce `current_slot_spans`，避免把 - row gap / 其他 request prefix 一起 reduce。 +4. `consume_prefix_with_current()` 只处理 `current_slot_spans`，避免把 + row gap / 其他 request prefix 一起纳入 suffix compose。 5. `consume()` 的 legacy full-materialize suffix 路径也改为使用 `current_slot_spans`，避免 bs>1 bounding suffix。 当前限制： -- MLA prefix spans 仍走现有 materialize + async all-reduce baseline；还没有 - 接入 spans-list TAI IPC 或 0SM CE。 +- MLA prefix spans 已不再走 async all-reduce fallback；当前要求 IPC/CE materialize + 生成 completion event，CP>1 下不可用则 fail-fast。 - index prefetcher 仍未改造，继续由 P3 处理。 **文件：** @@ -531,13 +553,13 @@ PYTHONPATH=python python -m pytest -q \ 1. `CpSharedKVMlaPrefetcher.maybe_create()` 不再因为 `batch_size != 1` 直接 skip。 2. create 阶段构造 `prefix_slot_spans/current_slot_spans`。 3. `start_next_layer_prefix()` 只 materialize prefix spans。 -4. `consume_prefix_with_current()` 只 fill/reduce current spans。 +4. `consume_prefix_with_current()` 只处理 current spans。 5. `consume()` 如果仍存在 legacy suffix path，必须使用 `current_slot_spans` 或 fail-fast，不能回到错误 bounding suffix。 **第一版允许：** -- prefix 使用现有 local materialize + all-reduce 或 SM IPC baseline。 +- prefix 使用现有 SM IPC 或 CE materialize baseline。 **第一版不允许：** @@ -567,8 +589,8 @@ RED 证据：旧代码在 `batch_size=2` 时以 - `current_slot_spans` - `prefix_page_count` - `current_page_count` -3. `start_next_layer_prefix()` 只 materialize/reduce index prefix spans。 -4. `consume_prefix_with_current()` 只 fill/reduce index current spans。 +3. `start_next_layer_prefix()` 只 materialize index prefix spans。 +4. `consume_prefix_with_current()` 只处理 index current spans。 5. `consume()` 的 legacy suffix 路径也改为 `current_slot_spans`， 不再使用 batch bounding suffix。 @@ -581,8 +603,8 @@ test_cp_shared_kv_runtime.py 当前限制： -- index prefetch 仍使用现有 materialize + async all-reduce baseline； - spans-list TAI IPC / 0SM CE 留给后续 P5/P6。 +- index prefetch 已不再走 async all-reduce fallback；当前要求 IPC/CE materialize + 生成 completion event，CP>1 下不可用则 fail-fast。 - active index layer / index skip 的 runtime hook 当前沿用已有 `nsa_backend.py` 调用点；本阶段没有修改 skip 参数语义。 @@ -637,11 +659,11 @@ test_cp_shared_kv_runtime.py - `MLATokenToKVPool.get_key_buffer_for_prefetch(layer_id, stream)` - `MLATokenToKVPool.get_index_k_with_scale_buffer_for_prefetch(layer_id, stream)` - 发现并修正一个 ordering 疏漏：P1-P3 后 prefix materialize 仍在 current stream 上执行，但 `start_next_layer_prefix()` 把 L2->L1 readiness wait 绑定到了 prefetch stream。这样只能保护后续 reduce，不能保护实际读取 L1 raw pages 的 materialize。 -- 修正策略：MLA/index 的 `start_next_layer_prefix()` 先取得 `current_stream`，把它传给 prefetch-safe getter，使 L2->L1 ready event 挂到实际读取 KV/index page 的 stream；随后仍保持 `prefetch_stream.wait_stream(current_stream)`，reduce 在 prefetch stream 上异步提交。 +- 修正策略：MLA/index 的 `start_next_layer_prefix()` 先取得 `current_stream`，把它传给 prefetch-safe getter，使 L2->L1 ready event 挂到实际读取 KV/index page 的 stream；随后仍保持 stream ordering，由 IPC/CE materialize event 表示 prefix ready。 - 没有新增 collective；没有把 forward stream 改成 CPU 同步等待。 - 单测锁住： - - `test_mla_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream` - - `test_index_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream` + - `test_mla_prefetch_waits_l2_l1_on_materialize_stream` + - `test_index_prefetch_waits_l2_l1_on_materialize_stream` ### P5. TAI SM IPC spans baseline @@ -995,3 +1017,79 @@ SYMM CE prefix + SM current: gpu_p50=0.521ms cpu_p50=0.546ms 2. 当前 PyTorch symmetric-memory prototype 没有优于 CUDA IPC CE；推测原因包括 symmetric memory backend/registration overhead、统一 window size 带来的空间浪费、以及仍使用 cudaMemcpyBatchAsync per-span submit。 3. 短期最有价值路径是继续用 IPC CE 做 prefix prefetch baseline，同时保留 symmetric window helper 作为后续 NCCL CE/NVSHMEM transport 的地址合同实验入口。 4. 生产接入前必须继续做 span coalescing、跨 layer descriptor 复用，以及更大 bs/不同 cached/extend sweep；不能只凭 symmetric window 概念直接替换现有路径。 + +--- + +## 10. FP8/RAGGED 与 index-skip prefetch 修正（2026-06-12） + +### 10.1 远端日志现象 + +最新 prefill 日志中 `index_prefetch` fallback 被刷屏： + +```text +[CP_SHARED_KV_FALLBACK][index_prefetch] + current_missing_prefetcher: 53136 + not_paged_topk: 1296 +``` + +排查结论： + +1. `current_missing_prefetcher` 大多来自无 cache / tiny extend / prefetcher 未创建的正常同步路径，不是运行时错误；consume 侧不应每层每 rank warning。 +2. `not_paged_topk` 是错误判定：FP8 sparse prefill 走 `TopkTransformMethod.RAGGED` 是预期行为，但 index/MLA prefix materialize 仍是 page-slot 语义，不能因为 top-k transform 不是 PAGED 就关闭 prefetch。 + +### 10.2 当前修正 + +1. `CpSharedKVIndexPrefetcher.maybe_create()` 对 RAGGED top-k 不再 fallback，而是继续创建 index prefetcher。 +2. `CpSharedKVMlaPrefetcher.maybe_create()` 对 RAGGED top-k 也继续创建 MLA prefetcher。 +3. `nsa_indexer._maybe_materialize_shared_index_buffer()` 不再对缺失 prefetcher 的正常同步路径打 warning；仅在 prefetcher 已存在但 consume miss 时保留 fallback warning。 +4. `nsa_backend.py` 的 FP8/RAGGED cache-hit compose 现在先尝试 `mla_prefetcher.consume_prefix_with_current()`，只有 prefetcher 存在但 miss 时才 warning 并回到同步 `materialize_prefix_and_reuse_current_kv_page_slots()`。 +5. attention-window 启动 index prefetch 时检查 `token_to_kv_pool.index_layer_to_slot`，index-skip/shared 层没有 index-cache slot 时只跳过 index prefetch，不影响 MLA prefetch。 + +### 10.3 验证 + +远端容器 `cjy-glm5-new` 已验证： + +```text +python -m py_compile \ + python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py \ + python/sglang/srt/layers/attention/nsa_backend.py + +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_fp8_ragged_mla_uses_page_slot_current_compose_for_cache_hit \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_mla_prefetch_create_allows_ragged_topk_transform \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVRuntimeHelpers::test_index_prefetch_create_allows_ragged_topk_transform \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py::TestCpSharedKVTaiMaterializeIntegration::test_attention_window_prefetch_skips_inactive_index_cache_layer +结果：4 passed +``` + +### 10.4 剩余风险 + +1. 旧进程日志不会反映本次降噪和 RAGGED prefetch 修正；需要重启 prefill 后再观察 warning 频率。 +2. RAGGED 当前层 materialize 的历史代码位置仍在 flashmla_sparse 分支内部；后续如果严格追求“当前层 materialize 完成后再启动下一层 prefetch”，需要把 RAGGED current compose 前移到 attention-window hook 之前。 +3. 本次只修正 create/consume 合同和单测；还需要 ETE 确认 FP8 cache-hit 下 index/MLA prefetch 命中率和精度不回退。 + +### 10.5 预期 skip 不再 warning（2026-06-12） + +远端新进程仍出现大量 fallback warning 后，继续按日志聚合定位： + +```text +[CP_SHARED_KV_FALLBACK][index_prefetch] reason=current_consume_miss 主要集中在 layer=0 +[CP_SHARED_KV_FALLBACK][mla_prefetch] reason=ragged_current_consume_miss 主要集中在 layer=0 +[CP_SHARED_KV_FALLBACK][index_prefetch|mla_prefetch] reason=prefix_not_page_aligned prefix_lens=[0] +``` + +结论：这些都不是异常 fallback： + +1. `layer=0` 没有上一层 attention-window 能提前发起本层 prefetch，consume miss 是预期同步路径。 +2. `prefix_lens=[0]` 表示 request 本身没有命中 cache，prefix prefetch 不应创建。 +3. 真实需要 warning 的场景应收窄为：非零 prefix 不 page-aligned、prefetcher 已创建且 `layer > start_layer` 仍 consume miss、metadata/slot 合同异常等。 + +修正合同： + +1. MLA/index prefetch create 遇到全零 prefix 时只走 debug skip，不再打 `[CP_SHARED_KV_FALLBACK]` warning。 +2. index partial-current prefetcher 在 `layer == token_to_kv_pool.start_layer` miss 时不 warning。 +3. FP8/RAGGED MLA prefetcher 在 `layer == start_layer` miss 时不 warning。 + +这类 skip 仍保留 debug 级别路径，方便需要时通过 debug 环境变量观察，但默认生产日志不应被无 cache/tiny/首层预期路径污染。 + +验证：远端 `cjy-glm5-new` 已通过 targeted regression（3 passed）以及完整 `test_cp_shared_kv_runtime.py`（156 passed, 21 warnings, 2 subtests passed）。 diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 668b50ce9..e3cfa8da2 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -8,8 +8,6 @@ from typing import Any, Optional import torch from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( - _all_reduce_materialized_buffer_async, - _all_reduce_materialized_buffer_range, build_batch_current_slot_spans, build_batch_prefix_slot_spans, cp_shared_kv_debug_enabled, @@ -34,13 +32,14 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( slot_range_to_token_slice, _raise_tai_ipc_materialize_required, _should_fail_fast_tai_ipc_materialize, + _try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into, + _try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into, _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into, _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into, _try_tai_ipc_materialize_paged_buffer_page_slot_spans_into, _try_tai_ipc_materialize_token_kv_page_slot_spans_into, ) from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split -from sglang.srt.layers.dp_attention import get_attention_cp_group from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout logger = logging.getLogger(__name__) @@ -215,32 +214,6 @@ def _materialize_local_paged_buffer_page_slot_spans_into( ) -def _all_reduce_materialized_buffer_ranges_async( - *, - dense_kv_cache: torch.Tensor, - row_slices: list[slice], - cp_size: int, - stream: torch.cuda.Stream, - nvtx_source: str, - nvtx_layer_id: int, - nvtx_cp_rank: int, -) -> torch.cuda.Event | None: - event: torch.cuda.Event | None = None - for rows in row_slices: - event = _all_reduce_materialized_buffer_async( - dense_kv_cache[rows], - cp_size=cp_size, - stream=stream, - nvtx_source=nvtx_source, - nvtx_layer_id=nvtx_layer_id, - nvtx_cp_rank=nvtx_cp_rank, - nvtx_rows=(rows.start, rows.stop), - ) - if event is None: - return None - return event - - def _record_event_on_stream(stream: torch.cuda.Stream) -> torch.cuda.Event: event = torch.cuda.Event() event.record(stream) @@ -277,7 +250,7 @@ class _PrefetchCpuTiming: start_max_ms: float = 0.0 get_total_ms: float = 0.0 materialize_total_ms: float = 0.0 - reduce_enqueue_total_ms: float = 0.0 + materialize_enqueue_total_ms: float = 0.0 consume_count: int = 0 consume_total_ms: float = 0.0 consume_wait_total_ms: float = 0.0 @@ -290,7 +263,7 @@ class _PrefetchCpuTiming: total_ms: float, get_ms: float, materialize_ms: float, - reduce_enqueue_ms: float, + materialize_enqueue_ms: float, ) -> None: if total_ms < 0.0: return @@ -299,7 +272,7 @@ class _PrefetchCpuTiming: self.start_max_ms = max(self.start_max_ms, total_ms) self.get_total_ms += max(get_ms, 0.0) self.materialize_total_ms += max(materialize_ms, 0.0) - self.reduce_enqueue_total_ms += max(reduce_enqueue_ms, 0.0) + self.materialize_enqueue_total_ms += max(materialize_enqueue_ms, 0.0) def record_consume( self, @@ -328,7 +301,7 @@ def _log_prefetch_cpu_start( total_ms: float, get_ms: float, materialize_ms: float, - reduce_enqueue_ms: float, + materialize_enqueue_ms: float, prefix_pages: int, total_slots: int, dense_units: int, @@ -338,14 +311,14 @@ def _log_prefetch_cpu_start( if cp_shared_kv_mla_prefetch_should_log_layer(layer_id): log_fn( "cpu_timing path=%s stage=start layer=%s total_ms=%.3f " - "get_ms=%.3f materialize_ms=%.3f reduce_enqueue_ms=%.3f " + "get_ms=%.3f materialize_ms=%.3f materialize_enqueue_ms=%.3f " "prefix_pages=%s total_slots=%s dense_units=%s", path, layer_id, total_ms, get_ms, materialize_ms, - reduce_enqueue_ms, + materialize_enqueue_ms, prefix_pages, total_slots, dense_units, @@ -354,14 +327,14 @@ def _log_prefetch_cpu_start( starts = timing.start_count log_fn( "cpu_summary path=%s starts=%s avg_start_ms=%.3f max_start_ms=%.3f " - "avg_get_ms=%.3f avg_materialize_ms=%.3f avg_reduce_enqueue_ms=%.3f", + "avg_get_ms=%.3f avg_materialize_ms=%.3f avg_materialize_enqueue_ms=%.3f", path, starts, timing.start_total_ms / starts, timing.start_max_ms, timing.get_total_ms / starts, timing.materialize_total_ms / starts, - timing.reduce_enqueue_total_ms / starts, + timing.materialize_enqueue_total_ms / starts, ) @@ -453,6 +426,7 @@ class CpSharedKVMlaPrefetcher: slot_sorted_logical_pages_by_row: torch.Tensor | None = None, slot_sorted_dense_pages_by_row: torch.Tensor | None = None, dense_num_pages: int, + slot_remap: Any | None = None, owned_prefix_pages: int = -1, owned_total_pages: int = -1, stream: Optional[torch.cuda.Stream] = None, @@ -465,6 +439,7 @@ class CpSharedKVMlaPrefetcher: self.slot_sorted_logical_pages_by_row = slot_sorted_logical_pages_by_row self.slot_sorted_dense_pages_by_row = slot_sorted_dense_pages_by_row self.dense_num_pages = dense_num_pages + self.slot_remap = slot_remap self.owned_prefix_pages = owned_prefix_pages self.owned_total_pages = owned_total_pages self.total_slots = int(slot_logical_pages.numel()) @@ -519,8 +494,11 @@ class CpSharedKVMlaPrefetcher: _prefetch_log("create_skip reason=not_in_seq_split") return None if not topk_transform_is_paged: - _prefetch_log("create_skip reason=not_paged_topk") - return None + # FP8 sparse prefill uses the RAGGED top-k transform by design. + # MLA prefix prefetch materializes page-slot KV rows before the + # final top-k offset semantics are consumed, so the RAGGED transform + # is still a valid prefetch target. + _prefetch_log("create_continue reason=ragged_topk") batch_size = int(getattr(forward_batch, "batch_size", 0)) if batch_size <= 0: _prefetch_log( @@ -553,11 +531,17 @@ class CpSharedKVMlaPrefetcher: for prefix_len in prefix_lens if prefix_len < 0 or prefix_len % page_size != 0 ] - if bad_prefix_lens or not any(prefix_len > 0 for prefix_len in prefix_lens): + if not any(prefix_len > 0 for prefix_len in prefix_lens): + _prefetch_log( + "create_skip reason=no_cache_prefix prefix_lens=%s page_size=%s", + prefix_lens, + page_size, + ) + return None + if bad_prefix_lens: _mla_prefetch_fallback_log( "prefix_not_page_aligned", - "prefix length is zero or not page-aligned. " - "prefix_lens=%s page_size=%s", + "non-zero prefix length is not page-aligned. prefix_lens=%s page_size=%s", prefix_lens, page_size, ) @@ -654,15 +638,6 @@ class CpSharedKVMlaPrefetcher: ) return None - cp_group = get_attention_cp_group() - if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: - _prefetch_log( - "create_skip reason=missing_pynccl cp_rank=%s cp_size=%s", - layout.cp_rank, - layout.cp_size, - ) - return None - prefetch_stream = stream if stream is not None else torch.cuda.Stream() create_cpu = _cpu_timing_start() get_cpu = _cpu_timing_start() @@ -738,6 +713,7 @@ class CpSharedKVMlaPrefetcher: slot_sorted_logical_pages_by_row=remap.slot_sorted_logical_pages_by_row, slot_sorted_dense_pages_by_row=remap.slot_sorted_dense_pages_by_row, dense_num_pages=remap.dense_num_pages, + slot_remap=remap, owned_prefix_pages=owned_prefix_pages, owned_total_pages=owned_total_pages, stream=prefetch_stream, @@ -779,12 +755,12 @@ class CpSharedKVMlaPrefetcher: self._log_layer(layer_id, "consume_miss layer=%s", layer_id) return None if handle.event is None: - self.launch_pending_reduce() + self.finalize_pending_materialize() handle = self.handles.get(layer_id) if handle is None or handle.event is None: self._log_layer( layer_id, - "consume_miss reason=prefix_reduce_not_ready layer=%s", + "consume_miss reason=prefix_materialize_not_ready layer=%s", layer_id, ) return None @@ -830,7 +806,11 @@ class CpSharedKVMlaPrefetcher: spans=suffix_spans, ) ) - if not materialized_suffix_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_kv_cache): + if ( + not materialized_suffix_by_ipc + and suffix_spans + and self.layout.cp_size > 1 + ): _raise_tai_ipc_materialize_required( "mla_prefetch_suffix_ipc_unavailable", cp_rank=self.layout.cp_rank, @@ -847,19 +827,6 @@ class CpSharedKVMlaPrefetcher: page_size=self.page_size, spans=suffix_spans, ) - if not materialized_suffix_by_ipc: - for suffix_rows in _slot_spans_to_token_slices( - self.page_size, suffix_spans - ): - _all_reduce_materialized_buffer_range( - dense_kv_cache, - self.layout.cp_size, - suffix_rows.start, - suffix_rows.stop, - nvtx_source="mla.consume_suffix", - nvtx_layer_id=layer_id, - nvtx_cp_rank=self.layout.cp_rank, - ) self._log_layer( layer_id, "consume_suffix_done layer=%s suffix_spans=%s", @@ -929,7 +896,7 @@ class CpSharedKVMlaPrefetcher: """Consume the prefetched prefix and append current-layer KV rows. This is the partial-current-reuse variant of :meth:`consume`: prefix - pages are already materialized/reduced by the prefetch stream, while + pages are already materialized by the prefetch stream, while current/suffix pages are not copied from the shared pool. Current locs in ``logical_locs`` are remapped to the appended current KV rows. """ @@ -954,12 +921,12 @@ class CpSharedKVMlaPrefetcher: self._log_layer(layer_id, "consume_prefix_current_miss layer=%s", layer_id) return None if handle.event is None: - self.launch_pending_reduce() + self.finalize_pending_materialize() handle = self.handles.get(layer_id) if handle is None or handle.event is None: self._log_layer( layer_id, - "consume_prefix_current_miss reason=prefix_reduce_not_ready layer=%s", + "consume_prefix_current_miss reason=prefix_materialize_not_ready layer=%s", layer_id, ) return None @@ -1022,7 +989,7 @@ class CpSharedKVMlaPrefetcher: if ( not current_materialized_by_ipc and self.current_slot_spans - and _should_fail_fast_tai_ipc_materialize(mixed_kv_cache) + and self.layout.cp_size > 1 ): _raise_tai_ipc_materialize_required( "mla_prefetch_current_ipc_unavailable", @@ -1031,19 +998,6 @@ class CpSharedKVMlaPrefetcher: spans=self.current_slot_spans, dense_shape=tuple(mixed_kv_cache.shape), ) - if not current_materialized_by_ipc: - for current_rows in _slot_spans_to_token_slices( - self.page_size, self.current_slot_spans - ): - _all_reduce_materialized_buffer_range( - mixed_kv_cache, - self.layout.cp_size, - current_rows.start, - current_rows.stop, - nvtx_source="mla.prefetch_current", - nvtx_layer_id=layer_id, - nvtx_cp_rank=self.layout.cp_rank, - ) remap_ms = _cpu_timing_ms(remap_cpu) total_ms = _cpu_timing_ms(consume_cpu) self._log_layer( @@ -1106,25 +1060,7 @@ class CpSharedKVMlaPrefetcher: start_cpu = _cpu_timing_start() current_stream = torch.cuda.current_stream() get_cpu = _cpu_timing_start() - try: - kv_cache = _prefetch_pool_get_key_buffer( - token_to_kv_pool=token_to_kv_pool, - layer_id=next_layer_id, - stream=current_stream, - path="mla", - ) - get_ms = _cpu_timing_ms(get_cpu) - except Exception: - logger.exception( - "Failed to get next-layer KV cache for CP shared KV MLA prefetch." - ) - self.disabled = True - self._log_next_layer( - next_layer_id, - "start_disable reason=get_kv_failed next_layer=%s", - next_layer_id, - ) - return + self.stream.wait_stream(current_stream) try: prefix_row_spans = _slot_spans_to_token_slices( @@ -1132,64 +1068,81 @@ class CpSharedKVMlaPrefetcher: ) prefix_rows = prefix_row_spans[0] if prefix_row_spans else slice(0, 0) materialize_cpu = _cpu_timing_start() - dense_kv_cache = kv_cache.new_zeros( - (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) - ) - self._log_next_layer( - next_layer_id, - "start_prefix_begin next_layer=%s prefix_slot_spans=%s " - "dense_rows=%s", - next_layer_id, - self.prefix_slot_spans, - int(dense_kv_cache.shape[0]), - ) - materialized_by_ipc = _try_tai_ipc_materialize_token_kv_page_slot_spans_into( - kv_cache=kv_cache, - dense_kv_cache=dense_kv_cache, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - page_size=self.page_size, - spans=self.prefix_slot_spans, - ) - if not materialized_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_kv_cache): - _raise_tai_ipc_materialize_required( - "mla_prefetch_prefix_ipc_unavailable", - cp_rank=self.layout.cp_rank, - cp_size=self.layout.cp_size, - spans=self.prefix_slot_spans, - dense_shape=tuple(dense_kv_cache.shape), + with torch.cuda.stream(self.stream): + kv_cache = _prefetch_pool_get_key_buffer( + token_to_kv_pool=token_to_kv_pool, + layer_id=next_layer_id, + stream=self.stream, + path="mla", ) - if not materialized_by_ipc: - _materialize_local_token_kv_page_slot_spans_into( - kv_cache=kv_cache, - dense_kv_cache=dense_kv_cache, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - page_size=self.page_size, - spans=self.prefix_slot_spans, + get_ms = _cpu_timing_ms(get_cpu) + dense_kv_cache = kv_cache.new_zeros( + (self.dense_num_pages * self.page_size, *kv_cache.shape[1:]) ) - materialize_ms = _cpu_timing_ms(materialize_cpu) - reduce_cpu = _cpu_timing_start() - if materialized_by_ipc: - event = _record_event_on_stream(current_stream) - else: - self.stream.wait_stream(current_stream) - with torch.cuda.stream(self.stream): - event = _all_reduce_materialized_buffer_ranges_async( + self._log_next_layer( + next_layer_id, + "start_prefix_begin next_layer=%s prefix_slot_spans=%s " + "dense_rows=%s stream=prefetch", + next_layer_id, + self.prefix_slot_spans, + int(dense_kv_cache.shape[0]), + ) + materialized_by_ce = ( + _try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into( + kv_cache=kv_cache, dense_kv_cache=dense_kv_cache, - row_slices=prefix_row_spans, - cp_size=self.layout.cp_size, - stream=self.stream, - nvtx_source="mla.prefetch_prefix", - nvtx_layer_id=next_layer_id, - nvtx_cp_rank=self.layout.cp_rank, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + spans=self.prefix_slot_spans, + slot_remap=self.slot_remap, ) - reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu) + ) + event = _record_event_on_stream(self.stream) if materialized_by_ce else None + + materialized_by_ipc = materialized_by_ce + if not materialized_by_ce: + current_stream.wait_stream(self.stream) + materialized_by_ipc = ( + _try_tai_ipc_materialize_token_kv_page_slot_spans_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + spans=self.prefix_slot_spans, + slot_remap=self.slot_remap, + ) + ) + if ( + not materialized_by_ipc + and self.prefix_slot_spans + and self.layout.cp_size > 1 + ): + _raise_tai_ipc_materialize_required( + "mla_prefetch_prefix_ipc_unavailable", + cp_rank=self.layout.cp_rank, + cp_size=self.layout.cp_size, + spans=self.prefix_slot_spans, + dense_shape=tuple(dense_kv_cache.shape), + ) + if not materialized_by_ipc: + _materialize_local_token_kv_page_slot_spans_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv_cache, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + page_size=self.page_size, + spans=self.prefix_slot_spans, + ) + event = _record_event_on_stream(current_stream) + materialize_ms = _cpu_timing_ms(materialize_cpu) + materialize_enqueue_ms = 0.0 if event is None: self.disabled = True logger.warning( "[CP_SHARED_KV_FALLBACK][mla_prefetch] " - "reason=async_reduce_unavailable layer_id=%s cp_rank=%s " + "reason=materialize_event_unavailable layer_id=%s cp_rank=%s " "cp_size=%s prefix_pages=%s total_slots=%s", next_layer_id, self.layout.cp_rank, @@ -1199,7 +1152,7 @@ class CpSharedKVMlaPrefetcher: ) self._log_next_layer( next_layer_id, - "start_disable reason=async_reduce_unavailable next_layer=%s", + "start_disable reason=materialize_event_unavailable next_layer=%s", next_layer_id, ) return @@ -1208,7 +1161,7 @@ class CpSharedKVMlaPrefetcher: total_ms=total_ms, get_ms=get_ms, materialize_ms=materialize_ms, - reduce_enqueue_ms=reduce_enqueue_ms, + materialize_enqueue_ms=materialize_enqueue_ms, ) _log_prefetch_cpu_start( log_fn=self._log, @@ -1219,7 +1172,7 @@ class CpSharedKVMlaPrefetcher: total_ms=total_ms, get_ms=get_ms, materialize_ms=materialize_ms, - reduce_enqueue_ms=reduce_enqueue_ms, + materialize_enqueue_ms=materialize_enqueue_ms, prefix_pages=self.prefix_pages, total_slots=self.total_slots, dense_units=int(dense_kv_cache.shape[0]), @@ -1246,74 +1199,44 @@ class CpSharedKVMlaPrefetcher: self._log_next_layer( next_layer_id, "start next_layer=%s prefix_pages=%s prefix_rows=%s dense_rows=%s " - "reduce_enqueued=True", + "materialize_enqueued=True", next_layer_id, self.prefix_pages, prefix_rows.stop - prefix_rows.start, int(dense_kv_cache.shape[0]), ) - def launch_pending_reduce(self) -> None: + def finalize_pending_materialize(self) -> None: handle = self.pending_attention_handle if handle is None: return if handle.event is not None: self._log_next_layer( handle.layer_id, - "start_prefix_reduce_skip reason=already_enqueued next_layer=%s", + "start_prefix_materialize_skip reason=already_enqueued next_layer=%s", handle.layer_id, ) return - prefix_row_spans = list(handle.prefix_row_spans or (handle.prefix_rows,)) try: - if _should_fail_fast_tai_ipc_materialize(handle.dense_kv_cache): - _raise_tai_ipc_materialize_required( - "mla_prefetch_deferred_prefix_ipc_unavailable", - cp_rank=self.layout.cp_rank, - cp_size=self.layout.cp_size, - dense_shape=tuple(handle.dense_kv_cache.shape), - layer_id=handle.layer_id, - ) - current_stream = torch.cuda.current_stream() - self.stream.wait_stream(current_stream) - with torch.cuda.stream(self.stream): - event = _all_reduce_materialized_buffer_ranges_async( - dense_kv_cache=handle.dense_kv_cache, - row_slices=prefix_row_spans, - cp_size=self.layout.cp_size, - stream=self.stream, - nvtx_source="mla.prefetch_prefix", - nvtx_layer_id=handle.layer_id, - nvtx_cp_rank=self.layout.cp_rank, - ) - if event is None: - self.disabled = True - self.handles.pop(handle.layer_id, None) - if self.pending_attention_handle is handle: - self.pending_attention_handle = None - self._log_next_layer( - handle.layer_id, - "start_disable reason=async_reduce_unavailable next_layer=%s", - handle.layer_id, - ) - return - handle.event = event - self._log_next_layer( - handle.layer_id, - "start_prefix_reduce_enqueued next_layer=%s row_spans=%s", - handle.layer_id, - [(rows.start, rows.stop) for rows in prefix_row_spans], - ) + _raise_tai_ipc_materialize_required( + "mla_prefetch_deferred_prefix_materialize_missing", + cp_rank=self.layout.cp_rank, + cp_size=self.layout.cp_size, + dense_shape=tuple(handle.dense_kv_cache.shape), + layer_id=handle.layer_id, + ) except Exception: - logger.exception("Failed to launch CP shared KV MLA prefix prefetch reduce.") + logger.exception( + "CP shared KV MLA prefix prefetch has no IPC/CE completion event." + ) self.disabled = True self.handles.pop(handle.layer_id, None) if self.pending_attention_handle is handle: self.pending_attention_handle = None self._log_next_layer( handle.layer_id, - "start_disable reason=reduce_launch_exception next_layer=%s", + "start_disable reason=materialize_event_missing next_layer=%s", handle.layer_id, ) return @@ -1366,6 +1289,7 @@ class CpSharedKVIndexPrefetcher: slot_sorted_logical_pages_by_row: torch.Tensor | None = None, slot_sorted_dense_pages_by_row: torch.Tensor | None = None, dense_num_pages: int, + slot_remap: Any | None = None, owned_prefix_pages: int = -1, owned_total_pages: int = -1, stream: Optional[torch.cuda.Stream] = None, @@ -1377,6 +1301,7 @@ class CpSharedKVIndexPrefetcher: self.slot_sorted_logical_pages_by_row = slot_sorted_logical_pages_by_row self.slot_sorted_dense_pages_by_row = slot_sorted_dense_pages_by_row self.dense_num_pages = dense_num_pages + self.slot_remap = slot_remap self.owned_prefix_pages = owned_prefix_pages self.owned_total_pages = owned_total_pages self.total_slots = int(slot_logical_pages.numel()) @@ -1448,11 +1373,11 @@ class CpSharedKVIndexPrefetcher: ) return None if not topk_transform_is_paged: - _index_prefetch_fallback_log( - "not_paged_topk", - "topk transform is not PAGED.", - ) - return None + # FP8 prefill uses the RAGGED top-k transform by design. Index + # prefix materialization is still page-slot based and remains a + # valid prefetch target; the transform method only affects the + # final top-k offset semantics after MQA logits are computed. + _prefetch_log("index_create_continue reason=ragged_topk") batch_size = int(getattr(forward_batch, "batch_size", 0)) if batch_size <= 0: _index_prefetch_fallback_log( @@ -1501,10 +1426,17 @@ class CpSharedKVIndexPrefetcher: for prefix_len in prefix_lens if prefix_len < 0 or prefix_len % page_size != 0 ] - if bad_prefix_lens or not any(prefix_len > 0 for prefix_len in prefix_lens): + if not any(prefix_len > 0 for prefix_len in prefix_lens): + _prefetch_log( + "index_create_skip reason=no_cache_prefix prefix_lens=%s page_size=%s", + prefix_lens, + page_size, + ) + return None + if bad_prefix_lens: _index_prefetch_fallback_log( "prefix_not_page_aligned", - "prefix length is zero or not page-aligned. prefix_lens=%s page_size=%s", + "non-zero prefix length is not page-aligned. prefix_lens=%s page_size=%s", prefix_lens, page_size, ) @@ -1605,16 +1537,6 @@ class CpSharedKVIndexPrefetcher: ) return None - cp_group = get_attention_cp_group() - if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: - _index_prefetch_fallback_log( - "missing_pynccl", - "pynccl communicator is unavailable. cp_rank=%s cp_size=%s", - layout.cp_rank, - layout.cp_size, - ) - return None - prefetch_stream = stream if stream is not None else torch.cuda.Stream() create_cpu = _cpu_timing_start() get_cpu = _cpu_timing_start() @@ -1690,6 +1612,7 @@ class CpSharedKVIndexPrefetcher: slot_sorted_logical_pages_by_row=remap.slot_sorted_logical_pages_by_row, slot_sorted_dense_pages_by_row=remap.slot_sorted_dense_pages_by_row, dense_num_pages=remap.dense_num_pages, + slot_remap=remap, owned_prefix_pages=owned_prefix_pages, owned_total_pages=owned_total_pages, stream=prefetch_stream, @@ -1730,12 +1653,12 @@ class CpSharedKVIndexPrefetcher: self._log_layer(layer_id, "index_consume_miss layer=%s", layer_id) return None if handle.event is None: - self.launch_pending_reduce() + self.finalize_pending_materialize() handle = self.handles.get(layer_id) if handle is None or handle.event is None: self._log_layer( layer_id, - "index_consume_miss reason=prefix_reduce_not_ready layer=%s", + "index_consume_miss reason=prefix_materialize_not_ready layer=%s", layer_id, ) return None @@ -1780,7 +1703,11 @@ class CpSharedKVIndexPrefetcher: spans=suffix_spans, ) ) - if not materialized_suffix_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_page_buffer): + if ( + not materialized_suffix_by_ipc + and suffix_spans + and self.layout.cp_size > 1 + ): _raise_tai_ipc_materialize_required( "index_prefetch_suffix_ipc_unavailable", cp_rank=self.layout.cp_rank, @@ -1796,17 +1723,6 @@ class CpSharedKVIndexPrefetcher: layout=self.layout, spans=suffix_spans, ) - if not materialized_suffix_by_ipc: - for suffix_rows in _slot_spans_to_page_slices(suffix_spans): - _all_reduce_materialized_buffer_range( - dense_page_buffer, - self.layout.cp_size, - suffix_rows.start, - suffix_rows.stop, - nvtx_source="index.consume_suffix", - nvtx_layer_id=layer_id, - nvtx_cp_rank=self.layout.cp_rank, - ) self._log_layer( layer_id, "index_consume_suffix_done layer=%s suffix_spans=%s", @@ -1895,13 +1811,13 @@ class CpSharedKVIndexPrefetcher: ) return None if handle.event is None: - self.launch_pending_reduce() + self.finalize_pending_materialize() handle = self.handles.get(layer_id) if handle is None or handle.event is None: self._log_layer( layer_id, "index_consume_prefix_current_miss " - "reason=prefix_reduce_not_ready layer=%s", + "reason=prefix_materialize_not_ready layer=%s", layer_id, ) return None @@ -1955,7 +1871,7 @@ class CpSharedKVIndexPrefetcher: if ( not current_materialized_by_ipc and self.current_slot_spans - and _should_fail_fast_tai_ipc_materialize(dense_page_buffer) + and self.layout.cp_size > 1 ): _raise_tai_ipc_materialize_required( "index_prefetch_current_ipc_unavailable", @@ -1964,17 +1880,6 @@ class CpSharedKVIndexPrefetcher: spans=self.current_slot_spans, dense_shape=tuple(dense_page_buffer.shape), ) - if not current_materialized_by_ipc: - for current_pages in _slot_spans_to_page_slices(self.current_slot_spans): - _all_reduce_materialized_buffer_range( - dense_page_buffer, - self.layout.cp_size, - current_pages.start, - current_pages.stop, - nvtx_source="index.prefetch_current", - nvtx_layer_id=layer_id, - nvtx_cp_rank=self.layout.cp_rank, - ) remap_ms = _cpu_timing_ms(remap_cpu) total_ms = _cpu_timing_ms(consume_cpu) self._log_layer( @@ -2039,87 +1944,83 @@ class CpSharedKVIndexPrefetcher: start_cpu = _cpu_timing_start() current_stream = torch.cuda.current_stream() get_cpu = _cpu_timing_start() - try: - page_buffer = _prefetch_pool_get_index_buffer( - token_to_kv_pool=token_to_kv_pool, - layer_id=next_layer_id, - stream=current_stream, - ) - get_ms = _cpu_timing_ms(get_cpu) - except Exception: - logger.exception( - "Failed to get next-layer index buffer for CP shared KV index prefetch." - ) - self.disabled = True - self._log_next_layer( - next_layer_id, - "index_start_disable reason=get_index_buffer_failed next_layer=%s", - next_layer_id, - ) - return + self.stream.wait_stream(current_stream) try: prefix_row_spans = _slot_spans_to_page_slices(self.prefix_slot_spans) prefix_rows = prefix_row_spans[0] if prefix_row_spans else slice(0, 0) materialize_cpu = _cpu_timing_start() - dense_page_buffer = page_buffer.new_zeros( - (self.dense_num_pages, *page_buffer.shape[1:]) - ) - self._log_next_layer( - next_layer_id, - "index_start_prefix_begin next_layer=%s prefix_slot_spans=%s " - "dense_pages=%s", - next_layer_id, - self.prefix_slot_spans, - int(dense_page_buffer.shape[0]), - ) - materialized_by_ipc = ( - _try_tai_ipc_materialize_paged_buffer_page_slot_spans_into( - page_buffer=page_buffer, - dense_page_buffer=dense_page_buffer, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - spans=self.prefix_slot_spans, + with torch.cuda.stream(self.stream): + page_buffer = _prefetch_pool_get_index_buffer( + token_to_kv_pool=token_to_kv_pool, + layer_id=next_layer_id, + stream=self.stream, ) - ) - if not materialized_by_ipc and _should_fail_fast_tai_ipc_materialize(dense_page_buffer): - _raise_tai_ipc_materialize_required( - "index_prefetch_prefix_ipc_unavailable", - cp_rank=self.layout.cp_rank, - cp_size=self.layout.cp_size, - spans=self.prefix_slot_spans, - dense_shape=tuple(dense_page_buffer.shape), + get_ms = _cpu_timing_ms(get_cpu) + dense_page_buffer = page_buffer.new_zeros( + (self.dense_num_pages, *page_buffer.shape[1:]) ) - if not materialized_by_ipc: - _materialize_local_paged_buffer_page_slot_spans_into( - page_buffer=page_buffer, - dense_page_buffer=dense_page_buffer, - slot_logical_pages=self.slot_logical_pages, - layout=self.layout, - spans=self.prefix_slot_spans, + self._log_next_layer( + next_layer_id, + "index_start_prefix_begin next_layer=%s prefix_slot_spans=%s " + "dense_pages=%s stream=prefetch", + next_layer_id, + self.prefix_slot_spans, + int(dense_page_buffer.shape[0]), ) - materialize_ms = _cpu_timing_ms(materialize_cpu) - reduce_cpu = _cpu_timing_start() - if materialized_by_ipc: - event = _record_event_on_stream(current_stream) - else: - self.stream.wait_stream(current_stream) - with torch.cuda.stream(self.stream): - event = _all_reduce_materialized_buffer_ranges_async( - dense_kv_cache=dense_page_buffer, - row_slices=prefix_row_spans, - cp_size=self.layout.cp_size, - stream=self.stream, - nvtx_source="index.prefetch_prefix", - nvtx_layer_id=next_layer_id, - nvtx_cp_rank=self.layout.cp_rank, + materialized_by_ce = ( + _try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + spans=self.prefix_slot_spans, + slot_remap=self.slot_remap, ) - reduce_enqueue_ms = _cpu_timing_ms(reduce_cpu) + ) + event = _record_event_on_stream(self.stream) if materialized_by_ce else None + + materialized_by_ipc = materialized_by_ce + if not materialized_by_ce: + current_stream.wait_stream(self.stream) + materialized_by_ipc = ( + _try_tai_ipc_materialize_paged_buffer_page_slot_spans_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + spans=self.prefix_slot_spans, + slot_remap=self.slot_remap, + ) + ) + if ( + not materialized_by_ipc + and self.prefix_slot_spans + and self.layout.cp_size > 1 + ): + _raise_tai_ipc_materialize_required( + "index_prefetch_prefix_ipc_unavailable", + cp_rank=self.layout.cp_rank, + cp_size=self.layout.cp_size, + spans=self.prefix_slot_spans, + dense_shape=tuple(dense_page_buffer.shape), + ) + if not materialized_by_ipc: + _materialize_local_paged_buffer_page_slot_spans_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + spans=self.prefix_slot_spans, + ) + event = _record_event_on_stream(current_stream) + materialize_ms = _cpu_timing_ms(materialize_cpu) + materialize_enqueue_ms = 0.0 if event is None: self.disabled = True _index_prefetch_fallback_log( - "async_reduce_unavailable", - "async reduce unavailable. layer_id=%s cp_rank=%s cp_size=%s " + "materialize_event_unavailable", + "materialize event unavailable. layer_id=%s cp_rank=%s cp_size=%s " "prefix_pages=%s total_slots=%s", next_layer_id, self.layout.cp_rank, @@ -2129,7 +2030,7 @@ class CpSharedKVIndexPrefetcher: ) self._log_next_layer( next_layer_id, - "index_start_disable reason=async_reduce_unavailable next_layer=%s", + "index_start_disable reason=materialize_event_unavailable next_layer=%s", next_layer_id, ) return @@ -2138,7 +2039,7 @@ class CpSharedKVIndexPrefetcher: total_ms=total_ms, get_ms=get_ms, materialize_ms=materialize_ms, - reduce_enqueue_ms=reduce_enqueue_ms, + materialize_enqueue_ms=materialize_enqueue_ms, ) _log_prefetch_cpu_start( log_fn=self._log, @@ -2149,7 +2050,7 @@ class CpSharedKVIndexPrefetcher: total_ms=total_ms, get_ms=get_ms, materialize_ms=materialize_ms, - reduce_enqueue_ms=reduce_enqueue_ms, + materialize_enqueue_ms=materialize_enqueue_ms, prefix_pages=self.prefix_pages, total_slots=self.total_slots, dense_units=int(dense_page_buffer.shape[0]), @@ -2176,67 +2077,35 @@ class CpSharedKVIndexPrefetcher: self._log_next_layer( next_layer_id, "index_start next_layer=%s prefix_pages=%s dense_pages=%s " - "reduce_enqueued=True", + "materialize_enqueued=True", next_layer_id, self.prefix_pages, int(dense_page_buffer.shape[0]), ) - def launch_pending_reduce(self) -> None: + def finalize_pending_materialize(self) -> None: handle = self.pending_attention_handle if handle is None: return if handle.event is not None: self._log_next_layer( handle.layer_id, - "index_start_prefix_reduce_skip reason=already_enqueued next_layer=%s", + "index_start_prefix_materialize_skip reason=already_enqueued next_layer=%s", handle.layer_id, ) return - prefix_row_spans = list(handle.prefix_row_spans or (handle.prefix_rows,)) try: - if _should_fail_fast_tai_ipc_materialize(handle.dense_page_buffer): - _raise_tai_ipc_materialize_required( - "index_prefetch_deferred_prefix_ipc_unavailable", - cp_rank=self.layout.cp_rank, - cp_size=self.layout.cp_size, - dense_shape=tuple(handle.dense_page_buffer.shape), - layer_id=handle.layer_id, - ) - current_stream = torch.cuda.current_stream() - self.stream.wait_stream(current_stream) - with torch.cuda.stream(self.stream): - event = _all_reduce_materialized_buffer_ranges_async( - dense_kv_cache=handle.dense_page_buffer, - row_slices=prefix_row_spans, - cp_size=self.layout.cp_size, - stream=self.stream, - nvtx_source="index.prefetch_prefix", - nvtx_layer_id=handle.layer_id, - nvtx_cp_rank=self.layout.cp_rank, - ) - if event is None: - self.disabled = True - self.handles.pop(handle.layer_id, None) - if self.pending_attention_handle is handle: - self.pending_attention_handle = None - self._log_next_layer( - handle.layer_id, - "index_start_disable reason=async_reduce_unavailable next_layer=%s", - handle.layer_id, - ) - return - handle.event = event - self._log_next_layer( - handle.layer_id, - "index_start_prefix_reduce_enqueued next_layer=%s row_spans=%s", - handle.layer_id, - [(rows.start, rows.stop) for rows in prefix_row_spans], - ) + _raise_tai_ipc_materialize_required( + "index_prefetch_deferred_prefix_materialize_missing", + cp_rank=self.layout.cp_rank, + cp_size=self.layout.cp_size, + dense_shape=tuple(handle.dense_page_buffer.shape), + layer_id=handle.layer_id, + ) except Exception: logger.exception( - "Failed to launch CP shared KV index prefix prefetch reduce." + "CP shared KV index prefix prefetch has no IPC/CE completion event." ) self.disabled = True self.handles.pop(handle.layer_id, None) @@ -2244,7 +2113,7 @@ class CpSharedKVIndexPrefetcher: self.pending_attention_handle = None self._log_next_layer( handle.layer_id, - "index_start_disable reason=reduce_launch_exception next_layer=%s", + "index_start_disable reason=materialize_event_missing next_layer=%s", handle.layer_id, ) return diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index c1fd9ab74..1008aa3ca 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -2679,6 +2679,253 @@ def _get_or_build_prefix_ipc_slot_descriptors( return descriptors +def _build_prefix_ipc_slot_descriptors_uncached( + *, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + spans: list[tuple[int, int]], + device: torch.device, + physical_page_capacity: int | None, +) -> _IpcPrefixSlotDescriptors: + flat_slot_logical_pages = _contiguous_for_tai( + slot_logical_pages.reshape(-1).to(device=device) + ) + slot_indices = _slot_spans_to_cuda_slot_indices( + spans, + total_slots=int(flat_slot_logical_pages.numel()), + device=device, + ) + if slot_indices.numel() == 0: + empty = torch.empty((0,), dtype=torch.long, device=device) + return _IpcPrefixSlotDescriptors(empty, empty, empty, empty) + slot_logical_pages_range = _contiguous_for_tai( + flat_slot_logical_pages.index_select(0, slot_indices) + ) + owner_ranks, src_page_indices = build_cp_shared_kv_ipc_page_descriptors( + slot_logical_pages_range, + layout, + physical_page_capacity=physical_page_capacity, + ) + return _IpcPrefixSlotDescriptors( + slot_indices=slot_indices.contiguous(), + owner_ranks=owner_ranks, + src_page_indices=src_page_indices, + dense_page_indices=(slot_indices + 1).to(torch.long).contiguous(), + ) + + +def _get_prefix_ipc_slot_descriptors( + *, + slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap | None, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + spans: list[tuple[int, int]], + device: torch.device, + physical_page_capacity: int | None, + cache_kind: str, +) -> _IpcPrefixSlotDescriptors: + if slot_remap is not None: + return _get_or_build_prefix_ipc_slot_descriptors( + slot_remap=slot_remap, + layout=layout, + spans=spans, + device=device, + physical_page_capacity=physical_page_capacity, + cache_kind=cache_kind, + ) + return _build_prefix_ipc_slot_descriptors_uncached( + slot_logical_pages=slot_logical_pages, + layout=layout, + spans=spans, + device=device, + physical_page_capacity=physical_page_capacity, + ) + + +def _filter_ipc_prefix_descriptors_for_ce( + descriptors: _IpcPrefixSlotDescriptors, +) -> _IpcPrefixSlotDescriptors: + """Drop invalid/sentinel slots before CE copy. + + The SM IPC materialize kernel can zero-fill invalid slots. The CE path is + intentionally just cudaMemcpyBatchAsync over valid spans; callers allocate + the dense destination with zeros, so invalid slots should be skipped instead + of submitted to the CE primitive. + """ + + if descriptors.owner_ranks.numel() == 0: + return descriptors + valid = ( + (descriptors.owner_ranks >= 0) + & (descriptors.src_page_indices >= 0) + & (descriptors.dense_page_indices >= 0) + ) + if bool(valid.all().item()): + return descriptors + return _IpcPrefixSlotDescriptors( + slot_indices=descriptors.slot_indices[valid].contiguous(), + owner_ranks=descriptors.owner_ranks[valid].contiguous(), + src_page_indices=descriptors.src_page_indices[valid].contiguous(), + dense_page_indices=descriptors.dense_page_indices[valid].contiguous(), + ) + + +def _try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into( + *, + kv_cache: torch.Tensor, + dense_kv_cache: torch.Tensor, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + page_size: int, + spans: list[tuple[int, int]], + slot_remap: SharedTokenKVSlotRemap | None = None, +) -> bool: + if not spans: + return True + if not dense_kv_cache.is_cuda or not dense_kv_cache.is_contiguous(): + _log_tai_ipc_materialize_fallback( + "ce_dense_token_tensor_unsupported", + "CP shared KV tai IPC CE token span materialize requires a contiguous " + "CUDA dense tensor; falling back to SM IPC or collective. " + "device=%s contiguous=%s dense_shape=%s spans=%s", + dense_kv_cache.device, + dense_kv_cache.is_contiguous(), + tuple(dense_kv_cache.shape), + spans, + limit=4, + ) + return False + + ipc_state = _get_or_open_tai_ipc_peer_ptrs(kv_cache, layout) + if ipc_state is None: + return False + kernels, peer_ptrs = ipc_state + ce_kernel = getattr(kernels, "materialize_cuda_ipc_peer_pages_slot_indices_ce", None) + if ce_kernel is None: + _log_tai_ipc_materialize_fallback( + "ce_kernel_missing", + "CP shared KV tai IPC CE token materialize is unavailable; " + "falling back to SM IPC or collective. Upgrade tai-kernel.", + limit=1, + ) + return False + + try: + descriptors = _get_prefix_ipc_slot_descriptors( + slot_remap=slot_remap, + slot_logical_pages=slot_logical_pages, + layout=layout, + spans=spans, + device=torch.device("cpu"), + physical_page_capacity=kv_cache.shape[0] // page_size, + cache_kind="token-ce", + ) + descriptors = _filter_ipc_prefix_descriptors_for_ce(descriptors) + if descriptors.owner_ranks.numel() == 0: + return True + ce_kernel( + peer_ptrs, + dense_kv_cache, + descriptors.owner_ranks, + descriptors.src_page_indices, + descriptors.dense_page_indices, + page_nbytes=_token_kv_page_nbytes(kv_cache, page_size), + ) + return True + except Exception as exc: + _log_tai_ipc_materialize_fallback( + "token_span_ce_kernel_failed", + "CP shared KV tai IPC CE token span materialize failed; falling back " + "to SM IPC or collective materialize. cp_rank=%s cp_size=%s spans=%s " + "page_size=%s kv_shape=%s dense_shape=%s error=%s", + layout.cp_rank, + layout.cp_size, + spans, + page_size, + tuple(kv_cache.shape), + tuple(dense_kv_cache.shape), + exc, + ) + return False + + +def _try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into( + *, + page_buffer: torch.Tensor, + dense_page_buffer: torch.Tensor, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + spans: list[tuple[int, int]], + slot_remap: SharedPagedBufferSlotRemap | None = None, +) -> bool: + if not spans: + return True + if not dense_page_buffer.is_cuda or not dense_page_buffer.is_contiguous(): + _log_tai_ipc_materialize_fallback( + "ce_dense_paged_tensor_unsupported", + "CP shared KV tai IPC CE paged span materialize requires a contiguous " + "CUDA dense tensor; falling back to SM IPC or collective. " + "device=%s contiguous=%s dense_shape=%s spans=%s", + dense_page_buffer.device, + dense_page_buffer.is_contiguous(), + tuple(dense_page_buffer.shape), + spans, + limit=4, + ) + return False + + ipc_state = _get_or_open_tai_ipc_peer_ptrs(page_buffer, layout) + if ipc_state is None: + return False + kernels, peer_ptrs = ipc_state + ce_kernel = getattr(kernels, "materialize_cuda_ipc_peer_pages_slot_indices_ce", None) + if ce_kernel is None: + _log_tai_ipc_materialize_fallback( + "ce_kernel_missing", + "CP shared KV tai IPC CE paged materialize is unavailable; " + "falling back to SM IPC or collective. Upgrade tai-kernel.", + limit=1, + ) + return False + + try: + descriptors = _get_prefix_ipc_slot_descriptors( + slot_remap=slot_remap, + slot_logical_pages=slot_logical_pages, + layout=layout, + spans=spans, + device=torch.device("cpu"), + physical_page_capacity=page_buffer.shape[0], + cache_kind="paged-ce", + ) + descriptors = _filter_ipc_prefix_descriptors_for_ce(descriptors) + if descriptors.owner_ranks.numel() == 0: + return True + ce_kernel( + peer_ptrs, + dense_page_buffer, + descriptors.owner_ranks, + descriptors.src_page_indices, + descriptors.dense_page_indices, + page_nbytes=_page_nbytes_from_page_tensor(page_buffer), + ) + return True + except Exception as exc: + _log_tai_ipc_materialize_fallback( + "paged_span_ce_kernel_failed", + "CP shared KV tai IPC CE paged span materialize failed; falling back " + "to SM IPC or collective materialize. cp_rank=%s cp_size=%s spans=%s " + "page_shape=%s dense_shape=%s error=%s", + layout.cp_rank, + layout.cp_size, + spans, + tuple(page_buffer.shape), + tuple(dense_page_buffer.shape), + exc, + ) + return False + + def _get_or_build_current_ipc_slot_descriptors( *, slot_remap: SharedTokenKVSlotRemap | SharedPagedBufferSlotRemap, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 6c06b9cb7..0de5e697c 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -20,7 +20,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( build_batch_prefix_slot_spans, cp_shared_kv_debug_enabled, cp_shared_kv_debug_log, - cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_log_enabled, cp_shared_kv_mla_prefetch_should_log_layer, @@ -477,14 +476,6 @@ def _log_cp_shared_kv_index_prefetch_fallback( ) -def _should_log_missing_index_prefetcher() -> bool: - # When MLA/index prefetch is disabled, a missing index prefetcher is the - # expected sync-materialize path, not a fast-path fallback. Still log consume - # misses when a prefetcher exists because those indicate an enabled prefetch - # pipeline failed to provide the requested layer/current buffer. - return cp_shared_kv_mla_prefetch_enabled() - - class BaseIndexerMetadata(ABC): @abstractmethod def get_seqlens_int32(self) -> torch.Tensor: @@ -798,25 +789,15 @@ class Indexer(MultiPlatformOp): ) if prefetched is not None: return prefetched - _log_cp_shared_kv_index_prefetch_fallback( - "current_consume_miss", - "prefetcher did not provide current index prefix buffer; " - "falling back to sync partial-current index compose. " - "layer=%s cp_rank=%s prefix_lens=%s extend_lens=%s " - "logical_page_table_shape=%s", - layer_id, - layout.cp_rank, - prefix_lens, - extend_lens, - tuple(logical_page_table.shape), - ) - else: - if _should_log_missing_index_prefetcher(): + if int(layer_id) > int( + getattr(forward_batch.token_to_kv_pool, "start_layer", 0) + ): _log_cp_shared_kv_index_prefetch_fallback( - "current_missing_prefetcher", - "index prefetcher is unavailable; falling back to sync " - "partial-current index compose. layer=%s cp_rank=%s " - "prefix_lens=%s extend_lens=%s logical_page_table_shape=%s", + "current_consume_miss", + "prefetcher did not provide current index prefix buffer; " + "falling back to sync partial-current index compose. " + "layer=%s cp_rank=%s prefix_lens=%s extend_lens=%s " + "logical_page_table_shape=%s", layer_id, layout.cp_rank, prefix_lens, @@ -901,16 +882,6 @@ class Indexer(MultiPlatformOp): layout.cp_rank, tuple(logical_page_table.shape), ) - else: - if _should_log_missing_index_prefetcher(): - _log_cp_shared_kv_index_prefetch_fallback( - "missing_prefetcher", - "index prefetcher is unavailable; falling back to sync paged " - "materialize. layer=%s cp_rank=%s logical_page_table_shape=%s", - layer_id, - layout.cp_rank, - tuple(logical_page_table.shape), - ) if cp_shared_kv_debug_enabled(): cp_shared_kv_debug_log( diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 1c1128e76..06a1259a2 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -76,7 +76,7 @@ from sglang.srt.layers.attention.utils import ( mla_quantize_and_rope_for_fp8, seqlens_expand_triton, ) -from sglang.srt.layers.dp_attention import get_attention_cp_group, get_attention_tp_size +from sglang.srt.layers.dp_attention import get_attention_tp_size from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import is_cuda, is_hip @@ -166,7 +166,7 @@ def _maybe_start_cp_shared_kv_attention_prefetch( This hook intentionally runs after current-layer index/MLA cache materialization has finished and immediately before the attention kernel. Earlier hooks can make the next-layer collective overlap current-layer KV - materialization/reduce instead of attention, which can serialize or contend + materialization instead of attention, which can serialize or contend with the current layer's required cache work. """ @@ -1061,17 +1061,11 @@ class NativeSparseAttnBackend( prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) extend_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) - try: - pynccl_available = ( - getattr(get_attention_cp_group(), "pynccl_comm", None) is not None - ) - except Exception: - pynccl_available = False cp_shared_kv_mla_prefetch_log( "create_result cp_rank=%s cp_size=%s uses_cp_shared_kv=%s " "forward_mode=%s topk_transform=%s batch_size=%s has_mla=%s " "has_index=%s prefix_lens=%s extend_lens=%s seq_lens=%s " - "real_pages=%s page_table_shape=%s pynccl=%s", + "real_pages=%s page_table_shape=%s", getattr(layout, "cp_rank", None), getattr(layout, "cp_size", None), getattr(forward_batch, "uses_cp_shared_kv", None), @@ -1095,7 +1089,6 @@ class NativeSparseAttnBackend( tuple(metadata.page_table_1.shape) if metadata.page_table_1 is not None else None, - pynccl_available, ) def _cal_indexer_k_start_end( @@ -2468,7 +2461,7 @@ class NativeSparseAttnBackend( forward_batch, "cp_shared_kv_index_prefetcher", None ) if index_prefetcher is not None: - index_prefetcher.launch_pending_reduce() + index_prefetcher.finalize_pending_materialize() try: if nsa_impl == "tilelang": @@ -2654,40 +2647,69 @@ class NativeSparseAttnBackend( ragged_current_req_id = ragged_current_req_id[ : int(current_locs_for_reuse.shape[0]) ] - slot_remap = get_or_build_shared_token_kv_slot_remap( - forward_batch, - kv_cache=kv_cache, - remap_logical_pages=metadata.real_page_table, - layout=forward_batch.cp_shared_kv_layout, - page_size=page_size, + prefetched_kv = None + mla_prefetcher = getattr( + forward_batch, "cp_shared_kv_mla_prefetcher", None ) - kv_cache, page_table_1_flattened = ( - materialize_prefix_and_reuse_current_kv_page_slots( + if mla_prefetcher is not None: + prefetched_kv = ( + mla_prefetcher.consume_prefix_with_current( + layer_id=layer.layer_id, + kv_cache=kv_cache, + logical_locs=page_table_1_flattened, + current_kv_cache=current_kv_cache, + current_locs=current_locs_for_reuse, + loc_req_id=logical_locs_row_ids, + current_req_id=ragged_current_req_id, + ) + ) + if prefetched_kv is not None: + kv_cache, page_table_1_flattened = prefetched_kv + else: + if ( + mla_prefetcher is not None + and int(layer.layer_id) > int(getattr(forward_batch.token_to_kv_pool, "start_layer", 0)) + ): + _log_cp_shared_kv_mla_prefetch_fallback( + "ragged_current_consume_miss", + "MLA prefetch fast path did not provide " + "RAGGED partial-current prefix; using " + "synchronous compose. cp_rank=%s layer=%s " + "prefix_lens=%s extend_lens=%s " + "current_rows=%s page_table_shape=%s", + forward_batch.cp_shared_kv_layout.cp_rank, + layer.layer_id, + [int(x) for x in prefix_lens_cpu], + [int(x) for x in extend_lens_cpu], + int(current_kv_cache.shape[0]), + tuple(page_table_1_flattened.shape), + ) + slot_remap = get_or_build_shared_token_kv_slot_remap( + forward_batch, kv_cache=kv_cache, - logical_locs=page_table_1_flattened, - current_kv_cache=current_kv_cache, - current_locs=current_locs_for_reuse, - slot_remap=slot_remap, + remap_logical_pages=metadata.real_page_table, layout=forward_batch.cp_shared_kv_layout, page_size=page_size, - prefix_pages=prefix_pages, - prefix_slot_spans=prefix_slot_spans, - current_slot_spans=current_slot_spans, - loc_req_id=logical_locs_row_ids, - current_req_id=ragged_current_req_id, - layer_id=layer.layer_id, - nvtx_source="mla.ragged_partial_current_sync", - current_page_writer_ranks=( - maybe_build_current_page_writer_ranks( - forward_batch=forward_batch, - prefix_lens_cpu=prefix_lens_cpu, - extend_lens_cpu=extend_lens_cpu, - page_size=page_size, - layout=forward_batch.cp_shared_kv_layout, - ) - ), + + ) + kv_cache, page_table_1_flattened = ( + materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=page_table_1_flattened, + current_kv_cache=current_kv_cache, + current_locs=current_locs_for_reuse, + slot_remap=slot_remap, + layout=forward_batch.cp_shared_kv_layout, + page_size=page_size, + prefix_pages=prefix_pages, + prefix_slot_spans=prefix_slot_spans, + current_slot_spans=current_slot_spans, + loc_req_id=logical_locs_row_ids, + current_req_id=ragged_current_req_id, + layer_id=layer.layer_id, + nvtx_source="mla.ragged_partial_current_sync", + ) ) - ) if envs.SGLANG_NSA_DEQUANT_ONLY_TOPK.get(): kv_cache = self._dequant_topk_inplace( kv_cache, @@ -2766,7 +2788,7 @@ class NativeSparseAttnBackend( ) finally: if mla_prefetcher is not None: - mla_prefetcher.launch_pending_reduce() + mla_prefetcher.finalize_pending_materialize() mla_prefetcher.wait_attention_window() index_prefetcher = getattr( forward_batch, "cp_shared_kv_index_prefetcher", None diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index bf8c25829..5adc4b04f 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -155,7 +155,7 @@ class _FakeExtendForwardMode: class TestCpSharedKVRuntimeHelpers(unittest.TestCase): - def test_mla_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream( + def test_mla_prefetch_waits_l2_l1_on_prefetch_stream_and_ipc_on_current_stream( self, ): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -182,6 +182,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def __exit__(self, exc_type, exc, tb): self.active_stream[0] = self.previous + class FakeEvent: + def __init__(self): + self.recorded = [] + + def record(self, stream): + self.recorded.append(stream.name) + class FakePool: start_layer = 0 page_size = 4 @@ -203,12 +210,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch_stream = FakeStream("prefetch") calls = [] - def record_materialize(**kwargs): - calls.append(("materialize", active_stream[0])) - - def record_reduce(buffer, cp_size, stream, **kwargs): - calls.append(("reduce", active_stream[0], stream.name)) - return object() + def record_ipc(**kwargs): + calls.append(("ipc", active_stream[0], kwargs["spans"])) + return True prefetcher = prefetch.CpSharedKVMlaPrefetcher( layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), @@ -227,13 +231,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): "stream", side_effect=lambda stream: FakeStreamContext(active_stream, stream), ), patch.object( - prefetch, - "materialize_local_token_kv_page_slots_into", - side_effect=record_materialize, + prefetch.torch.cuda, "Event", side_effect=FakeEvent ), patch.object( prefetch, - "_all_reduce_materialized_buffer_async", - side_effect=record_reduce, + "_try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into", + return_value=False, + create=True, + ), patch.object( + prefetch, + "_try_tai_ipc_materialize_token_kv_page_slot_spans_into", + side_effect=record_ipc, ): pool = FakePool() prefetcher.start_next_layer_prefix( @@ -241,22 +248,25 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): token_to_kv_pool=pool, ) - self.assertEqual(pool.prefetch_getter_streams, [(1, "current")]) + self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) self.assertEqual( calls, - [("materialize", "current"), ("reduce", "prefetch", "prefetch")], + [("ipc", "current", [(0, 2)])], ) self.assertEqual(prefetch_stream.waited, ["current"]) + self.assertEqual(current_stream.waited, ["prefetch"]) + self.assertEqual(prefetcher.handles[1].event.recorded, ["current"]) - prefetcher.launch_pending_reduce() + prefetcher.finalize_pending_materialize() self.assertEqual( calls, - [("materialize", "current"), ("reduce", "prefetch", "prefetch")], + [("ipc", "current", [(0, 2)])], ) self.assertEqual(prefetch_stream.waited, ["current"]) + self.assertEqual(current_stream.waited, ["prefetch"]) - def test_mla_prefetch_ipc_spans_skip_prefix_reduce(self): + def test_mla_prefetch_ipc_spans_skip_prefix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -275,6 +285,19 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def record(self, stream): self.recorded.append(stream.name) + class FakeStreamContext: + def __init__(self, active_stream, stream): + self.active_stream = active_stream + self.stream = stream + self.previous = None + + def __enter__(self): + self.previous = self.active_stream[0] + self.active_stream[0] = self.stream.name + + def __exit__(self, exc_type, exc, tb): + self.active_stream[0] = self.previous + class FakePool: start_layer = 0 page_size = 4 @@ -288,10 +311,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_stream = FakeStream("current") prefetch_stream = FakeStream("prefetch") + active_stream = ["current"] ipc_calls = [] def record_ipc(**kwargs): - ipc_calls.append(kwargs["spans"]) + ipc_calls.append((active_stream[0], kwargs["spans"])) return True prefetcher = prefetch.CpSharedKVMlaPrefetcher( @@ -308,28 +332,143 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream + ), patch.object( + prefetch.torch.cuda, + "stream", + side_effect=lambda stream: FakeStreamContext(active_stream, stream), ), patch.object( prefetch.torch.cuda, "Event", side_effect=FakeEvent + ), patch.object( + prefetch, + "_try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into", + return_value=False, + create=True, ), patch.object( prefetch, "_try_tai_ipc_materialize_token_kv_page_slot_spans_into", side_effect=record_ipc, - ), patch.object( - prefetch, - "_all_reduce_materialized_buffer_ranges_async", - side_effect=AssertionError("IPC prefix path must not all-reduce"), ): prefetcher.start_next_layer_prefix( next_layer_id=1, token_to_kv_pool=FakePool(), ) - self.assertEqual(ipc_calls, [[(0, 2), (4, 5)]]) + self.assertEqual(ipc_calls, [("current", [(0, 2), (4, 5)])]) handle = prefetcher.handles[1] self.assertEqual(handle.event.recorded, ["current"]) - self.assertEqual(prefetch_stream.waited, []) + self.assertEqual(prefetch_stream.waited, ["current"]) + self.assertEqual(current_stream.waited, ["prefetch"]) - def test_mla_prefetch_consume_ipc_suffix_skips_suffix_reduce(self): + def test_mla_prefetch_ce_spans_materialize_on_prefetch_stream(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeStream: + def __init__(self, name): + self.name = name + self.waited = [] + + def wait_stream(self, stream): + self.waited.append(stream.name) + + class FakeStreamContext: + def __init__(self, active_stream, stream): + self.active_stream = active_stream + self.stream = stream + self.previous = None + + def __enter__(self): + self.previous = self.active_stream[0] + self.active_stream[0] = self.stream.name + + def __exit__(self, exc_type, exc, tb): + self.active_stream[0] = self.previous + + class FakeEvent: + def __init__(self): + self.recorded = [] + + def record(self, stream): + self.recorded.append(stream.name) + + class FakeStreamContext: + def __init__(self, active_stream, stream): + self.active_stream = active_stream + self.stream = stream + self.previous = None + + def __enter__(self): + self.previous = self.active_stream[0] + self.active_stream[0] = self.stream.name + + def __exit__(self, exc_type, exc, tb): + self.active_stream[0] = self.previous + + class FakePool: + start_layer = 0 + page_size = 4 + kv_buffer = [object(), object(), object()] + + def __init__(self): + self.kv_cache = torch.zeros((64, 1, 2), dtype=torch.float32) + self.prefetch_getter_streams = [] + + def get_key_buffer_for_prefetch(self, layer_id, stream): + self.prefetch_getter_streams.append((layer_id, stream.name)) + return self.kv_cache + + active_stream = ["current"] + current_stream = FakeStream("current") + prefetch_stream = FakeStream("prefetch") + ce_calls = [] + + def record_ce(**kwargs): + ce_calls.append((active_stream[0], kwargs["spans"])) + return True + + prefetcher = prefetch.CpSharedKVMlaPrefetcher( + layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), + page_size=4, + prefix_pages=3, + prefix_slot_spans=[(0, 2), (4, 5)], + current_slot_spans=[(2, 4), (5, 6)], + slot_logical_pages=torch.tensor([0, 1, 2, 3, 8, 9], dtype=torch.int64), + page_inverse=torch.tensor([0, 1, 2, 3, -1, -1, -1, -1, 5, 6]), + dense_num_pages=7, + stream=prefetch_stream, + ) + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=current_stream + ), patch.object( + prefetch.torch.cuda, + "stream", + side_effect=lambda stream: FakeStreamContext(active_stream, stream), + ), patch.object( + prefetch.torch.cuda, "Event", side_effect=FakeEvent + ), patch.object( + prefetch, + "_try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into", + side_effect=record_ce, + create=True, + ), patch.object( + prefetch, + "_try_tai_ipc_materialize_token_kv_page_slot_spans_into", + side_effect=AssertionError("SM IPC prefix path must not run after CE"), + ): + pool = FakePool() + prefetcher.start_next_layer_prefix( + next_layer_id=1, + token_to_kv_pool=pool, + ) + + self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) + self.assertEqual(ce_calls, [("prefetch", [(0, 2), (4, 5)])]) + handle = prefetcher.handles[1] + self.assertEqual(handle.event.recorded, ["prefetch"]) + self.assertEqual(prefetch_stream.waited, ["current"]) + + def test_mla_prefetch_consume_ipc_suffix_skips_suffix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -376,11 +515,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "_materialize_local_token_kv_page_slot_spans_into", side_effect=AssertionError("IPC suffix path must not local materialize"), - ), patch.object( - prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=AssertionError("IPC suffix path must not all-reduce"), - ), patch.object( + ), patch.object( prefetch, "remap_logical_locs_to_slot_dense_locs_optimized", return_value=torch.tensor([12, 20], dtype=torch.int64), @@ -395,7 +530,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(ipc_calls, [[(2, 4), (5, 6)]]) self.assertEqual(current_stream.events, [fake_event]) - def test_index_prefetch_waits_l2_l1_on_materialize_stream_and_reduces_on_prefetch_stream( + def test_index_prefetch_waits_l2_l1_on_prefetch_stream_and_ipc_on_current_stream( self, ): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -422,6 +557,13 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def __exit__(self, exc_type, exc, tb): self.active_stream[0] = self.previous + class FakeEvent: + def __init__(self): + self.recorded = [] + + def record(self, stream): + self.recorded.append(stream.name) + class FakePool: start_layer = 0 page_size = 4 @@ -445,12 +587,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch_stream = FakeStream("prefetch") calls = [] - def record_materialize(**kwargs): - calls.append(("materialize", active_stream[0])) - - def record_reduce(buffer, cp_size, stream, **kwargs): - calls.append(("reduce", active_stream[0], stream.name)) - return object() + def record_ipc(**kwargs): + calls.append(("ipc", active_stream[0], kwargs["spans"])) + return True prefetcher = prefetch.CpSharedKVIndexPrefetcher( layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), @@ -468,13 +607,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): "stream", side_effect=lambda stream: FakeStreamContext(active_stream, stream), ), patch.object( - prefetch, - "materialize_local_paged_buffer_page_slots_into", - side_effect=record_materialize, + prefetch.torch.cuda, "Event", side_effect=FakeEvent ), patch.object( prefetch, - "_all_reduce_materialized_buffer_async", - side_effect=record_reduce, + "_try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into", + return_value=False, + create=True, + ), patch.object( + prefetch, + "_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into", + side_effect=record_ipc, ): pool = FakePool() prefetcher.start_next_layer_prefix( @@ -482,22 +624,25 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): token_to_kv_pool=pool, ) - self.assertEqual(pool.prefetch_getter_streams, [(1, "current")]) + self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) self.assertEqual( calls, - [("materialize", "current"), ("reduce", "prefetch", "prefetch")], + [("ipc", "current", [(0, 2)])], ) self.assertEqual(prefetch_stream.waited, ["current"]) + self.assertEqual(current_stream.waited, ["prefetch"]) + self.assertEqual(prefetcher.handles[1].event.recorded, ["current"]) - prefetcher.launch_pending_reduce() + prefetcher.finalize_pending_materialize() self.assertEqual( calls, - [("materialize", "current"), ("reduce", "prefetch", "prefetch")], + [("ipc", "current", [(0, 2)])], ) self.assertEqual(prefetch_stream.waited, ["current"]) + self.assertEqual(current_stream.waited, ["prefetch"]) - def test_index_prefetch_ipc_spans_skip_prefix_reduce(self): + def test_index_prefetch_ipc_spans_skip_prefix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -516,6 +661,19 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): def record(self, stream): self.recorded.append(stream.name) + class FakeStreamContext: + def __init__(self, active_stream, stream): + self.active_stream = active_stream + self.stream = stream + self.previous = None + + def __enter__(self): + self.previous = self.active_stream[0] + self.active_stream[0] = self.stream.name + + def __exit__(self, exc_type, exc, tb): + self.active_stream[0] = self.previous + class FakePool: start_layer = 0 page_size = 4 @@ -529,10 +687,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_stream = FakeStream("current") prefetch_stream = FakeStream("prefetch") + active_stream = ["current"] ipc_calls = [] def record_ipc(**kwargs): - ipc_calls.append(kwargs["spans"]) + ipc_calls.append((active_stream[0], kwargs["spans"])) return True prefetcher = prefetch.CpSharedKVIndexPrefetcher( @@ -548,28 +707,129 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream + ), patch.object( + prefetch.torch.cuda, + "stream", + side_effect=lambda stream: FakeStreamContext(active_stream, stream), ), patch.object( prefetch.torch.cuda, "Event", side_effect=FakeEvent + ), patch.object( + prefetch, + "_try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into", + return_value=False, + create=True, ), patch.object( prefetch, "_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into", side_effect=record_ipc, - ), patch.object( - prefetch, - "_all_reduce_materialized_buffer_ranges_async", - side_effect=AssertionError("IPC prefix path must not all-reduce"), ): prefetcher.start_next_layer_prefix( next_layer_id=1, token_to_kv_pool=FakePool(), ) - self.assertEqual(ipc_calls, [[(0, 2), (4, 5)]]) + self.assertEqual(ipc_calls, [("current", [(0, 2), (4, 5)])]) handle = prefetcher.handles[1] self.assertEqual(handle.event.recorded, ["current"]) - self.assertEqual(prefetch_stream.waited, []) + self.assertEqual(prefetch_stream.waited, ["current"]) + self.assertEqual(current_stream.waited, ["prefetch"]) - def test_index_prefetch_consume_ipc_suffix_skips_suffix_reduce(self): + def test_index_prefetch_ce_spans_materialize_on_prefetch_stream(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class FakeStream: + def __init__(self, name): + self.name = name + self.waited = [] + + def wait_stream(self, stream): + self.waited.append(stream.name) + + class FakeStreamContext: + def __init__(self, active_stream, stream): + self.active_stream = active_stream + self.stream = stream + self.previous = None + + def __enter__(self): + self.previous = self.active_stream[0] + self.active_stream[0] = self.stream.name + + def __exit__(self, exc_type, exc, tb): + self.active_stream[0] = self.previous + + class FakeEvent: + def __init__(self): + self.recorded = [] + + def record(self, stream): + self.recorded.append(stream.name) + + class FakePool: + start_layer = 0 + page_size = 4 + kv_buffer = [object(), object(), object()] + + def __init__(self): + self.page_buffer = torch.zeros((64, 3), dtype=torch.uint8) + self.prefetch_getter_streams = [] + + def get_index_k_with_scale_buffer_for_prefetch(self, layer_id, stream): + self.prefetch_getter_streams.append((layer_id, stream.name)) + return self.page_buffer + + active_stream = ["current"] + current_stream = FakeStream("current") + prefetch_stream = FakeStream("prefetch") + ce_calls = [] + + def record_ce(**kwargs): + ce_calls.append((active_stream[0], kwargs["spans"])) + return True + + prefetcher = prefetch.CpSharedKVIndexPrefetcher( + layout=CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0), + prefix_pages=3, + prefix_slot_spans=[(0, 2), (4, 5)], + current_slot_spans=[(2, 4), (5, 6)], + slot_logical_pages=torch.tensor([0, 1, 2, 3, 8, 9], dtype=torch.int64), + page_inverse=torch.tensor([0, 1, 2, 3, -1, -1, -1, -1, 5, 6]), + dense_num_pages=7, + stream=prefetch_stream, + ) + + with patch.object( + prefetch.torch.cuda, "current_stream", return_value=current_stream + ), patch.object( + prefetch.torch.cuda, + "stream", + side_effect=lambda stream: FakeStreamContext(active_stream, stream), + ), patch.object( + prefetch.torch.cuda, "Event", side_effect=FakeEvent + ), patch.object( + prefetch, + "_try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into", + side_effect=record_ce, + create=True, + ), patch.object( + prefetch, + "_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into", + side_effect=AssertionError("SM IPC prefix path must not run after CE"), + ): + pool = FakePool() + prefetcher.start_next_layer_prefix( + next_layer_id=1, + token_to_kv_pool=pool, + ) + + self.assertEqual(pool.prefetch_getter_streams, [(1, "prefetch")]) + self.assertEqual(ce_calls, [("prefetch", [(0, 2), (4, 5)])]) + handle = prefetcher.handles[1] + self.assertEqual(handle.event.recorded, ["prefetch"]) + self.assertEqual(prefetch_stream.waited, ["current"]) + + def test_index_prefetch_consume_ipc_suffix_skips_suffix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -615,11 +875,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "_materialize_local_paged_buffer_page_slot_spans_into", side_effect=AssertionError("IPC suffix path must not local materialize"), - ), patch.object( - prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=AssertionError("IPC suffix path must not all-reduce"), - ), patch.object( + ), patch.object( prefetch, "remap_logical_pages_to_slot_dense_pages", return_value=torch.tensor([[3, 6]], dtype=torch.int64), @@ -1434,9 +1690,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ).read_text() self.assertIn("[CP_SHARED_KV_FALLBACK][index_prefetch]", source) - self.assertIn("current_missing_prefetcher", source) self.assertIn("current_consume_miss", source) - self.assertIn("missing_prefetcher", source) + self.assertIn("consume_miss", source) + self.assertNotIn('"current_missing_prefetcher"', source) + self.assertNotIn('"missing_prefetcher"', source) def test_current_loc_remap_fast_path_args_only_for_current_only_extend(self): from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( @@ -1633,9 +1890,21 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): "materialize_prefix_and_reuse_current_kv_page_slots", ragged_source, ) + self.assertIn( + "mla_prefetcher.consume_prefix_with_current", + ragged_source, + "FP8/RAGGED cache-hit compose still has page-slot prefix pages and " + "must consume the same MLA prefetch path before sync fallback.", + ) self.assertIn("select_cp_current_valid_rows_for_reuse", ragged_source) self.assertIn("prefix_slot_spans=", ragged_source) self.assertIn("current_slot_spans=", ragged_source) + self.assertIn( + 'int(layer.layer_id) > int(getattr(forward_batch.token_to_kv_pool, "start_layer", 0))', + ragged_source, + "RAGGED layer-0 has no previous attention window to prefetch; " + "that expected miss must not warn.", + ) @unittest.skipIf(not torch.cuda.is_available(), "CUDA is required") def test_tai_current_slot_fill_sparse_page_self_test_passes_on_installed_kernel( @@ -2092,10 +2361,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "cp_shared_kv_mla_prefetch_min_async_extend_tokens", return_value=0, - ), patch.object( - prefetch, - "get_attention_cp_group", - return_value=SimpleNamespace(pynccl_comm=object()), ), patch.object( prefetch.torch.cuda, "Stream", return_value=stream ), patch.object( @@ -2117,7 +2382,80 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(prefetcher.prefix_page_count, 3) self.assertEqual(prefetcher.current_page_count, 3) - def test_mla_prefetch_batch_consume_reduces_exact_current_spans(self): + def test_mla_prefetch_create_allows_ragged_topk_transform(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + + class Mode: + def is_context_parallel_extend(self): + return True + + page_size = 4 + logical_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int64) + stream = object() + kv_cache = torch.zeros((128, 1, 1), dtype=torch.float32) + remap = SimpleNamespace( + slot_logical_pages=logical_pages.reshape(-1), + page_inverse=torch.zeros((1, 32), dtype=torch.int64), + slot_sorted_logical_pages_by_row=None, + slot_sorted_dense_pages_by_row=None, + dense_num_pages=5, + ) + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + hisparse_coordinator=None, + forward_mode=Mode(), + batch_size=1, + token_to_kv_pool=SimpleNamespace(page_size=page_size, start_layer=0), + cp_shared_kv_layout=SimpleNamespace(cp_size=2, cp_rank=0), + extend_prefix_lens_cpu=[8], + extend_seq_lens_cpu=[8], + ) + metadata = SimpleNamespace( + real_page_table=logical_pages, + page_table_1=torch.zeros((1, 4), dtype=torch.int32), + ) + + with patch.object( + prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True + ), patch.object( + prefetch, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + prefetch.torch.cuda, "is_available", return_value=True + ), patch.object( + prefetch, "_is_cuda_stream_capturing", return_value=False + ), patch.object( + prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True + ), patch.object( + prefetch, + "cp_shared_kv_mla_prefetch_min_prefix_pages", + return_value=0, + ), patch.object( + prefetch, + "cp_shared_kv_mla_prefetch_min_async_extend_tokens", + return_value=0, + ), patch.object( + prefetch.torch.cuda, "Stream", return_value=stream + ), patch.object( + prefetch, "_prefetch_pool_get_key_buffer", return_value=kv_cache + ), patch.object( + prefetch, + "get_or_build_shared_token_kv_slot_remap", + return_value=remap, + ), patch.object( + prefetch.logger, "warning" + ) as logger: + prefetcher = prefetch.CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=False, + ) + + self.assertIsNotNone(prefetcher) + self.assertEqual(prefetcher.prefix_slot_spans, [(0, 2)]) + self.assertEqual(prefetcher.current_slot_spans, [(2, 4)]) + logger.assert_not_called() + + def test_mla_prefetch_batch_consume_materializes_exact_current_spans(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -2179,7 +2517,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ], dtype=torch.int64, ) - reduce_ranges = [] + ipc_spans = [] class FakeCurrentStream: def __init__(self): @@ -2190,16 +2528,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_stream = FakeCurrentStream() - def record_range_reduce(buffer, cp_size, start_row, end_row, **kwargs): - reduce_ranges.append((start_row, end_row, kwargs.get("nvtx_source"))) - return buffer + def record_current_ipc(**kwargs): + ipc_spans.append(kwargs["spans"]) + return True with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream ), patch.object( prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=record_range_reduce, + "_try_tai_ipc_materialize_current_token_kv_page_slot_spans_into", + side_effect=record_current_ipc, ): mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current( layer_id=1, @@ -2218,13 +2556,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): mixed_locs.tolist(), [[4, 8, 12, 13, -1, -1], [20, 24, 25, -1, -1, -1]], ) - self.assertEqual( - reduce_ranges, - [ - (12, 16, "mla.prefetch_current"), - (24, 32, "mla.prefetch_current"), - ], - ) + self.assertEqual(ipc_spans, [[(2, 3), (5, 7)]]) def test_index_prefetch_create_batch_uses_exact_prefix_and_current_spans(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -2287,10 +2619,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "cp_shared_kv_mla_prefetch_min_async_extend_tokens", return_value=0, - ), patch.object( - prefetch, - "get_attention_cp_group", - return_value=SimpleNamespace(pynccl_comm=object()), ), patch.object( prefetch.torch.cuda, "Stream", return_value=stream ), patch.object( @@ -2312,7 +2640,90 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(prefetcher.prefix_page_count, 3) self.assertEqual(prefetcher.current_page_count, 3) - def test_index_prefetch_batch_consume_reduces_exact_current_spans(self): + def test_index_prefetch_create_allows_ragged_topk_transform(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + class Mode: + def is_context_parallel_extend(self): + return True + + page_size = 4 + logical_pages = torch.tensor( + [ + [1, 2, 5, 0], + [9, 11, 12, 13], + ], + dtype=torch.int64, + ) + stream = object() + page_buffer = torch.zeros((128, 32), dtype=torch.uint8) + remap = SimpleNamespace( + slot_logical_pages=logical_pages.reshape(-1), + page_inverse=torch.zeros((2, 32), dtype=torch.int64), + slot_sorted_logical_pages_by_row=None, + slot_sorted_dense_pages_by_row=None, + dense_pages=torch.tensor([[1, 2, 3, 0], [5, 6, 7, 8]], dtype=torch.int64), + dense_num_pages=9, + ) + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + hisparse_coordinator=None, + forward_mode=Mode(), + batch_size=2, + token_to_kv_pool=SimpleNamespace(page_size=page_size, start_layer=0), + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, cp_size=2, cp_rank=0 + ), + extend_prefix_lens_cpu=[8, 4], + extend_seq_lens_cpu=[2, 7], + ) + metadata = SimpleNamespace( + real_page_table=logical_pages, + page_table_1=torch.zeros((2, 4), dtype=torch.int32), + ) + + with patch.object( + prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True + ), patch.object( + prefetch, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + prefetch.torch.cuda, "is_available", return_value=True + ), patch.object( + prefetch, "_is_cuda_stream_capturing", return_value=False + ), patch.object( + prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True + ), patch.object( + prefetch, + "cp_shared_kv_mla_prefetch_min_prefix_pages", + return_value=0, + ), patch.object( + prefetch, + "cp_shared_kv_mla_prefetch_min_async_extend_tokens", + return_value=0, + ), patch.object( + prefetch.torch.cuda, "Stream", return_value=stream + ), patch.object( + prefetch, "_prefetch_pool_get_index_buffer", return_value=page_buffer + ), patch.object( + prefetch, + "get_or_build_shared_paged_buffer_slot_remap", + return_value=remap, + ), patch.object( + prefetch.logger, "warning" + ) as logger: + prefetcher = prefetch.CpSharedKVIndexPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=False, + ) + + self.assertIsNotNone(prefetcher) + self.assertEqual(prefetcher.prefix_slot_spans, [(0, 2), (4, 5)]) + self.assertEqual(prefetcher.current_slot_spans, [(2, 3), (5, 7)]) + logger.assert_not_called() + + def test_index_prefetch_batch_consume_materializes_exact_current_spans(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -2364,7 +2775,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_scale = torch.tensor([[1.25], [2.5], [3.5], [4.5]], dtype=torch.float32) current_locs = torch.tensor([20, 21, 44, 45], dtype=torch.int64) current_req_id = torch.tensor([0, 0, 1, 1], dtype=torch.int64) - reduce_ranges = [] + ipc_spans = [] class FakeCurrentStream: def __init__(self): @@ -2375,16 +2786,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_stream = FakeCurrentStream() - def record_range_reduce(buffer, cp_size, start_row, end_row, **kwargs): - reduce_ranges.append((start_row, end_row, kwargs.get("nvtx_source"))) - return buffer + def record_current_ipc(**kwargs): + ipc_spans.append(kwargs["spans"]) + return True with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream ), patch.object( prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=record_range_reduce, + "_try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into", + side_effect=record_current_ipc, ): mixed_buffer, dense_pages = prefetcher.consume_prefix_with_current( layer_id=1, @@ -2414,13 +2825,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): current_scale[2:].contiguous().view(torch.uint8).reshape(-1), ) ) - self.assertEqual( - reduce_ranges, - [ - (3, 4, "index.prefetch_current"), - (6, 8, "index.prefetch_current"), - ], - ) + self.assertEqual(ipc_spans, [[(2, 3), (5, 7)]]) def test_valid_page_mask_prevents_stale_rectangular_tail_remap(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime @@ -2960,11 +3365,11 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) prefetcher.handles[1] = handle prefetcher.pending_attention_handle = handle - range_calls = [] + ipc_spans = [] - def record_range_reduce(buffer, cp_size, start_row, end_row, **kwargs): - range_calls.append((start_row, end_row, kwargs.get("nvtx_source"))) - return buffer + def record_current_ipc(**kwargs): + ipc_spans.append(kwargs["spans"]) + return True with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream @@ -2974,8 +3379,8 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): side_effect=AssertionError("suffix materialize must not run"), ), patch.object( prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=record_range_reduce, + "_try_tai_ipc_materialize_current_token_kv_page_slot_spans_into", + side_effect=record_current_ipc, ): mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current( layer_id=1, @@ -2994,7 +3399,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): expected_kv[12:14] = current_kv self.assertTrue(torch.equal(mixed_kv, expected_kv)) self.assertEqual(mixed_locs.tolist(), [[4, 12], [13, 7], [-1, -1]]) - self.assertEqual(range_calls, [(12, 16, "mla.prefetch_current")]) + self.assertEqual(ipc_spans, [[(2, 3)]]) def test_mla_prefetch_consume_prefix_with_current_uses_ipc_without_all_reduce(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch @@ -3031,11 +3436,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "_try_tai_ipc_materialize_current_token_kv_page_slot_spans_into", return_value=True, - ), patch.object( - prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=AssertionError("prefetch current must use IPC"), - ): + ): mixed_kv, mixed_locs = prefetcher.consume_prefix_with_current( layer_id=1, kv_cache=torch.zeros((64, 1, 1), dtype=torch.float32), @@ -3087,7 +3488,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertIs(prefetcher.pending_attention_handle, handle) self.assertIs(prefetcher.handles[1], handle) - def test_mla_prefetch_attention_window_does_not_launch_pending_reduce(self): + def test_mla_prefetch_attention_window_does_not_finalize_pending_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -3119,10 +3520,10 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with patch.object( prefetch.torch.cuda, "current_stream", return_value=current_stream - ), patch.object(prefetcher, "launch_pending_reduce") as launch_pending_reduce: + ), patch.object(prefetcher, "finalize_pending_materialize") as finalize_pending_materialize: prefetcher.wait_attention_window() - launch_pending_reduce.assert_not_called() + finalize_pending_materialize.assert_not_called() self.assertEqual(current_stream.events, []) self.assertIs(prefetcher.pending_attention_handle, handle) self.assertIsNone(handle.event) @@ -3205,7 +3606,9 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with patch.object( prefetch.torch.cuda, "current_stream", return_value=FakeCurrentStream() ), patch.object( - prefetch, "_all_reduce_materialized_buffer_range", _identity_all_reduce + prefetch, + "_try_tai_ipc_materialize_paged_buffer_page_slot_spans_into", + return_value=True, ): dense_pages_buffer, dense_pages = prefetcher.consume( layer_id=1, @@ -3625,11 +4028,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "_try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into", return_value=True, - ), patch.object( - prefetch, - "_all_reduce_materialized_buffer_range", - side_effect=AssertionError("index prefetch current must use IPC"), - ): + ): mixed_buffer, dense_pages = prefetcher.consume_prefix_with_current( layer_id=1, logical_pages=torch.tensor([[1, 20]], dtype=torch.int64), @@ -3880,6 +4279,137 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(first.src_page_indices.tolist(), [1, 1, 2, 2]) self.assertEqual(first.dense_page_indices.tolist(), [1, 2, 3, 4]) + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_ipc_ce_token_prefix_uses_cpu_descriptors_and_skips_invalid_slots(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0) + kv_cache = torch.zeros((16, 1, 1), device="cuda", dtype=torch.float32) + dense_kv = torch.zeros((20, 1, 1), device="cuda", dtype=torch.float32) + calls = [] + + class FakeKernels: + @staticmethod + def materialize_cuda_ipc_peer_pages_slot_indices_ce( + peer_ptrs, + dst, + owner_ranks, + src_page_indices, + dst_page_indices, + *, + page_nbytes, + ): + calls.append( + ( + peer_ptrs.device.type, + owner_ranks.device.type, + src_page_indices.device.type, + dst_page_indices.device.type, + owner_ranks.tolist(), + src_page_indices.tolist(), + dst_page_indices.tolist(), + page_nbytes, + ) + ) + + with patch.object( + runtime, + "_get_or_open_tai_ipc_peer_ptrs", + return_value=(FakeKernels(), torch.tensor([123, 456], dtype=torch.int64)), + ): + ok = runtime._try_tai_ipc_ce_materialize_token_kv_page_slot_spans_into( + kv_cache=kv_cache, + dense_kv_cache=dense_kv, + slot_logical_pages=torch.tensor([0, 1, 2, 3], dtype=torch.int64), + layout=layout, + page_size=4, + spans=[(0, 4)], + ) + + self.assertTrue(ok) + self.assertEqual(len(calls), 1) + ( + peer_device, + owner_device, + src_device, + dst_device, + owners, + src_pages, + dst_pages, + page_nbytes, + ) = calls[0] + self.assertEqual((peer_device, owner_device, src_device, dst_device), ("cpu", "cpu", "cpu", "cpu")) + self.assertEqual(owners, [0, 1, 0]) + self.assertEqual(src_pages, [1, 1, 2]) + self.assertEqual(dst_pages, [2, 3, 4]) + self.assertEqual(page_nbytes, 4 * kv_cache.stride(0) * kv_cache.element_size()) + + @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA") + def test_ipc_ce_paged_prefix_uses_cpu_descriptors_and_skips_invalid_slots(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout + + layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=0) + page_buffer = torch.zeros((8, 16), device="cuda", dtype=torch.uint8) + dense_page_buffer = torch.zeros((5, 16), device="cuda", dtype=torch.uint8) + calls = [] + + class FakeKernels: + @staticmethod + def materialize_cuda_ipc_peer_pages_slot_indices_ce( + peer_ptrs, + dst, + owner_ranks, + src_page_indices, + dst_page_indices, + *, + page_nbytes, + ): + calls.append( + ( + peer_ptrs.device.type, + owner_ranks.device.type, + src_page_indices.device.type, + dst_page_indices.device.type, + owner_ranks.tolist(), + src_page_indices.tolist(), + dst_page_indices.tolist(), + page_nbytes, + ) + ) + + with patch.object( + runtime, + "_get_or_open_tai_ipc_peer_ptrs", + return_value=(FakeKernels(), torch.tensor([123, 456], dtype=torch.int64)), + ): + ok = runtime._try_tai_ipc_ce_materialize_paged_buffer_page_slot_spans_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=torch.tensor([0, 1, 2, 3], dtype=torch.int64), + layout=layout, + spans=[(0, 4)], + ) + + self.assertTrue(ok) + self.assertEqual(len(calls), 1) + ( + peer_device, + owner_device, + src_device, + dst_device, + owners, + src_pages, + dst_pages, + page_nbytes, + ) = calls[0] + self.assertEqual((peer_device, owner_device, src_device, dst_device), ("cpu", "cpu", "cpu", "cpu")) + self.assertEqual(owners, [0, 1, 0]) + self.assertEqual(src_pages, [1, 1, 2]) + self.assertEqual(dst_pages, [2, 3, 4]) + self.assertEqual(page_nbytes, page_buffer.stride(0) * page_buffer.element_size()) + def test_ipc_current_descriptors_are_cached_on_token_slot_remap(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -4185,6 +4715,56 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): ) self.assertIn("prefix_not_page_aligned", logger.call_args.args[1]) + def test_mla_and_index_prefetch_no_cache_prefix_skips_without_warning(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch + + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + hisparse_coordinator=None, + forward_mode=Mode(), + batch_size=1, + token_to_kv_pool=SimpleNamespace(page_size=64), + cp_shared_kv_layout=SimpleNamespace(cp_size=8, cp_rank=0), + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[64], + ) + metadata = SimpleNamespace( + real_page_table=torch.tensor([[0]], dtype=torch.int64), + page_table_1=torch.tensor([[0]], dtype=torch.int32), + ) + + with patch.object( + prefetch, "cp_shared_kv_mla_prefetch_enabled", return_value=True + ), patch.object( + prefetch, "cp_shared_kv_debug_enabled", return_value=False + ), patch.object( + prefetch.torch.cuda, "is_available", return_value=True + ), patch.object( + prefetch, "_is_cuda_stream_capturing", return_value=False + ), patch.object( + prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True + ), patch.object( + prefetch.logger, "warning" + ) as logger: + mla_result = prefetch.CpSharedKVMlaPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=False, + ) + index_result = prefetch.CpSharedKVIndexPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=False, + ) + + self.assertIsNone(mla_result) + self.assertIsNone(index_result) + logger.assert_not_called() + def test_mla_prefetch_min_prefix_pages_uses_cached_token_default_and_can_override( self, ): @@ -4310,10 +4890,6 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): prefetch, "_is_cuda_stream_capturing", return_value=False ), patch.object( prefetch, "is_nsa_prefill_cp_in_seq_split", return_value=True - ), patch.object( - prefetch, - "get_attention_cp_group", - return_value=SimpleNamespace(pynccl_comm=object()), ), patch.object( prefetch.torch.cuda, "Stream", return_value=stream ), patch.object( @@ -6671,6 +7247,33 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): self.assertEqual(index_prefetcher.calls, [(6, token_to_kv_pool)]) self.assertEqual(mla_prefetcher.calls, [(6, token_to_kv_pool)]) + def test_attention_window_prefetch_skips_inactive_index_cache_layer(self): + from sglang.srt.layers.attention import nsa_backend + + class FakePrefetcher: + def __init__(self): + self.calls = [] + + def start_next_layer_prefix(self, *, next_layer_id, token_to_kv_pool): + self.calls.append((next_layer_id, token_to_kv_pool)) + + token_to_kv_pool = SimpleNamespace(index_layer_to_slot={0: 0, 4: 1}) + index_prefetcher = FakePrefetcher() + mla_prefetcher = FakePrefetcher() + forward_batch = SimpleNamespace( + token_to_kv_pool=token_to_kv_pool, + cp_shared_kv_index_prefetcher=index_prefetcher, + cp_shared_kv_mla_prefetcher=mla_prefetcher, + cp_shared_kv_num_model_layers=12, + ) + + nsa_backend._maybe_start_cp_shared_kv_attention_prefetch( + forward_batch, layer_id=0 + ) + + self.assertEqual(index_prefetcher.calls, []) + self.assertEqual(mla_prefetcher.calls, [(1, token_to_kv_pool)]) + def test_index_prefetch_skips_when_current_layer_is_last(self): from sglang.srt.layers.attention.nsa import nsa_indexer from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( @@ -6818,6 +7421,123 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): logger.warning.assert_not_called() + def test_index_partial_current_missing_prefetcher_does_not_warn(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + class FakePool: + page_size = 4 + index_head_dim = 4 + + def get_index_k_with_scale_buffer(self, layer_id): + return torch.zeros((16, 32), dtype=torch.uint8) + + class FakeLayout: + page_size = 4 + cp_size = 1 + cp_rank = 0 + + fallback_buffer = torch.full((8, 32), 7, dtype=torch.uint8) + fallback_pages = torch.tensor([[0, 1]], dtype=torch.int64) + forward_batch = SimpleNamespace( + token_to_kv_pool=FakePool(), + uses_cp_shared_kv=True, + cp_shared_kv_layout=FakeLayout(), + cp_shared_kv_index_prefetcher=None, + extend_prefix_lens_cpu=[4], + extend_seq_lens_cpu=[4], + cp_local_out_cache_loc=torch.tensor([4, 5, 6, 7], dtype=torch.int64), + ) + logical_pages = torch.tensor([[0, 1]], dtype=torch.int64) + current_index_kv = ( + torch.zeros((4, 4), dtype=torch.uint8), + torch.zeros((4,), dtype=torch.float32), + ) + indexer = object.__new__(nsa_indexer.Indexer) + + with patch.object( + nsa_indexer, + "get_or_build_shared_paged_buffer_slot_remap", + return_value=object(), + ), patch.object( + nsa_indexer, + "materialize_prefix_and_reuse_current_index_page_slots", + return_value=(fallback_buffer, fallback_pages), + ), patch( + "sglang.srt.layers.attention.nsa.nsa_indexer.logger", + create=True, + ) as logger: + dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id=7, + logical_page_table=logical_pages, + current_index_kv=current_index_kv, + ) + + self.assertIs(dense_buffer, fallback_buffer) + self.assertIs(dense_pages, fallback_pages) + logger.warning.assert_not_called() + + def test_index_partial_current_first_layer_prefetch_miss_does_not_warn(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + class FakePool: + page_size = 4 + index_head_dim = 4 + start_layer = 0 + + def get_index_k_with_scale_buffer(self, layer_id): + return torch.zeros((16, 32), dtype=torch.uint8) + + class FakeLayout: + page_size = 4 + cp_size = 1 + cp_rank = 0 + + class MissingPrefetcher: + def consume_prefix_with_current(self, **kwargs): + return None + + fallback_buffer = torch.full((8, 32), 7, dtype=torch.uint8) + fallback_pages = torch.tensor([[0, 1]], dtype=torch.int64) + forward_batch = SimpleNamespace( + token_to_kv_pool=FakePool(), + uses_cp_shared_kv=True, + cp_shared_kv_layout=FakeLayout(), + cp_shared_kv_index_prefetcher=MissingPrefetcher(), + extend_prefix_lens_cpu=[4], + extend_seq_lens_cpu=[4], + cp_local_out_cache_loc=torch.tensor([4, 5, 6, 7], dtype=torch.int64), + ) + logical_pages = torch.tensor([[0, 1]], dtype=torch.int64) + current_index_kv = ( + torch.zeros((4, 4), dtype=torch.uint8), + torch.zeros((4,), dtype=torch.float32), + ) + indexer = object.__new__(nsa_indexer.Indexer) + + with patch.object( + nsa_indexer, + "get_or_build_shared_paged_buffer_slot_remap", + return_value=object(), + ), patch.object( + nsa_indexer, + "materialize_prefix_and_reuse_current_index_page_slots", + return_value=(fallback_buffer, fallback_pages), + ), patch( + "sglang.srt.layers.attention.nsa.nsa_indexer.logger", + create=True, + ) as logger: + dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id=0, + logical_page_table=logical_pages, + current_index_kv=current_index_kv, + ) + + self.assertIs(dense_buffer, fallback_buffer) + self.assertIs(dense_pages, fallback_pages) + logger.warning.assert_not_called() + def test_index_prefetch_create_skip_logs_fallback_when_enabled(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch