# NSA Prefill CP Compute Padding 实现计划

> **For agentic workers:** REQUIRED SUB-SKILL: Use `superpowers:subagent-driven-development` 或 `superpowers:executing-plans` 逐任务实现。本文记录当前代码依据、目标合同、分阶段实现和验证点。

**Goal:** 用 compute padding 消除 tiny extend 在 CP shared KV in-seq split 下产生 zero-lane/mostly-zero-lane 的 hang 风险，同时不把 dummy token 写入 radix/HiCache/KV cache。

**Architecture:** cache/radix/HiCache 继续以真实 page 为最小单位；forward compute 额外补 dummy query row，让每个 CP rank 至少拿到一个 suffix page 的 compute work。metadata 需要同时表达 compute split 与 valid/cache split，所有写 cache、current reuse、backup/load 只看 valid split，attention/last-token 的本地排布看 compute split。

**Tech Stack:** Python/SGLang NSA prefill CP metadata, CP shared KV direct write/current reuse, HiCache, unit tests under `test/registered/unit`.

---

## 1. 当前代码依据

### 1.1 tiny skip 目前只支持 bs=1

代码位置：

- `python/sglang/srt/layers/attention/nsa/utils.py`
  - `should_skip_cp_shared_kv_cp_split_for_short_page_extent()`
  - `can_cp_split()`

现状：

- `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` 要求 `len(extend_seq_lens_cpu) == 1`。
- bs>1 时，即使每个 request 都是 tiny extend，`can_cp_split()` 仍可能因为 flattened total token 足够而进入 CP split。
- 这会重新制造之前 bs=1 修过的 zero-lane / mostly-zero-lane 分布。

### 1.2 现有 batched plan 是 valid-token split，不是 compute split

代码位置：

- `CPSharedKVBatchPlan`
- `build_batch_page_aligned_in_seq_split_plan()`
- `build_page_aligned_in_seq_split_list()`

现状：

- 每个 request 独立 page-rounded，但 `request_split_lists` 仍按真实 valid token 计数。
- `extend=65,page=64,cp=8` 会得到类似 `[64, 1, 0, 0, 0, 0, 0, 0, ...]` 的 valid split。
- 这对 cache page contract 是正确的，但对 distributed compute 是不稳定的：大多数 rank 没有 query。

### 1.3 不能直接把 total_len 传大来“补齐”

`build_page_aligned_in_seq_split_list(total_len=512, extend_len=65, ...)` 当前会把 `padding_tokens = total_len - extend_len` 加到最后一个 segment，而不是把 8 个 page 分散到前 8 个 segment。

因此 compute padding 需要新的 split 构造逻辑，不能复用当前 `padding_tokens` 参数。

### 1.4 关键 consumer 目前默认 split rows == valid rows

受影响路径：

- `split_tensor_by_cp_batch_plan()`：要求 `tensor.shape[0] == sum(request_extend_lens)`。
- `cp_split_and_rebuild_position()`：bs>1 目前仍用 scalar `metadata.split_list/zigzag_index`，没有 batch-aware。
- `get_cp_shared_kv_local_out_cache_loc()`：按 valid `out_cache_loc` split，且 direct write 要求 local KV rows 与 local loc 数量一致。
- `DeepseekV2Model.forward_core()`：进入模型后对 `hidden_states/positions` 做 CP split。
- `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()`：要求 `k_nope/k_pe` rows 与 local loc rows 一致。
- `nsa_indexer.py::_store_cp_shared_local_index_k_cache()`：要求 local index KV rows 与 local loc rows 一致。
- `nsa_indexer.py::_get_topk_in_seq_cp_pair_batch()`：使用 `request_actual_seq_q_prev/next` 顺序消费 `q_fp8/weights`。
- `_in_seq_collect_last_token_batch()`：用 `request_last_token_owner/local_offset/rank_local_offsets` 从本地 hidden 中取 compact last hidden。

结论：compute padding 后必须区分 “本地 compute rows” 与 “本地 valid/cache rows”。不能把 dummy rows 直接送到 direct write/current reuse。

---

## 2. 目标合同

### 2.1 cache padding 与 compute padding 分离

对每个 request：

```text
valid_tokens = extend_len
valid_pages = ceil_div(valid_tokens, page_size)
valid_padded_tokens = valid_pages * page_size

compute_pages = max(valid_pages, cp_size)
compute_tokens = compute_pages * page_size
compute_padding_tokens = compute_tokens - valid_tokens
```

约束：

- radix/HiCache/device KV allocation 只使用 `valid_tokens/valid_pages`。
- compute padding 只存在于 forward local tensor 排布。
- dummy rows 不进入 out_cache_loc、direct write、backup、load、radix insert。

### 2.2 示例合同

`extend_len=65,page_size=64,cp_size=8`：

```text
valid_pages = 2
valid_padded_tokens = 128
compute_pages = 8
compute_tokens = 512

valid_split   = [64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
compute_split = [64,64,64,64,64,64,64,64,0,0,0,0,0,0,0,0]
```

CP rank local compute rows：

```text
rank0: segment0 + segment15 = 64 + 0
rank1: segment1 + segment14 = 64 + 0
...
rank7: segment7 + segment8 = 64 + 0
```

真实最后一个 token：

```text
token index = 64
owner rank = 1
local offset = 0
```

### 2.3 bs>1 合同

每个 request 独立 compute padding，不把 batch 当成一条长序列。

示例：

```text
extend_lens=[65, 1024], page_size=64, cp_size=8

req0:
  valid_pages=2, compute_pages=8
req1:
  valid_pages=16, compute_pages=16
```

每个 rank 的本地 tensor 是按 request 顺序拼接：

```text
rank_local = req0.local_compute_rows + req1.local_compute_rows + ...
```

---

## 3. Metadata 设计

### 3.1 保留 valid/cache split，新增 compute split

推荐字段：

```python
request_valid_split_lists: List[List[int]]
request_valid_rank_local_tokens: List[int]
request_valid_rank_local_offsets: List[int]
request_valid_padded_pages: List[int]
request_valid_padded_tokens: List[int]
request_valid_padding_tokens: List[int]

request_compute_split_lists: List[List[int]]
request_compute_rank_local_tokens: List[int]
request_compute_rank_local_offsets: List[int]
request_compute_padded_pages: List[int]
request_compute_padded_tokens: List[int]
request_compute_padding_tokens: List[int]

compute_padding_enabled: bool
```

兼容策略：

- 在 compute padding 模式下，`request_split_lists` 应明确作为 compute split 使用，因为它驱动 `cp_split_and_rebuild_*`、attention local q 长度、`max_rank_len`。
- valid/cache consumer 必须改为显式使用 `request_valid_split_lists`。
- bs=1 也应走同一套 batch plan 合同，避免 scalar path 继续隐藏 tiny 形状问题。

### 3.2 actual_seq_q/kv_len 语义

需要拆开：

- `request_actual_seq_q_prev/next`：给 attention/indexer 的 compute query row 数，应该来自 compute split。
- `request_valid_seq_q_prev/next`：给 direct write/current reuse/cache 写入的真实 row 数，来自 valid split。
- `request_kv_len_prev/next`：attention 的 KV 可见长度仍要按真实逻辑位置计算，不应把 dummy rows 当成已经写入 cache 的 KV。

注意：这部分是实现的核心风险点。attention 可以计算 dummy query，但 dummy query 不能增加真实 KV cache length。

---

## 4. 分阶段实现计划

### P1：planner 纯 CPU 合同

**Files:**

- Modify: `python/sglang/srt/layers/attention/nsa/utils.py`
- Test: `test/registered/unit/layers/test_nsa_cp_utils.py`

步骤：

1. 增加 compute padding split helper，例如：

```python
def build_page_aligned_compute_padding_split_list(
    *,
    extend_len: int,
    extend_prefix_len: int,
    page_size: int,
    cp_size: int,
) -> tuple[list[int], list[int], PageAlignedInSeqSplitInfo]:
    ...
```

返回：

- `valid_split_list`
- `compute_split_list`
- split info 内含 valid/compute padded pages/tokens

2. 扩展 `CPSharedKVBatchPlan`：

- 新增 valid split 字段。
- 新增 compute split 字段。
- `request_split_lists` 临时/兼容指向 compute split。

3. 更新 `build_batch_page_aligned_in_seq_split_plan()`：

- 每个 request 独立计算 valid/compute split。
- last-token owner 用 compute split + actual valid token count 计算。
- rank-local offsets 用 compute local tokens 计算。

4. 单测：

- `extend=65,page=64,cp=8`：
  - valid split 是 `[64,1,0,...]`
  - compute split 是前 8 段各 64。
  - last owner 是 rank1。
- `extend=100,page=64,cp=8`：
  - compute tokens 512，不是 1024。
- `extend=1024,page=64,cp=8`：
  - compute split 与 valid split 等价，不额外 padding。
- `extend_lens=[65,1024]`：
  - 每个 request 独立 padding，rank offsets 按 request 顺序累加。

### P2：CP split/rebuild 支持 compute rows

**Files:**

- Modify: `python/sglang/srt/layers/attention/nsa/utils.py`
- Test: `test/registered/unit/layers/test_nsa_cp_utils.py`

步骤：

1. 修改 `split_tensor_by_cp_batch_plan()`：

- 输入是 valid flattened rows 时，按 request 先 pad 到 compute rows，再按 compute split 切分。
- `mode="data"`：dummy rows 填 0。
- `mode="1d"`：需要由 caller 指定 pad 值；input_ids 默认 0。
- `mode="position"`：按 request 生成 dummy positions，建议用最后一个真实 position 继续递增，保证 RoPE 输入合法；dummy 输出后续不被收集/写 cache。

2. 修复 `cp_split_and_rebuild_position()`：

- bs>1 不能继续 scalar split。
- 统一调用 batch plan split helper。

3. 单测：

- `hidden_states` 输入 65 rows，rank1 输出 64 rows，其中第 1 row 是真实 token 64，其余 dummy 为 0。
- positions 65 rows 后 rank0-rank7 都有 64 positions。
- bs>1 时 request 边界不被打乱。

### P3：narrow last-token output 基于 compute offset

**Files:**

- Modify: `python/sglang/srt/layers/attention/nsa/utils.py`
- Test: `test/registered/unit/layers/test_nsa_cp_utils.py`

步骤：

1. `_get_in_seq_last_token_owner_and_offset()` 输入 compute split，`actual_token_count` 仍是真实 extend_len。
2. `_in_seq_collect_last_token_batch()` 使用 compute rank offsets。
3. full rerange/logprob/capture-hidden 在 compute padding 模式下先 fail-fast：

```text
[CP_SHARED_KV_FAIL_FAST][compute_padding_full_rerange_unsupported]
```

原因：full output 需要把 all-gather 后的 compute dummy rows trim 回 valid rows，不能复用当前 valid/full-rerange 假设。

### P4：valid-row selector，禁止 dummy rows 写 cache

**Files:**

- Modify: `python/sglang/srt/layers/attention/nsa/utils.py`
- Modify: `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
- Modify: `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- Test: `test/registered/unit/layers/test_nsa_cp_utils.py`

步骤：

1. 增加 helper：

```python
def select_cp_local_valid_rows_for_cache_write(
    forward_batch,
    local_tensor: torch.Tensor,
) -> torch.Tensor:
    ...
```

行为：

- 输入 local compute rows。
- 按 `request_valid_split_lists` + `request_compute_split_lists` 为当前 rank 选出真实 valid rows。
- 输出 rows 数必须等于 `get_cp_shared_kv_local_out_cache_loc()` 的 rows 数。

2. `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()` 前过滤 `k_nope/k_pe`。
3. `nsa_indexer.py::_store_cp_shared_local_index_k_cache()` 前过滤 local index key。
4. current reuse 不能再用 `key[:valid_current_rows]` 这种全局 prefix slice；需要同样用 valid selector 得到当前真实 rows。

单测：

- rank1 的 `extend=65` local compute rows=64，但 valid rows=1。
- direct write loc rows=1。
- 只有第 1 个真实 row 被写入，dummy 63 rows 不写。

### P5：index/top-k 与 attention metadata

**Files:**

- Modify: `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
- Modify: `python/sglang/srt/layers/attention/nsa/utils.py`
- Test: `test/registered/unit/layers/test_nsa_cp_utils.py`

步骤：

1. `request_actual_seq_q_prev/next` 使用 compute lengths，保证 attention/indexer 输入 rows 与本地 compute tensor 一致。
2. 新增 valid q metadata，供 top-k compact 真实 query：

```python
request_valid_seq_q_prev
request_valid_seq_q_next
request_valid_query_row_spans
```

3. `_get_topk_in_seq_cp_pair_batch()` 对 compute rows 输出同长度 result：

- valid rows 正常 top-k。
- dummy rows 填 `-1`。
- output_cursor 按 compute rows 前进。

4. `current_index_kv` reuse 使用 valid rows，不使用 dummy rows。

风险：

- 如果 attention kernel 不接受 dummy top-k `-1` rows，需要把 dummy top-k 映射到安全 page，而不是 `-1`。这一步必须用远端 CUDA/ETE 验证。

### P6：can_cp_split gate 与 bs>1 接入

**Files:**

- Modify: `python/sglang/srt/layers/attention/nsa/utils.py`
- Test: `test/registered/unit/layers/test_nsa_cp_utils.py`

步骤：

1. 删除/替换 `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` 的 bs=1-only tiny skip 逻辑。
2. `can_cp_split()` 对 CP shared KV 使用 per-request planner 判断：

- prefix 必须 page-aligned，否则 fail-fast。
- tiny request 不再 skip CP split，而是启用 compute padding。
- 仍要保证 `use_nsa`、`context_parallel_extend`、`nsa_enable_prefill_cp` 等原条件。

3. 对非 shared-KV 路径保持原逻辑。

### P7：HiCache/backup/load 不扩容到 compute padding

**Files:**

- Verify/possibly modify:
  - `python/sglang/srt/managers/schedule_batch.py`
  - `python/sglang/srt/model_executor/forward_batch_info.py`
  - `python/sglang/srt/mem_cache/hicache_controller.py`
  - `python/sglang/srt/managers/cache_controller.py`

检查点：

- `alloc_for_extend()` 仍按 valid `extend_num_tokens` 分配。
- `out_cache_loc` 仍只有 valid rows。
- HiCache per-layer backup batching 只提交 valid rows。
- load-back/prefetch 不因为 compute padding 请求更多 host/device slots。

### P8：验证与上线保护

**本地 CPU 单测：**

```bash
cd /root/sglang-work/sglang-dev
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/layers/test_nsa_cp_utils.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py
```

**远端 CUDA 验证：**

```bash
scp -o ControlMaster=no -o ControlPath=none \
  python/sglang/srt/layers/attention/nsa/utils.py \
  python/sglang/srt/layers/attention/nsa/nsa_indexer.py \
  python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py \
  g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/layers/attention/nsa/
```

实际同步时 `forward_mla.py` 路径要单独同步到：

```bash
g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/models/deepseek_common/attention_forward_methods/
```

远端先跑 targeted unit，再由用户启动 ETE：

```bash
ssh -o ControlMaster=no -o ControlPath=none g0034 \
  "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py'"
```

ETE 重点观察：

- repeated cache-hit `extend=65` 不再 hang。
- accept len 不因为 dummy rows 掉到 1。
- HiCache fallback/warning 不出现 dummy write。
- output len 不为 0。
- cache hit 后 direct write/current reuse 仍走 fast path。

---

## 5. 当前最关键风险

1. **attention 对 dummy query 的接受程度未验证。** 如果 top-k/attention kernel 不接受 dummy top-k `-1`，需要使用 safe page id 或 valid-query mask。
2. **current reuse 当前大量代码按 prefix slice。** compute padding 后必须改成按 local valid span select，否则会把 dummy rows 或错位 rows 当成 current KV。
3. **full rerange 暂不能直接复用。** compute padding 会让 all-gather 后 rows 包含 dummy，必须先 fail-fast 或实现 trim。
4. **positions padding 需要远端验证。** dummy position 不写 cache，但会进 q/projection/attention，必须保证不会触发 kernel shape/causal 边界问题。
5. **不能增加 cache capacity 需求。** 所有 allocation/backup/load 必须继续按 valid page，不按 compute page。

---

## 6. 推荐实施顺序

先做 P1-P3，锁住 metadata、split/rebuild 和 last-token narrow output；这部分可以主要靠 CPU/unit test。

然后做 P4-P5，处理 cache write/current reuse/top-k 的 valid-row 选择；这是 correctness 核心。

最后做 P6-P8，把 tiny skip gate 改为 compute padding 并做远端 ETE。

不要先改 `can_cp_split()` 放开 tiny bs>1，否则现有 consumer 仍会看到 zero/dummy 不一致，风险最高。

---

## 7. Implementation Ledger

### P1 complete: planner exposes valid/compute split side-by-side

Date: 2026-06-03

Implemented:

- `CPSharedKVBatchPlan` now exposes explicit valid/cache fields:
  - `request_valid_split_lists`
  - `request_valid_padded_pages`
  - `request_valid_padded_tokens`
  - `request_valid_padding_tokens`
  - `request_valid_rank_local_tokens`
  - `request_valid_rank_local_offsets`
  - `request_valid_actual_seq_q_prev/next`
- It also exposes compute-padding fields:
  - `request_compute_split_lists`
  - `request_compute_padded_pages`
  - `request_compute_padded_tokens`
  - `request_compute_padding_tokens`
  - `request_compute_padded_token_offsets`
  - `request_compute_rank_local_tokens`
  - `request_compute_rank_local_offsets`
  - `request_compute_actual_seq_q_prev/next`
  - `compute_padding_enabled`
- `NSAContextParallelMetadata` mirrors the new fields from the batch plan.
- Planner computes last-token owner/local offset using compute split plus real
  `extend_len`, so tiny suffixes can locate the real final token inside a
  padded local compute segment.

Important boundary:

- P1 does **not** switch runtime split helpers to compute rows.
- `request_split_lists`, `request_padded_pages`, `request_rank_local_tokens`,
  and `request_actual_seq_q_prev/next` intentionally keep the old valid-token
  semantics until P2 updates consumers.
- This avoids breaking direct write/current reuse/top-k paths before valid-row
  selection exists.

Verification:

- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py`
  - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py`
  - Result: `55 passed`.

Next:

- P2 should add compute-row split/rebuild support and then decide when
  `request_split_lists` aliases should move from valid split to compute split.

### P2 complete: split/rebuild can materialize local compute rows

Date: 2026-06-03

Implemented:

- `split_tensor_by_cp_batch_plan()` now uses `request_compute_split_lists` and
  `request_compute_padded_tokens` when `compute_padding_enabled=True`.
- Its input contract remains valid flattened rows.  The helper pads each
  request independently before CP segment splitting:
  - `mode="data"`: dummy rows are zeros.
  - `mode="1d"`: dummy token/id rows are zeros.
  - `mode="position"`: dummy positions continue from the last real position.
- `cp_split_and_rebuild_data()` and `cp_split_and_rebuild_1d()` now use the
  batch plan path whenever `nsa_cp_metadata.batch_plan` exists, including
  bs=1 plans used by compute-padding tests.
- `cp_split_and_rebuild_position()` is batch-plan aware and no longer forces
  scalar `metadata.split_list` when a batch plan exists.
- Added explicit fail-fast when padded request rows do not match the selected
  split total:
  `[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_target_len_mismatch]`.

Example verified:

- `extend=65,page_size=64,cp_size=8,cp_rank=1`
  - valid rows input: 65.
  - local compute rows output: 64.
  - first row is real token 64.
  - remaining 63 rows are dummy zeros for data/1d.
  - positions become `[40384, ..., 40447]`.

Important boundary:

- P2 still does not make dummy rows safe for cache writes/current reuse/top-k.
- `can_cp_split()` should not be changed to force tiny requests into CP split
  until P4/P5 provide valid-row selection and dummy-safe top-k/attention
  metadata.
- Existing metadata aliases (`request_split_lists`, `request_actual_seq_q_*`)
  still preserve prior valid-token semantics unless a caller explicitly uses
  the new compute fields through the split helper.

Verification:

- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py`
  - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
  - Result: `200 passed, 2 subtests passed`.

### P3 complete: narrow last-token collect uses compute offsets

Date: 2026-06-04

Implemented:

- `_in_seq_collect_last_token()` now dispatches to the batch-plan path whenever
  `nsa_cp_metadata.batch_plan` exists, including bs=1 compute-padding plans.
- `_in_seq_collect_last_token_batch()` now resolves last-token metadata from
  either `NSAContextParallelMetadata` or its `batch_plan`.
- When `compute_padding_enabled=True`, last-token collection uses
  `request_compute_rank_local_offsets` instead of valid/cache
  `request_rank_local_offsets`.
- Existing valid-offset behavior remains for non-compute-padding metadata.

Example verified:

- Single request `extend=65,page_size=64,cp_size=8,cp_rank=1`:
  - local compute rows: 64.
  - real last token is local row 0.
  - dummy row 1 is ignored.
- Batch request with two compute-padded suffixes:
  - valid offsets would select the wrong second request row.
  - compute offsets select the correct per-request last hidden.

Important boundary:

- P3 only fixes narrow output collection.
- Full rerange/logprob/capture-hidden still require a separate
  compute-to-valid trim implementation before they can be enabled under
  compute padding.
- Direct write/current reuse/top-k still need P4/P5 valid-row selection before
  tiny CP split can be opened in `can_cp_split()`.

Verification:

- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - Targeted P3 tests passed.
  - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py`
  - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
  - Result: `202 passed, 2 subtests passed`.

### P4 complete: cache/current paths select valid rows, not dummy compute rows

Date: 2026-06-04

Implemented:

- `split_tensor_by_cp_batch_plan()` now has an explicit `split_kind`:
  - `split_kind="compute"` keeps the P2 behavior and materializes compute
    padding rows.
  - `split_kind="valid"` uses valid/cache split metadata and never pads.
- `get_cp_shared_kv_local_out_cache_loc()` now uses the valid split whenever a
  batch plan exists, including bs=1 compute-padding plans.  This keeps
  `out_cache_loc` rows equal to real cache rows, not local compute rows.
- Added `select_cp_local_valid_rows_for_cache_write()`:
  - builds and caches per-forward-batch local row indices from
    `request_compute_split_lists`, `request_valid_split_lists`, and
    `request_zigzag_indices`;
  - fails fast if a caller passes rows that are not the expected local compute
    rows;
  - returns only real valid rows for cache/current consumers.
- `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()` filters `k_nope/k_pe`
  before MLA direct write.
- `nsa_indexer.py::_store_cp_shared_local_index_k_cache()` filters local index
  K before direct write.
- Index partial-current reuse no longer assumes `key[:valid_current_rows]`
  under compute padding:
  - current index preparation uses the valid-row selector and local valid
    `out_cache_loc`;
  - `_maybe_materialize_shared_index_buffer()` accepts local valid
    `current_index_kv` rows and local valid locs when compute padding is
    enabled.
- MLA partial-current reuse in `nsa_backend.py` similarly selects local valid
  rows and local valid locs before page-slot compose.

Example verified:

- `extend=65,page_size=64,cp_size=8,cp_rank=1`
  - local compute rows: 64.
  - local valid/cache rows: 1.
  - direct write writes only row 0 and loc page 2.
  - dummy rows 1..63 are not passed to MLA/index direct write or current
    compose.

Important boundary:

- P4 makes cache writes/current reuse dummy-safe.
- It does **not** make top-k/attention dummy rows fully safe yet; P5 still must
  ensure index/top-k output has a valid compute-row shape while dummy rows do
  not affect real outputs.
- It does **not** open `can_cp_split()` for tiny bs>1 yet; P6 remains blocked
  on P5.

Verification:

- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - RED observed first:
    - missing selector import failed collection;
    - direct write tests failed with `*_local_shape_mismatch`;
    - index partial-current compose failed because old code expected global
      valid rows.
  - Targeted P4 tests passed:
    - valid-row selector filters 64 compute rows to 1 valid row.
    - local `out_cache_loc` uses valid split under compute padding.
    - MLA/index direct write filters dummy rows.
    - index partial-current compose accepts local valid rows.
  - `python -m py_compile` passed for:
    - `python/sglang/srt/layers/attention/nsa/utils.py`
    - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py`
    - `python/sglang/srt/layers/attention/nsa_backend.py`
    - `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py`
    - `test/registered/unit/layers/test_nsa_cp_utils.py`
  - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
  - Result: `207 passed, 5 warnings, 2 subtests passed`.

### P5 complete: top-k/index consumes compute rows but compacts valid queries

Date: 2026-06-04

Implemented:

- `CPSharedKVBatchPlan` / `NSAContextParallelMetadata` now expose explicit
  top-k query-length aliases:
  - `request_compute_seq_q_prev/next`
  - `request_valid_seq_q_prev/next`
  - `request_valid_query_row_spans`
- `request_actual_seq_q_prev/next` now uses compute lengths in the batch plan.
  This supersedes the earlier P1/P2 boundary where actual-q aliases remained
  valid-token based.  Cache/page accounting aliases (`request_split_lists`,
  `request_padded_pages`) still remain valid-token based for compatibility.
- `_get_topk_in_seq_cp_pair()` routes any forward batch carrying a
  `batch_plan` into the batch-aware path, including bs=1 compute-padding plans.
  This avoids silently using the scalar path with compute-padded query tensors.
- `_get_topk_in_seq_cp_pair_batch()` now:
  - consumes local q/weight rows by compute segment length;
  - compacts only valid rows into `_get_topk_ragged_with_cp()`;
  - advances output offsets by compute rows;
  - fills dummy rows with `-1`;
  - computes `cp_index` from the valid tail, not from the padded compute tail.
- Existing bs>1 non-padding metadata remains compatible: if explicit valid-q
  aliases are absent and compute padding is disabled, valid q falls back to
  actual q.  Compute-padding metadata stays fail-fast when valid lengths are
  missing.

Example verified:

- `extend=65,page_size=64,cp_size=8,cp_rank=1`
  - local compute rows: 64.
  - local valid q rows: 1.
  - `_get_topk_ragged_with_cp()` receives only 1 compact row with
    `cp_index=[(0,64,65)]`.
  - returned top-k rows keep compute shape `(64, topk)`;
    row 0 is real and rows 1..63 are `-1`.

Important boundary:

- P5 makes top-k/index dummy-row safe at the Python metadata/dispatch layer.
- It does **not** open `can_cp_split()` for tiny bs>1 yet; P6 still owns the
  runtime gate change.
- CUDA/ETE still needs to prove the downstream attention kernel accepts dummy
  top-k rows filled with `-1`.  If not, dummy rows must map to a safe page id
  plus valid-query masking instead of `-1`.

Verification:

- Local:
  - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py`
- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - RED observed first:
    - compute-padded bs=1 batch-plan top-k initially used the scalar path and
      expected q rows to equal valid rows.
    - old bs>1 tests exposed missing valid-q fallback for non-padding metadata.
    - scalar current-index reuse exposed an accidentally inserted undefined
      `compute_padding_enabled` reference.
  - Targeted P5 tests:
    - `test_batch_plan_exposes_compute_padding_without_inflating_valid_cache_extent`
    - `test_indexer_in_seq_cp_pair_compute_padding_outputs_dummy_safe_rows`
    - plus the previous four failing batch/scalar top-k reuse tests.
  - Full unit slice:
    - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
    - Result: `208 passed, 5 warnings, 2 subtests passed`.

### P6 complete: tiny CP shared-KV requests enter compute-padding split

Date: 2026-06-04

Implemented:

- Removed the old effective tiny-suffix skip for CP shared KV:
  `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` is now a
  compatibility hook that validates the page-aligned contract and returns
  `False`.
- `can_cp_split()` now treats CP shared-KV in-seq split differently from the
  generic token-balanced path:
  - validates every request has matching extend/prefix metadata;
  - requires `token_to_kv_pool.page_size`;
  - fail-fasts on negative lengths and non-page-aligned prefixes;
  - allows positive tiny suffixes even when valid token count is below
    `cp_size`, because compute padding supplies one page of work per CP lane.
- `prepare_input_dp_with_cp_dsa()` now builds a `CPSharedKVBatchPlan` for
  **all** CP shared-KV requests, including bs=1.  This is required because bs=1
  tiny suffixes need the same valid/compute split contract as bs>1.
- `_build_batch_metadata_from_plan()` now separates scalar compatibility fields:
  - communication/rerange fields (`split_list`, `split_list_tensor`,
    `max_rank_len`, `per_rank_actual_token`, `reverse_split_len`,
    `cp_reverse_index`) use compute split when compute padding is enabled;
  - attention KV-length fields (`kv_len_prev/next`) still use valid split, so
    dummy rows do not extend visible KV length;
  - request-level cache fields remain valid-token based.
- NSA index partial-current preparation now selects valid current rows from
  `local_key` under compute padding, not from the already all-gathered key.

Example verified:

- `extend=65,page_size=64,cp_size=8,cp_rank=1`
  - `can_cp_split(...)=True`.
  - `prepare_input_dp_with_cp_dsa()` returns metadata with `batch_plan`.
  - scalar `metadata.split_list` is compute split
    `[64,64,64,64,64,64,64,64,0,0,0,0,0,0,0,0]`.
  - valid split remains `[64,1,0,...]`.
  - `max_rank_len=[64]*8`, while `kv_len_prev/next` remain valid-length based.

Important boundary:

- P6 opens the CP split gate, but still does not prove CUDA/ETE attention
  accepts dummy top-k `-1` rows.  That remains a P8 runtime validation item.
- Full rerange/logprob/capture-hidden are still not enabled for compute
  padding; narrow output remains the supported path.
- P7 still needs to verify HiCache allocation/backup/load never scales to
  compute padding rows.

Verification:

- Local:
  - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py`
- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - targeted P6 tests:
    - tiny single request `can_cp_split=True`;
    - tiny per-request bs>1 `can_cp_split=True`;
    - non-page-aligned prefix fail-fast;
    - bs=1 `prepare_input_dp_with_cp_dsa()` returns compute-padding
      batch-plan metadata;
    - P5 dummy-safe top-k remains green.
    - Result: `7 passed, 5 warnings`.
  - Added current-reuse regression:
    - `test_indexer_current_reuse_compute_padding_selects_local_key_not_gathered_key`
      constructs `local_key != gathered_key` and verifies compute-padding index
      current reuse quantizes the local valid row, not the all-gathered key.
    - Result: `1 passed, 5 warnings`.
  - full unit slice:
    - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
    - Result after regression test: `211 passed, 5 warnings, 2 subtests passed`.

### P7 complete: HiCache allocation/backup/load stay on valid/page extent

Date: 2026-06-04

Code audit:

- Scheduler allocation remains valid-token based:
  - `ScheduleBatch.prepare_for_extend()` sets `extend_num_tokens` from real
    `fill_ids[prefix:]`, not compute-padded rows.
  - `alloc_for_extend()` passes that same valid `extend_num_tokens` into
    `alloc_paged_token_slots_extend()`.
- CP shared-KV L1 owner-lane allocation remains valid-page based:
  - `build_in_seq_page_compute_owners()` returns one owner per real new page,
    including a tail page, and does **not** expand to `cp_size` compute pages.
  - For `extend=65,page_size=64,cp_size=8`, allocation owners are `[0, 1]`,
    output cache loc length is `65`, and no legacy allocation fallback is used.
- CP HiCache write reservation pads only to the physical tail page:
  - `HiCacheController.reserve_write_cp()` calls
    `pad_token_locs_to_page_boundary()`, whose contract is page-tail padding
    only and explicitly not CP-size padding.
  - For a 65-token logical span starting on a page boundary:
    `logical_len=65`, `padded_len=128`, `page_owners=[0,1]`; a rank owning
    the tail page reserves exactly one physical page (`64` slots), not
    `8*64` compute-padding slots.
- Per-layer backup uses reservation indices only:
  - `_submit_write_cp_layer_states()` concatenates
    `state.host_indices` and `state.physical_device_indices` created by
    `reserve_write_cp()`, so grouped bs>1 per-layer backup inherits the same
    valid/page extent.
  - No per-layer backup descriptor is built from local compute rows.
- CP HiCache load restores padded physical pages but exposes valid length:
  - `load_cp()` replays `page_owners` through `alloc_pages_with_owners()`.
  - It queues H2D only for this rank's saved `owned_positions/host_indices`.
  - It returns `visible_device_indices` trimmed to `valid_len`, so scheduler /
    radix-visible cache hit length does not become the padded physical length.

Important boundary:

- P7 did not require production code changes; the existing L1 allocation,
  HiCache reservation, per-layer backup, and load paths already separate
  compute padding from page-tail physical padding.
- The protected contract is now covered by regression tests.  Future changes
  must not pass compute-padded local rows into HiCache reservation or load.
- P8 still needs runtime CUDA/ETE validation that dummy top-k rows and
  compute-padded attention do not regress accept length or cause hangs.

Verification:

- Local:
  - `python -m py_compile test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
  - Local pytest is still blocked by missing local dependencies
    (`transformers` during `sglang.test.test_utils` import).
- Remote container `g0034:/sgl-workspace/sglang-tai`:
  - Added P7 regressions:
    - `test_alloc_extend_compute_owner_uses_valid_pages_not_compute_padding_pages`
    - `test_cp_hicache_write_reservation_uses_page_tail_not_compute_padding_extent`
    - `test_cp_hicache_load_returns_valid_visible_len_while_loading_owned_page_tail`
  - Targeted P7 result: `3 passed, 3 warnings`.
  - Full unit slice:
    - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py`
    - Result after P7 regressions: `214 passed, 5 warnings, 2 subtests passed`.
