Clarify the shared KV path toward compute-owned pages

Phase 4 is now documented as the page-aligned CP split step, leaving the existing shared KV layout untouched while establishing the page ownership invariant needed for direct writes. Phase 5 is documented as the follow-up compute-owner layout change that can route local MLA KV and NSA index writes into the owning rank's physical pool.

Constraint: Current Phase 2/3 layout is page-interleaved and still relies on compatibility materialize paths.
Rejected: Combine split alignment, allocator changes, and shard-aware attention into one phase | too broad to debug safely.
Confidence: high
Scope-risk: narrow
Directive: Do not start compute-owner KV layout work before preserving the Phase 4 page-aligned split invariant and fallback semantics.
Tested: git diff --cached --check
Not-tested: Runtime server startup; documentation-only commit.

docs/advanced_features/nsa_prefill_cp_phase1_narrow_output_collection.md

@@ -644,6 +644,15 @@ Phase 1 **不改变** 当前 KV cache 传输格式与 allocator 语义。
 
 ## 12. 后续阶段（本期不做）
 
+注意：以下是 Phase 1 当时的初始拆分。后续 shared KV 工作已经细化为：
+
+```text
+Phase 2: shared/sharded persistent KV pool
+Phase 3: current chunk reuse
+Phase 4: page-aligned CP token split
+Phase 5: compute-owner shared KV layout
+```
+
 ### Phase 2
 - prefill shard-KV 保存
 - 不再在每层 gather 完整 KV 再写本地 pool

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 先去掉 current chunk 的重复 materialize，Phase 4 再做 shard-aware attention/topk。
+   - 这是 Phase 2 的主要性能成本；Phase 3 先去掉 current chunk 的重复 materialize，Phase 4 先把 CP token split 对齐到 KV page，Phase 5 再改 compute-owner KV layout。
 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。
@@ -38,7 +38,8 @@ Phase 2 的目标是先扩大 **persistent KV cache pool 的逻辑容量**。它
 当前仍未完成、属于 Phase 3 或后续工作：
 
 - 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。
+- Phase 4：page-aligned CP token split，为 compute-owner KV layout 建立“一个 page 只由一个 CP compute rank 计算”的不变量。详见 `docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md`。
+- Phase 5：compute-owner shared KV layout，让 current KV/index 可以按 CP-local loc 直接写本 rank physical pool。详见 `docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md`。
 - `nsa_prefill_cp_mode=round-robin` 的 runtime wiring。
 - decode 侧 CP/shared KV。
 - Mooncake 之外的 PD transfer backend 完整支持。
@@ -774,13 +775,20 @@ mixed current/history 场景 current 直用，history 暂保留 Phase 2 compatib
 Phase 4 继续解决：
 
 ```text
-history/shared KV selected-page materialize
-NSA index/topk shard-aware
-owner-aware sparse attention
-distributed softmax/output reduce
+page-aligned CP token split
+确保 current KV page 不被多个 CP compute rank 同时切分
+为 compute-owner KV layout 提供 page ownership 不变量
 ```
 
-Phase 2 不应把后续 runtime workspace 问题隐藏起来。相反，Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来，方便 Phase 3/Phase 4 定位和优化。
+Phase 5 继续解决：
+
+```text
+compute-owner logical page allocation
+local MLA KV / NSA index K 直接写本 rank physical pool
+保持 scheduler/radix/PD transfer 的 logical loc 语义
+```
+
+更深层的 history selected-page materialize、shard-aware NSA topk、owner-aware sparse attention 和 distributed softmax/output reduce 在 Phase 5 之后继续拆分。Phase 2 不应把后续 runtime workspace 问题隐藏起来。相反，Phase 2 应通过日志和 metrics 把 full-view workspace 暴露出来，方便 Phase 3/4/5 定位和优化。
 
 ---
 

docs/advanced_features/nsa_prefill_cp_phase3_current_reuse_plan.md

@@ -2,7 +2,7 @@
 
 本文档定义 **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。
+Phase 3 的边界是：**复用当前 forward/chunk 中已经 CP all-gather + rerange 过的 current KV/index**，避免它们再次从 sharded persistent KV/index pool 里 materialize。后续拆分为：Phase 4 做 page-aligned CP token split，Phase 5 做 compute-owner shared KV layout；更深层的 history KV shard-aware topk、selected KV exchange、distributed sparse attention 继续后移。
 
 ---
 
@@ -91,7 +91,7 @@ history 仍走 Phase 2 compatibility materialize
 
 ### 1.3 非目标
 
-以下不属于 Phase 3，进入 Phase 4 或后续：
+以下不属于 Phase 3，进入 Phase 4/5 或后续：
 
 ```text
 - shard-aware NSA topk
@@ -372,9 +372,12 @@ metadata.page_table_1
 
 ```text
 SGLANG_CP_SHARED_KV_CURRENT_REUSE=0/1
+SGLANG_CP_SHARED_KV_NVTX=0/1
 ```
 
 建议初期默认关闭，验证稳定后再默认开启。
+`SGLANG_CP_SHARED_KV_NVTX=1` 只打开 shared-KV materialize 相关 NVTX range，
+用于 Nsight Systems 定位 Phase2/Phase3 materialize、local copy、all-reduce 开销。
 
 统计项：
 
@@ -483,7 +486,7 @@ history index materialize + current compact index
 -> topk result maintains logical loc semantics
 ```
 
-如果 topk loc semantic 不容易保证，允许保留 Phase 2 fallback，并把 mixed indexer 推迟到 Phase 4 前置任务。
+如果 topk loc semantic 不容易保证，允许保留 Phase 2 fallback，并把 mixed indexer 推迟到 Phase 5 之后的 shard-aware runtime 任务。
 
 ---
 
@@ -578,12 +581,14 @@ SGLANG_CP_SHARED_KV_CURRENT_REUSE=1
 ```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
 ```
 
 ---
@@ -610,17 +615,23 @@ MLA current-only + MLA mixed + NSA indexer current-only
 
 ---
 
-## 9. Phase 4 边界
+## 9. Phase 4 / Phase 5 边界
 
-Phase 4 继续解决 history/shared KV runtime 成本：
+Phase 4 先解决 CP split 与 KV page 不对齐的问题：
 
 ```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 的浪费
+- page-aligned in-seq-split
+- 保证一个 current KV page 只由一个 CP compute rank 计算
+- 不改变 shared KV persistent layout
 ```
 
-Phase 3 不隐藏 Phase 4 问题。Phase 3 只保证：**已经 gather 过的 current chunk 不再被 shared KV compatibility path 重复 gather/materialize。**
+Phase 5 再解决 persistent KV layout 与 compute owner 不一致的问题：
+
+```text
+- compute-owner logical page allocation
+- CP-local out_cache_loc
+- local MLA KV / NSA index K direct persistent write
+- 保持 materialize compatibility path 与 Mooncake PD transfer 正确
+```
+
+history selected-page materialize、shard-aware NSA topk、local topk + global merge、selected KV owner exchange、distributed sparse attention 继续后移。Phase 3 不隐藏这些问题。Phase 3 只保证：**已经 gather 过的 current chunk 不再被 shared KV compatibility path 重复 gather/materialize。**

docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md

@@ -0,0 +1,529 @@
+# NSA Prefill CP Phase 4 计划：Page-Aligned CP Token Split
+
+本文档定义 **Phase 4**：在不改变 shared KV persistent layout 的前提下，先把 NSA prefill CP 的 `in-seq-split` token 切分改造成 **page-aligned**。Phase 4 的目标是给后续 Phase 5 的 compute-owner KV layout 建立必要不变量：
+
+```text
+同一个 KV page 内的 current-token 只由一个 CP compute rank 负责计算。
+```
+
+当前 Phase 2/3 的 shared KV 仍使用 page-interleaved owner mapping。Phase 4 不改变 `CpSharedKVLayout`、allocator、PD transfer、materialize 策略；它只改变 CP compute token split 的边界。
+
+---
+
+## 0. 当前代码事实
+
+### 0.1 NSA CP split 当前是 token 平均切分
+
+主要路径：
+
+- `python/sglang/srt/models/deepseek_v2.py`
+  - `DeepseekV2ForCausalLM.forward()` 在 `can_cp_split(...)` 后设置 `forward_batch.nsa_cp_metadata`。
+  - `DeepseekV2Model.forward()` 在 `nsa_use_prefill_cp(...)` 时对 `hidden_states` 和 `positions` 做 CP split。
+- `python/sglang/srt/models/deepseek_nextn.py`
+  - nextn 路径也复用同一套 NSA CP metadata/split helper。
+- `python/sglang/srt/layers/attention/nsa/utils.py`
+  - `prepare_input_dp_with_cp_dsa(...)`
+  - `cp_split_and_rebuild_data(...)`
+  - `cp_split_and_rebuild_position(...)`
+  - `cp_all_gather_rerange_output(...)`
+
+当前 `in-seq-split` metadata 的核心逻辑：
+
+```python
+cp_segment_num = cp_size * 2
+seq_per_batch = kv_len // cp_segment_num
+split_list = seq_per_batch.repeat_interleave(cp_segment_num).int().tolist()
+remainder = kv_len % cp_segment_num
+if remainder > 0:
+    split_list[:remainder] = [x + 1 for x in split_list[:remainder]]
+```
+
+它只保证 token 数量大致均衡，不保证 segment 边界落在 KV page boundary。
+
+### 0.2 `in-seq-split` 的 zigzag compute owner
+
+当前 `in-seq-split` 把序列切成 `2 * cp_size` 个 segment，然后按 zigzag 分配：
+
+```text
+rank r owns:
+  segment r
+  segment (2 * cp_size - r - 1)
+```
+
+例如 `cp_size = 4`：
+
+```text
+original: block0 block1 block2 block3 block4 block5 block6 block7
+rank0  : block0 + block7
+rank1  : block1 + block6
+rank2  : block2 + block5
+rank3  : block3 + block4
+```
+
+Phase 4 保留这个 compute owner 规则，只把 `split_list` 从 token-level 改为 page-aware。
+
+### 0.3 shared KV persistent layout 当前仍是 page-interleaved
+
+当前 layout：
+
+- `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
+```
+
+该 owner 与 CP compute rank 没有关系。Phase 4 不改这里；Phase 5 再让 logical page 分配/owner 与 compute rank 对齐。
+
+---
+
+## 1. 为什么 Phase 4 要先做 page-aligned split
+
+Phase 5 想实现：
+
+```text
+local hidden -> local MLA KV / local NSA index K -> 直接写本 rank physical KV/index pool
+```
+
+这要求一个 page 的 current-token 不能同时由多个 CP rank 计算。否则 page-level shared KV layout 无法判断该 page 应该由哪个 rank 持久化。
+
+当前 token-average split 可能产生：
+
+```text
+page N:
+  token offset 0..31  -> rank A
+  token offset 32..63 -> rank B
+```
+
+这会导致后续 compute-owner KV layout 无法只用 page owner 表达真实写入关系。
+
+因此 Phase 4 的核心收益不是直接加速，而是建立后续布局改变的前置条件。
+
+---
+
+## 2. Phase 4 目标
+
+### 2.1 核心目标
+
+在 `--enable-nsa-prefill-cp-shared-kv` + `--nsa-prefill-cp-mode in-seq-split` 下，构造 page-aligned 的 `split_list`：
+
+```text
+sum(split_list) == extend token count
+
+对所有内部边界 boundary:
+  (request_relative_start + boundary) % page_size == 0
+```
+
+这里 `request_relative_start` 是当前 extend chunk 在 request KV 序列中的起点，即 `extend_prefix_lens_cpu[0]`。
+
+### 2.2 保留的行为
+
+Phase 4 不改变：
+
+- CP all-gather/rerange 语义。
+- `cp_reverse_index`、`reverse_split_len`、`per_rank_actual_token` 的用途。
+- attention/topk kernel 看到的 full-order runtime view。
+- persistent KV owner mapping。
+- PD transfer 的 page filtering。
+- materialize compatibility path。
+
+### 2.3 非目标
+
+Phase 4 不处理：
+
+- compute-owner KV layout。
+- local KV direct persistent write。
+- 删除 pre-write allgather。
+- shard-aware attention/topk。
+- `round-robin-split` 的 page-level layout 支持。
+- decode CP。
+- 非 NSA/MLA 路径。
+
+---
+
+## 3. Page-aligned split 规则
+
+### 3.1 MVP 约束
+
+Phase 4 第一版建议只支持：
+
+```text
+batch size == 1
+nsa_prefill_cp_mode == in-seq-split
+page_size > 1
+extend_prefix_lens_cpu[0] % page_size == 0
+```
+
+原因：
+
+1. 当前 `prepare_input_dp_with_cp_dsa(...)` 注释和实现本身就是 batch=1 语义。
+2. 如果 `extend_prefix_lens_cpu[0]` 不对齐 page，当前 chunk 的第一个 page 是历史 partial page。这个 page 的 owner 需要与前序 chunk 已经写入的 owner 保持一致。该约束更适合 Phase 5 处理。
+3. shared KV 当前主配置 `page_size=64`，chunked prefill 的主 chunk 通常是 page-aligned；非对齐场景可以先 fallback 到旧 token split。
+
+### 3.2 Page units
+
+在 MVP 约束成立时，当前 extend chunk 可以按 page unit 切：
+
+```python
+num_full_pages = extend_len // page_size
+tail_tokens = extend_len % page_size
+num_units = num_full_pages + (1 if tail_tokens > 0 else 0)
+```
+
+每个 unit 是不可切分的：
+
+```text
+full page unit: page_size tokens
+tail unit     : tail_tokens tokens
+```
+
+segment 边界只能落在 unit 边界。
+
+### 3.3 Segment 分配
+
+仍然构造 `cp_segment_num = cp_size * 2` 个 segment。
+
+推荐算法：
+
+```python
+base = num_units // cp_segment_num
+rem = num_units % cp_segment_num
+
+unit_counts[i] = base + (1 if i < rem else 0)
+```
+
+再把 unit count 转成 token count：
+
+```python
+split_list = []
+unit_cursor = 0
+for count in unit_counts:
+    token_count = 0
+    for _ in range(count):
+        if unit_cursor < num_full_pages:
+            token_count += page_size
+        else:
+            token_count += tail_tokens
+        unit_cursor += 1
+    split_list.append(token_count)
+```
+
+性质：
+
+```text
+sum(split_list) == extend_len
+所有非尾部内部边界 page aligned
+尾部 partial page 只属于最后一个拿到 tail unit 的 segment
+```
+
+### 3.4 Too-short fallback
+
+如果 page 数太少，很多 rank 会拿到 0 token。当前 CP path 对 zero-token rank 的通信/attention/kernel 行为需要谨慎验证。
+
+Phase 4 MVP 建议：
+
+```text
+如果 num_units < cp_size:
+  fallback 到旧 token-average split
+```
+
+更保守的 gate：
+
+```text
+如果 num_units < 2 * cp_size:
+  fallback 到旧 token-average split
+```
+
+推荐先使用保守 gate，跑通后再放宽。
+
+---
+
+## 4. Metadata 影响面
+
+`NSAContextParallelMetadata` 当前字段：
+
+```python
+split_list
+max_rank_len
+zigzag_index
+per_rank_actual_token
+reverse_split_len
+cp_reverse_index
+kv_len_prev / kv_len_next
+actual_seq_q_prev / actual_seq_q_next
+total_seq_lens
+```
+
+Phase 4 可以先复用这些字段，不强制加新字段。
+
+但为了 Phase 5 衔接，建议增加只读/debug 字段：
+
+```python
+page_aligned: bool
+page_size: int
+extend_prefix_len: int
+segment_page_starts: list[int]
+segment_page_ends: list[int]
+```
+
+这些字段的用途：
+
+- Debug/assert：确认 split boundary 没切 page。
+- Phase 5：从 segment/page position 推导 compute owner。
+- local `out_cache_loc` split：保证 local loc 对应完整 page run。
+
+---
+
+## 5. 需要改动的代码位置
+
+### 5.1 CP split metadata 构造
+
+文件：
+
+- `python/sglang/srt/layers/attention/nsa/utils.py`
+
+新增 helper：
+
+```python
+def build_page_aligned_in_seq_split_list(
+    *,
+    extend_len: int,
+    extend_prefix_len: int,
+    page_size: int,
+    cp_size: int,
+) -> tuple[list[int], PageAlignedSplitInfo | None]:
+    ...
+```
+
+`prepare_input_dp_with_cp_dsa(...)` 需要拿到：
+
+- `page_size`
+- `extend_prefix_len`
+- 是否启用 shared KV / page-aligned split
+
+当前函数签名只有：
+
+```python
+prepare_input_dp_with_cp_dsa(kv_len, cp_rank, cp_size, seqs_len)
+```
+
+因此 Phase 4 需要扩展签名，或直接传入 `forward_batch`。
+
+推荐签名：
+
+```python
+prepare_input_dp_with_cp_dsa(
+    kv_len,
+    cp_rank,
+    cp_size,
+    seqs_len,
+    *,
+    forward_batch=None,
+    page_size=None,
+)
+```
+
+这样对调用点侵入较小。
+
+### 5.2 调用点
+
+文件：
+
+- `python/sglang/srt/models/deepseek_v2.py`
+- `python/sglang/srt/models/deepseek_nextn.py`
+
+当前调用：
+
+```python
+forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa(
+    len(input_ids),
+    self.cp_rank,
+    self.cp_size,
+    forward_batch.seq_lens_cpu.tolist(),
+)
+```
+
+Phase 4 后需要传 `forward_batch` 或 `extend_prefix_lens_cpu`：
+
+```python
+forward_batch.nsa_cp_metadata = prepare_input_dp_with_cp_dsa(
+    len(input_ids),
+    self.cp_rank,
+    self.cp_size,
+    forward_batch.seq_lens_cpu.tolist(),
+    forward_batch=forward_batch,
+    page_size=forward_batch.token_to_kv_pool.page_size
+    if hasattr(forward_batch, "token_to_kv_pool")
+    else None,
+)
+```
+
+实际代码中 `forward_batch.token_to_kv_pool` 在 `ForwardBatch` 上可用，或从 model runner/server args 获取 `page_size`。
+
+### 5.3 Split application
+
+文件：
+
+- `python/sglang/srt/layers/attention/nsa/utils.py`
+
+`cp_split_and_rebuild_data(...)` 和 `cp_split_and_rebuild_position(...)` 可以不变，因为它们只消费 `split_list` 和 `zigzag_index`。
+
+### 5.4 Debug/assert
+
+新增 debug helper：
+
+```python
+def assert_page_aligned_split(
+    split_list: list[int],
+    *,
+    extend_prefix_len: int,
+    page_size: int,
+    extend_len: int,
+) -> None:
+    cursor = 0
+    for split in split_list[:-1]:
+        cursor += split
+        assert (extend_prefix_len + cursor) % page_size == 0
+    assert cursor + split_list[-1] == extend_len
+```
+
+该 assert 可以挂到已有 debug env：
+
+```text
+SGLANG_DEBUG_CP_SHARED_KV=1
+```
+
+正常推理路径不要引入 CUDA sync；这里全是 CPU metadata，成本可控。
+
+---
+
+## 6. Phase 4 数据流
+
+### 当前
+
+```text
+input_ids / hidden_states
+  |
+  |-- prepare_input_dp_with_cp_dsa()
+  |     split by token count
+  |
+  |-- cp_split_and_rebuild_data()
+        rank r gets segment r + segment mirror(r)
+        segment boundary may cut a KV page
+```
+
+### Phase 4
+
+```text
+input_ids / hidden_states
+  |
+  |-- prepare_input_dp_with_cp_dsa()
+  |     split by request-relative KV pages
+  |
+  |-- cp_split_and_rebuild_data()
+        rank r gets segment r + segment mirror(r)
+        segment boundary does not cut a KV page
+```
+
+Persistent KV at rest 仍然是 Phase 2/3 的 shared layout。
+
+---
+
+## 7. 测试计划
+
+### 7.1 Unit tests
+
+建议新增：
+
+```text
+test/registered/unit/attention/test_nsa_cp_page_aligned_split.py
+```
+
+覆盖：
+
+1. `extend_len=32768, page_size=64, cp_size=8`
+   - `sum(split_list) == 32768`
+   - 所有内部 cumulative boundary 是 64 的倍数。
+   - 每个 rank 两段 token 数接近。
+2. `extend_len=16384, page_size=64, cp_size=8`
+   - 同上。
+3. `extend_len=1000, page_size=64, cp_size=8`
+   - tail partial page 不被切分。
+4. `extend_prefix_len=128`
+   - boundary 检查使用 `prefix + boundary`。
+5. `extend_prefix_len=1`
+   - MVP fallback 到旧 split，或显式返回 `page_aligned=False`。
+6. too-short case
+   - `num_units < 2 * cp_size` fallback 或 `page_aligned=False`。
+
+### 7.2 Invariant tests
+
+对生成的 `split_list` 验证：
+
+```python
+cursor = 0
+for segment_len in split_list[:-1]:
+    cursor += segment_len
+    assert (extend_prefix_len + cursor) % page_size == 0
+```
+
+对 rank ownership 验证：
+
+```text
+rank r local token ranges:
+  segment r
+  segment 2*cp_size-r-1
+
+每个 local range 是 page-aligned range，除 request tail partial 外不切 page。
+```
+
+### 7.3 Runtime smoke
+
+使用当前 prefill 命令：
+
+```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
+```
+
+验证：
+
+- server 正常启动；
+- prefill health 正常；
+- router + decode 下短请求可返回；
+- debug 日志显示 `page_aligned=True`；
+- 与旧 split 对比输出没有明显异常。
+
+---
+
+## 8. Phase 4 完成标准
+
+Phase 4 完成时应满足：
+
+```text
+1. shared-KV + in-seq-split 下可以生成 page-aligned split_list。
+2. split_list 仍兼容现有 cp_split / cp_all_gather_rerange_output。
+3. 非 page-aligned prefix / too-short case 有明确 fallback 或 assert。
+4. 不改变 shared KV persistent layout。
+5. 不改变 PD transfer。
+6. 单测覆盖 split invariant。
+7. 启动测试不回归。
+```
+
+---
+
+## 9. Phase 4 到 Phase 5 的交付物
+
+Phase 4 需要为 Phase 5 提供以下不变量/metadata：
+
+```text
+page_aligned == True 时：
+  每个 current logical page 只有一个 CP compute owner。
+  owner 可由 segment index -> rank 规则推导。
+  local out_cache_loc 可以用同一 split_list 切出来，不会切 page。
+```
+
+Phase 5 将基于这个不变量，把 persistent KV page owner 改成 CP compute owner，并让 local KV/index 直接写本 rank physical pool。

docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md

@@ -0,0 +1,623 @@
+# 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 分配。
+
+如果命中到 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. 后续 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，调试面会过大。