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 dacd610c0..55c51376d 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,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 有下降 ``` diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d46ce03b4..9526820b5 100644 --- a/python/sglang/srt/environ.py +++ b/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) 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 d335d6f8d..6c49694c9 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 @@ -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, ) 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 6e8e4fb72..625f70d60 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 @@ -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) 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 d66d60683..46c2593e2 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 @@ -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