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.

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

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

@@ -67,12 +67,14 @@ from sglang.srt.distributed.parallel_state import get_pp_group
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.attention.nsa.utils import (
     cp_all_gather_rerange_output,
+    get_cp_shared_kv_batch_plan,
     get_cp_shared_kv_local_out_cache_loc,
     get_cp_shared_kv_local_physical_out_cache_loc,
     is_nsa_enable_prefill_cp,
     is_nsa_prefill_cp_in_seq_split,
     nsa_use_prefill_cp,
     raise_cp_shared_kv_direct_write_error,
+    select_cp_local_valid_rows_for_cache_write,
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.layers.communicator import ScatterMode
@@ -358,29 +360,57 @@ class Indexer(MultiPlatformOp):
                     f"logical_page_table_shape={tuple(logical_page_table.shape)} "
                     f"page_size={page_size}"
                 )
-            current_locs = forward_batch.out_cache_loc
-            valid_current_rows = current_extend_kv_rows_for_reuse(
-                forward_batch,
-                current_index_kv[0],
-                current_index_kv[1],
+            batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+            compute_padding_current = batch_plan is not None and bool(
+                getattr(batch_plan, "compute_padding_enabled", False)
             )
-            if valid_current_rows is None:
-                raise RuntimeError(
-                    "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
-                    "CP shared KV index partial-current compose received "
-                    "current_index_kv that does not satisfy current reuse "
-                    "metadata. "
-                    f"cp_rank={layout.cp_rank} layer_id={layer_id} "
-                    f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
-                    f"current_k_shape={tuple(current_index_kv[0].shape)} "
-                    f"current_scale_shape={tuple(current_index_kv[1].shape)} "
-                    f"out_cache_loc_shape={tuple(current_locs.shape)}"
+            if compute_padding_current:
+                current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+                if current_locs is None:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose requires local "
+                        "valid out_cache_loc when compute padding is enabled. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id}"
+                    )
+                valid_current_rows = int(current_locs.numel())
+                if (
+                    int(current_index_kv[0].shape[0]) != valid_current_rows
+                    or int(current_index_kv[1].shape[0]) != valid_current_rows
+                ):
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose received "
+                        "local current rows that do not match local valid locs. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id} "
+                        f"current_k_shape={tuple(current_index_kv[0].shape)} "
+                        f"current_scale_shape={tuple(current_index_kv[1].shape)} "
+                        f"current_locs_shape={tuple(current_locs.shape)}"
+                    )
+            else:
+                current_locs = forward_batch.out_cache_loc
+                valid_current_rows = current_extend_kv_rows_for_reuse(
+                    forward_batch,
+                    current_index_kv[0],
+                    current_index_kv[1],
+                )
+                if valid_current_rows is None:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] "
+                        "CP shared KV index partial-current compose received "
+                        "current_index_kv that does not satisfy current reuse "
+                        "metadata. "
+                        f"cp_rank={layout.cp_rank} layer_id={layer_id} "
+                        f"prefix_lens={prefix_lens} extend_lens={extend_lens} "
+                        f"current_k_shape={tuple(current_index_kv[0].shape)} "
+                        f"current_scale_shape={tuple(current_index_kv[1].shape)} "
+                        f"out_cache_loc_shape={tuple(current_locs.shape)}"
+                    )
+                current_locs = current_locs[:valid_current_rows]
+                current_index_kv = (
+                    current_index_kv[0][:valid_current_rows],
+                    current_index_kv[1][:valid_current_rows],
                 )
-            current_locs = current_locs[:valid_current_rows]
-            current_index_kv = (
-                current_index_kv[0][:valid_current_rows],
-                current_index_kv[1][:valid_current_rows],
-            )
             prefix_slot_span = None
             if len(prefix_lens_cpu) == 1:
                 prefix_pages = int(prefix_lens_cpu[0]) // page_size
@@ -1614,7 +1644,10 @@ class Indexer(MultiPlatformOp):
         current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
     ) -> torch.Tensor:
         assert forward_batch.nsa_cp_metadata is not None
-        if int(getattr(forward_batch.nsa_cp_metadata, "batch_size", 1) or 1) > 1:
+        if (
+            int(getattr(forward_batch.nsa_cp_metadata, "batch_size", 1) or 1) > 1
+            or get_cp_shared_kv_batch_plan(forward_batch) is not None
+        ):
             return self._get_topk_in_seq_cp_pair_batch(
                 forward_batch,
                 layer_id,
@@ -1648,9 +1681,8 @@ class Indexer(MultiPlatformOp):
         shared_index_buffer = None
         shared_block_tables = None
         current_index_kv_for_topk = current_index_kv
-        if current_index_kv is not None and not is_current_only_extend_batch(
-            forward_batch
-        ):
+        current_only_batch = is_current_only_extend_batch(forward_batch)
+        if current_index_kv is not None and not current_only_batch:
             current_index_kv_for_topk = None
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
@@ -1714,35 +1746,84 @@ class Indexer(MultiPlatformOp):
         cp_metadata = forward_batch.nsa_cp_metadata
         assert cp_metadata is not None
         batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1)
+        batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+        compute_padding_enabled = bool(
+            getattr(cp_metadata, "compute_padding_enabled", False)
+            or bool(getattr(batch_plan, "compute_padding_enabled", False))
+        )
+
+        def metadata_list(name: str, fallback_name: Optional[str] = None) -> List[int]:
+            values = getattr(cp_metadata, name, None)
+            if values is None and batch_plan is not None:
+                values = getattr(batch_plan, name, None)
+            if values is None and fallback_name is not None:
+                values = getattr(cp_metadata, fallback_name, None)
+                if values is None and batch_plan is not None:
+                    values = getattr(batch_plan, fallback_name, None)
+            return list(values or [])
 
         request_kv_len_prev = list(getattr(cp_metadata, "request_kv_len_prev", []) or [])
         request_kv_len_next = list(getattr(cp_metadata, "request_kv_len_next", []) or [])
-        request_actual_seq_q_prev = list(
-            getattr(cp_metadata, "request_actual_seq_q_prev", []) or []
+        if not request_kv_len_prev and batch_plan is not None:
+            request_kv_len_prev = list(
+                getattr(batch_plan, "request_kv_len_prev", []) or []
+            )
+        if not request_kv_len_next and batch_plan is not None:
+            request_kv_len_next = list(
+                getattr(batch_plan, "request_kv_len_next", []) or []
+            )
+        request_actual_seq_q_prev = metadata_list(
+            "request_compute_seq_q_prev"
+            if compute_padding_enabled
+            else "request_actual_seq_q_prev",
+            fallback_name="request_actual_seq_q_prev",
         )
-        request_actual_seq_q_next = list(
-            getattr(cp_metadata, "request_actual_seq_q_next", []) or []
+        request_actual_seq_q_next = metadata_list(
+            "request_compute_seq_q_next"
+            if compute_padding_enabled
+            else "request_actual_seq_q_next",
+            fallback_name="request_actual_seq_q_next",
         )
+        request_valid_seq_q_prev = metadata_list(
+            "request_valid_seq_q_prev",
+            fallback_name="request_valid_actual_seq_q_prev",
+        )
+        request_valid_seq_q_next = metadata_list(
+            "request_valid_seq_q_next",
+            fallback_name="request_valid_actual_seq_q_next",
+        )
+        if not compute_padding_enabled:
+            # Older bs>1 metadata did not have explicit valid-q aliases because
+            # actual q length was also the valid q length. Keep that path
+            # compatible while compute-padding remains fail-fast if valid
+            # lengths are missing.
+            if not request_valid_seq_q_prev:
+                request_valid_seq_q_prev = request_actual_seq_q_prev
+            if not request_valid_seq_q_next:
+                request_valid_seq_q_next = request_actual_seq_q_next
         if not (
             len(request_kv_len_prev) == batch_size
             and len(request_kv_len_next) == batch_size
             and len(request_actual_seq_q_prev) == batch_size
             and len(request_actual_seq_q_next) == batch_size
+            and len(request_valid_seq_q_prev) == batch_size
+            and len(request_valid_seq_q_next) == batch_size
         ):
             raise RuntimeError(
                 "[CP_SHARED_KV_FAIL_FAST][index_topk] "
                 "reason=batch_gt1_index_metadata_incomplete "
                 f"batch_size={batch_size} layer_id={layer_id} "
                 f"kv_prev={request_kv_len_prev} kv_next={request_kv_len_next} "
-                f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next}"
+                f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next} "
+                f"valid_q_prev={request_valid_seq_q_prev} "
+                f"valid_q_next={request_valid_seq_q_next}"
             )
 
         shared_index_buffer = None
         shared_block_tables = None
         current_index_kv_for_topk = current_index_kv
-        if current_index_kv is not None and not is_current_only_extend_batch(
-            forward_batch
-        ):
+        current_only_batch = is_current_only_extend_batch(forward_batch)
+        if current_index_kv is not None and not current_only_batch:
             current_index_kv_for_topk = None
             shared_block_tables = metadata.get_page_table_64()
             shared_index_buffer, shared_block_tables = (
@@ -1771,7 +1852,7 @@ class Indexer(MultiPlatformOp):
         compact_output_spans: List[Tuple[int, int]] = []
         current_only = (
             current_index_kv_for_topk is not None
-            and is_current_only_extend_batch(forward_batch)
+            and current_only_batch
         )
         page_table_1 = None if current_only else metadata.get_page_table_1()
 
@@ -1779,16 +1860,27 @@ class Indexer(MultiPlatformOp):
             *,
             req_id: int,
             segment_len: int,
+            valid_segment_len: int,
             kv_len: int,
         ) -> None:
             nonlocal cursor, output_cursor
             segment_len = int(segment_len)
+            valid_segment_len = int(valid_segment_len)
             kv_len = int(kv_len)
             q_segment = q_fp8[cursor : cursor + segment_len]
             weights_segment = weights[cursor : cursor + segment_len]
             cursor += segment_len
             if segment_len == 0:
                 return
+            if valid_segment_len < 0 or valid_segment_len > segment_len:
+                raise RuntimeError(
+                    "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                    "reason=batch_gt1_index_valid_q_len_mismatch "
+                    f"req_id={req_id} segment_len={segment_len} "
+                    f"valid_segment_len={valid_segment_len}"
+                )
+            if valid_segment_len == 0:
+                return
 
             seq_len = int(forward_batch.seq_lens_cpu[req_id].item())
             extend_seq_len = int(forward_batch.extend_seq_lens_cpu[req_id])
@@ -1807,13 +1899,13 @@ class Indexer(MultiPlatformOp):
                     logical_kv_limit = min(seq_len, int(page_table_1.shape[1]))
             valid_q_count = _compute_contiguous_valid_cp_query_count(
                 cp_kv_end=cp_kv_end,
-                actual_seq_q=segment_len,
+                actual_seq_q=valid_segment_len,
                 logical_kv_limit=logical_kv_limit,
             )
             if valid_q_count <= 0:
                 return
 
-            start_abs = cp_kv_end - segment_len
+            start_abs = cp_kv_end - valid_segment_len
             end_abs = start_abs + valid_q_count
             pre_chunk_offset = seq_len - extend_seq_len
             cp_index.append(
@@ -1831,12 +1923,14 @@ class Indexer(MultiPlatformOp):
             collect_segment(
                 req_id=req_id,
                 segment_len=request_actual_seq_q_prev[req_id],
+                valid_segment_len=request_valid_seq_q_prev[req_id],
                 kv_len=request_kv_len_prev[req_id],
             )
             output_cursor += int(request_actual_seq_q_prev[req_id])
             collect_segment(
                 req_id=req_id,
                 segment_len=request_actual_seq_q_next[req_id],
+                valid_segment_len=request_valid_seq_q_next[req_id],
                 kv_len=request_kv_len_next[req_id],
             )
             output_cursor += int(request_actual_seq_q_next[req_id])
@@ -2051,6 +2145,10 @@ class Indexer(MultiPlatformOp):
         local_out_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
         if local_out_loc is None:
             return False
+        local_key = select_cp_local_valid_rows_for_cache_write(
+            forward_batch,
+            local_key,
+        )
         if local_key.shape[0] != local_out_loc.numel():
             raise_cp_shared_kv_direct_write_error(
                 "index_local_shape_mismatch",
@@ -2239,10 +2337,41 @@ class Indexer(MultiPlatformOp):
 
         current_index_kv = None
         if self._can_reuse_current_index_kv(forward_batch):
-            valid_current_rows = current_extend_kv_rows_for_reuse(forward_batch, key)
-            if valid_current_rows is not None and key.shape[0] >= valid_current_rows:
+            batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+            compute_padding_current = batch_plan is not None and bool(
+                getattr(batch_plan, "compute_padding_enabled", False)
+            )
+            if compute_padding_current:
+                current_key = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch,
+                    local_key,
+                )
+                current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+                valid_current_rows = (
+                    int(current_locs.numel()) if current_locs is not None else None
+                )
+            else:
+                valid_current_rows = current_extend_kv_rows_for_reuse(
+                    forward_batch, key
+                )
+                current_key = (
+                    key[:valid_current_rows]
+                    if valid_current_rows is not None
+                    else None
+                )
+                current_locs = (
+                    forward_batch.out_cache_loc[:valid_current_rows]
+                    if valid_current_rows is not None
+                    else None
+                )
+            if (
+                valid_current_rows is not None
+                and current_key is not None
+                and current_locs is not None
+                and current_key.shape[0] == valid_current_rows
+            ):
                 current_k_fp8, current_k_scale = act_quant(
-                    key[:valid_current_rows].contiguous(),
+                    current_key.contiguous(),
                     self.block_size,
                     self.scale_fmt,
                 )
@@ -2258,9 +2387,7 @@ class Indexer(MultiPlatformOp):
                         if forward_batch.cp_shared_kv_layout is not None
                         else None,
                         layer_id,
-                        tensor_debug_summary(
-                            forward_batch.out_cache_loc[:valid_current_rows]
-                        ),
+                        tensor_debug_summary(current_locs),
                         tensor_debug_checksum(current_index_kv[0]),
                         tensor_debug_checksum(current_index_kv[1]),
                     )

python/sglang/srt/layers/attention/nsa/utils.py

@@ -320,6 +320,33 @@ class NSAContextParallelMetadata:
     flat_zigzag_index: List[int] = None
     flat_segment_request_ids: List[int] = None
     flat_segment_offsets: List[int] = None
+    compute_padding_enabled: bool = False
+    request_valid_split_lists: List[List[int]] = None
+    request_valid_segment_page_starts: List[List[int]] = None
+    request_valid_segment_page_ends: List[List[int]] = None
+    request_valid_padded_pages: List[int] = None
+    request_valid_padded_tokens: List[int] = None
+    request_valid_padding_tokens: List[int] = None
+    request_valid_rank_local_tokens: List[int] = None
+    request_valid_rank_local_offsets: List[int] = None
+    request_valid_actual_seq_q_prev: List[int] = None
+    request_valid_actual_seq_q_next: List[int] = None
+    request_valid_seq_q_prev: List[int] = None
+    request_valid_seq_q_next: List[int] = None
+    request_valid_query_row_spans: List[List[Tuple[int, int]]] = None
+    request_compute_split_lists: List[List[int]] = None
+    request_compute_segment_page_starts: List[List[int]] = None
+    request_compute_segment_page_ends: List[List[int]] = None
+    request_compute_padded_pages: List[int] = None
+    request_compute_padded_tokens: List[int] = None
+    request_compute_padding_tokens: List[int] = None
+    request_compute_padded_token_offsets: List[int] = None
+    request_compute_rank_local_tokens: List[int] = None
+    request_compute_rank_local_offsets: List[int] = None
+    request_compute_actual_seq_q_prev: List[int] = None
+    request_compute_actual_seq_q_next: List[int] = None
+    request_compute_seq_q_prev: List[int] = None
+    request_compute_seq_q_next: List[int] = None
     batch_plan: object = None
 
 
@@ -355,6 +382,33 @@ class CPSharedKVBatchPlan:
     flat_zigzag_index: List[int]
     flat_segment_request_ids: List[int]
     flat_segment_offsets: List[int]
+    compute_padding_enabled: bool
+    request_valid_split_lists: List[List[int]]
+    request_valid_segment_page_starts: List[List[int]]
+    request_valid_segment_page_ends: List[List[int]]
+    request_valid_padded_pages: List[int]
+    request_valid_padded_tokens: List[int]
+    request_valid_padding_tokens: List[int]
+    request_valid_rank_local_tokens: List[int]
+    request_valid_rank_local_offsets: List[int]
+    request_valid_actual_seq_q_prev: List[int]
+    request_valid_actual_seq_q_next: List[int]
+    request_valid_seq_q_prev: List[int]
+    request_valid_seq_q_next: List[int]
+    request_valid_query_row_spans: List[List[Tuple[int, int]]]
+    request_compute_split_lists: List[List[int]]
+    request_compute_segment_page_starts: List[List[int]]
+    request_compute_segment_page_ends: List[List[int]]
+    request_compute_padded_pages: List[int]
+    request_compute_padded_tokens: List[int]
+    request_compute_padding_tokens: List[int]
+    request_compute_padded_token_offsets: List[int]
+    request_compute_rank_local_tokens: List[int]
+    request_compute_rank_local_offsets: List[int]
+    request_compute_actual_seq_q_prev: List[int]
+    request_compute_actual_seq_q_next: List[int]
+    request_compute_seq_q_prev: List[int]
+    request_compute_seq_q_next: List[int]
 
 
 def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]:
@@ -378,6 +432,54 @@ def _prefix_offsets(lengths: List[int]) -> List[int]:
     return offsets
 
 
+def _build_full_page_unit_split(
+    *,
+    page_units: int,
+    extend_prefix_len: int,
+    page_size: int,
+    cp_size: int,
+) -> Tuple[List[int], List[int], List[int]]:
+    """Split full physical pages across in-seq CP segments.
+
+    This is the compute-side counterpart of
+    `build_page_aligned_in_seq_split_list`: every assigned page contributes a
+    full `page_size` rows because dummy compute rows fill any valid-token tail.
+    """
+
+    if page_units < 0:
+        raise ValueError(f"page_units must be non-negative, got {page_units}")
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+    if cp_size <= 0:
+        raise ValueError(f"cp_size must be positive, got {cp_size}")
+    if extend_prefix_len < 0 or extend_prefix_len % page_size != 0:
+        raise ValueError(
+            "extend_prefix_len must be non-negative and page-aligned, "
+            f"got extend_prefix_len={extend_prefix_len} page_size={page_size}"
+        )
+
+    cp_segment_num = cp_size * 2
+    base_units = page_units // cp_segment_num
+    remainder_units = page_units % cp_segment_num
+    unit_counts = [
+        base_units + (1 if i < remainder_units else 0)
+        for i in range(cp_segment_num)
+    ]
+
+    split_list: List[int] = []
+    segment_page_starts: List[int] = []
+    segment_page_ends: List[int] = []
+    unit_cursor = 0
+    base_page = extend_prefix_len // page_size
+    for unit_count in unit_counts:
+        segment_page_starts.append(base_page + unit_cursor)
+        unit_cursor += unit_count
+        segment_page_ends.append(base_page + unit_cursor)
+        split_list.append(unit_count * page_size)
+
+    return split_list, segment_page_starts, segment_page_ends
+
+
 def build_batch_page_aligned_in_seq_split_plan(
     *,
     extend_lens: List[int],
@@ -445,6 +547,29 @@ def build_batch_page_aligned_in_seq_split_plan(
     request_actual_seq_q_next: List[int] = []
     request_last_token_owner: List[int] = []
     request_last_token_local_offset: List[int] = []
+    request_valid_split_lists: List[List[int]] = []
+    request_valid_segment_page_starts: List[List[int]] = []
+    request_valid_segment_page_ends: List[List[int]] = []
+    request_valid_padded_pages: List[int] = []
+    request_valid_padded_tokens: List[int] = []
+    request_valid_padding_tokens: List[int] = []
+    request_valid_rank_local_tokens: List[int] = []
+    request_valid_actual_seq_q_prev: List[int] = []
+    request_valid_actual_seq_q_next: List[int] = []
+    request_valid_seq_q_prev: List[int] = []
+    request_valid_seq_q_next: List[int] = []
+    request_valid_query_row_spans: List[List[Tuple[int, int]]] = []
+    request_compute_split_lists: List[List[int]] = []
+    request_compute_segment_page_starts: List[List[int]] = []
+    request_compute_segment_page_ends: List[List[int]] = []
+    request_compute_padded_pages: List[int] = []
+    request_compute_padded_tokens: List[int] = []
+    request_compute_padding_tokens: List[int] = []
+    request_compute_rank_local_tokens: List[int] = []
+    request_compute_actual_seq_q_prev: List[int] = []
+    request_compute_actual_seq_q_next: List[int] = []
+    request_compute_seq_q_prev: List[int] = []
+    request_compute_seq_q_next: List[int] = []
     flat_split_list: List[int] = []
     flat_zigzag_index: List[int] = []
     flat_segment_request_ids: List[int] = []
@@ -468,8 +593,17 @@ def build_batch_page_aligned_in_seq_split_plan(
                 f"extend_len={extend_len} page_size={page_size}"
             )
 
+        compute_pages = max(split_info.extend_padded_pages, cp_size)
+        compute_split_list, compute_page_starts, compute_page_ends = (
+            _build_full_page_unit_split(
+                page_units=compute_pages,
+                extend_prefix_len=prefix_len,
+                page_size=page_size,
+                cp_size=cp_size,
+            )
+        )
         owner, local_offset = _get_in_seq_last_token_owner_and_offset(
-            split_list=split_list,
+            split_list=compute_split_list,
             cp_size=cp_size,
             actual_token_count=extend_len,
         )
@@ -479,6 +613,9 @@ def build_batch_page_aligned_in_seq_split_plan(
         rank_local_tokens = (
             split_list[cp_rank] + split_list[mirror_idx]
         )
+        compute_rank_local_tokens = (
+            compute_split_list[cp_rank] + compute_split_list[mirror_idx]
+        )
         split_prefix_list = [0] + prefix_sum_list[:-1]
 
         request_split_infos.append(split_info)
@@ -492,10 +629,38 @@ def build_batch_page_aligned_in_seq_split_plan(
         request_rank_local_tokens.append(rank_local_tokens)
         request_kv_len_prev.append(prefix_sum_list[cp_rank])
         request_kv_len_next.append(prefix_sum_list[mirror_idx])
-        request_actual_seq_q_prev.append(split_list[cp_rank])
-        request_actual_seq_q_next.append(split_list[mirror_idx])
+        request_actual_seq_q_prev.append(compute_split_list[cp_rank])
+        request_actual_seq_q_next.append(compute_split_list[mirror_idx])
         request_last_token_owner.append(owner)
         request_last_token_local_offset.append(local_offset)
+        request_valid_split_lists.append(split_list)
+        request_valid_segment_page_starts.append(split_info.segment_page_starts)
+        request_valid_segment_page_ends.append(split_info.segment_page_ends)
+        request_valid_padded_pages.append(split_info.extend_padded_pages)
+        request_valid_padded_tokens.append(split_info.extend_padded_tokens)
+        request_valid_padding_tokens.append(split_info.extend_padding_tokens)
+        request_valid_rank_local_tokens.append(rank_local_tokens)
+        request_valid_actual_seq_q_prev.append(split_list[cp_rank])
+        request_valid_actual_seq_q_next.append(split_list[mirror_idx])
+        request_valid_seq_q_prev.append(split_list[cp_rank])
+        request_valid_seq_q_next.append(split_list[mirror_idx])
+        request_valid_query_row_spans.append(
+            [
+                (0, split_list[cp_rank]),
+                (compute_split_list[cp_rank], split_list[mirror_idx]),
+            ]
+        )
+        request_compute_split_lists.append(compute_split_list)
+        request_compute_segment_page_starts.append(compute_page_starts)
+        request_compute_segment_page_ends.append(compute_page_ends)
+        request_compute_padded_pages.append(compute_pages)
+        request_compute_padded_tokens.append(compute_pages * page_size)
+        request_compute_padding_tokens.append(compute_pages * page_size - extend_len)
+        request_compute_rank_local_tokens.append(compute_rank_local_tokens)
+        request_compute_actual_seq_q_prev.append(compute_split_list[cp_rank])
+        request_compute_actual_seq_q_next.append(compute_split_list[mirror_idx])
+        request_compute_seq_q_prev.append(compute_split_list[cp_rank])
+        request_compute_seq_q_next.append(compute_split_list[mirror_idx])
         flat_split_list.extend(split_list)
         segment_base = req_id * cp_segment_num
         flat_zigzag_index.extend(segment_base + idx for idx in zigzag_index)
@@ -533,6 +698,42 @@ def build_batch_page_aligned_in_seq_split_plan(
         flat_zigzag_index=flat_zigzag_index,
         flat_segment_request_ids=flat_segment_request_ids,
         flat_segment_offsets=flat_segment_offsets,
+        compute_padding_enabled=any(
+            compute_tokens != valid_tokens
+            for compute_tokens, valid_tokens in zip(
+                request_compute_padded_tokens, request_valid_padded_tokens
+            )
+        ),
+        request_valid_split_lists=request_valid_split_lists,
+        request_valid_segment_page_starts=request_valid_segment_page_starts,
+        request_valid_segment_page_ends=request_valid_segment_page_ends,
+        request_valid_padded_pages=request_valid_padded_pages,
+        request_valid_padded_tokens=request_valid_padded_tokens,
+        request_valid_padding_tokens=request_valid_padding_tokens,
+        request_valid_rank_local_tokens=request_valid_rank_local_tokens,
+        request_valid_rank_local_offsets=_prefix_offsets(request_valid_rank_local_tokens),
+        request_valid_actual_seq_q_prev=request_valid_actual_seq_q_prev,
+        request_valid_actual_seq_q_next=request_valid_actual_seq_q_next,
+        request_valid_seq_q_prev=request_valid_seq_q_prev,
+        request_valid_seq_q_next=request_valid_seq_q_next,
+        request_valid_query_row_spans=request_valid_query_row_spans,
+        request_compute_split_lists=request_compute_split_lists,
+        request_compute_segment_page_starts=request_compute_segment_page_starts,
+        request_compute_segment_page_ends=request_compute_segment_page_ends,
+        request_compute_padded_pages=request_compute_padded_pages,
+        request_compute_padded_tokens=request_compute_padded_tokens,
+        request_compute_padding_tokens=request_compute_padding_tokens,
+        request_compute_padded_token_offsets=_prefix_offsets(
+            request_compute_padded_tokens
+        ),
+        request_compute_rank_local_tokens=request_compute_rank_local_tokens,
+        request_compute_rank_local_offsets=_prefix_offsets(
+            request_compute_rank_local_tokens
+        ),
+        request_compute_actual_seq_q_prev=request_compute_actual_seq_q_prev,
+        request_compute_actual_seq_q_next=request_compute_actual_seq_q_next,
+        request_compute_seq_q_prev=request_compute_seq_q_prev,
+        request_compute_seq_q_next=request_compute_seq_q_next,
     )
 
 
@@ -555,25 +756,43 @@ def split_tensor_by_cp_batch_plan(
     plan,
     *,
     mode: str = "data",
+    split_kind: str = "compute",
 ) -> torch.Tensor:
     """Split a flattened batch tensor by per-request in-seq CP plan.
 
     `mode` is kept explicit for future shape-specific kernels.  The current
     CPU/Python planner path splits along dim0 for 1d, position, and data views.
+    `split_kind="compute"` materializes padded compute rows.  Cache writes must
+    use `split_kind="valid"` so dummy compute rows never receive cache locs.
     """
 
     if mode not in ("1d", "data", "position"):
         raise ValueError(f"unsupported CP batch split mode={mode!r}")
+    if split_kind not in ("compute", "valid"):
+        raise ValueError(f"unsupported CP batch split_kind={split_kind!r}")
 
     request_extend_lens = getattr(plan, "request_extend_lens", None)
-    request_split_lists = getattr(plan, "request_split_lists", None)
+    compute_padding_enabled = bool(getattr(plan, "compute_padding_enabled", False))
+    if split_kind == "valid":
+        request_split_lists = getattr(
+            plan, "request_valid_split_lists", None
+        ) or getattr(plan, "request_split_lists", None)
+        request_target_lens = request_extend_lens
+    elif compute_padding_enabled:
+        request_split_lists = getattr(plan, "request_compute_split_lists", None)
+        request_target_lens = getattr(plan, "request_compute_padded_tokens", None)
+    else:
+        request_split_lists = getattr(plan, "request_split_lists", None)
+        request_target_lens = request_extend_lens
     request_zigzag_indices = getattr(plan, "request_zigzag_indices", None)
     batch_size = int(getattr(plan, "batch_size", 1) or 1)
     if (
         request_extend_lens is None
+        or request_target_lens is None
         or request_split_lists is None
         or request_zigzag_indices is None
         or len(request_extend_lens) != batch_size
+        or len(request_target_lens) != batch_size
         or len(request_split_lists) != batch_size
         or len(request_zigzag_indices) != batch_size
     ):
@@ -591,9 +810,28 @@ def split_tensor_by_cp_batch_plan(
 
     local_chunks = []
     request_tensors = torch.split(tensor, [int(x) for x in request_extend_lens], dim=0)
-    for req_tensor, split_list, zigzag_index in zip(
-        request_tensors, request_split_lists, request_zigzag_indices
+    for req_id, (req_tensor, target_len, split_list, zigzag_index) in enumerate(
+        zip(
+            request_tensors,
+            request_target_lens,
+            request_split_lists,
+            request_zigzag_indices,
+        )
     ):
+        req_tensor = _pad_cp_request_tensor_for_split(
+            req_tensor,
+            target_len=int(target_len),
+            mode=mode,
+            req_id=req_id,
+        )
+        split_total = sum(int(x) for x in split_list)
+        if split_total != int(req_tensor.shape[0]):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_target_len_mismatch] "
+                "request split rows must equal padded request rows. "
+                f"req_id={req_id} split_total={split_total} "
+                f"request_rows={int(req_tensor.shape[0])} mode={mode}"
+            )
         req_segments = list(torch.split(req_tensor, [int(x) for x in split_list], dim=0))
         local_chunks.extend(req_segments[int(index)] for index in zigzag_index)
 
@@ -602,6 +840,150 @@ def split_tensor_by_cp_batch_plan(
     return torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:])
 
 
+def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.device):
+    cached = getattr(forward_batch, "cp_local_valid_row_indices_for_cache_write", None)
+    cached_expected_rows = getattr(
+        forward_batch, "cp_local_valid_compute_rows_for_cache_write", None
+    )
+    if cached is not None and cached_expected_rows is not None:
+        if cached.device == device:
+            return cached, int(cached_expected_rows)
+
+    batch_size = int(getattr(plan, "batch_size", 1) or 1)
+    request_compute_split_lists = getattr(plan, "request_compute_split_lists", None)
+    request_valid_split_lists = getattr(
+        plan, "request_valid_split_lists", None
+    ) or getattr(plan, "request_split_lists", None)
+    request_zigzag_indices = getattr(plan, "request_zigzag_indices", None)
+    if (
+        request_compute_split_lists is None
+        or request_valid_split_lists is None
+        or request_zigzag_indices is None
+        or len(request_compute_split_lists) != batch_size
+        or len(request_valid_split_lists) != batch_size
+        or len(request_zigzag_indices) != batch_size
+    ):
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
+            "CP shared-KV cache writes require valid and compute split metadata "
+            "when compute padding is enabled."
+        )
+
+    chunks = []
+    local_cursor = 0
+    for req_id in range(batch_size):
+        compute_split = request_compute_split_lists[req_id]
+        valid_split = request_valid_split_lists[req_id]
+        if len(compute_split) != len(valid_split):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
+                "valid and compute split metadata disagree. "
+                f"req_id={req_id} compute_segments={len(compute_split)} "
+                f"valid_segments={len(valid_split)}"
+            )
+        for segment_id in request_zigzag_indices[req_id]:
+            segment_id = int(segment_id)
+            compute_len = int(compute_split[segment_id])
+            valid_len = int(valid_split[segment_id])
+            if valid_len > compute_len:
+                raise RuntimeError(
+                    "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] "
+                    "valid segment is longer than compute segment. "
+                    f"req_id={req_id} segment_id={segment_id} "
+                    f"valid_len={valid_len} compute_len={compute_len}"
+                )
+            if valid_len > 0:
+                chunks.append(
+                    torch.arange(
+                        local_cursor,
+                        local_cursor + valid_len,
+                        device=device,
+                        dtype=torch.long,
+                    )
+                )
+            local_cursor += compute_len
+
+    if chunks:
+        indices = torch.cat(chunks, dim=0)
+    else:
+        indices = torch.empty((0,), device=device, dtype=torch.long)
+    forward_batch.cp_local_valid_row_indices_for_cache_write = indices
+    forward_batch.cp_local_valid_compute_rows_for_cache_write = local_cursor
+    return indices, local_cursor
+
+
+def select_cp_local_valid_rows_for_cache_write(
+    forward_batch,
+    local_tensor: torch.Tensor,
+) -> torch.Tensor:
+    """Drop compute-padding rows before writing CP shared KV into persistent cache."""
+
+    plan = get_cp_shared_kv_batch_plan(forward_batch)
+    if plan is None or not bool(getattr(plan, "compute_padding_enabled", False)):
+        return local_tensor
+
+    indices, expected_compute_rows = _get_cp_local_valid_row_indices_cache(
+        forward_batch,
+        plan,
+        local_tensor.device,
+    )
+    local_rows = int(local_tensor.shape[0])
+    if local_rows != expected_compute_rows:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_shape] "
+            "CP shared-KV cache write tensor must contain local compute rows "
+            "before valid-row selection. "
+            f"local_rows={local_rows} expected_compute_rows={expected_compute_rows} "
+            f"valid_rows={int(indices.numel())}"
+        )
+    if indices.numel() == local_rows:
+        return local_tensor
+    if indices.numel() == 0:
+        return local_tensor.new_empty((0, *local_tensor.shape[1:]))
+    return local_tensor.index_select(0, indices)
+
+
+def _pad_cp_request_tensor_for_split(
+    tensor: torch.Tensor,
+    *,
+    target_len: int,
+    mode: str,
+    req_id: int,
+) -> torch.Tensor:
+    current_len = int(tensor.shape[0])
+    if target_len < current_len:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_len_mismatch] "
+            "target split length is shorter than request valid rows. "
+            f"req_id={req_id} target_len={target_len} current_len={current_len}"
+        )
+    pad_len = target_len - current_len
+    if pad_len == 0:
+        return tensor
+    if current_len <= 0:
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_empty_request] "
+            f"cannot compute-pad an empty request. req_id={req_id}"
+        )
+
+    if mode == "position":
+        if tensor.dim() != 1:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_position_padding_rank] "
+                "position compute padding expects a 1D position tensor. "
+                f"req_id={req_id} shape={tuple(tensor.shape)}"
+            )
+        start = tensor[-1] + 1
+        padding = torch.arange(
+            pad_len,
+            device=tensor.device,
+            dtype=tensor.dtype,
+        ) + start
+    else:
+        padding = tensor.new_zeros((pad_len, *tensor.shape[1:]))
+    return torch.cat([tensor, padding], dim=0)
+
+
 def build_flat_page_owner_plan(plan) -> List[int]:
     from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
         build_in_seq_page_compute_owners,
@@ -761,11 +1143,16 @@ def _build_in_seq_split_for_forward_batch(
 
 
 def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
+    communication_split_lists = (
+        plan.request_compute_split_lists
+        if plan.compute_padding_enabled
+        else plan.request_split_lists
+    )
     per_rank_actual_token = []
     for rank in range(plan.cp_size):
         rank_tokens = 0
         mirror = plan.cp_size * 2 - rank - 1
-        for split_list in plan.request_split_lists:
+        for split_list in communication_split_lists:
             rank_tokens += split_list[rank] + split_list[mirror]
         per_rank_actual_token.append(rank_tokens)
     max_rank_token = max(per_rank_actual_token) if per_rank_actual_token else 0
@@ -773,31 +1160,58 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
 
     # Scalar fields remain populated for compatibility, but scalar-only
     # consumers must not use them when batch_size > 1.
-    first_split = plan.request_split_lists[0] if plan.request_split_lists else []
+    first_split = communication_split_lists[0] if communication_split_lists else []
+    first_valid_split = (
+        plan.request_valid_split_lists[0]
+        if plan.request_valid_split_lists
+        else plan.request_split_lists[0]
+        if plan.request_split_lists
+        else []
+    )
     first_info = plan.request_split_infos[0] if plan.request_split_infos else None
     first_zigzag = plan.request_zigzag_indices[0] if plan.request_zigzag_indices else []
-    first_prefix_sum = list(accumulate(first_split))
-    first_kv_len_prev = first_prefix_sum[plan.cp_rank] if first_prefix_sum else 0
+    first_valid_prefix_sum = list(accumulate(first_valid_split))
+    first_kv_len_prev = (
+        first_valid_prefix_sum[plan.cp_rank] if first_valid_prefix_sum else 0
+    )
     first_mirror = plan.cp_size * 2 - plan.cp_rank - 1
-    first_kv_len_next = first_prefix_sum[first_mirror] if first_prefix_sum else 0
+    first_kv_len_next = (
+        first_valid_prefix_sum[first_mirror] if first_valid_prefix_sum else 0
+    )
     first_actual_seq_q_prev = first_split[plan.cp_rank] if first_split else 0
     first_actual_seq_q_next = first_split[first_mirror] if first_split else 0
+    flat_communication_split_list = [
+        token_count
+        for split_list in communication_split_lists
+        for token_count in split_list
+    ]
+    first_reverse_split_len = [
+        element
+        for i in range(plan.cp_size)
+        for element in (first_split[i], first_split[plan.cp_size * 2 - i - 1])
+    ]
+    first_cp_reverse_index = (
+        list(range(0, plan.cp_size * 2, 2))
+        + list(range(plan.cp_size * 2 - 1, 0, -2))
+        if first_split
+        else []
+    )
 
     return NSAContextParallelMetadata(
         split_list=first_split,
         split_list_tensor=torch.tensor(
-            plan.flat_split_list, device="cuda", dtype=torch.int32
+            flat_communication_split_list, device="cuda", dtype=torch.int32
         ),
         split_prefix_tensor=torch.tensor(
-            [0] + list(accumulate(plan.flat_split_list))[:-1],
+            [0] + list(accumulate(flat_communication_split_list))[:-1],
             device="cuda",
             dtype=torch.int32,
         ),
         max_rank_len=max_rank_len,
         zigzag_index=first_zigzag,
         per_rank_actual_token=per_rank_actual_token,
-        reverse_split_len=None,
-        cp_reverse_index=None,
+        reverse_split_len=first_reverse_split_len,
+        cp_reverse_index=first_cp_reverse_index,
         kv_len_prev=first_kv_len_prev,
         kv_len_next=first_kv_len_next,
         actual_seq_q_prev=first_actual_seq_q_prev,
@@ -892,6 +1306,33 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan):
         flat_zigzag_index=plan.flat_zigzag_index,
         flat_segment_request_ids=plan.flat_segment_request_ids,
         flat_segment_offsets=plan.flat_segment_offsets,
+        compute_padding_enabled=plan.compute_padding_enabled,
+        request_valid_split_lists=plan.request_valid_split_lists,
+        request_valid_segment_page_starts=plan.request_valid_segment_page_starts,
+        request_valid_segment_page_ends=plan.request_valid_segment_page_ends,
+        request_valid_padded_pages=plan.request_valid_padded_pages,
+        request_valid_padded_tokens=plan.request_valid_padded_tokens,
+        request_valid_padding_tokens=plan.request_valid_padding_tokens,
+        request_valid_rank_local_tokens=plan.request_valid_rank_local_tokens,
+        request_valid_rank_local_offsets=plan.request_valid_rank_local_offsets,
+        request_valid_actual_seq_q_prev=plan.request_valid_actual_seq_q_prev,
+        request_valid_actual_seq_q_next=plan.request_valid_actual_seq_q_next,
+        request_valid_seq_q_prev=plan.request_valid_seq_q_prev,
+        request_valid_seq_q_next=plan.request_valid_seq_q_next,
+        request_valid_query_row_spans=plan.request_valid_query_row_spans,
+        request_compute_split_lists=plan.request_compute_split_lists,
+        request_compute_segment_page_starts=plan.request_compute_segment_page_starts,
+        request_compute_segment_page_ends=plan.request_compute_segment_page_ends,
+        request_compute_padded_pages=plan.request_compute_padded_pages,
+        request_compute_padded_tokens=plan.request_compute_padded_tokens,
+        request_compute_padding_tokens=plan.request_compute_padding_tokens,
+        request_compute_padded_token_offsets=plan.request_compute_padded_token_offsets,
+        request_compute_rank_local_tokens=plan.request_compute_rank_local_tokens,
+        request_compute_rank_local_offsets=plan.request_compute_rank_local_offsets,
+        request_compute_actual_seq_q_prev=plan.request_compute_actual_seq_q_prev,
+        request_compute_actual_seq_q_next=plan.request_compute_actual_seq_q_next,
+        request_compute_seq_q_prev=plan.request_compute_seq_q_prev,
+        request_compute_seq_q_next=plan.request_compute_seq_q_next,
         batch_plan=plan,
     )
 
@@ -914,20 +1355,26 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent(
     forward_batch: "ForwardBatch",
     cp_size: int,
 ) -> bool:
-    """Avoid in-seq CP split when a shared-KV suffix has too few pages.
+    """Compatibility hook for the old tiny-suffix skip gate.
 
-    CP shared KV is page-owned.  A cache-hit suffix with fewer physical pages
-    than CP lanes creates mostly-zero in-seq segments; the distributed NSA path
-    has repeatedly shown hangs on that shape.  Keep the page cache contract by
-    running those tiny suffixes without NSA in-seq CP split instead of falling
-    back to token-balanced page-splitting.
+    Compute padding now handles suffixes with fewer physical pages than CP
+    lanes, so this function validates the page-aligned shared-KV contract and
+    no longer requests a skip.
     """
 
-    if (
-        forward_batch is None
-        or not getattr(forward_batch, "uses_cp_shared_kv", False)
-        or cp_size <= 1
-    ):
+    if not _is_cp_shared_kv_forward_batch(forward_batch) or cp_size <= 1:
+        return False
+    _validate_cp_shared_kv_cp_split_plan_inputs(forward_batch, cp_size)
+    return False
+
+
+def _validate_cp_shared_kv_cp_split_plan_inputs(
+    forward_batch: "ForwardBatch",
+    cp_size: int,
+) -> bool:
+    if not _is_cp_shared_kv_forward_batch(forward_batch):
+        return False
+    if cp_size <= 1:
         return False
 
     extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
@@ -936,32 +1383,55 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent(
     page_size = int(getattr(token_to_kv_pool, "page_size", 0) or 0)
     if (
         extend_seq_lens_cpu is None
-        or len(extend_seq_lens_cpu) != 1
         or extend_prefix_lens_cpu is None
-        or len(extend_prefix_lens_cpu) != 1
         or page_size <= 0
     ):
-        return False
-
-    extend_len = int(extend_seq_lens_cpu[0])
-    if extend_len <= 0:
-        return False
-
-    prefix_len = int(extend_prefix_lens_cpu[0])
-    if prefix_len < 0:
-        return False
-    if prefix_len % page_size != 0:
         raise RuntimeError(
-            "[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] "
-            "CP shared KV NSA in-seq split requires a page-aligned prefix. "
-            "The radix/HiCache match path should floor cache hits to the "
-            "previous page boundary before CP split planning. "
-            f"prefix_len={prefix_len} extend_len={extend_len} "
+            "[CP_SHARED_KV_FAIL_FAST][cp_split_missing_page_plan_inputs] "
+            "CP shared KV NSA in-seq split requires extend lengths, prefix "
+            "lengths, and token_to_kv_pool.page_size before planning. "
+            f"extend_seq_lens_cpu={extend_seq_lens_cpu} "
+            f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} "
             f"page_size={page_size} cp_size={cp_size}"
         )
+    if len(extend_seq_lens_cpu) != len(extend_prefix_lens_cpu):
+        raise RuntimeError(
+            "[CP_SHARED_KV_FAIL_FAST][cp_split_length_mismatch] "
+            "CP shared KV NSA in-seq split requires one prefix length per "
+            "extend length. "
+            f"extend_seq_lens_cpu={extend_seq_lens_cpu} "
+            f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} "
+            f"cp_size={cp_size}"
+        )
+    if len(extend_seq_lens_cpu) == 0:
+        return False
 
-    padded_pages = ceil_div(extend_len, page_size)
-    return padded_pages < cp_size
+    has_extend = False
+    for req_id, (extend_len_raw, prefix_len_raw) in enumerate(
+        zip(extend_seq_lens_cpu, extend_prefix_lens_cpu)
+    ):
+        extend_len = int(extend_len_raw)
+        prefix_len = int(prefix_len_raw)
+        if extend_len < 0 or prefix_len < 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][cp_split_negative_len] "
+                "CP shared KV NSA in-seq split received a negative length. "
+                f"req_id={req_id} prefix_len={prefix_len} "
+                f"extend_len={extend_len} page_size={page_size} "
+                f"cp_size={cp_size}"
+            )
+        if prefix_len % page_size != 0:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] "
+                "CP shared KV NSA in-seq split requires a page-aligned prefix. "
+                "The radix/HiCache match path should floor cache hits to the "
+                "previous page boundary before CP split planning. "
+                f"req_id={req_id} prefix_len={prefix_len} "
+                f"extend_len={extend_len} page_size={page_size} "
+                f"cp_size={cp_size}"
+            )
+        has_extend = has_extend or extend_len > 0
+    return has_extend
 
 
 def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
@@ -977,18 +1447,28 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
         # the seq data needs to be divided and recombined at twice the size of cp_size.
         if should_use_replicated_compute_for_short_radix_hit(forward_batch, cp_size):
             return False
-        if should_skip_cp_shared_kv_cp_split_for_short_page_extent(
-            forward_batch, cp_size
-        ):
-            return False
-        cur_cp_seq_len = seq_len // (cp_size * 2)
+        if _is_cp_shared_kv_forward_batch(forward_batch):
+            cur_cp_seq_len = (
+                1
+                if _validate_cp_shared_kv_cp_split_plan_inputs(
+                    forward_batch, cp_size
+                )
+                else 0
+            )
+        else:
+            cur_cp_seq_len = seq_len // (cp_size * 2)
+    extend_token_count = sum(forward_batch.extend_seq_lens_cpu)
+    if _is_cp_shared_kv_forward_batch(forward_batch):
+        min_extend_token_count = 1
+    else:
+        min_extend_token_count = cp_size
     if (
         cur_cp_seq_len != 0
         and cp_size > 1
         and use_nsa
         and forward_batch.forward_mode.is_context_parallel_extend()
         and is_nsa_enable_prefill_cp()
-        and sum(forward_batch.extend_seq_lens_cpu) >= cp_size
+        and extend_token_count >= min_extend_token_count
     ):
         return True
     else:
@@ -1005,7 +1485,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
         return nsa_cp_round_robin_split_data(input_)
 
     metadata = forward_batch.nsa_cp_metadata
-    if getattr(metadata, "batch_size", 1) > 1:
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None:
         return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
 
     input_list = list(
@@ -1027,7 +1507,7 @@ def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor):
         return nsa_cp_round_robin_split_data(input_)
 
     metadata = forward_batch.nsa_cp_metadata
-    if getattr(metadata, "batch_size", 1) > 1:
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None:
         return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_)
 
     input_list = list(
@@ -1135,7 +1615,7 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
             "nsa_prefill_cp_mode is not in-seq-split",
         )
     batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
-    if batch_plan is not None and int(getattr(batch_plan, "batch_size", 1) or 1) > 1:
+    if batch_plan is not None:
         split_tokens = sum(int(x) for x in getattr(batch_plan, "request_extend_lens", []))
         mismatch_reason = "batch_split_out_cache_len_mismatch"
     else:
@@ -1150,10 +1630,18 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
             out_cache_tokens,
         )
 
-    local_out_cache_loc = cp_split_and_rebuild_1d(
-        forward_batch,
-        out_cache_loc.contiguous(),
-    )
+    if batch_plan is not None:
+        local_out_cache_loc = split_tensor_by_cp_batch_plan(
+            out_cache_loc.contiguous(),
+            batch_plan,
+            mode="1d",
+            split_kind="valid",
+        )
+    else:
+        local_out_cache_loc = cp_split_and_rebuild_1d(
+            forward_batch,
+            out_cache_loc.contiguous(),
+        )
     if local_out_cache_loc.numel() == 0:
         forward_batch.cp_local_out_cache_loc = local_out_cache_loc
         return local_out_cache_loc
@@ -1231,6 +1719,13 @@ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
         )
         return nsa_cp_round_robin_split_data(positions)
 
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None:
+        return split_tensor_by_cp_batch_plan(
+            positions,
+            get_cp_shared_kv_batch_plan(forward_batch),
+            mode="position",
+        )
+
     position_id_list = list(
         torch.split(positions, forward_batch.nsa_cp_metadata.split_list, dim=-1)
     )
@@ -1855,11 +2350,10 @@ def prepare_input_dp_with_cp_dsa(
         and getattr(forward_batch, "uses_cp_shared_kv", False)
         and getattr(forward_batch, "extend_seq_lens_cpu", None) is not None
         and getattr(forward_batch, "extend_prefix_lens_cpu", None) is not None
-        and len(forward_batch.extend_seq_lens_cpu) > 1
     ):
         if page_size is None:
             raise RuntimeError(
-                "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_page_size] "
+                "[CP_SHARED_KV_FAIL_FAST][cp_shared_missing_page_size] "
                 "CP shared-KV batch planning requires token_to_kv_pool.page_size"
             )
         batch_plan = build_batch_page_aligned_in_seq_split_plan(
@@ -2098,7 +2592,7 @@ def _in_seq_collect_last_token(
     cp_rank = get_attention_cp_rank()
     bs = len(forward_batch.extend_seq_lens_cpu)
     metadata = getattr(forward_batch, "nsa_cp_metadata", None)
-    if bs > 1:
+    if get_cp_shared_kv_batch_plan(forward_batch) is not None or bs > 1:
         return _in_seq_collect_last_token_batch(hidden_states, metadata, cp_size, cp_rank, bs)
 
     owner = 0
@@ -2134,9 +2628,25 @@ def _in_seq_collect_last_token_batch(
             "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_metadata] "
             "CP in-seq bs>1 last-token collection requires batch metadata."
         )
-    owners = getattr(metadata, "request_last_token_owner", None)
-    local_offsets = getattr(metadata, "request_last_token_local_offset", None)
-    rank_offsets = getattr(metadata, "request_rank_local_offsets", None)
+    plan = getattr(metadata, "batch_plan", None)
+    owners = _get_cp_last_token_metadata_list(
+        metadata, plan, "request_last_token_owner"
+    )
+    local_offsets = _get_cp_last_token_metadata_list(
+        metadata, plan, "request_last_token_local_offset"
+    )
+    compute_padding_enabled = bool(
+        getattr(metadata, "compute_padding_enabled", False)
+        or bool(getattr(plan, "compute_padding_enabled", False))
+    )
+    if compute_padding_enabled:
+        rank_offsets = _get_cp_last_token_metadata_list(
+            metadata, plan, "request_compute_rank_local_offsets"
+        )
+    else:
+        rank_offsets = _get_cp_last_token_metadata_list(
+            metadata, plan, "request_rank_local_offsets"
+        )
     if (
         owners is None
         or local_offsets is None
@@ -2181,3 +2691,16 @@ def _in_seq_collect_last_token_batch(
         dtype=torch.long,
     )
     return gathered.index_select(0, gather_indices)
+
+
+def _get_cp_last_token_metadata_list(
+    metadata: NSAContextParallelMetadata,
+    plan,
+    field_name: str,
+):
+    value = getattr(metadata, field_name, None)
+    if value is not None:
+        return value
+    if plan is not None:
+        return getattr(plan, field_name, None)
+    return None

python/sglang/srt/layers/attention/nsa_backend.py

@@ -51,11 +51,14 @@ from sglang.srt.layers.attention.nsa.transform_index import (
 from sglang.srt.layers.attention.nsa.utils import (
     can_nsa_prefill_cp_round_robin_split,
     compute_nsa_seqlens,
+    get_cp_shared_kv_batch_plan,
+    get_cp_shared_kv_local_out_cache_loc,
     is_nsa_enable_prefill_cp,
     nsa_cp_round_robin_split_data,
     nsa_cp_round_robin_split_q_seqs,
     nsa_use_prefill_cp,
     pad_nsa_cache_seqlens,
+    select_cp_local_valid_rows_for_cache_write,
 )
 from sglang.srt.layers.attention.utils import (
     concat_mla_absorb_q_general,
@@ -1807,12 +1810,41 @@ class NativeSparseAttnBackend(
             mla_prefetcher = getattr(
                 forward_batch, "cp_shared_kv_mla_prefetcher", None
             )
-            current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse(
-                forward_batch,
-                k,
-                k_rope,
+            batch_plan = get_cp_shared_kv_batch_plan(forward_batch)
+            compute_padding_current = batch_plan is not None and bool(
+                getattr(batch_plan, "compute_padding_enabled", False)
             )
-            can_reuse_current_kv = current_kv_rows_for_reuse is not None
+            current_locs_for_reuse = None
+            current_k_nope = None
+            current_k_rope = None
+            if compute_padding_current:
+                assert k is not None and k_rope is not None
+                current_k_nope = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch, k
+                )
+                current_k_rope = select_cp_local_valid_rows_for_cache_write(
+                    forward_batch, k_rope
+                )
+                current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc(
+                    forward_batch
+                )
+                current_kv_rows_for_reuse = (
+                    int(current_locs_for_reuse.numel())
+                    if current_locs_for_reuse is not None
+                    else None
+                )
+                can_reuse_current_kv = (
+                    current_kv_rows_for_reuse is not None
+                    and int(current_k_nope.shape[0]) == current_kv_rows_for_reuse
+                    and int(current_k_rope.shape[0]) == current_kv_rows_for_reuse
+                )
+            else:
+                current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse(
+                    forward_batch,
+                    k,
+                    k_rope,
+                )
+                can_reuse_current_kv = current_kv_rows_for_reuse is not None
             if cp_shared_kv_mla_prefetch_log_enabled():
                 if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id):
                     prefix_lens_cpu = getattr(
@@ -1850,8 +1882,15 @@ class NativeSparseAttnBackend(
                 assert k is not None and k_rope is not None
                 assert current_kv_rows_for_reuse is not None
                 valid_current_rows = int(current_kv_rows_for_reuse)
-                current_k_nope = k[:valid_current_rows]
-                current_k_rope = k_rope[:valid_current_rows]
+                if not compute_padding_current:
+                    current_k_nope = k[:valid_current_rows]
+                    current_k_rope = k_rope[:valid_current_rows]
+                    current_locs_for_reuse = forward_batch.out_cache_loc[
+                        :valid_current_rows
+                    ]
+                assert current_k_nope is not None
+                assert current_k_rope is not None
+                assert current_locs_for_reuse is not None
                 if is_packed_fp8_mla_kv_cache(kv_cache):
                     current_kv_cache = pack_current_mla_kv_for_reuse(
                         current_k_nope,
@@ -1860,9 +1899,6 @@ class NativeSparseAttnBackend(
                     )
                 else:
                     current_kv_cache = _cat([current_k_nope, current_k_rope], dim=-1)
-                current_locs_for_reuse = forward_batch.out_cache_loc[
-                    :valid_current_rows
-                ]
                 logical_page_table_1 = page_table_1
                 current_remap_page_size, current_remap_logical_page_capacity = (
                     current_loc_remap_fast_path_args(forward_batch)

python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py

@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     log_cp_draft_shared_kv_debug,
     nsa_use_prefill_cp,
     raise_cp_shared_kv_direct_write_error,
+    select_cp_local_valid_rows_for_cache_write,
 )
 from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     should_reuse_current_extend_kv,
@@ -557,6 +558,8 @@ class DeepseekMLAForwardMixin:
         local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
         if local_out_cache_loc is None:
             return False
+        k_nope = select_cp_local_valid_rows_for_cache_write(forward_batch, k_nope)
+        k_pe = select_cp_local_valid_rows_for_cache_write(forward_batch, k_pe)
         if (
             k_nope.shape[0] != local_out_cache_loc.numel()
             or k_pe.shape[0] != local_out_cache_loc.numel()

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     build_batch_page_aligned_in_seq_split_plan,
     build_page_aligned_cache_extent,
     _get_in_seq_last_token_owner_and_offset,
+    _build_batch_metadata_from_plan,
     build_page_aligned_in_seq_split_list,
     build_token_balanced_in_seq_split_list,
     can_cp_split,
@@ -19,6 +20,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     cp_collect_last_token_hidden,
     cp_split_and_rebuild_1d,
     cp_split_and_rebuild_data,
+    cp_split_and_rebuild_position,
     _torch_batch_in_seq_all_gather_rerange,
     get_cp_shared_kv_batch_plan,
     get_cp_shared_kv_local_out_cache_loc,
@@ -26,6 +28,7 @@ from sglang.srt.layers.attention.nsa.utils import (
     get_cp_local_embedding_padded_token_count,
     pad_cp_local_input_ids_for_embedding,
     prepare_input_dp_with_cp_dsa,
+    select_cp_local_valid_rows_for_cache_write,
     split_tensor_by_cp_batch_plan,
     split_in_seq_cp_local_pair,
 )
@@ -242,7 +245,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             split_list, extend_prefix_len=54464, extend_len=256, page_size=64
         )
 
-    def test_can_cp_split_skips_cp_when_radix_hit_suffix_has_too_few_pages(self):
+    def test_can_cp_split_uses_compute_padding_for_short_radix_hit_suffix(self):
         class Mode:
             def is_context_parallel_extend(self):
                 return True
@@ -265,9 +268,11 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 return_value=True,
             ),
         ):
-            self.assertFalse(can_cp_split(128, 8, True, forward_batch))
+            self.assertTrue(can_cp_split(128, 8, True, forward_batch))
 
-    def test_can_cp_split_skips_cp_when_page_units_do_not_cover_all_lanes(self):
+    def test_can_cp_split_uses_compute_padding_when_page_units_do_not_cover_all_lanes(
+        self,
+    ):
         class Mode:
             def is_context_parallel_extend(self):
                 return True
@@ -290,9 +295,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 return_value=True,
             ),
         ):
-            self.assertFalse(can_cp_split(256, 8, True, forward_batch))
+            self.assertTrue(can_cp_split(256, 8, True, forward_batch))
 
-    def test_can_cp_split_skips_current_only_when_page_units_do_not_cover_all_lanes(
+    def test_can_cp_split_uses_compute_padding_for_current_only_one_page_suffix(
         self,
     ):
         class Mode:
@@ -317,7 +322,34 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 return_value=True,
             ),
         ):
-            self.assertFalse(can_cp_split(64, 8, True, forward_batch))
+            self.assertTrue(can_cp_split(64, 8, True, forward_batch))
+
+    def test_can_cp_split_uses_compute_padding_per_request_for_batched_tiny_suffix(
+        self,
+    ):
+        class Mode:
+            def is_context_parallel_extend(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            extend_seq_lens_cpu=[65, 64],
+            extend_prefix_lens_cpu=[54464, 8192],
+            token_to_kv_pool=SimpleNamespace(page_size=64),
+            forward_mode=Mode(),
+        )
+
+        with (
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_enable_prefill_cp",
+                return_value=True,
+            ),
+        ):
+            self.assertTrue(can_cp_split(129, 8, True, forward_batch))
 
     def test_can_cp_split_fails_on_non_page_aligned_cp_shared_prefix(self):
         class Mode:
@@ -445,6 +477,70 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertEqual(plan.flat_segment_request_ids, [0, 0, 0, 0, 1, 1, 1, 1])
         self.assertEqual(plan.flat_segment_offsets, [0, 4, 4, 4, 0, 4, 8, 9])
 
+    def test_batch_plan_exposes_compute_padding_without_inflating_valid_cache_extent(
+        self,
+    ):
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[40320],
+            page_size=64,
+            cp_size=8,
+            cp_rank=1,
+        )
+
+        self.assertTrue(plan.compute_padding_enabled)
+        self.assertEqual(
+            plan.request_valid_split_lists,
+            [[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
+        )
+        self.assertEqual(
+            plan.request_compute_split_lists,
+            [[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]],
+        )
+        self.assertEqual(plan.request_valid_padded_pages, [2])
+        self.assertEqual(plan.request_valid_padded_tokens, [128])
+        self.assertEqual(plan.request_compute_padded_pages, [8])
+        self.assertEqual(plan.request_compute_padded_tokens, [512])
+        self.assertEqual(plan.request_compute_padding_tokens, [447])
+        self.assertEqual(plan.request_compute_rank_local_tokens, [64])
+        self.assertEqual(plan.request_compute_rank_local_offsets, [0])
+        self.assertEqual(plan.request_valid_rank_local_tokens, [1])
+        self.assertEqual(plan.request_valid_rank_local_offsets, [0])
+        self.assertEqual(plan.request_last_token_owner, [1])
+        self.assertEqual(plan.request_last_token_local_offset, [0])
+
+        # Compatibility aliases for cache/page accounting stay valid-token
+        # based. Query-length metadata is split separately below: attention and
+        # top-k consume compute rows, cache/current paths consume valid rows.
+        self.assertEqual(plan.request_split_lists, plan.request_valid_split_lists)
+        self.assertEqual(plan.request_padded_pages, plan.request_valid_padded_pages)
+        self.assertEqual(plan.request_actual_seq_q_prev, [64])
+        self.assertEqual(plan.request_actual_seq_q_next, [0])
+        self.assertEqual(plan.request_valid_seq_q_prev, [1])
+        self.assertEqual(plan.request_valid_seq_q_next, [0])
+        self.assertEqual(plan.request_compute_seq_q_prev, [64])
+        self.assertEqual(plan.request_compute_seq_q_next, [0])
+
+    def test_batch_plan_compute_padding_is_per_request_not_batch_total(self):
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65, 1024],
+            prefix_lens=[40320, 8192],
+            page_size=64,
+            cp_size=8,
+            cp_rank=0,
+        )
+
+        self.assertTrue(plan.compute_padding_enabled)
+        self.assertEqual(plan.request_valid_padded_pages, [2, 16])
+        self.assertEqual(plan.request_compute_padded_pages, [8, 16])
+        self.assertEqual(plan.request_compute_padded_tokens, [512, 1024])
+        self.assertEqual(plan.request_compute_padding_tokens, [447, 0])
+        self.assertEqual(plan.request_compute_rank_local_tokens, [64, 128])
+        self.assertEqual(plan.request_compute_rank_local_offsets, [0, 64])
+        self.assertEqual(plan.request_valid_rank_local_tokens, [64, 128])
+        self.assertEqual(plan.request_valid_rank_local_offsets, [0, 64])
+        self.assertEqual(plan.request_last_token_owner, [1, 0])
+
     def test_batch_plan_stable_helpers_split_and_build_page_owner_plan(self):
         import torch
 
@@ -461,12 +557,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertEqual(build_flat_page_owner_plan(plan), [0, 0, 1, 1])
 
         local_1d = split_tensor_by_cp_batch_plan(torch.arange(13), plan, mode="1d")
-        self.assertEqual(local_1d.tolist(), [8, 9, 10, 11, 12])
+        self.assertEqual(local_1d.tolist(), [0, 0, 0, 0, 8, 9, 10, 11, 12, 0, 0, 0])
 
         local_data = split_tensor_by_cp_batch_plan(
             torch.arange(13 * 2).view(13, 2), plan, mode="data"
         )
-        self.assertEqual(local_data[:, 0].tolist(), list(range(16, 26, 2)))
+        self.assertEqual(
+            local_data[:, 0].tolist(),
+            [0, 0, 0, 0, 16, 18, 20, 22, 24, 0, 0, 0],
+        )
 
     def test_collect_last_token_hidden_uses_batch_owner_metadata(self):
         import torch
@@ -506,6 +605,91 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(collected.tolist(), [[13.0], [99.0]])
 
+    def test_collect_last_token_hidden_uses_compute_padding_for_single_request(self):
+        import torch
+
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[40320],
+            page_size=64,
+            cp_size=8,
+            cp_rank=1,
+        )
+        hidden_states = torch.zeros((64, 1), dtype=torch.float32)
+        hidden_states[0] = 123.0
+        hidden_states[1] = 999.0
+        forward_batch = SimpleNamespace(
+            extend_seq_lens_cpu=[65],
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+            ),
+        )
+
+        def fake_all_gather(output, local_last):
+            self.assertEqual(local_last.tolist(), [[123.0]])
+            output.zero_()
+            output[1] = local_last[0]
+
+        with (
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
+                return_value=1,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor",
+                side_effect=fake_all_gather,
+            ),
+        ):
+            collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 8)
+
+        self.assertEqual(collected.tolist(), [[123.0]])
+
+    def test_collect_last_token_hidden_uses_compute_rank_offsets_for_batch(self):
+        import torch
+
+        hidden_states = torch.zeros((8, 1), dtype=torch.float32)
+        hidden_states[0] = 10.0
+        hidden_states[1] = 99.0
+        hidden_states[4] = 20.0
+        forward_batch = SimpleNamespace(
+            extend_seq_lens_cpu=[5, 5],
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=2,
+                request_last_token_owner=[1, 1],
+                request_last_token_local_offset=[0, 0],
+                request_rank_local_offsets=[0, 1],
+                request_compute_rank_local_offsets=[0, 4],
+                compute_padding_enabled=True,
+            ),
+        )
+
+        def fake_all_gather(output, local_last):
+            self.assertEqual(local_last.tolist(), [[10.0], [20.0]])
+            output.copy_(torch.tensor([[0.0], [0.0], [10.0], [20.0]]))
+
+        with (
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+                return_value=False,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.get_attention_cp_rank",
+                return_value=1,
+            ),
+            patch(
+                "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor",
+                side_effect=fake_all_gather,
+            ),
+        ):
+            collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2)
+
+        self.assertEqual(collected.tolist(), [[10.0], [20.0]])
+
     def test_collect_last_token_hidden_fails_fast_without_batch_owner_metadata(self):
         import torch
 
@@ -577,6 +761,51 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 page_size=64,
             )
 
+    def test_cp_shared_kv_prepare_uses_batch_plan_for_bs1_compute_padding(self):
+        class Mode:
+            def is_context_parallel_extend(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            extend_seq_lens_cpu=[65],
+            extend_prefix_lens_cpu=[0],
+            token_to_kv_pool=SimpleNamespace(page_size=64),
+            forward_mode=Mode(),
+        )
+
+        with patch(
+            "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+            return_value=False,
+        ):
+            metadata = prepare_input_dp_with_cp_dsa(
+                65,
+                cp_rank=1,
+                cp_size=8,
+                seqs_len=[65],
+                forward_batch=forward_batch,
+                page_size=64,
+            )
+
+        self.assertIsNotNone(metadata.batch_plan)
+        self.assertEqual(metadata.batch_size, 1)
+        self.assertTrue(metadata.compute_padding_enabled)
+        self.assertEqual(
+            metadata.request_valid_split_lists,
+            [[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
+        )
+        self.assertEqual(
+            metadata.request_compute_split_lists,
+            [[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]],
+        )
+        self.assertEqual(metadata.split_list, metadata.request_compute_split_lists[0])
+        self.assertEqual(metadata.max_rank_len, [64] * 8)
+        self.assertEqual(metadata.per_rank_actual_token, [64] * 8)
+        self.assertEqual(metadata.actual_seq_q_prev, 64)
+        self.assertEqual(metadata.actual_seq_q_next, 0)
+        self.assertEqual(metadata.request_valid_seq_q_prev, [1])
+        self.assertEqual(metadata.request_valid_seq_q_next, [0])
+
     def test_cp_shared_kv_all_gather_rejects_round_robin_mode(self):
         import torch
 
@@ -792,6 +1021,34 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2)))
 
+    def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self):
+        import torch
+
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[40320],
+            page_size=64,
+            cp_size=8,
+            cp_rank=1,
+        )
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+            )
+        )
+        tensor = torch.arange(65 * 2, dtype=torch.float32).view(65, 2)
+
+        with patch(
+            "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+            return_value=False,
+        ):
+            local = cp_split_and_rebuild_data(forward_batch, tensor)
+
+        self.assertEqual(local.shape, (64, 2))
+        self.assertEqual(local[0].tolist(), [128.0, 129.0])
+        self.assertTrue(torch.equal(local[1:], torch.zeros((63, 2))))
+
     def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self):
         import torch
 
@@ -812,6 +1069,85 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         self.assertEqual(local.tolist(), [8, 9, 10, 11, 12])
 
+    def test_cp_split_and_rebuild_1d_uses_zero_compute_padding_rows(self):
+        import torch
+
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[40320],
+            page_size=64,
+            cp_size=8,
+            cp_rank=1,
+        )
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+            )
+        )
+
+        with patch(
+            "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+            return_value=False,
+        ):
+            local = cp_split_and_rebuild_1d(forward_batch, torch.arange(65))
+
+        self.assertEqual(local.shape, (64,))
+        self.assertEqual(local[0].item(), 64)
+        self.assertEqual(local[1:].tolist(), [0] * 63)
+
+    def test_select_cp_local_valid_rows_filters_compute_padding_rows(self):
+        import torch
+
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=64,
+            cp_size=8,
+            cp_rank=1,
+        )
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+            )
+        )
+        local_compute_rows = torch.full((64, 2), -1.0)
+        local_compute_rows[0] = torch.tensor([50.0, 51.0])
+
+        selected = select_cp_local_valid_rows_for_cache_write(
+            forward_batch, local_compute_rows
+        )
+
+        self.assertEqual(selected.tolist(), [[50.0, 51.0]])
+
+    def test_cp_split_and_rebuild_position_is_batch_aware_and_compute_padded(self):
+        import torch
+
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[40320],
+            page_size=64,
+            cp_size=8,
+            cp_rank=1,
+        )
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+            )
+        )
+        positions = torch.arange(40320, 40385, dtype=torch.int32)
+
+        with patch(
+            "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+            return_value=False,
+        ):
+            local = cp_split_and_rebuild_position(forward_batch, positions)
+
+        self.assertEqual(local.shape, (64,))
+        self.assertEqual(local.tolist(), list(range(40384, 40448)))
+
     def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self):
         from types import SimpleNamespace
 
@@ -952,6 +1288,48 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             list(range(2 * page_size, 3 * page_size)) + [4 * page_size],
         )
 
+    def test_local_out_cache_loc_uses_valid_rows_under_compute_padding(self):
+        import torch
+
+        page_size = 64
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=page_size,
+            cp_size=8,
+            cp_rank=1,
+        )
+        out_cache_loc = torch.cat(
+            [
+                torch.arange(page_size, 2 * page_size),
+                torch.tensor([2 * page_size]),
+            ]
+        ).to(torch.int64)
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=8,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+        )
+
+        with patch(
+            "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split",
+            return_value=False,
+        ):
+            local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+
+        self.assertEqual(local_locs.tolist(), [2 * page_size])
+
     def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self):
         import torch
         from types import SimpleNamespace
@@ -1135,6 +1513,67 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                     act_quant=None,
                 )
 
+    def test_indexer_direct_write_filters_compute_padding_rows(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        page_size = 64
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=page_size,
+            cp_size=8,
+            cp_rank=1,
+        )
+        out_cache_loc = torch.cat(
+            [
+                torch.arange(page_size, 2 * page_size),
+                torch.tensor([2 * page_size]),
+            ]
+        ).to(torch.int64)
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=8,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+            token_to_kv_pool=SimpleNamespace(page_size=page_size),
+        )
+        indexer = object.__new__(Indexer)
+        indexer.nsa_enable_prefill_cp = True
+        calls = []
+
+        def fake_store_index_k_cache(**kwargs):
+            calls.append(kwargs)
+
+        indexer._store_index_k_cache = fake_store_index_k_cache
+        local_key = torch.full((64, 2), -1.0)
+        local_key[0] = torch.tensor([9.0, 10.0])
+
+        with patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True):
+            stored = Indexer._store_cp_shared_local_index_k_cache(
+                indexer,
+                forward_batch,
+                layer_id=0,
+                local_key=local_key,
+                act_quant=None,
+            )
+
+        self.assertTrue(stored)
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0]["key"].tolist(), [[9.0, 10.0]])
+
     def test_mla_direct_write_fails_fast_on_local_shape_mismatch(self):
         import torch
 
@@ -1161,6 +1600,292 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                     k_pe=torch.empty((2, 8)),
                 )
 
+    def test_mla_direct_write_filters_compute_padding_rows(self):
+        import torch
+
+        from sglang.srt.models.deepseek_common.attention_forward_methods import (
+            forward_mla,
+        )
+
+        page_size = 64
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=page_size,
+            cp_size=8,
+            cp_rank=1,
+        )
+        out_cache_loc = torch.cat(
+            [
+                torch.arange(page_size, 2 * page_size),
+                torch.tensor([2 * page_size]),
+            ]
+        ).to(torch.int64)
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=8,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+            token_to_kv_pool=SimpleNamespace(page_size=page_size),
+        )
+        mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=0))
+        calls = []
+
+        def fake_tai_store(**kwargs):
+            calls.append(kwargs)
+            return True
+
+        k_nope = torch.full((64, 2), -1.0)
+        k_nope[0] = torch.tensor([1.0, 2.0])
+        k_pe = torch.full((64, 2), -1.0)
+        k_pe[0] = torch.tensor([3.0, 4.0])
+
+        with patch.object(forward_mla, "try_tai_fused_mla_store", fake_tai_store):
+            stored = (
+                forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv(
+                    mla,
+                    forward_batch,
+                    k_nope=k_nope,
+                    k_pe=k_pe,
+                )
+            )
+
+        self.assertTrue(stored)
+        self.assertEqual(len(calls), 1)
+        self.assertEqual(calls[0]["k_nope"].tolist(), [[1.0, 2.0]])
+        self.assertEqual(calls[0]["k_rope"].tolist(), [[3.0, 4.0]])
+        self.assertEqual(calls[0]["logical_locs"].tolist(), [2 * page_size])
+
+    def test_index_partial_current_compose_accepts_local_valid_compute_padding_rows(
+        self,
+    ):
+        import torch
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        page_size = 64
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=page_size,
+            cp_size=8,
+            cp_rank=1,
+        )
+
+        class FakePool:
+            page_size = 64
+            index_head_dim = 2
+
+            def get_index_k_with_scale_buffer(self, layer_id):
+                return torch.zeros((4, 3), dtype=torch.float32)
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=8,
+                cp_rank=1,
+            ),
+            cp_shared_kv_index_prefetcher=None,
+            token_to_kv_pool=FakePool(),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                batch_size=1,
+                batch_plan=plan,
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=torch.cat(
+                [
+                    torch.arange(page_size, 2 * page_size),
+                    torch.tensor([2 * page_size]),
+                ]
+            ).to(torch.int64),
+            extend_prefix_lens_cpu=[page_size],
+            extend_seq_lens_cpu=[65],
+        )
+        logical_page_table = torch.tensor([[1, 2, 3]], dtype=torch.int32)
+        current_index_kv = (
+            torch.tensor([[7.0, 8.0]], dtype=torch.float32),
+            torch.tensor([[0.5]], dtype=torch.float32),
+        )
+        materialize_calls = []
+        expected_buffer = torch.ones((3, 3), dtype=torch.float32)
+        expected_pages = torch.tensor([[0, 1, 2]], dtype=torch.int32)
+        indexer = object.__new__(Indexer)
+
+        def fake_materialize(**kwargs):
+            materialize_calls.append(kwargs)
+            return expected_buffer, expected_pages
+
+        with patch.object(
+            nsa_indexer,
+            "get_or_build_shared_paged_buffer_slot_remap",
+            return_value=torch.tensor([0, 1, 2], dtype=torch.int64),
+        ), patch.object(
+            nsa_indexer,
+            "materialize_prefix_and_reuse_current_index_page_slots",
+            side_effect=fake_materialize,
+        ):
+            dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer(
+                forward_batch,
+                layer_id=0,
+                logical_page_table=logical_page_table,
+                current_index_kv=current_index_kv,
+            )
+
+        self.assertIs(dense_buffer, expected_buffer)
+        self.assertIs(dense_pages, expected_pages)
+        self.assertEqual(len(materialize_calls), 1)
+        self.assertIs(materialize_calls[0]["current_index_k"], current_index_kv[0])
+        self.assertEqual(materialize_calls[0]["current_locs"].tolist(), [2 * page_size])
+
+    def test_indexer_current_reuse_compute_padding_selects_local_key_not_gathered_key(
+        self,
+    ):
+        import torch
+        import types
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        page_size = 64
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=page_size,
+            cp_size=8,
+            cp_rank=1,
+        )
+        metadata_obj = _build_batch_metadata_from_plan(plan)
+
+        class Mode:
+            def is_extend_without_speculative(self):
+                return True
+
+            def is_decode_or_idle(self):
+                return False
+
+            def is_target_verify(self):
+                return False
+
+            def is_draft_extend(self, include_v2=False):
+                return False
+
+            def is_context_parallel_extend(self):
+                return True
+
+        class AttnBackend:
+            def get_indexer_metadata(self, layer_id, forward_batch):
+                return object()
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=8,
+                cp_rank=1,
+            ),
+            token_to_kv_pool=SimpleNamespace(page_size=page_size),
+            nsa_cp_metadata=metadata_obj,
+            out_cache_loc=torch.cat(
+                [
+                    torch.arange(page_size, 2 * page_size),
+                    torch.tensor([2 * page_size]),
+                ]
+            ).to(torch.int64),
+            extend_prefix_lens_cpu=[0],
+            extend_seq_lens_cpu=[65],
+            seq_lens_cpu=torch.tensor([65], dtype=torch.int64),
+            forward_mode=Mode(),
+            attn_backend=AttnBackend(),
+            hisparse_coordinator=None,
+        )
+
+        local_key = torch.full((64, 2), -1.0, dtype=torch.float32)
+        local_key[0] = torch.tensor([11.0, 12.0])
+        gathered_key = torch.full((64, 2), 99.0, dtype=torch.float32)
+        gathered_key[0] = torch.tensor([101.0, 102.0])
+        query = torch.zeros((64, 2), dtype=torch.float32)
+        act_quant_inputs = []
+        topk_current_index_kv = []
+
+        def fake_act_quant(tensor, block_size, scale_fmt):
+            act_quant_inputs.append(tensor.detach().clone())
+            return tensor.detach().clone(), torch.ones(
+                (int(tensor.shape[0]), 1), dtype=torch.float32
+            )
+
+        fake_triton_kernel = types.ModuleType(
+            "sglang.srt.layers.attention.nsa.triton_kernel"
+        )
+        fake_triton_kernel.act_quant = fake_act_quant
+
+        indexer = object.__new__(Indexer)
+        indexer.alt_stream = None
+        indexer.nsa_enable_prefill_cp = True
+        indexer.index_topk = 2
+        indexer.block_size = 64
+        indexer.scale_fmt = None
+        indexer._get_q_k_bf16 = (
+            lambda *args, **kwargs: (query, gathered_key, local_key)
+        )
+        indexer._store_cp_shared_local_index_k_cache = lambda **kwargs: True
+        indexer._can_reuse_current_index_kv = lambda forward_batch: True
+        indexer._get_logits_head_gate = (
+            lambda x_for_gate, q_scale: torch.zeros((64, 1), dtype=torch.float32)
+        )
+
+        def fake_topk(*args, **kwargs):
+            topk_current_index_kv.append(kwargs["current_index_kv"])
+            return torch.zeros((64, 2), dtype=torch.int32)
+
+        indexer._get_topk_in_seq_cp_pair = fake_topk
+
+        with (
+            patch.dict(
+                sys.modules,
+                {
+                    "sglang.srt.layers.attention.nsa.triton_kernel": fake_triton_kernel
+                },
+            ),
+            patch.object(nsa_indexer, "_is_cuda", True),
+            patch.object(nsa_indexer, "_is_hip", False),
+            patch.object(nsa_indexer, "_is_npu", False),
+            patch.object(
+                nsa_indexer,
+                "is_nsa_prefill_cp_in_seq_split",
+                return_value=True,
+            ),
+        ):
+            result = Indexer.forward_cuda(
+                indexer,
+                x=torch.zeros((64, 2), dtype=torch.float32),
+                q_lora=torch.zeros((64, 2), dtype=torch.float32),
+                positions=torch.arange(64, dtype=torch.int64),
+                forward_batch=forward_batch,
+                layer_id=0,
+                return_indices=True,
+            )
+
+        self.assertEqual(result.shape, (64, 2))
+        self.assertEqual(len(act_quant_inputs), 2)
+        self.assertEqual(act_quant_inputs[1].tolist(), [[11.0, 12.0]])
+        self.assertNotEqual(act_quant_inputs[1].tolist(), [[101.0, 102.0]])
+        self.assertEqual(len(topk_current_index_kv), 1)
+        self.assertEqual(topk_current_index_kv[0][0].tolist(), [[11.0, 12.0]])
+
     def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self):
         import torch
         from types import SimpleNamespace
@@ -1416,6 +2141,113 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
         )
 
+    def test_indexer_in_seq_cp_pair_compute_padding_outputs_dummy_safe_rows(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        page_size = 64
+        plan = build_batch_page_aligned_in_seq_split_plan(
+            extend_lens=[65],
+            prefix_lens=[0],
+            page_size=page_size,
+            cp_size=8,
+            cp_rank=1,
+        )
+        metadata_obj = _build_batch_metadata_from_plan(plan)
+        indexer = object.__new__(Indexer)
+        indexer.index_topk = 2
+        logical_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        materialized_index = torch.tensor([11], dtype=torch.int32)
+        dense_pages = torch.tensor([[1, 2]], dtype=torch.int32)
+        materialize_calls = []
+        topk_calls = []
+
+        class Metadata:
+            def get_page_table_64(self):
+                return logical_pages
+
+            def get_page_table_1(self):
+                return torch.empty((1, 65), dtype=torch.int32)
+
+        def fake_materialize(forward_batch, layer_id, logical_page_table):
+            materialize_calls.append((layer_id, logical_page_table))
+            return materialized_index, dense_pages
+
+        def fake_get_topk(
+            forward_batch,
+            layer_id,
+            q_fp8,
+            weights,
+            metadata,
+            kv_len,
+            actual_seq_q,
+            cp_index=None,
+            current_index_kv=None,
+            shared_index_buffer=None,
+            shared_block_tables=None,
+            actual_seq_q_tensor=None,
+            actual_seq_q_cu_tensor=None,
+            batch_idx=0,
+        ):
+            topk_calls.append(
+                {
+                    "actual_seq_q": actual_seq_q,
+                    "cp_index": cp_index,
+                    "q": q_fp8.flatten().tolist(),
+                    "weights": weights.flatten().tolist(),
+                    "shared_index_buffer": shared_index_buffer,
+                    "shared_block_tables": shared_block_tables,
+                }
+            )
+            rows = int(q_fp8.shape[0])
+            return (
+                torch.arange(1, rows + 1, dtype=torch.int32)
+                .view(rows, 1)
+                .repeat(1, 2)
+            )
+
+        indexer._maybe_materialize_shared_index_buffer = fake_materialize
+        indexer._get_topk_ragged_with_cp = fake_get_topk
+
+        forward_batch = SimpleNamespace(
+            batch_size=1,
+            forward_mode=SimpleNamespace(
+                is_extend_without_speculative=lambda: True,
+            ),
+            extend_prefix_lens_cpu=[0],
+            extend_seq_lens_cpu=[65],
+            seq_lens_cpu=torch.tensor([65], dtype=torch.int64),
+            nsa_cp_metadata=metadata_obj,
+        )
+        q_fp8 = torch.arange(64, dtype=torch.float32).view(64, 1)
+        weights = (torch.arange(64, dtype=torch.float32) + 100).view(64, 1)
+
+        result = Indexer._get_topk_in_seq_cp_pair(
+            indexer,
+            forward_batch,
+            layer_id=7,
+            q_fp8=q_fp8,
+            weights=weights,
+            metadata=Metadata(),
+            current_index_kv=None,
+        )
+
+        self.assertEqual(len(materialize_calls), 1)
+        self.assertIs(materialize_calls[0][1], logical_pages)
+        self.assertEqual(len(topk_calls), 1)
+        self.assertEqual(topk_calls[0]["actual_seq_q"], 1)
+        self.assertEqual(topk_calls[0]["cp_index"], [(0, 64, 65)])
+        self.assertEqual(topk_calls[0]["q"], [0.0])
+        self.assertEqual(topk_calls[0]["weights"], [100.0])
+        self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
+        self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages)
+        self.assertEqual(result.shape, (64, 2))
+        self.assertEqual(result[0].tolist(), [1, 1])
+        self.assertTrue(
+            torch.equal(result[1:], torch.full((63, 2), -1, dtype=torch.int32))
+        )
+
     def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self):
         import torch
 

test/registered/unit/mem_cache/test_cp_shared_kv_layout.py

@@ -297,6 +297,173 @@ class TestCPSharedPagedAllocator(CustomTestCase):
 
         self.assertEqual(owners, [0, 1])
 
+    def test_alloc_extend_compute_owner_uses_valid_pages_not_compute_padding_pages(
+        self,
+    ):
+        from types import SimpleNamespace
+
+        from sglang.srt.mem_cache import common
+
+        page_size = 64
+
+        class FakeAllocator:
+            def __init__(self):
+                self.page_size = page_size
+                self.cp_size = 8
+                self.owner_calls = []
+                self.extend_num_tokens = []
+
+            def alloc_extend_compute_owner(
+                self,
+                _prefix_lens,
+                _prefix_lens_cpu,
+                _seq_lens,
+                _seq_lens_cpu,
+                _last_loc,
+                extend_num_tokens,
+                page_compute_owners,
+            ):
+                self.extend_num_tokens.append(int(extend_num_tokens))
+                self.owner_calls.append(list(page_compute_owners))
+                return torch.arange(
+                    1024, 1024 + int(extend_num_tokens), dtype=torch.int64
+                )
+
+            def alloc_extend(self, *_args, **_kwargs):
+                raise AssertionError("legacy allocation should not be used")
+
+        class FakeTreeCache:
+            def __init__(self):
+                self.token_to_kv_pool_allocator = FakeAllocator()
+
+            def is_chunk_cache(self):
+                return False
+
+            def evict(self, *_args, **_kwargs):
+                raise AssertionError("eviction should not be needed")
+
+        server_args = SimpleNamespace(
+            enable_nsa_prefill_cp_shared_kv=True,
+            enable_nsa_prefill_context_parallel=True,
+            nsa_prefill_cp_mode="in-seq-split",
+        )
+        tree_cache = FakeTreeCache()
+
+        with patch.object(common, "get_global_server_args", return_value=server_args):
+            out_cache_loc = common.alloc_paged_token_slots_extend(
+                tree_cache=tree_cache,
+                prefix_lens=torch.tensor([0], dtype=torch.int64),
+                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                seq_lens=torch.tensor([65], dtype=torch.int64),
+                seq_lens_cpu=torch.tensor([65], dtype=torch.int64),
+                last_loc=torch.tensor([-1], dtype=torch.int64),
+                extend_num_tokens=65,
+            )
+
+        self.assertEqual(out_cache_loc.numel(), 65)
+        self.assertEqual(tree_cache.token_to_kv_pool_allocator.extend_num_tokens, [65])
+        self.assertEqual(tree_cache.token_to_kv_pool_allocator.owner_calls, [[0, 1]])
+
+    def test_cp_hicache_write_reservation_uses_page_tail_not_compute_padding_extent(
+        self,
+    ):
+        from sglang.srt.managers.cache_controller import HiCacheController
+
+        page_size = 64
+
+        class FakeHostPool:
+            def __init__(self):
+                self.alloc_sizes = []
+
+            def alloc_contiguous_preferred(self, need_size):
+                self.alloc_sizes.append(int(need_size))
+                return torch.arange(1024, 1024 + int(need_size), dtype=torch.int64)
+
+            def alloc(self, need_size):
+                return self.alloc_contiguous_preferred(need_size)
+
+            def free(self, _indices):
+                raise AssertionError("reservation should not roll back")
+
+        controller = HiCacheController.__new__(HiCacheController)
+        controller.page_size = page_size
+        controller.cp_shared_kv_layout = CpSharedKVLayout(
+            page_size=page_size, cp_size=8, cp_rank=1
+        )
+        controller.mem_pool_host = FakeHostPool()
+        controller.draft_mem_pool_host = None
+        controller.draft_mem_pool_device = None
+
+        reservation = controller.reserve_write_cp(
+            torch.arange(page_size, page_size + 65, dtype=torch.int64),
+            node_id=123,
+        )
+
+        self.assertEqual(controller.mem_pool_host.alloc_sizes, [page_size])
+        self.assertEqual(reservation.metadata.logical_len, 65)
+        self.assertEqual(reservation.metadata.padded_len, page_size * 2)
+        self.assertEqual(reservation.metadata.page_owners.tolist(), [0, 1])
+        self.assertEqual(reservation.host_indices.numel(), page_size)
+        self.assertEqual(reservation.physical_device_indices.numel(), page_size)
+        self.assertEqual(
+            reservation.metadata.owned_positions.tolist(),
+            list(range(64, 128)),
+        )
+
+    def test_cp_hicache_load_returns_valid_visible_len_while_loading_owned_page_tail(
+        self,
+    ):
+        from types import SimpleNamespace
+
+        from sglang.srt.managers.cache_controller import HiCacheController
+        from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata
+
+        page_size = 64
+
+        class FakeDeviceAllocator:
+            def __init__(self):
+                self.owner_calls = []
+                self.freed = []
+
+            def alloc_pages_with_owners(self, page_owners):
+                self.owner_calls.append(list(page_owners))
+                return torch.arange(page_size, page_size * 3, dtype=torch.int64)
+
+            def free(self, indices):
+                self.freed.append(indices.clone())
+
+        controller = HiCacheController.__new__(HiCacheController)
+        controller.page_size = page_size
+        controller.cp_shared_kv_layout = CpSharedKVLayout(
+            page_size=page_size, cp_size=8, cp_rank=1
+        )
+        controller.mem_pool_device_allocator = FakeDeviceAllocator()
+        controller.load_queue = []
+        controller.draft_load_queue = []
+        controller.draft_mem_pool_host = None
+        controller.draft_mem_pool_device = None
+
+        metadata = CpHiCacheNodeMetadata(
+            logical_len=65,
+            padded_len=page_size * 2,
+            owned_positions=torch.arange(page_size, page_size * 2, dtype=torch.int64),
+            host_indices=torch.arange(1024, 1024 + page_size, dtype=torch.int64),
+            page_owners=torch.tensor([0, 1], dtype=torch.int8),
+            page_size=page_size,
+        )
+        node = SimpleNamespace(cp_hicache=metadata, host_len=65, id=321)
+
+        visible_device_indices = controller.load_cp([node], node_id=321)
+
+        self.assertEqual(controller.mem_pool_device_allocator.owner_calls, [[0, 1]])
+        self.assertEqual(controller.mem_pool_device_allocator.freed, [])
+        self.assertEqual(visible_device_indices.numel(), 65)
+        self.assertEqual(visible_device_indices.tolist(), list(range(64, 129)))
+        self.assertEqual(len(controller.load_queue), 1)
+        load_op = controller.load_queue[0]
+        self.assertEqual(load_op.host_indices.tolist(), list(range(1024, 1024 + 64)))
+        self.assertEqual(load_op.device_indices.tolist(), list(range(64, 128)))
+
     def test_compute_owner_page_assignment_allows_radix_hit_suffix_with_one_page_per_rank(
         self,
     ):