Define the current-reuse boundary for shared KV Phase 3

Phase 2 expanded prefill CP persistent KV capacity by sharding KV at rest, but its compatibility runtime can still rebuild full-view KV/index buffers from sharded storage. This documents a narrower Phase 3 target: reuse current chunk KV/index that was already CP all-gathered and reranged, while leaving deeper history shard-aware work to Phase 4. The new plan separates current-only and mixed current/history cases, defines fallback requirements to the Phase 2 full-view compatibility path, and records validation criteria before any runtime implementation starts. Constraint: Phase 3 must reduce avoidable duplicate materialization without rewriting NSA topk or sparse attention kernels Constraint: History KV/index access remains Phase 2 compatibility materialization until Phase 4 Rejected: Start with shard-aware topk/distributed attention | too broad for Phase 3 and harder to validate against the current Phase 2 baseline Rejected: Remove all shared-KV runtime materialization immediately | history/prefix/chunked-prefill paths still need a compatibility fallback Confidence: high Scope-risk: narrow Directive: Keep Phase 3 scoped to current chunk reuse; put selected-history and shard-aware algorithms in Phase 4 unless the design is explicitly revised Tested: git diff --check for the two documentation files Not-tested: Runtime behavior; documentation-only change
2026-04-26 05:57:33 +08:00
parent f8fca72635
commit 5af232e9de
2 changed files with 640 additions and 5 deletions
--- 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 定位和优化。

 ---

--- a/docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md
+++ 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。**