# NSA Prefill CP Phase 6: reuse prev/next index materialize

Phase 6 是一个小范围性能优化阶段，目标是在不改变 NSA topk 语义的前提下，减少 `in-seq-split` CP shared KV 路径里重复的 NSA index materialize。

## 背景

在 `nsa_prefill_cp_mode=in-seq-split` 下，一个 CP rank 本地 query 被拆成两个段：

```text
prev segment + next segment
```

两个段的 causal 可见 KV 长度不同，因此 topk 计算仍然需要分别执行：

```text
topk(prev, kv_len_prev)
topk(next, kv_len_next)
```

但它们读取的是同一层、同一 batch 的 NSA index K/scale 和同一份 request page table。Phase 5 后，persistent index cache 已经按 compute-owner shard 写入；read path 仍通过 compatibility materialize 得到 dense full-view index buffer。

原路径在 prev/next 两次 topk 前各调用一次：

```text
_maybe_materialize_shared_index_buffer()
  -> materialize_shared_paged_buffer()
  -> local copy + CP all-reduce
```

这会导致同一份 index K/scale 被 materialize 两次。

## 目标

将 `in-seq-split` CP pair 路径从：

```text
materialize index for prev
materialize index for next
materialize MLA KV for attention
```

改为：

```text
materialize index once for prev+next
materialize MLA KV for attention
```

Phase 6 不改变：

- persistent KV/index layout；
- `topk` 语义；
- prev/next 两段的 causal range；
- MLA KV materialize；
- PD transfer；
- radix cache 逻辑。

## 实现

代码位置：

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

新增 `_get_topk_in_seq_cp_pair(...)`，负责：

1. 根据 `forward_batch.nsa_cp_metadata.actual_seq_q_prev/next` 拆分 `q_fp8` 和 `weights`。
2. 当 `current_index_kv is None` 时，对 `metadata.get_page_table_64()` 调用一次 `_maybe_materialize_shared_index_buffer(...)`。
3. 将同一个 `shared_index_buffer` 和 `shared_block_tables` 传给 prev/next 两次 `_get_topk_ragged_with_cp(...)`。
4. 当 `current_index_kv` 可复用时，不读取 page table、不 materialize，保持 Phase 3 current reuse 行为。

`_get_topk_ragged_with_cp(...)` 增加可选参数：

```python
shared_index_buffer
shared_block_tables
```

如果两个参数同时提供，则直接使用这份 dense full-view index buffer 和 remapped block table；否则保留原有内部 materialize 行为。

## 验证

新增 CPU 级单元测试覆盖：

1. prev/next pair 在没有 `current_index_kv` 时只 materialize 一次，并且两次 topk 共用同一份 materialized index/block table。
2. 存在 `current_index_kv` 时不触发 page table 读取和 materialize。

运行：

```bash
PYTHONPATH=python python3 -m pytest \
  test/registered/unit/layers/test_nsa_cp_utils.py \
  test/registered/unit/mem_cache/test_cp_shared_kv_layout.py -q
```

## 2026-05-31 update: index partial/current reuse

The index path now follows the same target-model partial-current contract as MLA:

```text
page-aligned cached prefix index pages
+ valid current index K/scale rows copied into their slot-dense suffix pages
+ page-tail slack left zero and invisible through valid seq/page-table lengths
```

Implementation notes:

- `Indexer._can_reuse_current_index_kv()` uses `should_reuse_current_extend_kv()` instead of the old current-only gate, so target cache-hit suffixes can reuse freshly computed current index K/scale.
- Current-only batches still use the compact `current_index_kv` path and avoid page-table materialization.
- Partial cache-hit batches materialize the prefix once for the prev/next in-seq split pair, then pass the composed dense index buffer to both topk calls.
- If the async index prefetcher has a ready prefix handle, `CpSharedKVIndexPrefetcher.consume_prefix_with_current()` consumes it and only fills current rows into the suffix page slots.
- If no prefetch handle exists, `materialize_prefix_and_reuse_current_index_page_slots()` synchronously materializes only prefix slots, all-reduces that prefix range, and fills current rows locally.
- `forward_cuda()` quantizes only `valid_current_rows = extend_seq_lens_cpu[0]` current rows. This preserves current reuse when `out_cache_loc` is physically padded for page alignment.

Debug contract correction:

- The prefetch consume path must remap `logical_pages` through the slot `page_inverse` directly. It must not validate these logical pages against `dense_page_buffer.shape[0]` as if the dense buffer were the physical pool capacity; that debug-only check can reject valid high logical page ids that are intentionally packed into a small slot-dense buffer.

Verification:

```bash
# g0034 / sglang-glm5-dev-2
cd /sgl-workspace/sglang-tai
PYTHONPATH=python python -m pytest -q \
  test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py
# 77 passed, 5 warnings, 2 subtests passed
```
