From 46be97adc0d230c9c6588d3ee34fe9769652b994 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Mon, 1 Jun 2026 01:09:43 +0800 Subject: [PATCH] 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. --- ...nsa_prefill_cp_eagle_mtp_shared_kv_plan.md | 109 ++++++++++++++++++ ...fill_cp_phase8_mla_prefix_prefetch_plan.md | 19 +++ .../srt/layers/attention/nsa/nsa_indexer.py | 42 ++----- .../srt/layers/attention/nsa_backend.py | 46 ++++++-- .../attention_forward_methods/forward_mla.py | 44 ------- .../mem_cache/test_cp_shared_kv_runtime.py | 31 ++++- 6 files changed, 206 insertions(+), 85 deletions(-) diff --git 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 index 85b42a3bf..05dd8fd6b 100644 --- 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，不能静默降级。 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 70218741f..b80e1d47f 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 @@ -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. diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 8386f84b8..2b50a3a50 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 ( 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( diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index 9fa434007..54f5c6476 100644 --- 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 ) 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 c0df53737..9ba806756 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 @@ -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: 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 53dc3eca9..8af754fdd 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 @@ -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