Keep Phase8 prefetch on the deferred-consume path

Phase8 only gains useful overlap when the next-layer MLA prefix prefetch is allowed to run until the next layer actually consumes the prefetched buffer. The old wait-after-attention switch let runtime configuration collapse the optimization back into current-layer tail latency, so the prefetch path now has one wait policy and the documentation records the implemented behavior.

Constraint: Phase8 should keep the production environment surface minimal while preserving the existing enable and debug-log knobs
Rejected: SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION | it reintroduced current-layer synchronous waiting and made profiling behavior depend on a nonessential policy knob
Confidence: medium
Scope-risk: narrow
Directive: Do not add another Phase8 wait policy knob without first proving the added policy improves end-to-end prefill latency under CP shared KV
Tested: Python AST parse for touched Python files
Tested: git diff --check
Not-tested: Full pytest and remote server integration were not run in this commit

docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md

@@ -1,8 +1,36 @@
 # NSA Prefill CP Phase 8: MLA 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 前的同步阻塞路径里移出，并尽量与当前层 attention compute 重叠。
+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 时。
 
-本阶段只做 MLA KV prefix prefetch。暂时不做 index K/scale prefetch，不引入 `SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND` 这类选择型环境变量。
+本阶段只做 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 策略开关。
+
+---
+
+
+## 0. 当前实现状态
+
+Phase 8 已按 deferred-consume 策略实现：
+
+```text
+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
+```
+
+保留的环境变量：
+
+```text
+SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1   # 生产开关，默认关闭
+SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1      # 调试日志，只打印 probe layer，默认关闭
+```
+
+已移除的策略开关：
+
+```text
+SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION
+```
+
+移除原因：attention 返回前立刻 wait 会把 prefetch 重新变成当前层同步尾延迟，抵消主要 overlap 收益。Phase 8 现在固定使用“下一层消费前等待”的单一策略，减少运行时分支和环境变量读取。
 
 ---
 
@@ -184,12 +212,13 @@ consume(prefetched prefix)
 materialize(current/suffix only)
 ```
 
-并把下一层 prefix 的 materialize/all-reduce 放到当前 attention compute 窗口内：
+并把下一层 prefix 的 materialize/all-reduce 提前到当前层 attention 前启动，允许它持续到下一层消费前完成：
 
 ```text
-current layer attention compute
+current/next layer compute window
   overlaps with
 next layer prefix materialize + async CP all-reduce
+next layer consume waits only if prefetch has not finished
 ```
 
 ### 3.3 非目标
@@ -278,26 +307,25 @@ chunked/radix-hit 下 prefix 往往占大头。
 ```text
 Layer L attention:
   consume Layer L prefetched prefix if available
+    -> wait event only when the prefetched buffer is actually needed
   materialize Layer L suffix/current pages
+  start async prefetch for Layer L+1 prefix MLA KV
   run Layer L attention
+  return without waiting for Layer L+1 prefetch
 
-  before Layer L attention kernel:
-    start async prefetch for Layer L+1 prefix MLA KV
-
-  after Layer L attention kernel and before returning:
-    wait Layer L+1 prefetch complete
-
-Layer L MLP / EP:
-  no outstanding CP shared KV prefetch collective
+Layer L+1 attention:
+  consume Layer L+1 prefetched prefix
+    -> wait event here if async prefetch has not completed yet
+  materialize Layer L+1 suffix/current pages
 ```
 
-注意这个策略是 **attention-bounded**：
+注意这个策略是 **one-layer deferred-consume**：
 
 ```text
-prefetch 只允许与当前 attention compute 重叠。
-不能跨出 attention backend 返回边界。
-不能与后面的 prepare_mlp / EP / MoE A2A 竞争不可控带宽。
-不能改变 CP collective 的跨层顺序。
+prefetch 最多提前一层。
+等待点固定在下一层 consume，而不是当前层 attention 返回前。
+不跨 forward/batch 保存 dense KV。
+不改变 CP collective 的 rank 间顺序；所有 rank 对同一 next layer 同步启动同一类 collective。
 ```
 
 ### 5.3 为什么不是“复用上一 chunk dense KV”
@@ -362,23 +390,28 @@ Phase 8 不能改变 dense page id，否则现有 `page_table_1` / `topk_indices
 
 ## 7. Runtime 设计
 
-### 7.1 新环境变量
+### 7.1 环境变量
 
-只新增一个环境变量：
+生产路径只需要一个启用开关：
 
 ```text
 SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1
 ```
 
-默认关闭。
+默认关闭。调试时可额外开启：
 
-不新增：
+```text
+SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1
+```
+
+不新增/不保留：
 
 ```text
 SGLANG_CP_SHARED_KV_LAYER_PREFETCH_KIND
+SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION
 ```
 
-因为 Phase 8 v1 只有 MLA path，没有 index/index+mla 可选组合。
+因为 Phase 8 v1 只有 MLA path，没有 index/index+mla 可选组合；wait 策略固定为 deferred-consume，避免 attention 返回前 wait 抵消 overlap。
 
 ### 7.2 ForwardBatch 挂载 prefetcher
 
@@ -422,7 +455,7 @@ CpSharedKVMlaPrefetcher:
   - 管理 layer_id -> handle
   - consume 当前 layer prefix handle
   - start 下一 layer prefix prefetch
-  - wait attention window 内已启动的 prefetch
+  - 记录 deferred wait 状态；真实 wait 在下一层 consume
   - fallback 到同步 materialize
 
 CpSharedKVMlaPrefetchHandle:
@@ -531,7 +564,7 @@ disable MLA prefetch and fallback sync materialize
 
 ---
 
-## 8. Attention-bounded 调度
+## 8. Deferred-consume 调度
 
 ### 8.1 启动时机
 
@@ -550,8 +583,8 @@ disable MLA prefetch and fallback sync materialize
 2. 当前层 materialize/reuse 决策已经完成；
 3. attention kernel 即将开始；
 4. 此时启动下一层 prefix prefetch；
-5. 当前 attention kernel 提供 overlap window；
-6. attention 返回前 wait prefetch event。
+5. 当前 attention 以及后续到下一层 consume 前的计算提供 overlap window；
+6. attention 返回前不 wait；下一层 consume prefetched KV 时 wait event。
 ```
 
 ### 8.2 伪代码
@@ -576,17 +609,19 @@ if shared_kv_paged_path:
 attn_output = run_attention_kernel(...)
 
 if prefetcher is not None:
-    prefetcher.wait_attention_window()
+    prefetcher.wait_attention_window()  # historical name; now only records deferred wait state/log
 
 return attn_output
 ```
 
-`wait_attention_window()` 必须在 `forward_extend(...)` 返回前调用，保证：
+`wait_attention_window()` 仍在 `forward_extend(...)` 的 finally 中调用，但当前实现不做同步等待，只保留 deferred 日志/状态检查入口。真正的同步点在下一层 `consume(...)`：
 
 ```text
-没有 outstanding prefetch collective 跨入 prepare_mlp / EP / 下一层 prepare。
+torch.cuda.current_stream().wait_event(prefetch_event)
 ```
 
+这样避免 attention 返回前 wait 把 prefetch 收益吃掉。
+
 ---
 
 ## 9. Correctness 约束
@@ -792,7 +827,7 @@ allocate dense buffer
 prefetch prefix range
 record event
 consume handle
-wait attention window
+defer wait until consume
 fallback
 ```
 
@@ -815,7 +850,7 @@ try consume prefetched prefix for current layer
 materialize current/suffix into prefetched buffer
 start next layer prefix prefetch
 run attention
-wait prefetch before return
+do not wait before return; consume waits before use
 ```
 
 ### Step 5: 单元测试
@@ -850,7 +885,7 @@ router:  g0034
 3. 重复请求触发 radix hit；
 4. 检查输出质量；
 5. 检查没有 collective hang；
-6. profile 确认 `cp_shared_kv.materialize.token` 同步时间下降或被 attention overlap。
+6. profile 确认 `cp_shared_kv.materialize.token` 同步时间下降或被后续计算 overlap，并能看到下一层 consume 前的必要 wait。
 
 ---
 
@@ -865,23 +900,24 @@ async CP all-reduce 插入层间路径，最大风险是不同 rank collective
 ```text
 gate 必须全 rank 一致
 prefetch 只在 attention backend 内启动
-forward_extend 返回前 wait
-不跨入 prepare_mlp / EP
+每个 batch 最多 outstanding 一个 next-layer prefetch
+下一层 consume 前强制 wait event
+不跨 forward/batch 保存 dense KV
 ```
 
 ### 13.2 带宽竞争
 
 prefetch 会消耗 NVLink/NCCL 带宽。
 
-Phase 8 暂时不做 bandwidth 控制，但通过 attention-bounded wait 限制影响范围：
+Phase 8 暂时不做 bandwidth 控制。deferred-consume 策略会扩大 overlap 窗口，但也可能让 prefetch 与 MLP/EP/下一层 prepare 存在带宽竞争：
 
 ```text
-只与 current attention compute overlap
-不与 EP/MoE A2A overlap
-不与下一层 prepare_attn overlap
+最多 one-layer-ahead outstanding prefetch
+下一层 consume 前必须 wait
+不做多层预取
 ```
 
-如果 profiling 显示 attention 本身也被明显拖慢，再进入后续 Phase 做 bandwidth throttle/page budget。
+如果 profiling 显示带宽竞争明显，再进入后续 Phase 做 bandwidth throttle/page budget。
 
 ### 13.3 显存增加
 
@@ -925,7 +961,7 @@ Performance：
 
 ```text
 prefetch on 时，MLA KV prefix materialize 的同步阻塞减少
-Nsight 中可看到 next-layer prefix materialize/all-reduce 与 current attention 有 overlap
+Nsight 中可看到 next-layer prefix materialize/all-reduce 被提前，并在下一层 consume 前只剩必要 wait
 总 TTFT / chunked prefill latency 有下降
 ```
 

python/sglang/srt/environ.py

@@ -207,7 +207,6 @@ class Envs:
     SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False)
     SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False)
-    SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)

python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py

@@ -14,7 +14,6 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     cp_shared_kv_mla_prefetch_enabled,
     cp_shared_kv_mla_prefetch_log,
     cp_shared_kv_mla_prefetch_should_log_layer,
-    cp_shared_kv_mla_prefetch_wait_after_attention_enabled,
     filter_locs_mappable_to_physical_pool,
     materialize_local_token_kv_page_slots_into,
     remap_logical_locs_to_slot_dense_locs_optimized,
@@ -393,24 +392,12 @@ class CpSharedKVMlaPrefetcher:
         )
 
     def wait_attention_window(self) -> None:
-        if not cp_shared_kv_mla_prefetch_wait_after_attention_enabled():
-            handle = self.pending_attention_handle
-            if handle is not None:
-                self._log_next_layer(
-                    handle.layer_id,
-                    "attention_wait_deferred next_layer=%s",
-                    handle.layer_id,
-                )
-            return
-
         handle = self.pending_attention_handle
-        self.pending_attention_handle = None
         if handle is None:
             return
-        torch.cuda.current_stream().wait_event(handle.event)
         self._log_next_layer(
             handle.layer_id,
-            "attention_wait next_layer=%s",
+            "attention_wait_deferred next_layer=%s",
             handle.layer_id,
         )
 

python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py

@@ -37,10 +37,6 @@ def cp_shared_kv_mla_prefetch_log_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get()
 
 
-def cp_shared_kv_mla_prefetch_wait_after_attention_enabled() -> bool:
-    return envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.get()
-
-
 def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None:
     if cp_shared_kv_mla_prefetch_log_enabled():
         logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args)

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -257,21 +257,16 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
             cp_shared_kv_mla_prefetch_log_enabled,
             cp_shared_kv_mla_prefetch_should_log_layer,
-            cp_shared_kv_mla_prefetch_wait_after_attention_enabled,
         )
 
         envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.clear()
-        envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.clear()
         self.assertFalse(cp_shared_kv_mla_prefetch_log_enabled())
-        self.assertFalse(cp_shared_kv_mla_prefetch_wait_after_attention_enabled())
         self.assertFalse(cp_shared_kv_mla_prefetch_should_log_layer(1))
         self.assertTrue(cp_shared_kv_mla_prefetch_should_log_layer(2))
         self.assertFalse(cp_shared_kv_mla_prefetch_should_log_layer(3))
 
         with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True):
             self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled())
-        with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_WAIT_AFTER_ATTENTION.override(True):
-            self.assertTrue(cp_shared_kv_mla_prefetch_wait_after_attention_enabled())
 
     def test_token_range_materialize_uses_tai_kernel_when_enabled(self):
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime