Enable CP shared-KV compute padding without inflating cache state

Tiny extend requests can leave most CP lanes without query work, which has been tied to hangs and accept-length regressions. This change introduces a dual valid/compute metadata contract: forward paths may materialize compute-padded rows, while cache, current reuse, direct write, HiCache backup, and load remain valid/page based. The implementation keeps radix/HiCache/device allocation on real page extents, filters dummy compute rows before MLA/index cache writes and current reuse, makes top-k/index consume compute rows while compacting valid rows, and opens tiny CP shared-KV in-seq split through compute padding. The accompanying plan document records the contract and P1-P7 evidence. Constraint: CP shared KV and HiCache must stay page-granular; dummy compute rows must not allocate, write, backup, or load KV cache. Constraint: Avoid silent fallback and avoid adding collectives on hot paths. Rejected: Pad cache allocations to cp_size pages | would waste KV capacity and pollute radix/HiCache state. Rejected: Keep tiny suffixes out of CP split | preserves the zero-lane behavior that compute padding is meant to remove. Confidence: medium Scope-risk: broad Directive: Do not route compute-padded dummy rows into out_cache_loc, current reuse, HiCache reservation, or backup descriptors; keep valid/cache metadata explicit. Tested: Remote g0034 container targeted P7 tests: 3 passed, 3 warnings. Tested: Remote g0034 container 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 => 214 passed, 5 warnings, 2 subtests passed. Tested: Local py_compile for touched P7 test file. Not-tested: Latest CUDA/ETE traffic validation for dummy top-k rows, accept len, output len, and detokenizer hang behavior.
2026-06-04 01:34:35 +08:00
parent b3913046b6
commit 3e3f1b776b
7 changed files with 2681 additions and 124 deletions
--- a/docs/advanced_features/nsa_prefill_cp_compute_padding_plan_zh.md
+++ b/docs/advanced_features/nsa_prefill_cp_compute_padding_plan_zh.md
@@ -0,0 +1,869 @@
+# 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`.