# NSA Prefill CP Phase 5 计划：Compute-Owner Shared KV Layout

本文档定义 **Phase 5**：在 Phase 4 已经把 `in-seq-split` 改成 page-aligned split 后，把 shared KV persistent layout 从当前的 page-interleaved owner 改造成 **CP compute-owner layout**。

Phase 5 的核心目标：

```text
local hidden -> local MLA KV / local NSA index K -> 直接写本 rank physical KV/index pool
```

避免当前 shared KV 路径中为了 persistent write 先把 current KV/index all-gather 成 full view，再过滤 owner pages 写入。

---

## 0. 与前序 Phase 的边界

### Phase 2/3 当前状态

当前 shared KV 已经实现：

```text
logical KV loc/page:
  scheduler / radix / req_to_token / PD transfer 统一可见

physical KV loc/page:
  每个 CP rank 本地实际 KV pool

owner mapping:
  owner = (logical_page - 1) % cp_size
```

当前 runtime compatibility path 仍可能 materialize full-view KV/index：

```text
每 rank copy 本地 owner pages，其余填 0
CP group all-reduce
得到完整 dense view
```

### Phase 4 前置条件

Phase 5 依赖 Phase 4 的不变量：

```text
page_aligned == True 时，同一个 current KV page 只由一个 CP compute rank 计算。
```

如果 Phase 4 fallback 到旧 token-average split，Phase 5 的 direct-write path 必须禁用。

---

## 1. 当前代码事实

### 1.1 当前 layout 是 logical page modulo owner

文件：

- `python/sglang/srt/mem_cache/cp_shared_kv_layout.py`

当前规则：

```python
owner = (logical_page - 1) % cp_size
physical_page = (logical_page - 1) // cp_size + 1
```

这与 CP compute rank 无关。

### 1.2 当前 allocator 只扩大 logical capacity

文件：

- `python/sglang/srt/mem_cache/allocator.py`
- `python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py`

`CPSharedPagedTokenToKVPoolAllocator` 只要求：

```python
logical_size == physical_size * cp_size
```

然后复用 `PagedTokenToKVPoolAllocator` 的 free page 顺序。它不会根据 CP split owner 分配 logical pages。

### 1.3 当前 MLA persistent write 依赖 full KV

文件：

- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
- `python/sglang/srt/models/deepseek_v2.py`
- `python/sglang/srt/layers/attention/nsa_backend.py`

当前 MLA path：

```text
local hidden
  -> local latent_cache / k_nope / k_pe
  -> rebuild_cp_kv_cache()
       CP all-gather + rerange into full current KV
  -> nsa_backend.forward_extend()
       _maybe_filter_shared_mla_kv_write()
       filter logical pages by current layout owner
       write owned physical pages
```

`_maybe_filter_shared_mla_kv_write(...)` 当前要求：

```python
k.shape[0] == cache_loc.numel()
k_rope.shape[0] == cache_loc.numel()
```

也就是它只接受 full-order KV，不接受 CP-local KV。

### 1.4 当前 NSA index persistent write 也依赖 full key

文件：

- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`

当前 indexer path：

```text
local hidden
  -> local index key
  -> _get_q_k_bf16()
       CP all-gather + rerange index key
  -> _store_index_k_cache()
       _filter_shared_index_write()
       write owned physical index K/scale pages
```

`_filter_shared_index_write(...)` 当前同样要求：

```python
key.shape[0] == forward_batch.out_cache_loc.numel()
```

---

## 2. Phase 5 目标

### 2.1 核心目标

在 Phase 4 `page_aligned=True` 的 batch 上：

```text
1. logical page 的 owner 与 CP compute rank 一致。
2. 本 rank local KV/index key 对应的 local out_cache_loc 全部 owned_by_this_rank。
3. MLA KV 直接用 local KV + local physical loc 写入 persistent KV pool。
4. NSA index K/scale 直接用 local key + local physical loc 写入 index pool。
5. scheduler/radix/req_to_token 仍保存 CP group 统一 logical loc/page。
6. decode 侧仍接收 full layout。
```

### 2.2 性能目标

Phase 5 要去掉的是 **persistent write 前不必要的 current KV/index full gather 依赖**。

注意：如果 attention/topk 计算仍需要 full current KV/index，计算路径上的 all-gather 可以暂时保留。Phase 5 的第一步不承诺彻底消除所有 attention runtime communication。

更深层目标，例如 shard-aware topk、selected KV exchange、distributed sparse attention，可以放到后续 Phase。

### 2.3 非目标

Phase 5 不强制实现：

- `round-robin-split` compute-owner page layout。
- decode CP。
- 完全取消 materialize compatibility path。
- shard-aware distributed attention softmax/output reduce。
- 非 Mooncake transfer backend 完整优化。

---

## 3. 推荐设计：allocation-aware modulo owner

### 3.1 为什么不优先做 explicit owner table

可选方案 A：新增 owner table。

```text
logical_page -> owner_rank
logical_page -> physical_page
```

优点：

- 最灵活。
- 可表达任意 owner 策略。

缺点：

- materialize、PD transfer、debug、allocator 都要读 owner table。
- table 生命周期要与 radix/cache eviction/free 同步。
- 多 rank 必须一致维护 table，错误面更大。

### 3.2 推荐方案：保留 modulo layout 公式，但改变 logical page 分配

当前 `CpSharedKVLayout` 的公式其实可以继续使用：

```python
owner = (logical_page - 1) % cp_size
physical_page = (logical_page - 1) // cp_size + 1
```

Phase 5 不一定要把公式改掉，而是让 allocator 在分配 logical pages 时保证：

```text
需要由 compute rank r 写入的 request page
  -> 分配一个满足 (logical_page - 1) % cp_size == r 的 logical_page
```

这样：

```text
logical page id 本身编码 owner
existing logical_to_physical mapping 继续成立
materialize_shared_* 基本可以继续复用
PD transfer filter 基本可以继续复用
```

这仍然是 KV layout 改变，因为 logical page 的分配策略从“顺序 page-interleaved”变成“按 CP compute owner lane 分配”。

---

## 4. 关键不变量

Phase 5 direct-write path 只能在以下条件成立时启用：

```text
forward_batch.uses_cp_shared_kv == True
forward_batch.nsa_cp_metadata.page_aligned == True
nsa_prefill_cp_mode == in-seq-split
batch size == 1
page_size > 1
```

并且对本 rank：

```python
local_out_cache_loc = cp_split_out_cache_loc_1d(
    forward_batch.out_cache_loc,
    forward_batch.nsa_cp_metadata,
)

assert all(layout.owned_by_this_rank(local_out_cache_loc[local_out_cache_loc > 0]))
```

如果 assert 不成立，必须 fallback 到 Phase 2/3 兼容路径，不能静默写错。

---

## 5. 需要改动的模块

### 5.1 CP-aware logical page allocation

文件：

- `python/sglang/srt/mem_cache/allocator.py`
- `python/sglang/srt/mem_cache/common.py`
- `python/sglang/srt/managers/schedule_batch.py`

当前 `alloc_extend(...)` 只从统一 `free_pages` 顺序取新 page。

Phase 5 需要新增一条 shared-KV CP-aware allocation path：

```text
输入：
  prefix_len
  extend_len
  page_size
  nsa_cp_metadata.page_aligned segment/page ownership

输出：
  full-order logical out_cache_loc
  其中每个 newly allocated page id 满足 owner(logical_page) == compute_owner(page)
```

推荐实现方式：

```text
allocator 维护 cp_size 条 logical page free lane：
  lane r: pages where (page - 1) % cp_size == r

分配 page 时：
  根据 request page position 找 compute owner rank r
  从 lane r 取下一个 logical page
  写入 full-order out_cache_loc 对应 page 的 token locs
```

注意：所有 CP ranks 必须运行相同 deterministic allocation，得到相同 logical `out_cache_loc`。

### 5.2 local out_cache_loc split

文件：

- `python/sglang/srt/layers/attention/nsa/utils.py`
- `python/sglang/srt/model_executor/forward_batch_info.py`

当前 `cp_split_and_rebuild_data(...)` 针对 2D hidden，不适合 1D loc tensor。

新增 helper：

```python
def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor) -> torch.Tensor:
    if is_nsa_prefill_cp_round_robin_split():
        return nsa_cp_round_robin_split_data(input_)

    input_list = list(torch.split(input_, forward_batch.nsa_cp_metadata.split_list, dim=0))
    return torch.cat(
        [input_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index],
        dim=0,
    )
```

并在 `ForwardBatch` 上缓存：

```python
forward_batch.cp_local_out_cache_loc
```

避免每层重复 split。

### 5.3 MLA KV direct write

文件：

- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
- `python/sglang/srt/models/deepseek_v2.py`
- `python/sglang/srt/layers/attention/nsa_backend.py`

当前 `forward_absorb_prepare(...)` 在 `rebuild_cp_kv_cache(...)` 前有 local `k_nope/k_pe`。

Phase 5 需要保留 local KV：

```text
local_k_nope
local_k_pe
```

然后：

```text
local_k_nope/local_k_pe + cp_local_out_cache_loc
  -> layout.logical_locs_to_physical(...)
  -> token_to_kv_pool.set_mla_kv_buffer(...)
```

如果 attention 计算仍需要 full current KV：

```text
local KV
  -> direct persistent write
  -> cp_all_gather_rerange_output for attention runtime
```

这样 Phase 5 第一版先解耦 persistent write 和 attention compute。

### 5.4 NSA index K/scale direct write

文件：

- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`

当前 `_get_q_k_bf16(...)` 在 allgather 前已经有 local `key`。

Phase 5 需要：

```text
local key + cp_local_out_cache_loc
  -> act_quant / fused store
  -> set_index_k_scale_buffer(...)
```

如果 topk 计算仍需要 gathered full key，保留计算用 allgather，但不再用 gathered key 做 persistent write。

### 5.5 materialize compatibility path

文件：

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

如果采用 allocation-aware modulo owner，现有 materialize path 可以继续使用：

```python
layout.owned_pages_mask(...)
layout.logical_pages_to_physical(...)
```

因为 logical page id 已经按 compute owner lane 分配。

Phase 5 只需要增加 debug assertion：

```text
page_aligned direct-write path 写入后，
materialize_shared_token_kv_buffer(...) 读出的 current KV
与 no-direct-write baseline 对齐。
```

### 5.6 Mooncake PD transfer

文件：

- `python/sglang/srt/disaggregation/utils.py`
- `python/sglang/srt/disaggregation/mooncake/conn.py`

现有 shared KV transfer：

```text
filter_kv_pages_for_cp_shared_kv(...)
  -> layout.filter_owned_pages_np(...)
  -> returns physical source pages + request logical positions
```

如果 Phase 5 采用 allocation-aware modulo owner，这条路径可以继续工作：

```text
compute-owner page 被分配到对应 modulo owner lane
layout.filter_owned_pages_np 会选中正确 prefill CP rank 的 pages
logical_page_positions 仍用于 decode dst page selection
```

需要补充测试确认：

```text
prefill source physical pages 与 decode dst pages 按 request positions 对齐。
```

---

## 6. Phase 5 数据流

### 当前 Phase 2/3

```text
local hidden
  |
  +-- local MLA KV / local index K
  |
  +-- CP all-gather + rerange current KV/index
  |
  +-- filter by page-interleaved owner
  |
  +-- write owned physical pages
  |
  +-- attention/topk may materialize full-view shared KV/index
```

### Phase 5 direct-write

```text
local hidden
  |
  +-- local MLA KV / local index K
  |
  +-- cp_local_out_cache_loc
  |     pages already allocated from this compute rank owner lane
  |
  +-- logical_to_physical
  |
  +-- direct write local physical KV/index pool
  |
  +-- attention/topk compute path keeps existing gather/materialize as needed
```

---

## 7. Correctness risks

### 7.1 Prefix partial page

如果 current chunk 从 page 中间开始：

```text
extend_prefix_len % page_size != 0
```

则第一个 current page 已经包含历史 tokens。这个 page 的 owner 必须与历史写入 owner 一致。

Phase 5 MVP 建议：

```text
只有 extend_prefix_len % page_size == 0 时启用 direct-write layout。
否则 fallback 到 Phase 2/3 path。
```

后续可以支持 partial page continuation，但需要 allocator 保存 page owner continuation 语义。

### 7.2 Radix/cache hit

radix prefix 命中会让 current extend 从已有 logical pages 之后开始。只要 `extend_prefix_len` page-aligned，new pages 可以按 compute owner lane 分配。

Phase 4/5 对 radix-hit short suffix 放宽 too-short gate：

- 当 `extend_prefix_len > 0` 且 prefix page-aligned，并且 current suffix page 数至少为 `cp_size` 时，即使 page 数小于 `2 * cp_size`，仍生成 page-aligned split / compute-owner page owner list；未覆盖的第二段 zigzag segment 长度为 0。
- 当前 page-aligned contract 已废弃 short suffix 的 replicated compute 特例。suffix page 数小于 `cp_size` 时，compute-owner page allocation 仍返回实际 page 的 owner list，CP split 使用 zero-token segments 表示没有 page 的 rank。这样保持 owner-lane / HiCache / materialize 的物理 page pattern，不再因为短 suffix 回退到 legacy allocation。

如果命中到 partial page，MVP fallback。

### 7.3 Page owner lane free/evict

allocator 需要保证：

```text
free(logical_page) 回到其 modulo owner lane
evict/release 不破坏 lane 去重
所有 CP ranks 对 logical free state 一致
```

否则不同 rank 可能得到不同 logical loc，导致 req_to_token/radix/PD 全部错位。

### 7.4 Current allgather 仍可能存在

Phase 5 第一版即使 direct-write 成功，也可能仍保留计算用 allgather：

```text
local KV -> direct write
local KV -> allgather for attention/topk compute
```

这不是 Phase 5 失败。Phase 5 主要先消除 persistent write 对 full gathered KV 的依赖，并建立 compute-owner layout。

---

## 8. 测试计划

### 8.1 Unit tests

建议新增/修改：

```text
test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
test/registered/unit/attention/test_nsa_cp_compute_owner_layout.py
```

覆盖：

1. owner-lane page allocation：
   - compute owner 为 `r` 的 request pages 分配到 `(logical_page - 1) % cp_size == r`。
2. `cp_split_and_rebuild_1d(...)`：
   - 对 `out_cache_loc` 的 split 顺序与 hidden split 一致。
3. local loc ownership：
   - `layout.owned_by_this_rank(cp_local_out_cache_loc)` 全 True。
4. MLA direct-write shape：
   - local KV token 数等于 local loc token 数。
5. NSA index direct-write shape：
   - local key token 数等于 local loc token 数。
6. materialize compatibility：
   - direct-write 后 `materialize_shared_token_kv_buffer(...)` 读出的 dense KV 与 baseline 一致。
7. PD filtering：
   - `filter_kv_pages_for_cp_shared_kv(...)` 返回的 physical pages 与 request positions 对齐。

### 8.2 Runtime smoke

使用 prefill CP + decode + router：

```bash
--enable-nsa-prefill-context-parallel
--nsa-prefill-cp-mode in-seq-split
--enable-nsa-prefill-cp-shared-kv
--page-size 64
--chunked-prefill-size 16384
```

检查：

- prefill 正常启动；
- decode 正常启动；
- router health 正常；
- 短 prompt 与长 prompt 均能返回正常内容；
- debug 下确认：
  - `page_aligned=True`
  - `cp_local_out_cache_loc` owned by current rank
  - direct-write KV checksum 与 compatibility materialize checksum 对齐。

### 8.3 性能验证

NVTX 建议区分：

```text
cp_shared_kv.write.mla.direct
cp_shared_kv.write.index.direct
cp_shared_kv.materialize.token.local_copy
cp_shared_kv.materialize.token.all_reduce
cp_shared_kv.materialize.paged.local_copy
cp_shared_kv.materialize.paged.all_reduce
```

Phase 5 预期：

```text
persistent write 前的 full-KV filter/copy 开销下降
materialize compatibility 开销不一定消失
```

---

## 9. Phase 5 完成标准

Phase 5 完成时应满足：

```text
1. Phase 4 page_aligned batch 可以启用 compute-owner logical page allocation。
2. CP-local out_cache_loc 与 local KV/index key token 数一致。
3. MLA KV 可以直接写本 rank physical KV pool。
4. NSA index K/scale 可以直接写本 rank physical index pool。
5. req_to_token/radix 仍保存 full logical loc。
6. materialize compatibility path 仍能正确读出 full-view KV/index。
7. Mooncake PD transfer 仍能把 prefill CP shards 传给 decode full KV。
8. 非 page-aligned / unsupported case 有明确 fallback。
```

## 10. 当前实现切入点

第一版实现按 **allocation-aware modulo owner** 落地：

```text
mem_cache/cp_shared_kv_compute_owner.py
  根据 Phase 4 page-aligned in-seq split 规则生成 current page -> compute owner。

mem_cache/allocator.py
  CPSharedPagedTokenToKVPoolAllocator.alloc_extend_compute_owner(...)
  从对应 modulo owner lane 选择 logical page。

mem_cache/common.py
  在 shared KV + in-seq-split + 单请求 page-aligned 场景调用 compute-owner allocation；
  lane 不足或不满足 gate 时 fallback 到旧 allocator。

layers/attention/nsa/utils.py
  cp_split_and_rebuild_1d(...)
  get_cp_shared_kv_local_out_cache_loc(...)
  生成并缓存本 rank local out_cache_loc，且只有 owner 校验通过才启用 direct write。

forward_mla.py / nsa_indexer.py
  MLA KV 与 NSA index K/scale 在 all-gather 前用 local KV/key + local physical loc 直接写 persistent pool；
  attention/topk 计算路径仍保留原有 all-gather。
  NSA index direct-write 只在 nsa_use_prefill_cp(...) 为 true 时尝试，避免 warmup/短 batch/
  decode 等没有 nsa_cp_metadata 的非 CP 阶段刷 missing_metadata。
```

当前仍保留 fallback：

```text
- 非 page_aligned batch；
- split_list 与 out_cache_loc 长度不一致（例如 padding 场景）；
- local logical loc 不属于当前 cp_rank；
- local KV/index key token 数与 local loc 数不一致；
- compute-owner lane 分配失败。
```

compute-owner lane 分配失败通常不是总 free page 不足，而是 modulo owner lane
不均衡：例如历史 legacy allocation / radix eviction 释放了足够总页数，但某个
`(logical_page - 1) % cp_size == r` lane 的 free page 不够。当前实现会在
compute-owner allocation 前按 owner lane deficit 主动触发 radix cache eviction，
尽量释放对应 lane 的 logical page；只有多次 eviction 后仍不足才 fallback，并在
fallback log 中打印 `required_by_owner / available_by_owner / deficit_by_owner`。

这些 fallback 会通过 logger 每次触发都提示，不能按 reason 去重。PD warmup
可能先触发和真实请求相同的 fallback reason；如果只提示一次，会隐藏后续真实
请求仍在 fallback 的问题：

```text
CP shared KV compute-owner allocation fallback (...)
CP shared KV direct-write fallback (...)
```

---

## 11. 后续 Phase 候选

Phase 5 后，如果仍然慢，下一步应集中在 runtime compute path：

```text
- selected history page materialize，而不是 full dense materialize；
- local topk + global topk merge；
- owner-to-requester selected KV exchange；
- distributed sparse attention；
- 避免 attention compute 需要 full current KV allgather。
```

这些不建议混入 Phase 5，否则同时改 split、allocator、write、PD、attention kernel，调试面会过大。
