# 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
```

当前 page-aligned contract 已去掉 short-suffix fallback：

```text
如果 prefix page-aligned 且 extend_len > 0:
  使用 page-aligned split
  每个物理 page unit 只归属一个 zigzag segment
  没有 page 的 segment 长度为 0
  out_cache_loc 仍只包含 valid tokens，不追加 fake padding tokens
```

原因是 CP 本身不适合无 cache 的短序列；但长 agent 上下文里 radix cache hit 很频繁，命中后 current suffix 可能只有少量 page。如果继续因为 suffix 太短 fallback，会导致 compute-owner allocation/write 在高频 radix-hit 路径失效。当前目标约束是：prefix 必须 page-aligned，实际 suffix page 不被切开。即使 `num_units < cp_size`，也保持 page-aligned CP split，使用 zero-token segment 表示没有 page 的 rank，由 shared-KV write filter 只写本 rank 拥有的 page。

---

## 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
   - cache-miss: `num_units < 2 * cp_size` fallback 或 `page_aligned=False`。
   - radix-hit 且 prefix page-aligned: `num_units < 2 * cp_size` 也允许；没有 page 的 segment 使用 zero-token segment，compute-owner allocation 仍可用。

### 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。
