# 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
SGLANG_CP_SHARED_KV_NVTX=0/1
SGLANG_DEBUG_SORT_NVTX=0/1
SGLANG_DEBUG_MOE_SORT_NVTX=0/1
```

建议初期默认关闭，验证稳定后再默认开启。
`SGLANG_CP_SHARED_KV_NVTX=1` 只打开 shared-KV materialize 相关 NVTX range，
`SGLANG_DEBUG_SORT_NVTX=1` 只标记 CP shared KV current remap 与 KV allocator free-page merge 中可能触发 ATen sort kernel 的 `torch.sort` callsite。MoE EP preprocess sort 标记被拆到 `SGLANG_DEBUG_MOE_SORT_NVTX=1`，因为 MoE sort 调用频率很高，常规 profile 不应默认开启，避免 Nsight/CPU 侧被动态 NVTX range 拖慢。两个开关默认关闭，并在模块导入时缓存，避免热路径反复读取环境变量。
用于 Nsight Systems 定位 Phase2/Phase3 materialize、local copy、all-reduce 开销。

统计项：

```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
SGLANG_CP_SHARED_KV_NVTX=1
```

性能压测必须关闭：

```text
SGLANG_DEBUG_CP_SHARED_KV=0
SGLANG_CP_SHARED_KV_NVTX=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。**
