Align CP shared-KV prefetch with the attention overlap window

Launching CP shared-KV prefetch from MLA prepare or the indexer can make next-layer prefix work overlap current-layer MQA/materialization instead of the attention window. Centralize the launch in the NSA backend after current-layer materialization and before attention, and leave the early indexer hook inert so the call site cannot regress silently. The accompanying notes capture the draft-as-forward-layer follow-up plan and the latest OOM diagnosis: the observed 178945-token failure matches the CP in-seq MQA logits allocation, so the follow-up fix is q-dimension chunking in _get_topk_ragged_with_cp(), not a max-prefetch-size gate. Constraint: CP shared-KV prefetch must avoid overlapping current-layer MQA/materialization and must not add silent fallback behavior. Rejected: Limit maximum prefetch size | hides the CP logits peak and can reduce cache/prefetch effectiveness. Rejected: Keep prepare/indexer launch sites | they place next-layer collectives in the wrong overlap window. Confidence: medium Scope-risk: moderate Directive: Do not reintroduce early CP prefetch launch without checking Nsight overlap and CP MQA memory peaks. Tested: Local git diff --check and py_compile for touched Python files. Tested: Remote container py_compile plus targeted pytest: 3 passed, 5 warnings. Not-tested: Full ETE under production traffic after this commit. Not-tested: CP in-seq MQA logits chunking; documented as follow-up.
2026-06-01 01:09:43 +08:00
parent 4342de0463
commit 46be97adc0
6 changed files with 206 additions and 85 deletions
--- a/docs/advanced_features/nsa_prefill_cp_eagle_mtp_shared_kv_plan.md
+++ b/docs/advanced_features/nsa_prefill_cp_eagle_mtp_shared_kv_plan.md
@@ -472,3 +472,112 @@ CP draft shared KV fallback: unsupported draft architecture
 4. draft persistent KV 使用 CP shared-KV physical shard，而不是每 rank full KV。
 5. prefill->decode draft KV transfer 正常。
 6. fallback 路径安全、可观测，非 GLM-5 draft 不误走新路径。
+
+## 2026-06-01：Draft layer as forward-layer in prefetch/backup pipeline plan
+
+### 背景
+
+当前 EAGLE / NextN draft 已经作为 target HiCache 的 shadow payload 接入了一部分生命周期：
+
+- `CacheController.attach_draft_pool()` 会把 draft device/host pool 绑定到 target controller，并注册同一个 `LayerDoneCounter`。
+- draft decoder layer end 会通过 `notify_layer_end_for_backup()` 进入 `CacheController.on_layer_end(layer_id, source="draft")`，再提交 draft D2H backup。
+- `CacheController.start_loading()` 已经会把 target load queue 与 draft load queue 一起处理。
+
+但 CP shared-KV 的 index/MLA prefetch 当前仍把 draft 排除在外：
+
+- `nsa_backend.init_forward_metadata()` 对 `cp_shared_kv_is_draft_input(forward_batch)` 显式禁用 prefetcher 创建。
+- `cp_shared_kv_should_prefetch_next_layer()` 对 draft input 返回 `False`。
+- draft/NextN 只有一个 executable layer，`layer_id=0`，不能直接复用 target 的 `next_layer` prefetch 语义。
+
+### 设计目标
+
+把 draft 层视为 target forward 后的一个逻辑 forward layer：
+
+```text
+target layer 0
+...
+target layer N-1
+draft layer 0
+```
+
+但第一阶段不重构全局 event/counter timeline；先让 draft layer 拥有独立的 same-layer prefetch 合同：
+
+```text
+draft metadata ready
+  -> start draft current-layer index/MLA prefix prefetch
+  -> draft layer 0 consumes prefix + current KV
+  -> draft layer 0 end submits async backup
+```
+
+### 分阶段计划
+
+#### P0：锁住当前 draft 合同
+
+- 测试 draft input 不应走 target-style next-layer prefetch。
+- 测试 draft layer id 只允许 `0`。
+- 测试 draft backup notifier 仍通过 `source="draft"` 触发。
+- 保留 fail-fast，不允许静默 fallback 到 target prefetch 路径。
+
+#### P1：新增 draft same-layer prefetch API
+
+新增独立语义，避免误用 target `start_next_layer_prefix()`：
+
+```python
+start_current_layer_prefix(layer_id=0, token_to_kv_pool=draft_pool)
+```
+
+约束：
+
+- 只允许 draft / NextN path 使用。
+- 只允许 `layer_id == 0`。
+- metadata 必须已经构造完成。
+- 不满足条件时明确 warning/fail-fast，不做静默慢 fallback。
+
+#### P2：把 draft prefetch 启动点放到 draft metadata ready 后、attention 前
+
+目标顺序：
+
+```text
+draft init_forward_metadata
+pre-attention hook:
+  start index prefix prefetch
+  start MLA prefix prefetch
+draft indexer / attention consume:
+  wait pending event if needed
+```
+
+这和 target 当前的 attention-front prefetch 位置保持一致，但语义是 current-layer，不是 next-layer。
+
+#### P3：先不重构 backup 主流程
+
+当前 draft backup 已经在 draft layer end 触发。第一阶段只验证：
+
+- 是否确实 per-layer async 提交；
+- 是否还有 catch-up fallback；
+- final ack 是否仍等待 target + draft payload 完成。
+
+若发现 draft backup 延迟到 request 结束，再单独修复。
+
+#### P4：后续再评估 target-last-layer -> draft-layer bridge
+
+更激进的优化是让 target 最后一层提前发起 draft layer 0 prefetch。但这需要更早构造 draft metadata，而 draft forward batch 当前依赖 target 输出 hidden / verified token，风险较高。
+
+暂不作为第一阶段实现项。
+
+#### P5：后续再统一虚拟 layer timeline
+
+理想模型是：
+
+```text
+event 0..N-1: target layers
+event N:      draft layer 0
+```
+
+但这会影响 `LayerDoneCounter`、draft H2D wait、backup ack、target/draft strong-sync 等多个路径。只有当 P1-P3 证明收益不足或现有 event 绑定导致 correctness/perf 问题时再做。
+
+### 风险记录
+
+- 直接打开 draft prefetcher 可能复现之前 EAGLE hang；draft 没有 next-layer，必须单独建 current-layer API。
+- draft metadata 生成时间晚于 target metadata，target last-layer 提前 prefetch 需要额外桥接。
+- index/MLA prefix materialize 会占 SM，短 draft 层可能出现启动开销大于 overlap 收益，需要 benchmark/trace 验证。
+- 任何 fallback 必须 warning/fail-fast，不能静默降级。
--- 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
@@ -1028,3 +1028,22 @@ Fallback：
 debug enabled 时自动回到现有同步 materialize
 pynccl unavailable 时自动回到现有同步 materialize
 ```
+
+## 2026-06-01：attention-front prefetch OOM observation
+
+Remote log:
+
+- `/mnt/beegfs/cjy/log/sglang_cp_hicache_20260531_161740.log`
+- Failure time: `2026-05-31 16:21:21`
+- Stack: target prefill `eagle_worker.forward_target_extend()` -> `DeepseekV2AttentionMLA.forward_prepare()` -> `Indexer._get_topk_in_seq_cp_pair()` -> `_get_topk_ragged_with_cp()` -> `deep_gemm.fp8_mqa_logits()`.
+- Error: CUDA OOM while allocating `7.58 GiB`; GPU0 had `6.67 GiB` free, process used `132.84 GiB`.
+- Immediately preceding successful batch included very large requests, e.g. `#new-token: 140672`, `#cached-token: 13824`, `#inflight-req: 2`; next failed request prepared backup for `logical_len=178945`.
+
+Confirmed root cause direction:
+
+- The OOM is not a host-cache capacity error; it is a CUDA activation/temp-buffer peak inside current-layer NSA index MQA/topk.
+- CP in-seq `_get_topk_ragged_with_cp()` calls `deep_gemm.fp8_mqa_logits()` without the q-dimension chunking guard that the non-CP ragged path already has.
+- For `logical_len=178945`, CP0's zigzag tail segment is about `11184` query rows against about `178945` KV rows, so the fp32 logits allocation is about `7.45 GiB`, matching the `7.58 GiB` OOM.
+- Do not solve this by limiting max prefetch size. The follow-up fix should add chunked MQA logits handling to the CP in-seq path.
+
+Do not re-debug from scratch before checking this section.
--- 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 (
    cp_shared_kv_mla_prefetch_log,
    cp_shared_kv_mla_prefetch_log_enabled,
    cp_shared_kv_mla_prefetch_should_log_layer,
-    cp_shared_kv_should_prefetch_next_layer,
    filter_owned_logical_locs,
    get_or_build_shared_paged_buffer_slot_remap,
    is_current_only_extend_batch,
@@ -478,36 +477,16 @@ class Indexer(MultiPlatformOp):
        forward_batch: ForwardBatch,
        layer_id: int,
    ) -> None:
-        next_layer_id = layer_id + 1
-        index_prefetcher = getattr(
-            forward_batch, "cp_shared_kv_index_prefetcher", None
-        )
-        if (
-            cp_shared_kv_mla_prefetch_log_enabled()
-            and cp_shared_kv_mla_prefetch_should_log_layer(next_layer_id)
-        ):
-            layout = getattr(forward_batch, "cp_shared_kv_layout", None)
-            seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None)
-            cp_shared_kv_mla_prefetch_log(
-                "index_start_request cp_rank=%s cp_size=%s layer=%s "
-                "next_layer=%s has_index=%s uses_cp_shared_kv=%s "
-                "seq_lens_cpu_len=%s",
-                getattr(layout, "cp_rank", None),
-                getattr(layout, "cp_size", None),
-                layer_id,
-                next_layer_id,
-                index_prefetcher is not None,
-                getattr(forward_batch, "uses_cp_shared_kv", None),
-                len(seq_lens_cpu) if seq_lens_cpu is not None else None,
-            )
-        if index_prefetcher is None:
-            return
-        if not cp_shared_kv_should_prefetch_next_layer(forward_batch, layer_id):
-            return
-        index_prefetcher.start_next_layer_prefix(
-            next_layer_id=next_layer_id,
-            token_to_kv_pool=forward_batch.token_to_kv_pool,
-        )
+        """Deprecated early prefetch hook.
+
+        Next-layer CP shared-KV prefetch must be launched by the attention
+        backend after current-layer index/MLA materialization is complete and
+        immediately before the attention kernel.  Starting here overlaps the
+        next-layer collective with current-layer MLA materialization/all-reduce,
+        which is the wrong dependency window.
+        """
+
+        return

    def _filter_shared_index_write(
        self,
@@ -1901,7 +1880,6 @@ 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(
--- a/python/sglang/srt/layers/attention/nsa_backend.py
+++ b/python/sglang/srt/layers/attention/nsa_backend.py
@@ -98,6 +98,43 @@ global_workspace_buffer = None
 _USE_FUSED_METADATA_COPY = envs.SGLANG_USE_FUSED_METADATA_COPY.get() and not _is_hip


+def _maybe_start_cp_shared_kv_attention_prefetch(
+    forward_batch: ForwardBatch,
+    layer_id: int,
+) -> None:
+    """Launch next-layer CP shared-KV prefetch in the attention overlap window.
+
+    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
+    with the current layer's required cache work.
+    """
+
+    token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
+    if token_to_kv_pool is None:
+        return
+    if not cp_shared_kv_should_prefetch_next_layer(forward_batch, layer_id):
+        return
+
+    next_layer_id = int(layer_id) + 1
+    index_prefetcher = getattr(
+        forward_batch, "cp_shared_kv_index_prefetcher", None
+    )
+    if index_prefetcher is not None:
+        index_prefetcher.start_next_layer_prefix(
+            next_layer_id=next_layer_id,
+            token_to_kv_pool=token_to_kv_pool,
+        )
+
+    mla_prefetcher = getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None)
+    if mla_prefetcher is not None:
+        mla_prefetcher.start_next_layer_prefix(
+            next_layer_id=next_layer_id,
+            token_to_kv_pool=token_to_kv_pool,
+        )
+
+
@dataclass(frozen=True)
 class NSAFlashMLAMetadata:
    """Metadata only needed by FlashMLA"""
@@ -2070,16 +2107,11 @@ class NativeSparseAttnBackend(
                        else None,
                        tuple(out_cache_loc.shape) if out_cache_loc is not None else None,
                    )
-            if mla_prefetcher is not None and cp_shared_kv_should_prefetch_next_layer(
-                forward_batch, layer.layer_id
-            ):
-                mla_prefetcher.start_next_layer_prefix(
-                    next_layer_id=layer.layer_id + 1,
-                    token_to_kv_pool=forward_batch.token_to_kv_pool,
-                )
        else:
            mla_prefetcher = None

+        _maybe_start_cp_shared_kv_attention_prefetch(forward_batch, layer.layer_id)
+
        index_prefetcher = getattr(
            forward_batch, "cp_shared_kv_index_prefetcher", None
        )
--- 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
@@ -14,7 +14,6 @@ from sglang.srt.layers.attention.nsa.utils import (
    nsa_use_prefill_cp,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
-    cp_shared_kv_should_prefetch_next_layer,
    should_reuse_current_extend_kv,
    try_tai_fused_mla_store,
 )
@@ -95,47 +94,6 @@ class DeepseekMLAForwardMixin:
            get_global_server_args().flashinfer_mla_disable_ragged
        )

-    def _maybe_start_cp_shared_next_layer_prefetch(
-        self: DeepseekV2AttentionMLA,
-        forward_batch: ForwardBatch,
-    ) -> None:
-        """Launch next-layer CP shared-KV prefix prefetch before rank-skewed work.
-
-        The original Phase8 hook launched prefetch from the indexer/backend after
-        current-layer MQA/materialize work.  Those paths can be rank-skewed in
-        in-seq CP, so rank 0 may enqueue the next-layer collective early while
-        other ranks reach it too late to overlap.  Starting from layer prepare
-        keeps the launch point before the current layer's per-rank MQA/topk and
-        materialize imbalance.  Existing later hooks remain as fallbacks and
-        become no-ops via the prefetchers' already-started guard.
-        """
-
-        token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
-        if token_to_kv_pool is None:
-            return
-
-        if not cp_shared_kv_should_prefetch_next_layer(
-            forward_batch, self.layer_id
-        ):
-            return
-
-        next_layer_id = int(self.layer_id) + 1
-        index_prefetcher = getattr(
-            forward_batch, "cp_shared_kv_index_prefetcher", None
-        )
-        if index_prefetcher is not None:
-            index_prefetcher.start_next_layer_prefix(
-                next_layer_id=next_layer_id,
-                token_to_kv_pool=token_to_kv_pool,
-            )
-
-        mla_prefetcher = getattr(forward_batch, "cp_shared_kv_mla_prefetcher", None)
-        if mla_prefetcher is not None:
-            mla_prefetcher.start_next_layer_prefix(
-                next_layer_id=next_layer_id,
-                token_to_kv_pool=token_to_kv_pool,
-            )
-
    def forward_absorb_prepare(
        self: DeepseekV2AttentionMLA,
        positions: torch.Tensor,
@@ -146,8 +104,6 @@ class DeepseekMLAForwardMixin:
    ):
        from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode

-        self._maybe_start_cp_shared_next_layer_prefetch(forward_batch)
-
        q_lora = None
        topk_indices = None
        if self.q_lora_rank is not None:
--- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
+++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
@@ -3101,7 +3101,7 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
        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):
+    def test_indexer_does_not_start_next_layer_prefetch_before_attention_window(self):
        from sglang.srt.layers.attention.nsa import nsa_indexer

        class FakePrefetcher:
@@ -3121,7 +3121,34 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):

        indexer._maybe_start_next_layer_index_prefetch(forward_batch, layer_id=11)

-        self.assertEqual(fake_prefetcher.calls, [(12, token_to_kv_pool)])
+        self.assertEqual(fake_prefetcher.calls, [])
+
+    def test_attention_window_prefetch_starts_index_and_mla_next_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 = object()
+        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=5
+        )
+
+        self.assertEqual(index_prefetcher.calls, [(6, token_to_kv_pool)])
+        self.assertEqual(mla_prefetcher.calls, [(6, token_to_kv_pool)])

    def test_index_prefetch_skips_when_current_layer_is_last(self):
        from sglang.srt.layers.attention.nsa import nsa_indexer