From 43ad2fe52d8129a981b13d51f587ff4d780f2a92 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Wed, 6 May 2026 05:27:43 +0800 Subject: [PATCH] Reduce CP shared KV overhead without changing ownership semantics The shared-KV path now keeps more CP metadata on-device and reuses physical out-cache locations across MLA and NSA index writes, so each layer avoids repeating logical-to-physical remaps. The in-seq CP all-gather rerange path now delegates to tai-kernel when available and falls back to the existing torch split/cat path with an explicit log. This also extends the Phase8 prefetch machinery to cover shared KV materialization metadata and keeps debug/fallback behavior gated so the fast path is not polluted by diagnostic checks. Constraint: Custom CP kernels must live in tai-kernel and be imported lazily from SGLang Constraint: Decode does not use CP; these changes target NSA prefill CP in-seq-split shared KV Rejected: Recompute physical local cache locations separately for MLA and index writes | repeats the same remap work every layer Rejected: Keep the in-seq rerange Triton code inline in SGLang | duplicates kernel ownership and blocks tai-kernel reuse Confidence: medium Scope-risk: moderate Directive: Keep CP collective ordering identical across ranks; do not add rank-local fallback decisions inside shared KV materialize paths Tested: Remote g0034 container py_compile for modified SGLang/tai-kernel files; remote pytest test/registered/unit/layers/test_nsa_cp_utils.py passed with 24 tests Not-tested: Full multi-node GLM5 prefill/decode throughput after the final commit boundary --- ...fill_cp_phase8_mla_prefix_prefetch_plan.md | 82 +++- .../attention/nsa/cp_shared_kv_prefetch.py | 441 ++++++++++++++++++ .../attention/nsa/cp_shared_kv_runtime.py | 99 +++- .../srt/layers/attention/nsa/nsa_indexer.py | 83 +++- .../sglang/srt/layers/attention/nsa/utils.py | 180 ++++++- .../srt/layers/attention/nsa_backend.py | 16 + .../srt/model_executor/forward_batch_info.py | 2 + .../attention_forward_methods/forward_mla.py | 7 +- .../unit/layers/test_nsa_cp_utils.py | 40 ++ .../mem_cache/test_cp_shared_kv_runtime.py | 248 ++++++++++ 10 files changed, 1152 insertions(+), 46 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md index 55c51376d..f95ed9fe8 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md @@ -1,8 +1,8 @@ -# NSA Prefill CP Phase 8: MLA prefix prefetch +# NSA Prefill CP Phase 8: MLA / index prefix prefetch -Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 chunked prefill / radix cache hit 场景引入 **MLA KV one-layer-ahead prefix prefetch**，把历史 prefix 的 shared KV materialize 从当前层 attention 前的同步阻塞路径里移出，并把等待点延迟到下一层真正消费 prefetched KV 时。 +Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 chunked prefill / radix cache hit 场景引入 **MLA KV 与 NSA index K/scale 的 one-layer-ahead prefix prefetch**，把历史 prefix 的 shared KV/index materialize 从当前层同步阻塞路径里移出，并把等待点延迟到下一层真正消费 prefetched buffer 时。 -本阶段只做 MLA KV prefix prefetch。暂时不做 index K/scale prefetch，不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量，也不保留 `SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION` 这类 wait 策略开关。 +本阶段不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量，也不保留 `SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION` 这类 wait 策略开关。`SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=1` 现在同时打开 MLA KV prefix prefetch 和 NSA index K/scale prefix prefetch。 --- @@ -12,15 +12,16 @@ Phase 8 的目标是在不改变 Phase 2-7 shared KV 语义的前提下，为 ch Phase 8 已按 deferred-consume 策略实现： ```text +Layer L index/topk 后启动 Layer L+1 prefix index K/scale prefetch Layer L attention 前启动 Layer L+1 prefix MLA KV prefetch -Layer L forward_extend 返回时不等待该 prefetch -Layer L+1 consume prefetched KV 时 wait event，再补齐 suffix/current pages +Layer L attention 返回时不强制等待 prefetch +Layer L+1 consume prefetched buffer 时 wait event，再补齐 suffix/current pages ``` 保留的环境变量： ```text -SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 # 生产开关，默认关闭 +SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 # 生产开关，默认关闭；同时控制 MLA KV 与 index K/scale prefetch SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1 # 调试日志，只打印 probe layer，默认关闭 ``` @@ -48,13 +49,13 @@ owner-sharded physical MLA KV on each CP rank -> existing NSA attention kernel ``` -在 chunked prefill 的第二个 chunk 及之后，或者 radix cache 命中时，`extend_prefix_len > 0`。这部分 prefix KV 已经在之前的 chunk/request 中写入 persistent KV pool；当前层 attention 仍会同步 materialize 整个可见 KV，包括历史 prefix 和当前 suffix。 +在 chunked prefill 的第二个 chunk 及之后，或者 radix cache 命中时，`extend_prefix_len > 0`。这部分 prefix KV 与 index K/scale 已经在之前的 chunk/request 中写入 persistent pool；当前层 index/topk 和 attention 仍会同步 materialize 整个可见范围，包括历史 prefix 和当前 suffix。 Phase 8 的机会是： ```text 历史 prefix 部分已经存在，可以提前为下一层 materialize。 -当前 suffix 部分必须等下一层 prepare 写入后才能 materialize。 +当前 suffix 部分必须等下一层 prepare/index write 写入后才能 materialize。 ``` 因此 Phase 8 不做跨 chunk 的 dense KV cache 复用，而是做 **每次 forward 内、相邻 layer 之间的一层提前预取**。 @@ -173,12 +174,12 @@ _get_topk_in_seq_cp_pair(...) -> materialize_shared_paged_buffer(...) ``` -Phase 8 暂时不碰 index path。原因： +Index K/scale prefetch 在 MLA KV prefetch 之后补齐。当前约束： 1. index materialize 已经完成 Phase 6 的一次合并； 2. indexer 在 MLA attention 之前执行，调度窗口不同； -3. 本阶段目标是先验证 MLA prefix prefetch 是否能显著隐藏最大块的 KV materialize； -4. 避免同时改变 topk/index 和 attention KV 两条路径，降低正确性风险。 +3. 当前实现先落地 MLA KV，再补 index K/scale；两者共用同一个 page-aligned slot layout； +4. index K/scale prefetch 只覆盖 PAGED topk 路径，RAGGED 仍不进入本阶段。 --- @@ -192,9 +193,16 @@ Phase 8 暂时不碰 index path。原因： Layer L attention 计算期间， 提前为 Layer L+1 materialize prefix MLA KV。 +Layer L topk/indexer 完成后， +提前为 Layer L+1 materialize prefix index K/scale。 + Layer L+1 attention 前， 复用已经 materialize 完成的 prefix dense KV， 只同步补齐 current/suffix pages。 + +Layer L+1 topk/indexer 前， +复用已经 materialize 完成的 prefix dense index buffer， +只同步补齐 current/suffix pages。 ``` ### 3.2 性能目标 @@ -225,7 +233,6 @@ next layer consume waits only if prefetch has not finished Phase 8 不做： -- 不做 NSA index K/scale prefetch； - 不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND`； - 不做 bandwidth throttle/page budget； - 不做多层 dense KV 常驻缓存； @@ -576,7 +583,7 @@ disable MLA prefetch and fallback sync materialize 2. 太早会和当前层 index materialize/topk 抢带宽； 3. 太早会增加 NCCL collective 顺序风险。 -推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动： +MLA KV 推荐在 `NativeSparseAttnBackend.forward_extend(...)` 中启动： ```text 1. 当前层需要的 kv_cache/page_table_1 已经准备好； @@ -614,6 +621,19 @@ if prefetcher is not None: return attn_output ``` +Index K/scale 的启动点在 `Indexer.forward_cuda(...)` 末尾： + +```python +topk_result = run_current_layer_topk(...) +index_prefetcher.start_next_layer_prefix( + next_layer_id=layer_id + 1, + token_to_kv_pool=forward_batch.token_to_kv_pool, +) +return topk_result +``` + +当前层 `_maybe_materialize_shared_index_buffer(...)` 优先消费 `cp_shared_kv_index_prefetcher.consume(...)`；miss 时回退原来的 `materialize_shared_paged_buffer(...)`，因此关闭环境变量或 prefetch 条件不满足时语义不变。 + `wait_attention_window()` 仍在 `forward_extend(...)` 的 finally 中调用，但当前实现不做同步等待，只保留 deferred 日志/状态检查入口。真正的同步点在下一层 `consume(...)`： ```text @@ -790,9 +810,10 @@ Phase 8 的 local copy/remap helper 应尽量复用 Phase 7 的 tai materialize ```text SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False) forward_batch.cp_shared_kv_mla_prefetcher +forward_batch.cp_shared_kv_index_prefetcher ``` -### Step 2: 拆分 token materialize helper +### Step 2: 拆分 token/page materialize helper 修改： @@ -803,6 +824,8 @@ forward_batch.cp_shared_kv_mla_prefetcher ```text slot remap builder range local token KV materialize into existing dense buffer +range local paged buffer materialize into existing dense buffer +logical page -> slot dense page remap sync range all-reduce async all-reduce wrapper ``` @@ -813,7 +836,7 @@ async all-reduce wrapper full materialize 结果 == prefix materialize + suffix materialize 结果 ``` -### Step 3: 新增 MLA prefetcher +### Step 3: 新增 MLA / index prefetcher 新增： @@ -822,7 +845,7 @@ full materialize 结果 == prefix materialize + suffix materialize 结果 职责： ```text -gate +common gate allocate dense buffer prefetch prefix range record event @@ -831,7 +854,15 @@ defer wait until consume fallback ``` -### Step 4: 接入 `nsa_backend.forward_extend` +Index K/scale prefetch 与 MLA prefetch 使用相同 slot-layout page table： + +```text +real_page_table.reshape(-1)[slot] -> dense page slot + 1 +``` + +区别是 MLA KV materialize 的单位是 token rows，需要 `page_size` 展开；index K/scale materialize 的单位是 page rows，prefix/suffix range 直接对应 dense page 行。 + +### Step 4: 接入 `nsa_backend.forward_extend` / `nsa_indexer` 修改： @@ -840,7 +871,7 @@ fallback 接入位置： ```text -shared KV PAGED path 的 MLA materialize 分支。 +shared KV PAGED path 的 MLA materialize 分支，以及 NSA indexer 的 `_maybe_materialize_shared_index_buffer(...)`。 ``` 逻辑： @@ -853,6 +884,15 @@ run attention do not wait before return; consume waits before use ``` +Index K/scale 的启动点在当前层 topk/indexer 完成后： + +```text +Layer L writes current layer index K/scale +Layer L materializes/uses current layer index K/scale for topk +Layer L starts Layer L+1 prefix index K/scale prefetch +Layer L+1 indexer consumes prefetched prefix and fills suffix +``` + ### Step 5: 单元测试 修改： @@ -861,12 +901,13 @@ do not wait before return; consume waits before use 新增覆盖： -1. prefix/suffix range materialize 拼接结果等价于 full materialize； +1. token/paged prefix/suffix range materialize 拼接结果等价于 full materialize； 2. prefix page-aligned gate； 3. prefix_len=0 不 prefetch； 4. non-PAGED / batch_size>1 / debug enabled fallback； 5. consume handle layer mismatch fallback； 6. started async handle 必须 wait。 +7. index materialize 优先消费 prefetched buffer，miss 后回退 full materialize； ### Step 6: 远端集成验证 @@ -885,7 +926,7 @@ router: g0034 3. 重复请求触发 radix hit； 4. 检查输出质量； 5. 检查没有 collective hang； -6. profile 确认 `cp_shared_kv.materialize.token` 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。 +6. profile 确认 `cp_shared_kv.materialize.token` / index paged materialize 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。 --- @@ -972,4 +1013,3 @@ Fallback： debug enabled 时自动回到现有同步 materialize pynccl unavailable 时自动回到现有同步 materialize ``` - 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 6c49694c9..5100f4e7b 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 @@ -10,13 +10,19 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( _all_reduce_materialized_buffer_async, _all_reduce_materialized_buffer_range, build_shared_token_kv_slot_remap, + build_slot_page_inverse_optimized, + build_slot_page_remap, cp_shared_kv_debug_enabled, cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, cp_shared_kv_mla_prefetch_should_log_layer, filter_locs_mappable_to_physical_pool, + filter_pages_mappable_to_physical_pool, + materialize_local_paged_buffer_page_slots_into, materialize_local_token_kv_page_slots_into, + remap_logical_pages_to_slot_dense_pages, remap_logical_locs_to_slot_dense_locs_optimized, + slot_range_to_page_slice, slot_range_to_token_slice, ) from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split @@ -30,6 +36,14 @@ def _prefetch_log(message: str, *args) -> None: cp_shared_kv_mla_prefetch_log(message, *args) +def _index_prefetch_fallback_log(reason: str, message: str, *args) -> None: + logger.info( + "CP shared KV index prefetch fallback (%s): " + message, + reason, + *args, + ) + + def _is_cuda_stream_capturing() -> bool: if not torch.cuda.is_available(): return False @@ -46,6 +60,13 @@ class CpSharedKVMlaPrefetchHandle: event: torch.cuda.Event +@dataclass +class CpSharedKVIndexPrefetchHandle: + layer_id: int + dense_page_buffer: torch.Tensor + event: torch.cuda.Event + + class CpSharedKVMlaPrefetcher: """One-layer-ahead MLA prefix materialize prefetch for CP shared KV. @@ -416,3 +437,423 @@ class CpSharedKVMlaPrefetcher: self.layout.cp_size, *args, ) + + +class CpSharedKVIndexPrefetcher: + """One-layer-ahead NSA index K/scale prefix materialize prefetch. + + The index buffer is page-granular, unlike MLA KV which is token-granular. + It still uses the same slot-layout page table as MLA prefetch so the next + layer can consume the prefetched prefix and synchronously fill only current + suffix pages before running fp8 MQA/topk. + """ + + def __init__( + self, + *, + layout: CpSharedKVLayout, + prefix_pages: int, + slot_logical_pages: torch.Tensor, + page_inverse: torch.Tensor, + dense_num_pages: int, + ) -> None: + self.layout = layout + self.prefix_pages = prefix_pages + self.slot_logical_pages = slot_logical_pages + self.page_inverse = page_inverse + self.dense_num_pages = dense_num_pages + self.total_slots = int(slot_logical_pages.numel()) + self.stream = torch.cuda.Stream() + self.handles: dict[int, CpSharedKVIndexPrefetchHandle] = {} + self.pending_attention_handle: Optional[CpSharedKVIndexPrefetchHandle] = None + self.disabled = False + + @classmethod + def maybe_create( + cls, + *, + forward_batch: Any, + metadata: Any, + topk_transform_is_paged: bool, + ) -> Optional["CpSharedKVIndexPrefetcher"]: + if not cp_shared_kv_mla_prefetch_enabled(): + return None + if cp_shared_kv_debug_enabled(): + _index_prefetch_fallback_log( + "debug_enabled", + "SGLANG_DEBUG_CP_SHARED_KV is enabled.", + ) + return None + if not torch.cuda.is_available() or _is_cuda_stream_capturing(): + _index_prefetch_fallback_log( + "cuda_unavailable_or_stream_capturing", + "CUDA is unavailable or the current stream is capturing. cuda_available=%s", + torch.cuda.is_available(), + ) + return None + if not getattr(forward_batch, "uses_cp_shared_kv", False): + _index_prefetch_fallback_log( + "not_cp_shared_kv", + "forward batch is not using CP shared KV.", + ) + return None + if getattr(forward_batch, "hisparse_coordinator", None) is not None: + _index_prefetch_fallback_log( + "hisparse", + "HiSparse coordinator is active.", + ) + return None + forward_mode = getattr(forward_batch, "forward_mode", None) + if forward_mode is None or not forward_mode.is_context_parallel_extend(): + _index_prefetch_fallback_log( + "not_context_parallel_extend", + "forward mode is not context-parallel extend. forward_mode=%s", + forward_mode, + ) + return None + if not is_nsa_prefill_cp_in_seq_split(): + _index_prefetch_fallback_log( + "not_in_seq_split", + "NSA prefill CP mode is not in-seq-split.", + ) + return None + if not topk_transform_is_paged: + _index_prefetch_fallback_log( + "not_paged_topk", + "topk transform is not PAGED.", + ) + return None + if int(getattr(forward_batch, "batch_size", 0)) != 1: + _index_prefetch_fallback_log( + "batch_size", + "batch size is not supported. batch_size=%s", + getattr(forward_batch, "batch_size", None), + ) + return None + + token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None) + if token_to_kv_pool is None: + _index_prefetch_fallback_log( + "missing_token_to_kv_pool", + "forward batch has no token_to_kv_pool.", + ) + return None + if getattr(token_to_kv_pool, "layer_transfer_counter", None) is not None: + _index_prefetch_fallback_log( + "layer_transfer_active", + "layer transfer is active.", + ) + return None + + layout = getattr(forward_batch, "cp_shared_kv_layout", None) + if layout is None: + _index_prefetch_fallback_log( + "missing_layout", + "forward batch has no CP shared KV layout.", + ) + return None + + extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None) + if extend_prefix_lens_cpu is None or len(extend_prefix_lens_cpu) != 1: + _index_prefetch_fallback_log( + "bad_prefix_lens_metadata", + "extend_prefix_lens_cpu is missing or not single-batch. value=%s", + extend_prefix_lens_cpu, + ) + return None + page_size = int(getattr(token_to_kv_pool, "page_size", 1)) + if page_size <= 1: + _index_prefetch_fallback_log( + "bad_page_size", + "page size is not supported. page_size=%s", + page_size, + ) + return None + extend_prefix_len = int(extend_prefix_lens_cpu[0]) + if extend_prefix_len <= 0 or extend_prefix_len % page_size != 0: + _index_prefetch_fallback_log( + "prefix_not_page_aligned", + "prefix length is zero or not page-aligned. prefix_len=%s page_size=%s", + extend_prefix_len, + page_size, + ) + return None + prefix_pages = extend_prefix_len // page_size + + real_page_table = getattr(metadata, "real_page_table", None) + page_table_1 = getattr(metadata, "page_table_1", None) + if real_page_table is None or page_table_1 is None: + _index_prefetch_fallback_log( + "missing_page_tables", + "metadata is missing real_page_table or page_table_1.", + ) + return None + if prefix_pages <= 0 or prefix_pages > int(real_page_table.numel()): + _index_prefetch_fallback_log( + "prefix_pages_out_of_range", + "prefix pages are outside real page table. prefix_pages=%s real_pages=%s", + prefix_pages, + int(real_page_table.numel()), + ) + 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 + + try: + first_layer_id = int(getattr(token_to_kv_pool, "start_layer", 0)) + page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer( + layer_id=first_layer_id + ) + remap_logical_pages = filter_pages_mappable_to_physical_pool( + logical_pages=real_page_table, + layout=layout, + physical_page_capacity=page_buffer.shape[0], + ) + slot_logical_pages, _ = build_slot_page_remap(remap_logical_pages) + logical_page_capacity = max(int(page_buffer.shape[0]) - 1, 0) * ( + layout.cp_size + ) + 1 + page_inverse = build_slot_page_inverse_optimized( + slot_logical_pages, + logical_page_capacity=logical_page_capacity, + ) + except Exception as exc: + _index_prefetch_fallback_log( + "init_exception", + "failed to initialize index prefetcher; falling back to sync materialize. error=%s", + exc, + ) + logger.exception("Failed to initialize CP shared KV index prefetcher.") + return None + + _prefetch_log( + "index_create cp_rank=%s cp_size=%s prefix_pages=%s total_slots=%s dense_pages=%s page_size=%s", + layout.cp_rank, + layout.cp_size, + prefix_pages, + int(slot_logical_pages.numel()), + int(slot_logical_pages.numel()) + 1, + page_size, + ) + + return cls( + layout=layout, + prefix_pages=prefix_pages, + slot_logical_pages=slot_logical_pages, + page_inverse=page_inverse, + dense_num_pages=int(slot_logical_pages.numel()) + 1, + ) + + def _layer_in_pool(self, token_to_kv_pool: Any, layer_id: int) -> bool: + start_layer = int(getattr(token_to_kv_pool, "start_layer", 0)) + kv_buffer = getattr(token_to_kv_pool, "kv_buffer", None) + if kv_buffer is None: + return layer_id >= start_layer + return start_layer <= layer_id < start_layer + len(kv_buffer) + + def consume( + self, + *, + layer_id: int, + page_buffer: torch.Tensor, + logical_pages: torch.Tensor, + ) -> Optional[tuple[torch.Tensor, torch.Tensor]]: + if self.disabled: + self._log_layer( + layer_id, + "index_consume_skip reason=disabled layer=%s", + layer_id, + ) + return None + + handle = self.handles.pop(layer_id, None) + if handle is None: + self._log_layer(layer_id, "index_consume_miss layer=%s", layer_id) + return None + if self.pending_attention_handle is handle: + self.pending_attention_handle = None + if handle.layer_id != layer_id: + self.disabled = True + self._log_layer( + layer_id, + "index_consume_skip reason=layer_mismatch expected=%s actual=%s", + layer_id, + handle.layer_id, + ) + return None + + torch.cuda.current_stream().wait_event(handle.event) + dense_page_buffer = handle.dense_page_buffer + suffix_slots = self.total_slots - self.prefix_pages + + if self.prefix_pages < self.total_slots: + materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + start_slot=self.prefix_pages, + end_slot=self.total_slots, + ) + suffix_rows = slot_range_to_page_slice( + self.prefix_pages, + self.total_slots, + ) + _all_reduce_materialized_buffer_range( + dense_page_buffer, + self.layout.cp_size, + suffix_rows.start, + suffix_rows.stop, + ) + + self._log_layer( + layer_id, + "index_consume_hit layer=%s prefix_pages=%s suffix_slots=%s dense_pages=%s", + layer_id, + self.prefix_pages, + suffix_slots, + int(dense_page_buffer.shape[0]), + ) + + logical_pages = filter_pages_mappable_to_physical_pool( + logical_pages=logical_pages, + layout=self.layout, + physical_page_capacity=page_buffer.shape[0], + ) + dense_pages = remap_logical_pages_to_slot_dense_pages( + logical_pages, + page_inverse=self.page_inverse, + ) + return dense_page_buffer, dense_pages + + def start_next_layer_prefix( + self, + *, + next_layer_id: int, + token_to_kv_pool: Any, + ) -> None: + if self.disabled: + self._log_next_layer( + next_layer_id, + "index_start_skip reason=disabled next_layer=%s", + next_layer_id, + ) + return + if next_layer_id in self.handles: + self._log_next_layer( + next_layer_id, + "index_start_skip reason=already_started next_layer=%s", + next_layer_id, + ) + return + if not self._layer_in_pool(token_to_kv_pool, next_layer_id): + self._log_next_layer( + next_layer_id, + "index_start_skip reason=layer_out_of_pool next_layer=%s", + next_layer_id, + ) + return + + try: + page_buffer = token_to_kv_pool.get_index_k_with_scale_buffer( + layer_id=next_layer_id + ) + 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 + + current_stream = torch.cuda.current_stream() + self.stream.wait_stream(current_stream) + try: + with torch.cuda.stream(self.stream): + dense_page_buffer = page_buffer.new_zeros( + (self.dense_num_pages, *page_buffer.shape[1:]) + ) + materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=self.slot_logical_pages, + layout=self.layout, + start_slot=0, + end_slot=self.prefix_pages, + ) + prefix_rows = slot_range_to_page_slice(0, self.prefix_pages) + event = _all_reduce_materialized_buffer_async( + dense_page_buffer[prefix_rows], + cp_size=self.layout.cp_size, + stream=self.stream, + ) + if event is None: + self.disabled = True + self._log_next_layer( + next_layer_id, + "index_start_disable reason=async_reduce_unavailable next_layer=%s", + next_layer_id, + ) + return + except Exception: + logger.exception("Failed to start CP shared KV index prefix prefetch.") + self.disabled = True + self._log_next_layer( + next_layer_id, + "index_start_disable reason=start_exception next_layer=%s", + next_layer_id, + ) + return + + handle = CpSharedKVIndexPrefetchHandle( + layer_id=next_layer_id, + dense_page_buffer=dense_page_buffer, + event=event, + ) + self.handles[next_layer_id] = handle + self.pending_attention_handle = handle + self._log_next_layer( + next_layer_id, + "index_start next_layer=%s prefix_pages=%s dense_pages=%s", + next_layer_id, + self.prefix_pages, + int(dense_page_buffer.shape[0]), + ) + + def wait_attention_window(self) -> None: + handle = self.pending_attention_handle + if handle is None: + return + self._log_next_layer( + handle.layer_id, + "index_attention_wait_deferred next_layer=%s", + handle.layer_id, + ) + + def _log_layer(self, layer_id: int, message: str, *args) -> None: + if cp_shared_kv_mla_prefetch_should_log_layer(layer_id): + self._log(message, *args) + + def _log_next_layer(self, next_layer_id: int, message: str, *args) -> None: + if cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id): + self._log(message, *args) + + def _log(self, message: str, *args) -> None: + _prefetch_log( + "cp_rank=%s cp_size=%s " + message, + self.layout.cp_rank, + self.layout.cp_size, + *args, + ) 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 4fda36787..78bfe146c 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 @@ -773,6 +773,43 @@ def remap_logical_locs_to_slot_dense_locs( return torch.where(mapped, dense_values, dense_locs) +def remap_logical_pages_to_slot_dense_pages( + logical_pages: torch.Tensor, + page_inverse: torch.Tensor, +) -> torch.Tensor: + """Map logical page ids through a fixed-shape slot page inverse. + + This is the paged-buffer equivalent of + `remap_logical_locs_to_slot_dense_locs`. Page 0 remains the dummy dense page + 0; negative sentinels and pages absent from `page_inverse` remain `-1`. + """ + + dense_pages_out = torch.full_like(logical_pages, -1) + if logical_pages.numel() == 0 or page_inverse.numel() == 0: + return dense_pages_out + + pages_long = logical_pages.to(torch.long) + valid_pages = pages_long >= 0 + safe_pages = torch.where(valid_pages, pages_long, torch.zeros_like(pages_long)) + pages_in_range = safe_pages < page_inverse.numel() + clamped_pages = torch.clamp(safe_pages, max=page_inverse.numel() - 1) + dense_pages = page_inverse[clamped_pages] + mapped = valid_pages & pages_in_range & (dense_pages >= 0) + + if cp_shared_kv_debug_enabled() and torch.any( + valid_pages & pages_in_range & (dense_pages < 0) + ): + missing_pages = pages_long[valid_pages & pages_in_range & (dense_pages < 0)] + raise RuntimeError( + "CP shared KV slot remap got logical pages outside remap_logical_pages. " + f"missing_page_min={int(missing_pages.min().item())} " + f"missing_page_max={int(missing_pages.max().item())} " + f"logical_pages={tensor_debug_summary(logical_pages)}" + ) + + return torch.where(mapped, dense_pages.to(logical_pages.dtype), dense_pages_out) + + def build_shared_token_kv_slot_remap( kv_cache: torch.Tensor, logical_locs: torch.Tensor | None, @@ -1261,7 +1298,50 @@ def materialize_local_paged_buffer_page_slots( if slot_logical_pages.numel() == 0: return dense_page_buffer - logical_pages = slot_logical_pages.reshape(-1).to(torch.long) + materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=dense_page_buffer, + slot_logical_pages=slot_logical_pages, + layout=layout, + start_slot=0, + end_slot=int(slot_logical_pages.numel()), + ) + return dense_page_buffer + + +def materialize_local_paged_buffer_page_slots_into( + page_buffer: torch.Tensor, + dense_page_buffer: torch.Tensor, + slot_logical_pages: torch.Tensor, + layout: CpSharedKVLayout, + start_slot: int, + end_slot: int | None = None, +) -> None: + """Materialize a slot range into an existing dense paged buffer. + + Dense page 0 is the dummy page; slot `i` writes dense page `i + 1`. + This is used by layer prefetch: prefix pages are written asynchronously, + then only suffix pages are filled synchronously when the layer is reached. + """ + + flat_slot_logical_pages = slot_logical_pages.reshape(-1) + total_slots = int(flat_slot_logical_pages.numel()) + if end_slot is None: + end_slot = total_slots + if start_slot < 0 or end_slot < start_slot or end_slot > total_slots: + raise ValueError( + "Invalid CP shared KV paged slot materialize range: " + f"start_slot={start_slot} end_slot={end_slot} total_slots={total_slots}" + ) + if start_slot == end_slot: + return + if dense_page_buffer.shape[0] < total_slots + 1: + raise ValueError( + "CP shared KV dense paged buffer is too small for slot materialize: " + f"dense_pages={dense_page_buffer.shape[0]} expected_at_least={total_slots + 1}" + ) + + logical_pages = flat_slot_logical_pages[start_slot:end_slot].to(torch.long) owned_mask = layout.owned_pages_mask(logical_pages) physical_pages = layout.logical_pages_to_physical(logical_pages).to(torch.long) safe_physical_pages = torch.where( @@ -1272,8 +1352,21 @@ def materialize_local_paged_buffer_page_slots( gathered = page_buffer[safe_physical_pages] owned_view = owned_mask.view(-1, *([1] * (gathered.ndim - 1))) zero = torch.zeros((), dtype=page_buffer.dtype, device=page_buffer.device) - dense_page_buffer[1:].copy_(torch.where(owned_view, gathered, zero)) - return dense_page_buffer + dense_page_buffer[start_slot + 1 : end_slot + 1].copy_( + torch.where(owned_view, gathered, zero) + ) + + +def slot_range_to_page_slice( + start_slot: int, + end_slot: int, +) -> slice: + if start_slot < 0 or end_slot < start_slot: + raise ValueError( + "Invalid CP shared KV slot page slice range: " + f"start_slot={start_slot} end_slot={end_slot}" + ) + return slice(start_slot + 1, end_slot + 1) def paged_copy_debug_checksum( diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 7d63d9d81..12b64efde 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1,6 +1,7 @@ from __future__ import annotations import contextlib +import logging from abc import ABC, abstractmethod from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple @@ -34,6 +35,7 @@ _is_cuda = is_cuda() _is_hip = is_hip() _is_npu = is_npu() _is_fp8_fnuz = is_fp8_fnuz() +logger = logging.getLogger(__name__) if _is_cuda: try: import deep_gemm @@ -53,6 +55,7 @@ from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.nsa.utils import ( cp_all_gather_rerange_output, get_cp_shared_kv_local_out_cache_loc, + get_cp_shared_kv_local_physical_out_cache_loc, is_nsa_enable_prefill_cp, is_nsa_prefill_cp_in_seq_split, log_cp_shared_kv_direct_write_fallback, @@ -102,6 +105,24 @@ def _compute_contiguous_valid_cp_query_count( return max(0, min(actual_seq_q, valid_count)) +def _log_cp_shared_kv_index_prefetch_fallback( + reason: str, + message: str, + *args, +) -> None: + """Log every index prefetch fallback event. + + Warmup and real requests can hit the same reason independently; do not + dedupe here or later real fallbacks become invisible during profiling. + """ + + logger.info( + "CP shared KV index prefetch fallback (%s): " + message, + reason, + *args, + ) + + class BaseIndexerMetadata(ABC): @abstractmethod def get_seqlens_int32(self) -> torch.Tensor: @@ -282,6 +303,39 @@ class Indexer(MultiPlatformOp): assert forward_batch.cp_shared_kv_layout is not None layout = forward_batch.cp_shared_kv_layout + index_prefetcher = getattr( + forward_batch, "cp_shared_kv_index_prefetcher", None + ) + if index_prefetcher is not None: + prefetched = index_prefetcher.consume( + layer_id=layer_id, + page_buffer=index_buffer, + logical_pages=logical_page_table, + ) + if prefetched is not None: + if cp_shared_kv_debug_enabled(): + cp_shared_kv_debug_log( + "index_materialize_prefetch_hit", + "NSA index materialize prefetch hit cp_rank=%s layer=%s dense_pages=%s buffer_ck=%s", + layout.cp_rank, + layer_id, + tensor_debug_summary(prefetched[1]), + tensor_debug_checksum(prefetched[0]), + ) + return prefetched + start_layer = int(getattr(forward_batch.token_to_kv_pool, "start_layer", 0)) + if int(layer_id) > start_layer: + _log_cp_shared_kv_index_prefetch_fallback( + "consume_miss", + "prefetcher did not provide layer buffer; falling back to " + "sync paged materialize. layer=%s start_layer=%s cp_rank=%s " + "logical_page_table_shape=%s", + layer_id, + start_layer, + layout.cp_rank, + tuple(logical_page_table.shape), + ) + if cp_shared_kv_debug_enabled(): cp_shared_kv_debug_log( "index_materialize_call", @@ -303,9 +357,24 @@ class Indexer(MultiPlatformOp): layer_id, tensor_debug_summary(dense_pages), tensor_debug_checksum(materialized), - ) + ) return materialized, dense_pages + def _maybe_start_next_layer_index_prefetch( + self, + forward_batch: ForwardBatch, + layer_id: int, + ) -> None: + index_prefetcher = getattr( + forward_batch, "cp_shared_kv_index_prefetcher", None + ) + if index_prefetcher is None: + return + index_prefetcher.start_next_layer_prefix( + next_layer_id=layer_id + 1, + token_to_kv_pool=forward_batch.token_to_kv_pool, + ) + def _filter_shared_index_write( self, forward_batch: ForwardBatch, @@ -1355,12 +1424,9 @@ class Indexer(MultiPlatformOp): if local_out_loc.numel() == 0: return True - assert forward_batch.cp_shared_kv_layout is not None - physical_out_loc = ( - forward_batch.cp_shared_kv_layout.logical_locs_to_physical( - local_out_loc - ).contiguous() - ) + physical_out_loc = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) + if physical_out_loc is None: + return False self._store_index_k_cache( forward_batch=forward_batch, layer_id=layer_id, @@ -1579,7 +1645,7 @@ class Indexer(MultiPlatformOp): forward_batch.nsa_cp_metadata is not None and is_nsa_prefill_cp_in_seq_split() ): - return self._get_topk_in_seq_cp_pair( + topk_result = self._get_topk_in_seq_cp_pair( forward_batch, layer_id, q_fp8, @@ -1604,6 +1670,7 @@ class Indexer(MultiPlatformOp): topk=self.index_topk, layer_id=layer_id, ) + self._maybe_start_next_layer_index_prefetch(forward_batch, layer_id) return topk_result def forward_npu( diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 2996c5293..5a1bc6298 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -167,6 +167,8 @@ class PageAlignedInSeqSplitInfo: @dataclass class NSAContextParallelMetadata: split_list: List[int] = None + split_list_tensor: torch.Tensor = None + split_prefix_tensor: torch.Tensor = None max_rank_len: List[int] = None zigzag_index: List[int] = None per_rank_actual_token: List[int] = None @@ -529,6 +531,44 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"): return local_out_cache_loc +def get_cp_shared_kv_local_physical_out_cache_loc(forward_batch: "ForwardBatch"): + """Return cached physical rows for this CP rank's shared-KV direct writes. + + `get_cp_shared_kv_local_out_cache_loc` returns logical shared-KV locs because + radix/scheduler/PD metadata are expressed in the global logical address + space. Persistent writes into this rank's compact physical pool need the + same locs remapped through `CpSharedKVLayout`. The logical loc tensor is + batch-scoped, not layer-scoped, so cache the physical remap on ForwardBatch + and reuse it for both MLA KV and NSA index K/scale writes across layers. + """ + + cached = getattr(forward_batch, "cp_local_physical_out_cache_loc", None) + if cached is not None: + return cached + + local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch) + if local_out_cache_loc is None: + return None + + if local_out_cache_loc.numel() == 0: + forward_batch.cp_local_physical_out_cache_loc = local_out_cache_loc + return local_out_cache_loc + + layout = getattr(forward_batch, "cp_shared_kv_layout", None) + if layout is None: + log_cp_shared_kv_direct_write_fallback( + "missing_layout", + "cp_shared_kv_layout is missing while building physical out_cache_loc", + ) + return None + + physical_out_cache_loc = layout.logical_locs_to_physical( + local_out_cache_loc + ).contiguous() + forward_batch.cp_local_physical_out_cache_loc = physical_out_cache_loc + return physical_out_cache_loc + + def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor): if is_nsa_prefill_cp_round_robin_split(): cp_size = get_attention_cp_size() @@ -655,7 +695,7 @@ def nsa_use_prefill_cp(forward_batch, nsa_enable_prefill_cp=None): return False -def cp_attn_tp_all_gather_reorganazied_into_tensor( +def _cp_attn_tp_all_gather_padded_tensor( input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op ): """ @@ -683,11 +723,14 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor( get_attention_cp_group().cp_all_gather_into_tensor_async( input_tensor_all, input_, stream_op ) - # step3 + return input_tensor_all + + +def _trim_cp_rank_padding_after_all_gather(input_tensor_all, forward_batch): outputs_list_max = list( torch.split(input_tensor_all, forward_batch.nsa_cp_metadata.max_rank_len, dim=0) ) - outputs = torch.cat( + return torch.cat( [ outputs_list_max[index][:per_rank_len] for index, per_rank_len in enumerate( @@ -696,9 +739,114 @@ def cp_attn_tp_all_gather_reorganazied_into_tensor( ], dim=0, ) + + +def cp_attn_tp_all_gather_reorganazied_into_tensor( + input_: torch.Tensor, total_len, attn_tp_size, forward_batch, stream_op +): + input_tensor_all = _cp_attn_tp_all_gather_padded_tensor( + input_, total_len, attn_tp_size, forward_batch, stream_op + ) + # step3 + outputs = _trim_cp_rank_padding_after_all_gather(input_tensor_all, forward_batch) return outputs +def _log_tai_in_seq_rerange_fallback(reason: str, message: str, *args) -> None: + logger.info( + "CP NSA in-seq all-gather rerange tai-kernel fallback (%s): " + message, + reason, + *args, + ) + + +def _try_tai_in_seq_all_gather_rerange( + input_tensor_all: torch.Tensor, + forward_batch: "ForwardBatch", + *, + hidden_size: int, + cp_size: int, +) -> torch.Tensor | None: + metadata = forward_batch.nsa_cp_metadata + split_lens = getattr(metadata, "split_list_tensor", None) + split_prefix = getattr(metadata, "split_prefix_tensor", None) + split_list = getattr(metadata, "split_list", None) + max_rank_len = getattr(metadata, "max_rank_len", None) + if split_lens is None or split_prefix is None: + _log_tai_in_seq_rerange_fallback( + "missing_metadata", + "split_list_tensor or split_prefix_tensor is missing", + ) + return None + if split_list is None or max_rank_len is None: + _log_tai_in_seq_rerange_fallback( + "missing_cpu_metadata", + "split_list or max_rank_len is missing", + ) + return None + if not input_tensor_all.is_cuda or input_tensor_all.dim() != 2: + return None + + total_tokens = int(sum(split_list)) + if total_tokens == 0: + return input_tensor_all.new_empty((0, hidden_size)) + + max_segment_len = int(max(split_list)) + max_rank_token = int(max_rank_len[0]) if len(max_rank_len) > 0 else 0 + if max_segment_len <= 0 or max_rank_token <= 0: + return None + if input_tensor_all.shape[0] < max_rank_token * cp_size: + return None + + try: + from tai_kernel.nsa_prefill import in_seq_all_gather_rerange + except Exception as exc: + _log_tai_in_seq_rerange_fallback( + "import_failed", + "failed to import tai_kernel.nsa_prefill.in_seq_all_gather_rerange: %s", + exc, + ) + return None + + try: + return in_seq_all_gather_rerange( + input_tensor_all, + split_lens, + split_prefix, + total_tokens=total_tokens, + hidden_size=hidden_size, + max_segment_len=max_segment_len, + max_rank_token=max_rank_token, + cp_size=cp_size, + ) + except Exception as exc: + _log_tai_in_seq_rerange_fallback( + "kernel_failed", + "tai-kernel in-seq rerange failed: %s", + exc, + ) + return None + + +def _torch_in_seq_all_gather_rerange( + input_tensor_all: torch.Tensor, + forward_batch: "ForwardBatch", + hidden_size: int, +) -> torch.Tensor: + output_tensor = _trim_cp_rank_padding_after_all_gather( + input_tensor_all, forward_batch + ) + outputs_list = list( + torch.split( + output_tensor, forward_batch.nsa_cp_metadata.reverse_split_len, dim=0 + ) + ) + output_tensor = torch.cat( + [outputs_list[i] for i in forward_batch.nsa_cp_metadata.cp_reverse_index], dim=0 + ) + return output_tensor.view(-1, hidden_size) + + def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream): """ # for in-seq-split @@ -746,23 +894,26 @@ def cp_all_gather_rerange_output(input_tensor, cp_size, forward_batch, stream): return output_tensor bs_seq_len, hidden_size = input_tensor.shape - output_tensor = cp_attn_tp_all_gather_reorganazied_into_tensor( + input_tensor_all = _cp_attn_tp_all_gather_padded_tensor( input_tensor, forward_batch.nsa_cp_metadata.total_seq_lens, cp_size, forward_batch, stream, ) - outputs_list = list( - torch.split( - output_tensor, forward_batch.nsa_cp_metadata.reverse_split_len, dim=0 - ) + output_tensor = _try_tai_in_seq_all_gather_rerange( + input_tensor_all, + forward_batch, + hidden_size=hidden_size, + cp_size=cp_size, ) - output_tensor = torch.cat( - [outputs_list[i] for i in forward_batch.nsa_cp_metadata.cp_reverse_index], dim=0 + if output_tensor is not None: + return output_tensor + return _torch_in_seq_all_gather_rerange( + input_tensor_all, + forward_batch, + hidden_size, ) - output_tensor = output_tensor.view(-1, hidden_size) - return output_tensor def calculate_cp_seq_idx(cp_chunks_len, seqs_len): @@ -908,6 +1059,7 @@ def prepare_input_dp_with_cp_dsa( ) ) prefix_sum_list = list(accumulate(split_list)) + split_prefix_list = [0] + prefix_sum_list[:-1] # TODO Support multi-batch-cp-split, multi-batch-cp support has accuracy issues # cp_seq_index = calculate_cp_seq_idx(split_list[:], seqs_len[:]) @@ -932,6 +1084,10 @@ def prepare_input_dp_with_cp_dsa( nsa_cp_metadata = NSAContextParallelMetadata( split_list=split_list, + split_list_tensor=torch.tensor(split_list, device="cuda", dtype=torch.int32), + split_prefix_tensor=torch.tensor( + split_prefix_list, device="cuda", dtype=torch.int32 + ), max_rank_len=max_rank_len, zigzag_index=zigzag_index, per_rank_actual_token=per_rank_actual_token, diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index fd7495faf..cdb909a00 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -10,6 +10,7 @@ from sglang.srt.configs.model_config import get_nsa_index_topk, is_deepseek_nsa from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import ( + CpSharedKVIndexPrefetcher, CpSharedKVMlaPrefetcher, ) from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( @@ -603,6 +604,7 @@ class NativeSparseAttnBackend( batch_size = forward_batch.batch_size device = forward_batch.seq_lens.device forward_batch.cp_shared_kv_mla_prefetcher = None + forward_batch.cp_shared_kv_index_prefetcher = None if forward_batch.forward_mode.is_target_verify(): draft_token_num = self.speculative_num_draft_tokens @@ -884,6 +886,15 @@ class NativeSparseAttnBackend( ), ) ) + forward_batch.cp_shared_kv_index_prefetcher = ( + CpSharedKVIndexPrefetcher.maybe_create( + forward_batch=forward_batch, + metadata=metadata, + topk_transform_is_paged=( + topk_transform_method == TopkTransformMethod.PAGED + ), + ) + ) def _cal_indexer_k_start_end( self, @@ -1825,6 +1836,11 @@ class NativeSparseAttnBackend( finally: if mla_prefetcher is not None: mla_prefetcher.wait_attention_window() + index_prefetcher = getattr( + forward_batch, "cp_shared_kv_index_prefetcher", None + ) + if index_prefetcher is not None: + index_prefetcher.wait_attention_window() return attn_output diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index b2c949414..1ca1a0a95 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -423,8 +423,10 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin): uses_cp_shared_kv: bool = False cp_shared_kv_layout: Optional[CpSharedKVLayout] = None cp_local_out_cache_loc: Optional[torch.Tensor] = None + cp_local_physical_out_cache_loc: Optional[torch.Tensor] = None cp_shared_mla_direct_write_done: bool = False cp_shared_kv_mla_prefetcher: Optional[Any] = None + cp_shared_kv_index_prefetcher: Optional[Any] = None # For hidden states before normal return_hidden_states_before_norm: bool = False diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index b4e2ebbe7..9ee93cbaa 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py @@ -8,6 +8,7 @@ from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cud from sglang.srt.layers import deep_gemm_wrapper from sglang.srt.layers.attention.nsa.utils import ( get_cp_shared_kv_local_out_cache_loc, + get_cp_shared_kv_local_physical_out_cache_loc, log_cp_shared_kv_direct_write_fallback, nsa_use_prefill_cp, ) @@ -569,9 +570,11 @@ class DeepseekMLAForwardMixin: ): return True - physical_out_cache_loc = ( - layout.logical_locs_to_physical(local_out_cache_loc).contiguous() + physical_out_cache_loc = get_cp_shared_kv_local_physical_out_cache_loc( + forward_batch ) + if physical_out_cache_loc is None: + return False forward_batch.token_to_kv_pool.set_mla_kv_buffer( self.attn_mqa, physical_out_cache_loc, diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 76e6f0487..6d5bdddee 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -11,6 +11,7 @@ from sglang.srt.layers.attention.nsa.utils import ( can_cp_split, cp_split_and_rebuild_1d, get_cp_shared_kv_local_out_cache_loc, + get_cp_shared_kv_local_physical_out_cache_loc, split_in_seq_cp_local_pair, ) from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -335,6 +336,45 @@ class TestNSAInSeqCPUtils(unittest.TestCase): + list(range(6 * page_size, 7 * page_size)), ) + def test_local_physical_out_cache_loc_is_cached(self): + import torch + from types import SimpleNamespace + + page_size = 4 + segment_pages = [1, 2, 3, 4, 8, 7, 6, 5] + out_cache_loc = torch.cat( + [ + torch.arange(page * page_size, (page + 1) * page_size) + for page in segment_pages + ] + ) + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, + cp_size=4, + cp_rank=1, + ), + nsa_cp_metadata=NSAContextParallelMetadata( + split_list=[page_size] * 8, + zigzag_index=[1, 6], + page_aligned=True, + page_size=page_size, + extend_prefix_len=0, + ), + out_cache_loc=out_cache_loc, + ) + + physical_locs = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) + second_read = get_cp_shared_kv_local_physical_out_cache_loc(forward_batch) + + self.assertIs(physical_locs, second_read) + self.assertEqual( + physical_locs.tolist(), + list(range(1 * page_size, 2 * page_size)) + + list(range(2 * page_size, 3 * page_size)), + ) + def test_local_out_cache_loc_falls_back_when_owner_mismatch(self): import torch from types import SimpleNamespace 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 03feb4675..ad44e50df 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 @@ -1214,6 +1214,254 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1])) self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[2])) + def test_materialize_local_paged_buffer_page_slots_into_matches_full_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=64, cp_size=2, cp_rank=0) + page_buffer = torch.arange(0, 6 * 3, dtype=torch.float32).view(6, 3) + slot_logical_pages = torch.tensor([1, 2, 3, 4, 0, -1], dtype=torch.int64) + + full = runtime.materialize_local_paged_buffer_page_slots( + page_buffer=page_buffer, + slot_logical_pages=slot_logical_pages, + layout=layout, + ) + split = page_buffer.new_full(full.shape, -7) + split[0].zero_() + + runtime.materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=split, + slot_logical_pages=slot_logical_pages, + layout=layout, + start_slot=0, + end_slot=3, + ) + runtime.materialize_local_paged_buffer_page_slots_into( + page_buffer=page_buffer, + dense_page_buffer=split, + slot_logical_pages=slot_logical_pages, + layout=layout, + start_slot=3, + end_slot=slot_logical_pages.numel(), + ) + + self.assertTrue(torch.equal(split, full)) + + def test_remap_logical_pages_to_slot_dense_pages_preserves_sentinels(self): + from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime + + slot_logical_pages = torch.tensor([1, 2, 0, 4, -1, 5], dtype=torch.int64) + page_inverse = runtime.build_slot_page_inverse( + slot_logical_pages, + logical_page_capacity=8, + ) + logical_pages = torch.tensor([[0, 1, 5, -1, 3, 9]], dtype=torch.int32) + + dense_pages = runtime.remap_logical_pages_to_slot_dense_pages( + logical_pages, + page_inverse=page_inverse, + ) + + self.assertEqual(dense_pages.tolist(), [[0, 1, 6, -1, -1, -1]]) + + def test_index_materialize_uses_prefetched_buffer_before_fallback(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + class FakePool: + def __init__(self): + self.index_buffer = torch.arange(0, 12, dtype=torch.float32).view(4, 3) + + def get_index_k_with_scale_buffer(self, layer_id): + self.layer_id = layer_id + return self.index_buffer + + class FakePrefetcher: + def __init__(self): + self.calls = [] + self.dense_buffer = torch.full((3, 3), 5.0) + self.dense_pages = torch.tensor([[1, 2]], dtype=torch.int32) + + def consume(self, *, layer_id, page_buffer, logical_pages): + self.calls.append((layer_id, page_buffer, logical_pages)) + return self.dense_buffer, self.dense_pages + + fake_pool = FakePool() + fake_prefetcher = FakePrefetcher() + forward_batch = SimpleNamespace( + token_to_kv_pool=fake_pool, + uses_cp_shared_kv=True, + cp_shared_kv_layout=object(), + cp_shared_kv_index_prefetcher=fake_prefetcher, + ) + logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) + indexer = object.__new__(nsa_indexer.Indexer) + + with patch.object( + nsa_indexer, + "materialize_shared_paged_buffer", + side_effect=AssertionError("prefetch hit must bypass full materialize"), + ): + dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id=7, + logical_page_table=logical_pages, + ) + + self.assertIs(dense_buffer, fake_prefetcher.dense_buffer) + self.assertIs(dense_pages, fake_prefetcher.dense_pages) + self.assertEqual(fake_prefetcher.calls[0][0], 7) + self.assertIs(fake_prefetcher.calls[0][1], fake_pool.index_buffer) + self.assertIs(fake_prefetcher.calls[0][2], logical_pages) + + def test_index_prefetch_start_targets_next_layer(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + 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 = object() + fake_prefetcher = FakePrefetcher() + forward_batch = SimpleNamespace( + token_to_kv_pool=token_to_kv_pool, + cp_shared_kv_index_prefetcher=fake_prefetcher, + ) + indexer = object.__new__(nsa_indexer.Indexer) + + indexer._maybe_start_next_layer_index_prefetch(forward_batch, layer_id=11) + + self.assertEqual(fake_prefetcher.calls, [(12, token_to_kv_pool)]) + + def test_index_prefetch_consume_miss_logs_fallback_after_first_layer(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + class FakePool: + start_layer = 0 + + def __init__(self): + self.index_buffer = torch.arange(0, 12, dtype=torch.float32).view(4, 3) + + def get_index_k_with_scale_buffer(self, layer_id): + return self.index_buffer + + class FakeLayout: + cp_rank = 3 + + class MissingPrefetcher: + def consume(self, *, layer_id, page_buffer, logical_pages): + return None + + fallback_buffer = torch.full((3, 3), 9.0) + fallback_pages = torch.tensor([[1, 2]], dtype=torch.int32) + forward_batch = SimpleNamespace( + token_to_kv_pool=FakePool(), + uses_cp_shared_kv=True, + cp_shared_kv_layout=FakeLayout(), + cp_shared_kv_index_prefetcher=MissingPrefetcher(), + ) + logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) + indexer = object.__new__(nsa_indexer.Indexer) + + with patch.object( + nsa_indexer, + "materialize_shared_paged_buffer", + 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, + ) + + self.assertIs(dense_buffer, fallback_buffer) + self.assertIs(dense_pages, fallback_pages) + logger.info.assert_called_once() + self.assertIn( + "CP shared KV index prefetch fallback", + logger.info.call_args.args[0], + ) + self.assertIn("consume_miss", logger.info.call_args.args[1]) + + def test_index_prefetch_first_layer_miss_does_not_log_fallback(self): + from sglang.srt.layers.attention.nsa import nsa_indexer + + class FakePool: + start_layer = 0 + + def __init__(self): + self.index_buffer = torch.arange(0, 12, dtype=torch.float32).view(4, 3) + + def get_index_k_with_scale_buffer(self, layer_id): + return self.index_buffer + + class FakeLayout: + cp_rank = 0 + + class MissingPrefetcher: + def consume(self, *, layer_id, page_buffer, logical_pages): + return None + + fallback_buffer = torch.full((3, 3), 9.0) + fallback_pages = torch.tensor([[1, 2]], dtype=torch.int32) + forward_batch = SimpleNamespace( + token_to_kv_pool=FakePool(), + uses_cp_shared_kv=True, + cp_shared_kv_layout=FakeLayout(), + cp_shared_kv_index_prefetcher=MissingPrefetcher(), + ) + logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) + indexer = object.__new__(nsa_indexer.Indexer) + + with patch.object( + nsa_indexer, + "materialize_shared_paged_buffer", + return_value=(fallback_buffer, fallback_pages), + ), patch( + "sglang.srt.layers.attention.nsa.nsa_indexer.logger", + create=True, + ) as logger: + indexer._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id=0, + logical_page_table=logical_pages, + ) + + logger.info.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 + + 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=False + ), patch.object( + prefetch.logger, "info" + ) as logger: + result = prefetch.CpSharedKVIndexPrefetcher.maybe_create( + forward_batch=SimpleNamespace(), + metadata=SimpleNamespace(), + topk_transform_is_paged=True, + ) + + self.assertIsNone(result) + logger.assert_called_once() + self.assertIn( + "CP shared KV index prefetch fallback", + logger.call_args.args[0], + ) + self.assertIn("cuda_unavailable_or_stream_capturing", logger.call_args.args[1]) + if __name__ == "__main__": unittest.main()