diff --git a/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md b/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md index 5e5250cd5..f1bf528f2 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md +++ b/docs/advanced_features/nsa_prefill_cp_phase2_shared_kv_design.md @@ -26,7 +26,7 @@ Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它 3. **attention runtime 仍走 full-view compatibility**: - 现有 NSA attention/indexer kernel 仍假设可读 full KV/index view。 - shared KV 下运行时会 materialize dense full-view buffer:每 rank 拷贝本地 owner pages,其余填 0,然后 CP group all-reduce 得到完整 runtime view。 - - 这是 Phase 2 的主要性能成本;Phase 3 才会改成 shard-aware attention/topk。 + - 这是 Phase 2 的主要性能成本;Phase 3 先去掉 current chunk 的重复 materialize,Phase 4 再做 shard-aware attention/topk。 4. **Mooncake PD transfer 已支持 prefill CP shards -> decode full KV**: - prefill CP rank 发送本 rank owner physical pages。 - transfer chunk 携带 `logical_page_positions`,decode 侧用这些位置选择 full-layout dst pages。 @@ -37,8 +37,8 @@ Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它 当前仍未完成、属于 Phase 3 或后续工作: -- runtime 不再 materialize full/maxlen KV。 -- NSA index/topk/attention 的 shard-aware 计算与 global topk merge。 +- Phase 3:复用已经 CP all-gather + rerange 的 current chunk KV/index,避免 shared KV compatibility path 重复 materialize current chunk。详见 `docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md`。 +- Phase 4:history/shared KV 的 selected-page materialize、NSA index/topk/attention 的 shard-aware 计算与 global topk merge。 - `nsa_prefill_cp_mode=round-robin` 的 runtime wiring。 - decode 侧 CP/shared KV。 - Mooncake 之外的 PD transfer backend 完整支持。 @@ -765,13 +765,22 @@ attention runtime: may use full-view materialization compatibility path Phase 3 继续解决: ```text -attention runtime 不再 materialize full/maxlen KV +复用 current chunk 已经 CP all-gather + rerange 的 KV/index +current-only 场景跳过 shared KV/index full-view materialize +mixed current/history 场景 current 直用,history 暂保留 Phase 2 compatibility materialize +保留所有 fast path 到 Phase 2 full-view compatibility 的 fallback +``` + +Phase 4 继续解决: + +```text +history/shared KV selected-page materialize NSA index/topk shard-aware owner-aware sparse attention distributed softmax/output reduce ``` -Phase 2 不应把 Phase 3 的 runtime workspace 问题隐藏起来。相反,Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来,方便 Phase 3 定位和优化。 +Phase 2 不应把后续 runtime workspace 问题隐藏起来。相反,Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来,方便 Phase 3/Phase 4 定位和优化。 --- diff --git a/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md b/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md new file mode 100644 index 000000000..091c93aa8 --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md @@ -0,0 +1,626 @@ +# NSA Prefill CP Phase 3 计划:复用 Current Chunk,减少 Shared KV 重复 Materialize + +本文档定义 **Phase 3**:在 Phase 2 已经实现 CP shared/sharded persistent KV 的基础上,减少 shared KV compatibility path 引入的 **可避免重复 KV/index materialize**。 + +Phase 3 的边界是:**复用当前 forward/chunk 中已经 CP all-gather + rerange 过的 current KV/index**,避免它们再次从 sharded persistent KV/index pool 里 materialize。更深层的 history KV shard-aware topk、selected KV exchange、distributed sparse attention 进入 Phase 4。 + +--- + +## 0. 当前状态与问题 + +Phase 2 当前语义: + +```text +persistent KV at rest: CP-sharded/shared +scheduler/radix/req_to_token: logical loc/page +attention/index runtime: full-view compatibility materialize +``` + +当前性能问题集中在两条 runtime compatibility 路径: + +1. **NSA indexer topk 前 materialize full index buffer**: + - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py::_maybe_materialize_shared_index_buffer` + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py::materialize_shared_paged_buffer` +2. **MLA sparse attention 前 materialize full KV buffer**: + - `python/sglang/srt/layers/attention/nsa_backend.py::forward_extend` + - `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py::materialize_shared_token_kv_buffer` + +这两条路径都会构造 dense full-view buffer: + +```text +每个 CP rank 拷贝自己 owner 的 physical pages,其余填 0 +然后 CP group all-reduce 得到完整 dense view +``` + +但在 prefill CP 中,**current chunk 的 KV/index 在写入 persistent pool 之前已经 gather 过一次**: + +- MLA current KV:`DeepseekV2AttentionMLA.rebuild_cp_kv_cache()` 调用 `cp_all_gather_rerange_output()`。 +- NSA index current K:`NativeSparseIndexer._get_q_k_bf16()` 中调用 `cp_all_gather_rerange_output()`。 + +因此 current chunk 这部分数据的当前路径是: + +```text +current KV/index 先 CP all-gather + rerange +-> shared KV 写入时只写 owner shard +-> attention/topk 前又从 persistent shard materialize + all-reduce +``` + +Phase 3 要去掉的就是这部分 **current chunk 重复 materialize**。 + +--- + +## 1. Phase 3 目标 + +### 1.1 核心目标 + +Phase 3 只优化 **Phase 2 compatibility path 中对 current chunk 的重复读/通信**: + +```text +如果 attention/topk 要读的 loc 属于当前 forward/chunk 的 out_cache_loc: + 直接复用已经 CP all-gather 后的 current tensor + 不再从 persistent shared KV/index pool materialize + +如果 attention/topk 要读的 loc 属于 history: + 暂时继续走 Phase 2 full-view compatibility materialize +``` + +这里的 history 包括: + +```text +cached prefix +radix cache hit +chunked prefill 前序 chunk +已写入 persistent KV pool 的历史 tokens +``` + +### 1.2 预期收益 + +current-only prefill 场景: + +```text +MLA KV materialize all-reduce: 0 +NSA index materialize all-reduce: 0 +``` + +mixed current/history 场景: + +```text +current chunk 不再进入 shared KV/index materialize +history 仍走 Phase 2 compatibility materialize +``` + +### 1.3 非目标 + +以下不属于 Phase 3,进入 Phase 4 或后续: + +```text +- shard-aware NSA topk +- local topk + global topk merge +- selected-page history materialize / selected-page global union +- owner-to-requester KV exchange +- distributed sparse attention / distributed softmax/output reduce +- current chunk owner-routed all_to_all write +- nsa_prefill_cp_mode=round-robin runtime wiring +- decode CP/shared KV +- 非 Mooncake PD backend 完整支持 +``` + +--- + +## 2. 当前数据流 + +### 2.1 MLA attention 当前路径 + +```text +forward_absorb_prepare + | + |-- compute current k_nope / k_pe + | + |-- rebuild_cp_kv_cache() + | | + | +-- CP all_gather + rerange current MLA KV + | + v +forward_absorb_core -> attn_mqa -> nsa_backend.forward_extend + | + |-- save_kv_cache + | | + | +-- _maybe_filter_shared_mla_kv_write() + | +-- only owner logical locs are written to physical KV pool + | + |-- materialize_shared_token_kv_buffer() + | + +-- materialize dense full-view KV from persistent shared KV + +-- CP all-reduce dense KV buffer + +-- run flashmla_sparse/tilelang with remapped page_table_1 +``` + +当前冗余点:`k_nope/k_pe` 已经是 CP all-gather 后的 current chunk KV,但 attention 仍可能通过 `materialize_shared_token_kv_buffer()` 再从 persistent shard 重建 current chunk KV。 + +### 2.2 NSA indexer 当前路径 + +```text +nsa_indexer.forward_cuda + | + |-- _get_q_k_bf16() + | | + | +-- CP all_gather + rerange current index K + | + |-- _store_index_k_cache() + | | + | +-- only owner logical locs are written to physical index K/scale pool + | + |-- _get_topk_paged() / _get_topk_ragged_with_cp() + | + +-- _maybe_materialize_shared_index_buffer() + +-- materialize dense full-view index K/scale from persistent shared index pool + +-- CP all-reduce dense index buffer + +-- compute logits/topk +``` + +当前冗余点:index K 已经是 CP all-gather 后的 current chunk K,但 topk 前仍可能从 persistent shard 重建 current chunk index K/scale。 + +--- + +## 3. Phase 3 目标数据流 + +### 3.1 MLA current-only fast path + +适用条件: + +```text +forward_batch.uses_cp_shared_kv == True +forward mode 是 prefill/extend CP path +page_table_1/topk logical locs 全部属于 forward_batch.out_cache_loc +无 cached prefix / previous chunk / radix history loc 被当前 attention 使用 +``` + +目标路径: + +```text +current gathered k_nope/k_pe + | + +-- build compact current KV buffer + | + +-- remap page_table_1 logical loc -> compact current row id + | + +-- run existing flashmla_sparse/tilelang kernel +``` + +跳过: + +```text +materialize_shared_token_kv_buffer() +CP all-reduce dense KV buffer +``` + +### 3.2 MLA mixed current/history path + +适用条件: + +```text +page_table_1/topk logical locs 同时包含 current loc 和 history loc +``` + +目标路径: + +```text +current locs: + use gathered k_nope/k_pe directly + +history locs: + keep Phase 2 materialize_shared_token_kv_buffer compatibility path + but exclude current locs from the remap set when it is safe and group-consistent + +combined compact KV: + concat(history_compact_kv, current_compact_kv) + remap page_table_1: + history loc -> history compact row + current loc -> current compact row + history_offset + run existing flashmla_sparse/tilelang kernel +``` + +注意:Phase 3 mixed path 不做 selected-page global union。history 的 materialize 仍以 Phase 2 compatibility 语义为准;只保证 current chunk 不被重复 materialize。 + +### 3.3 NSA indexer current-only fast path + +适用条件: + +```text +indexer topk 需要访问的 K 全部来自 current out_cache_loc +``` + +目标路径: + +```text +current gathered index K + | + +-- quantize to current fp8 K/scale buffer + | + +-- build compact current index page buffer + | + +-- remap logical page table -> compact current pages + | + +-- compute logits/topk with existing deep_gemm path where possible +``` + +跳过: + +```text +_maybe_materialize_shared_index_buffer() +materialize_shared_paged_buffer() +CP all-reduce dense index buffer +``` + +MVP 可以接受 current index K 被额外 quantize 一次,以换取避免 full index materialize/all-reduce。后续如果需要再优化 fused store,使 current quantized result 同时用于 persistent write 和 topk。 + +### 3.4 NSA indexer mixed current/history path + +适用条件: + +```text +topk 需要同时访问 current index K 和 history index K +``` + +目标路径: + +```text +current index: + use current gathered key -> quantized compact current index buffer + +history index: + keep Phase 2 materialize_shared_paged_buffer compatibility path + +combined index: + concat(history_index_buffer, current_index_buffer) + remap logical positions/pages to combined compact ids + compute logits/topk + map topk result back to global logical loc semantics +``` + +这一步比 MLA mixed path 风险更高,因为 topk 输出必须保持 global logical loc/page_table 语义。实施顺序应晚于 MLA current-only 和 MLA mixed。 + +--- + +## 4. 关键接口设计 + +### 4.1 Current loc remap helper + +建议新增到: + +```text +python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +``` + +候选接口: + +```python +def build_current_loc_remap( + query_locs: torch.Tensor, + current_locs: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return (is_current_mask, compact_row_ids). + + query_locs: + logical locs from topk/page_table. May contain -1 sentinel. + + current_locs: + forward_batch.out_cache_loc. Row order must match the gathered current + KV/index tensor order. + + compact_row_ids: + For entries where is_current_mask is true, row id into current compact + tensor. Other entries are unspecified or -1. + """ +``` + +Implementation requirements: + +```text +- 不能使用 Python set/CPU membership 检查热路径数据 +- 支持 -1 sentinel +- 支持 1D/2D topk/page_table tensor +- 不做 CUDA -> CPU sync +- current_locs 不假设一定连续;可通过 sort + searchsorted 建映射 +``` + +### 4.2 History remap set + +MLA mixed path 中如果要排除 current loc,需要构造 group-consistent history remap set: + +```text +metadata.page_table_1 + -> remove locs that belong to current_locs + -> history_remap_locs +``` + +约束: + +```text +所有 CP rank 必须得到相同 shape/内容的 history_remap_locs, +否则 materialize_shared_token_kv_buffer 内部 all-reduce buffer shape 会不一致。 +``` + +如果无法证明 group-consistent,则 fallback 到 Phase 2 full materialize。 + +### 4.3 Current KV reuse context + +需要把 current gathered tensor 显式传递到 shared KV fast path,避免隐式依赖局部变量。 + +候选方案: + +```text +方案 A:只在 nsa_backend.forward_extend 内使用已有 k/k_rope 参数 + - MLA 最简单 + - 不需要扩展 ForwardBatch + +方案 B:ForwardBatch/metadata 增加 current_mla_kv/current_index_k context + - indexer 更方便 + - 需要注意 batch/layer 生命周期 + +推荐: + MLA 先用方案 A;indexer 再引入最小 scoped context,避免过早污染 ForwardBatch。 +``` + +--- + +## 5. 实施顺序 + +### Step 0:观测与开关 + +新增轻量开关/统计: + +```text +SGLANG_CP_SHARED_KV_CURRENT_REUSE=0/1 +``` + +建议初期默认关闭,验证稳定后再默认开启。 + +统计项: + +```text +cp_shared_kv_current_reuse_mla_current_only_hit +cp_shared_kv_current_reuse_mla_mixed_hit +cp_shared_kv_current_reuse_index_current_only_hit +cp_shared_kv_current_reuse_index_mixed_hit +cp_shared_kv_current_reuse_full_fallback +cp_shared_kv_current_reuse_skipped_current_tokens +cp_shared_kv_current_reuse_history_tokens +``` + +debug 细节仍受 `SGLANG_DEBUG_CP_SHARED_KV=1` 控制。 + +### Step 1:current/history loc helper + 单元测试 + +新增并测试: + +```text +build_current_loc_remap +split_current_history_locs +remap_mixed_locs_to_compact_rows +``` + +测试覆盖: + +```text +- current_locs 连续 +- current_locs 非连续 +- query_locs 含 -1 +- query_locs 全 current +- query_locs 全 history +- query_locs current/history 混合 +- duplicate locs/topk rows +``` + +### Step 2:MLA current-only fast path + +在 `nsa_backend.forward_extend()` 中,`materialize_shared_token_kv_buffer()` 前增加 fast path: + +```text +if shared KV enabled and all page_table_1 locs are current: + build compact current KV from k/k_rope + remap page_table_1 + skip materialize_shared_token_kv_buffer +else: + fallback Phase 2 path +``` + +验收: + +```text +- first chunk/no-prefix 场景命中 current-only fast path +- 不调用 materialize_shared_token_kv_buffer +- 输出与 Phase 2 full-view baseline 对齐 +``` + +### Step 3:MLA mixed current/history path + +在 current-only 正确后加入 mixed: + +```text +if page_table_1 contains current + history: + current locs use k/k_rope + history locs use Phase 2 materialize + concat compact buffers + remap page_table_1 +``` + +如果 history remap set 无法保证 group-consistent,则 fallback Phase 2 full materialize。 + +验收: + +```text +- chunked prefill 第二个及之后 chunk 可命中 mixed path +- current chunk tokens 不进入 history materialize +- 输出与 Phase 2 full-view baseline 对齐 +``` + +### Step 4:NSA indexer current-only fast path + +在 `_get_topk_paged()` / `_get_topk_ragged_with_cp()` 前建立 current-only index path: + +```text +current gathered index K + -> quantize current compact K/scale + -> remap block table/page table + -> compute logits/topk +``` + +验收: + +```text +- current-only indexer 不调用 materialize_shared_paged_buffer +- topk logical loc 与 Phase 2 full-view baseline 对齐 +``` + +### Step 5:NSA indexer mixed current/history path + +最后处理 mixed indexer: + +```text +history index materialize + current compact index +-> combined index buffer +-> topk result maintains logical loc semantics +``` + +如果 topk loc semantic 不容易保证,允许保留 Phase 2 fallback,并把 mixed indexer 推迟到 Phase 4 前置任务。 + +--- + +## 6. Fallback 策略 + +Phase 3 所有 fast path 必须有无损 fallback: + +```text +fallback = Phase 2 full-view compatibility materialize +``` + +触发 fallback 的情况: + +```text +- 开关关闭 +- return_logprob / hidden capture / graph capture 路径不满足约束 +- current loc remap 出现不完整覆盖 +- current/history split 结果不能保证所有 CP rank shape 一致 +- topk transform method / backend 不在已验证组合内 +- debug validation 发现非法 loc 或 remap out-of-bound +``` + +fallback 不能改变输出语义,只影响性能。 + +--- + +## 7. 验证计划 + +### 7.1 单元测试 + +新增/扩展测试文件建议: + +```text +test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +test/registered/unit/layers/test_nsa_cp_current_reuse.py +``` + +覆盖: + +```text +- current loc remap +- mixed loc remap +- compact current KV row mapping +- history remap set 排除 current loc +- -1 sentinel 传播 +- fallback 条件 +``` + +### 7.2 小规模集成测试 + +对比组合: + +```text +A. shared KV Phase 2 full materialize baseline +B. shared KV Phase 3 current reuse enabled +``` + +场景: + +```text +- prompt 64 / 512 / 4096 +- chunked-prefill-size 16384 +- prefix cache miss +- prefix cache hit +- chunked prefill 多 chunk +- max_tokens 小值,避免 decode 差异掩盖 prefill 问题 +``` + +检查: + +```text +- 输出 token/logits/topk loc 与 baseline 对齐 +- current-only 场景 materialize count 为 0 +- mixed 场景 skipped_current_tokens > 0 +- no fallback unexpected +``` + +### 7.3 生产启动命令验证 + +主验证组合: + +```text +--enable-nsa-prefill-context-parallel +--nsa-prefill-cp-mode in-seq-split +--attn-cp-size 8 +--enable-nsa-prefill-cp-shared-kv +SGLANG_CP_SHARED_KV_CURRENT_REUSE=1 +``` + +保留 debug 诊断组合: + +```text +SGLANG_DEBUG_CP_SHARED_KV=1 +SGLANG_CP_SHARED_KV_CURRENT_REUSE=1 +``` + +性能压测必须关闭: + +```text +SGLANG_DEBUG_CP_SHARED_KV=0 +``` + +--- + +## 8. 接受标准 + +Phase 3 完成标准: + +```text +1. Phase 2 shared KV correctness 不回退。 +2. current-only MLA attention 不再调用 materialize_shared_token_kv_buffer。 +3. current-only NSA indexer 不再调用 materialize_shared_paged_buffer。 +4. mixed MLA path 至少能避免 current chunk 重复 materialize,history 可继续走 Phase 2 full-view。 +5. 所有 fast path 都有 Phase 2 fallback。 +6. debug/metrics 能区分:current-only hit、mixed hit、fallback、skipped current tokens、history tokens。 +7. 生产路径默认不执行 expensive CUDA sync debug 检查。 +``` + +如果 Step 5 的 NSA indexer mixed path 风险过高,可以作为 Phase 3 的 stretch goal;Phase 3 最小可交付为: + +```text +MLA current-only + MLA mixed + NSA indexer current-only +``` + +--- + +## 9. Phase 4 边界 + +Phase 4 继续解决 history/shared KV runtime 成本: + +```text +- history selected-page materialize +- shard-aware NSA topk +- local topk + global merge +- selected KV owner exchange +- distributed sparse attention +- owner-routed current write,消除写入前 current chunk all-gather 后再过滤 owner 的浪费 +``` + +Phase 3 不隐藏 Phase 4 问题。Phase 3 只保证:**已经 gather 过的 current chunk 不再被 shared KV compatibility path 重复 gather/materialize。**