Batch CP shared-KV index work for bs>1 fast paths

The bs>1 path needs index top-k, shared-index prepare, current-index compact, and current-slot compose to consume flattened batch descriptors instead of falling back to per-request or per-segment Python/Torch work. This change wires SGLang to the new TAI batch prepare kernels, keeps fallback explicit, and records the remaining HiCache/load-backup gaps in the bs>1 workstream docs.

Constraint: CP shared-KV bs>1 must reuse fast paths rather than adding slow batch-only fallbacks

Constraint: No new collective operations were introduced

Rejected: Leave current-only cp_index as Python slice/cat | it keeps per-segment overhead in the short-extend bs>1 case

Rejected: Infer max segment lengths from CUDA descriptor tensors | .item() would add CPU synchronization on the hot path

Confidence: medium

Scope-risk: moderate

Directive: Do not remove the explicit fallback warnings without verifying the corresponding TAI symbols are present in production

Tested: local py_compile for touched SGLang files

Tested: remote g0034 test_nsa_cp_utils.py passed, 53 tests

Tested: remote g0034 test_fill_current_index_page_slots_uses_tai_kernel_when_available passed

Not-tested: full ETE bs>1 traffic with HiCache load/backup and draft/EAGLE enabled

docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md

@@ -651,6 +651,15 @@ Phase 0 -> Phase 2
 5. **workspace pressure：** bs>1 materialize buffer 更大，必须测显存和临时 buffer，不只看 latency。
 6. **collective 顺序：** 所有 CP rank 必须按完全相同的 layer/request 顺序执行 collective。
 7. **fallback 可见性：** 开发阶段优先 fail-fast；生产 fallback 必须 warning 且限频。
+8. **W2 allocator CPU descriptor overhead：** 当前 W2 correctness path 仍在
+   allocation 热路径中用 Python list/tensor 临时对象构造 batch owner descriptor：
+   `extend_lens` / `extend_prefix_lens` list、`flat_page_compute_owners`
+   list、`positions_by_owner` list，以及每个 owner lane 的 `position_tensor`。
+   这仍复用了 owner free/release bucket 本地记账，比旧全量 scan/sort 路径轻，
+   但在线上大量 200-2000 token 短 extend 叠 batch 时可能成为 CPU hot path。
+   后续应让 W1 batch plan 直接携带 page-owner descriptor，W2 消费
+   tensor/固定 buffer，并用 kernel 或固定 workspace 完成 selected-pages
+   gather/scatter，避免每次 allocation 重建 Python descriptor。
 
 ---
 

docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md

@@ -253,6 +253,44 @@ out_cache_loc: torch.Tensor  # flattened request order
 - CP shared-KV bs>1 进入 allocator 时不再出现 `multi_batch` fallback。
 - W3 direct write 可以信任 `out_cache_loc` 的 owner-lane 合同。
 
+### 当前实现状态与性能风险
+
+PR #10 已把 W2 正确性路径接入：
+
+- bs>1 会按 request 独立调用 in-seq page owner planner；
+- owner list 按 request order flatten；
+- `alloc_extend_compute_owner()` 删除 bs=1 hard guard；
+- allocator 复用已有 `alloc_extend_naive()` 生成 flattened `out_cache_loc`；
+- supported CP shared-KV case 不再走 `multi_batch` legacy fallback。
+
+但当前实现仍有一段 Python/control-path descriptor overhead，主要包括：
+
+1. `alloc_paged_token_slots_extend()` 每次 allocation 从 CPU tensor `.item()` 构造
+   `extend_lens` / `extend_prefix_lens` Python lists。
+2. `build_batch_in_seq_page_compute_owners()` 用 Python loop 生成
+   `flat_page_compute_owners: List[int]`。
+3. `_select_compute_owner_pages()` 再用 Python loop 构造
+   `required_by_owner` / `positions_by_owner`。
+4. 每个 owner lane 会临时构造 `position_tensor`，再把 lane pages scatter 回
+   `selected_pages`。
+
+这比旧的全量 free-page scan / `torch.isin` / sort-merge 路径轻，因为 allocator
+仍使用 `_owner_free_pages` / `_owner_release_pages` 本地 bucket 记账；但它还不是最终
+descriptor 化 fast path。在线上大量 200-2000 token 短 extend 叠 batch 时，allocation
+频率高，这段 Python list/tensor descriptor 构造可能重新成为 CPU hot path。
+
+后续优化方向：
+
+```text
+W1 batch plan 直接携带 page-owner descriptor
+  -> W2 allocator 消费 tensor/descriptor 而不是重新 Python 推导
+  -> selected_pages gather/scatter 用固定 buffer 或 kernel 完成
+  -> 避免每 tick/request 反复创建 Python lists 和临时 position tensors
+```
+
+这个风险不影响当前 W2 correctness，但需要在 bs>1 ETE/perf 阶段用 micro benchmark
+和线上 trace 单独验证。
+
 ---
 
 ## 6. W3：CP split/rebuild + direct write fast path
@@ -662,3 +700,48 @@ runtime / kernel 都消费 CPSharedKVBatchPlan descriptors
    - 设计 variable-length descriptor benchmark，覆盖 bf16/fp8、page_first_direct、random/owner-lane pages。
 
 第一批完成后，再开始 W3/W4/W6/W7 draft 接入。
+
+---
+
+## 15. 2026-06-03 W5/W4 kernel 接入状态更新
+
+### 已接入并验证
+
+1. **NSA index top-k 的 batch cp_index 调用已从 per request/segment MQA top-k 收敛为单次调用。**
+   - SGLang 路径：`python/sglang/srt/layers/attention/nsa/nsa_indexer.py::_get_topk_in_seq_cp_pair_batch`。
+   - 语义：多个 request 的 prev/next CP segment 先 compact 成一个 `cp_index` descriptor，再调用一次 `_get_topk_ragged_with_cp`。
+   - 已有测试：`test_indexer_in_seq_cp_pair_batch_*` 与 `test_indexer_ragged_cp_index_current_batch_does_not_materialize`。
+
+2. **index partial/current compose 的 current-slot fill 已接入 TAI kernel。**
+   - SGLang 路径：`cp_shared_kv_runtime.fill_current_index_page_slots`。
+   - TAI kernel：`tai_kernel.nsa_prefill.cp_shared_kv_materialize.fill_current_index_page_slots`。
+   - 输入是 flatten 后的 current rows，可一次处理 bs>1 的 current suffix；不再在 CUDA 上默认走 PyTorch advanced indexing。
+   - fallback 仍为显式 warning：`[CP_SHARED_KV_FALLBACK][tai_materialize] reason=index_current_fill_*`。
+   - 远端 CUDA smoke：通过 64/page slot compose 的 K/scale 字节布局检查。
+
+3. **shared-index batch cp_index 的 K/S + range prepare 已接入 TAI batch descriptor kernel。**
+   - SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_index_batch`。
+   - TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_kv_and_range_batch`。
+   - descriptor：`batch_indices, kv_lens, q_starts, q_lens, k_bases, q_bases, total_kv_len, total_q_count, max_kv_len, max_q_len`。
+   - 关键约束：`max_kv_len/max_q_len` 由 Python planner 传入，避免在 wrapper 内对 CUDA tensor 做 `.max().item()` 同步。
+   - 输出合同：最终仍是按 segment concat 的 dense K buffer / scale buffer；`ks` 为 segment 的 K base；`ke_offset` 为 segment-local causal end offset，调用端计算 `ke = ks + ke_offset`。
+   - 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。
+
+### 仍需继续处理的 kernel/runtime 缺口
+
+1. **current-only cp_index 的 K/S compact copy 已接入 TAI batch kernel。**
+   - SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_current_index_batch`。
+   - TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_current_kv_and_range_batch`。
+   - 语义：从 flatten current rows 按 segment `current_base/current_len` 拷到 concat K/S buffer，同时生成 `ks/ke_offset`。
+   - 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。
+
+2. **HiCache per-layer D2H backup 仍按 reservation/node 提交。**
+   - `scheduler._prepare_hicache_write_backups_before_forward()` 仍 `for req in batch.reqs: prepare_write_backup_for_req(req)`。
+   - `cache_controller.submit_write_cp_layer()` 每个 reservation 在每层调用 host pool backup。
+   - 底层 TAI direct transfer kernel 接收 flatten indices，本身可以吃 batched descriptor；但 runtime 还没有把多个 request/reservation 合并成一个 layer-level descriptor。
+
+3. **L2->L1 load/prefetch 需要继续确认实际 queue 合并粒度。**
+   - `cache_controller.load_cp()` / `start_loading()` 已比 backup 更接近 batched queue，但仍需用 ETE/NVTX 确认 layer-level descriptor 是否足够大，是否存在 per-node launch。
+
+4. **MLA current/partial-current 和 direct store 已有 flatten kernel 基础，但还需用 bs>1 ETE 验证没有 fallback hot path。**
+   - 特别关注 FP8/BF16 两种 dtype、draft/EAGLE 路径、prefetch 关闭时 full/current reuse 是否仍启用。

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

@@ -577,6 +577,26 @@ def _load_tai_index_mqa_prepare_kernel():
         return None
 
 
+@lru_cache(maxsize=1)
+def _load_tai_index_mqa_prepare_batch_kernel():
+    try:
+        from tai_kernel.nsa_prefill import prepare_cp_mqa_kv_and_range_batch
+
+        return prepare_cp_mqa_kv_and_range_batch
+    except Exception:
+        return None
+
+
+@lru_cache(maxsize=1)
+def _load_tai_index_mqa_current_prepare_batch_kernel():
+    try:
+        from tai_kernel.nsa_prefill import prepare_cp_mqa_current_kv_and_range_batch
+
+        return prepare_cp_mqa_current_kv_and_range_batch
+    except Exception:
+        return None
+
+
 @lru_cache(maxsize=1)
 def _load_tai_index_mqa_range_kernel():
     try:
@@ -800,6 +820,214 @@ def try_tai_prepare_cp_mqa_index(
         return None
 
 
+def try_tai_prepare_cp_mqa_current_index_batch(
+    *,
+    current_index_k: torch.Tensor,
+    current_index_scale: torch.Tensor,
+    current_bases: torch.Tensor,
+    kv_lens: torch.Tensor,
+    q_starts: torch.Tensor,
+    q_lens: torch.Tensor,
+    k_bases: torch.Tensor,
+    q_bases: torch.Tensor,
+    total_kv_len: int,
+    total_q_count: int,
+    max_kv_len: int,
+    max_q_len: int,
+    index_head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
+    """Try the TAI batched current-index compact + MQA range prepare path."""
+
+    if not cp_shared_kv_tai_index_mqa_prepare_enabled():
+        _log_tai_index_mqa_prepare_fallback(
+            "current_batch_env_disabled",
+            "CP shared KV TAI batched current-index prepare fast path is disabled; "
+            "falling back to Python slice/cat. segments=%s total_kv_len=%s "
+            "total_q_count=%s index_head_dim=%s",
+            int(current_bases.numel()),
+            total_kv_len,
+            total_q_count,
+            index_head_dim,
+            limit=1,
+        )
+        return None
+
+    kernel = _load_tai_index_mqa_current_prepare_batch_kernel()
+    if kernel is None:
+        _log_tai_index_mqa_prepare_fallback(
+            "current_batch_kernel_missing",
+            "CP shared KV TAI batched current-index prepare kernel is unavailable; "
+            "falling back to Python slice/cat. segments=%s total_kv_len=%s "
+            "total_q_count=%s index_head_dim=%s",
+            int(current_bases.numel()),
+            total_kv_len,
+            total_q_count,
+            index_head_dim,
+            limit=1,
+        )
+        return None
+
+    if index_head_dim != 128:
+        _log_tai_index_mqa_prepare_fallback(
+            "current_batch_unsupported_layout",
+            "CP shared KV TAI batched current-index prepare supports "
+            "index_head_dim=128 only; falling back to Python slice/cat. "
+            "index_head_dim=%s",
+            index_head_dim,
+            limit=4,
+        )
+        return None
+
+    try:
+        return kernel(
+            _contiguous_for_tai(current_index_k),
+            _contiguous_for_tai(current_index_scale),
+            _contiguous_for_tai(current_bases),
+            _contiguous_for_tai(kv_lens),
+            _contiguous_for_tai(q_starts),
+            _contiguous_for_tai(q_lens),
+            _contiguous_for_tai(k_bases),
+            _contiguous_for_tai(q_bases),
+            total_kv_len=int(total_kv_len),
+            total_q_count=int(total_q_count),
+            max_kv_len=int(max_kv_len),
+            max_q_len=int(max_q_len),
+            index_head_dim=int(index_head_dim),
+        )
+    except Exception as exc:
+        _log_tai_index_mqa_prepare_fallback(
+            "current_batch_kernel_failed",
+            "CP shared KV TAI batched current-index prepare failed; falling back "
+            "to Python slice/cat. error=%s segments=%s total_kv_len=%s "
+            "total_q_count=%s",
+            exc,
+            int(current_bases.numel()),
+            total_kv_len,
+            total_q_count,
+        )
+        return None
+
+
+def try_tai_prepare_cp_mqa_index_batch(
+    *,
+    index_buffer: torch.Tensor,
+    block_tables: torch.Tensor,
+    batch_indices: torch.Tensor,
+    kv_lens: torch.Tensor,
+    q_starts: torch.Tensor,
+    q_lens: torch.Tensor,
+    k_bases: torch.Tensor,
+    q_bases: torch.Tensor,
+    total_kv_len: int,
+    total_q_count: int,
+    max_kv_len: int,
+    max_q_len: int,
+    page_size: int,
+    index_head_dim: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None:
+    """Try the TAI batched GetK/GetS + MQA range prepare path.
+
+    This is the bs>1 / multi-CP-segment counterpart of
+    :func:`try_tai_prepare_cp_mqa_index`: all segment descriptors are flattened
+    and one kernel family prepares the concatenated K/S buffer and ks/ke ranges.
+    """
+
+    if not cp_shared_kv_tai_index_mqa_prepare_enabled():
+        _log_tai_index_mqa_prepare_fallback(
+            "batch_env_disabled",
+            "CP shared KV TAI batched index MQA prepare fast path is disabled; "
+            "falling back to per-segment GetK/GetS. segments=%s total_kv_len=%s "
+            "total_q_count=%s page_size=%s index_head_dim=%s",
+            int(batch_indices.numel()),
+            total_kv_len,
+            total_q_count,
+            page_size,
+            index_head_dim,
+            limit=1,
+        )
+        return None
+
+    kernel = _load_tai_index_mqa_prepare_batch_kernel()
+    if kernel is None:
+        _log_tai_index_mqa_prepare_fallback(
+            "batch_kernel_missing",
+            "CP shared KV TAI batched index MQA prepare kernel is unavailable; "
+            "falling back to per-segment GetK/GetS. segments=%s total_kv_len=%s "
+            "total_q_count=%s page_size=%s index_head_dim=%s",
+            int(batch_indices.numel()),
+            total_kv_len,
+            total_q_count,
+            page_size,
+            index_head_dim,
+            limit=1,
+        )
+        return None
+
+    if index_buffer.dtype != torch.uint8:
+        _log_tai_index_mqa_prepare_fallback(
+            "batch_unsupported_dtype",
+            "CP shared KV TAI batched index MQA prepare requires uint8 page "
+            "buffer; falling back to per-segment GetK/GetS. dtype=%s",
+            index_buffer.dtype,
+            limit=4,
+        )
+        return None
+    if not index_buffer.is_contiguous() or not block_tables.is_contiguous():
+        _log_tai_index_mqa_prepare_fallback(
+            "batch_non_contiguous",
+            "CP shared KV TAI batched index MQA prepare requires contiguous "
+            "index buffer and block tables; falling back to per-segment GetK/GetS. "
+            "index_shape=%s index_stride=%s table_shape=%s table_stride=%s",
+            tuple(index_buffer.shape),
+            tuple(index_buffer.stride()),
+            tuple(block_tables.shape),
+            tuple(block_tables.stride()),
+            limit=4,
+        )
+        return None
+    if index_head_dim != 128 or page_size != 64:
+        _log_tai_index_mqa_prepare_fallback(
+            "batch_unsupported_layout",
+            "CP shared KV TAI batched index MQA prepare supports page_size=64 "
+            "and index_head_dim=128 only; falling back to per-segment GetK/GetS. "
+            "page_size=%s index_head_dim=%s",
+            page_size,
+            index_head_dim,
+            limit=4,
+        )
+        return None
+
+    try:
+        return kernel(
+            index_buffer,
+            _contiguous_for_tai(block_tables),
+            _contiguous_for_tai(batch_indices),
+            _contiguous_for_tai(kv_lens),
+            _contiguous_for_tai(q_starts),
+            _contiguous_for_tai(q_lens),
+            _contiguous_for_tai(k_bases),
+            _contiguous_for_tai(q_bases),
+            total_kv_len=int(total_kv_len),
+            total_q_count=int(total_q_count),
+            max_kv_len=int(max_kv_len),
+            max_q_len=int(max_q_len),
+            page_size=int(page_size),
+            index_head_dim=int(index_head_dim),
+        )
+    except Exception as exc:
+        _log_tai_index_mqa_prepare_fallback(
+            "batch_kernel_failed",
+            "CP shared KV TAI batched index MQA prepare failed; falling back to "
+            "per-segment GetK/GetS. error=%s segments=%s total_kv_len=%s "
+            "total_q_count=%s",
+            exc,
+            int(batch_indices.numel()),
+            total_kv_len,
+            total_q_count,
+        )
+        return None
+
+
 def try_tai_prepare_cp_mqa_range(
     *,
     valid_q_count: int,
@@ -1222,6 +1450,64 @@ def fill_current_kv_page_slots_and_remap_locs(
     return dense_kv_cache, mixed_locs, current_mask
 
 
+def _try_tai_fill_current_index_page_slots(
+    *,
+    dense_page_buffer: torch.Tensor,
+    current_index_k: torch.Tensor,
+    current_index_scale: torch.Tensor,
+    current_locs: torch.Tensor,
+    page_inverse: torch.Tensor,
+    page_size: int,
+    index_head_dim: int,
+) -> torch.Tensor | None:
+    if not _tai_materialize_runtime_enabled():
+        _log_tai_materialize_runtime_disabled("fill_current_index_page_slots")
+        return None
+
+    kernels = _load_tai_materialize_kernels()
+    if kernels is None:
+        return None
+    fill_kernel = getattr(kernels, "fill_current_index_page_slots", None)
+    if fill_kernel is None:
+        _log_tai_materialize_fallback(
+            "index_fill_current_missing",
+            "CP shared KV tai index current-slot fill kernel is unavailable; "
+            "falling back to torch reference. Upgrade tai-kernel to keep this "
+            "hot path off PyTorch. page_size=%s index_head_dim=%s current_rows=%s "
+            "dense_pages=%s",
+            page_size,
+            index_head_dim,
+            int(current_locs.numel()),
+            int(dense_page_buffer.shape[0]),
+            limit=1,
+        )
+        return None
+
+    try:
+        return fill_kernel(
+            _contiguous_for_tai(dense_page_buffer),
+            _contiguous_for_tai(current_index_k),
+            _contiguous_for_tai(current_index_scale),
+            _contiguous_for_tai(current_locs.reshape(-1)),
+            _contiguous_for_tai(page_inverse),
+            page_size=int(page_size),
+            index_head_dim=int(index_head_dim),
+        )
+    except Exception as exc:
+        _log_tai_materialize_fallback(
+            "index_fill_current_failed",
+            "CP shared KV tai index current-slot fill failed; falling back to "
+            "torch reference. error=%s page_size=%s index_head_dim=%s "
+            "current_rows=%s dense_pages=%s",
+            exc,
+            page_size,
+            index_head_dim,
+            int(current_locs.numel()),
+            int(dense_page_buffer.shape[0]),
+        )
+        return None
+
+
 def fill_current_index_page_slots(
     *,
     dense_page_buffer: torch.Tensor,
@@ -1244,6 +1530,17 @@ def fill_current_index_page_slots(
     current_rows = int(current_locs.numel())
     if current_rows == 0:
         return dense_page_buffer
+    tai_result = _try_tai_fill_current_index_page_slots(
+        dense_page_buffer=dense_page_buffer,
+        current_index_k=current_index_k,
+        current_index_scale=current_index_scale,
+        current_locs=current_locs,
+        page_inverse=page_inverse,
+        page_size=page_size,
+        index_head_dim=index_head_dim,
+    )
+    if tai_result is not None:
+        return tai_result
     if dense_page_buffer.is_cuda:
         _log_tai_materialize_fallback(
             "index_current_fill_torch_reference_cuda",

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

@@ -32,7 +32,9 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import (
     should_reuse_current_extend_kv,
     tensor_debug_checksum,
     tensor_debug_summary,
+    try_tai_prepare_cp_mqa_current_index_batch,
     try_tai_prepare_cp_mqa_index,
+    try_tai_prepare_cp_mqa_index_batch,
     try_tai_prepare_cp_mqa_range,
 )
 from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor
@@ -1139,8 +1141,15 @@ class Indexer(MultiPlatformOp):
         actual_seq_q_list = []
         batch_idx_list = []
 
-        if current_index_kv is not None and cp_index is not None:
-            current_index_kv = None
+        if (
+            current_index_kv is not None
+            and cp_index is not None
+            and (shared_index_buffer is not None or shared_block_tables is not None)
+        ):
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                "reason=batch_cp_index_ambiguous_current_and_shared_buffer"
+            )
         if current_index_kv is None:
             if shared_index_buffer is not None or shared_block_tables is not None:
                 if shared_index_buffer is None or shared_block_tables is None:
@@ -1177,32 +1186,234 @@ class Indexer(MultiPlatformOp):
             and forward_batch.extend_seq_lens_cpu is not None
         )
         if cp_index is not None:
-            # TODO Multi-batch support has accuracy issues
-            for batch_idx, start_seq_position, end_seq_position in cp_index:
+            current_req_offsets: Optional[List[int]] = None
+            if current_index_kv is not None:
+                current_req_offsets = []
+                current_cursor = 0
+                for extend_len in forward_batch.extend_seq_lens_cpu:
+                    current_req_offsets.append(current_cursor)
+                    current_cursor += int(extend_len)
+
+            segment_records: List[Tuple[int, int, int, int, int, int, int]] = []
+            batch_idx_list = []
+            kv_lens_list = []
+            q_starts_list = []
+            q_lens_list = []
+            k_bases_list = []
+            q_bases_list = []
+            k_cursor = 0
+            q_cursor = 0
+            for raw_batch_idx, start_seq_position, end_seq_position in cp_index:
+                batch_idx = int(raw_batch_idx)
                 pre_chunk_offset = (
                     forward_batch.seq_lens_cpu[batch_idx].item()
                     - forward_batch.extend_seq_lens_cpu[batch_idx]
                 )
                 start_seq_position += pre_chunk_offset
                 end_seq_position += pre_chunk_offset
-                if offset == 0 and batch_idx != 0:
-                    offset += forward_batch.extend_seq_lens_cpu[batch_idx - 1]
-                k_fp8 = index_buf_accessor.GetK.execute(
-                    forward_batch.token_to_kv_pool,
-                    index_buffer,
-                    seq_len=end_seq_position,
-                    page_indices=block_tables[batch_idx],
-                )
-                k_scale = index_buf_accessor.GetS.execute(
-                    forward_batch.token_to_kv_pool,
-                    index_buffer,
-                    seq_len=end_seq_position,
-                    page_indices=block_tables[batch_idx],
+                if end_seq_position < start_seq_position:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                        "reason=batch_cp_index_bad_segment "
+                        f"batch_idx={batch_idx} start={start_seq_position} "
+                        f"end={end_seq_position}"
+                    )
+                extend_seq_len = int(end_seq_position - start_seq_position)
+                kv_len_i = int(end_seq_position)
+                segment_records.append(
+                    (
+                        batch_idx,
+                        int(start_seq_position),
+                        int(end_seq_position),
+                        extend_seq_len,
+                        kv_len_i,
+                        k_cursor,
+                        q_cursor,
+                    )
                 )
+                batch_idx_list.append(batch_idx)
+                kv_lens_list.append(kv_len_i)
+                q_starts_list.append(int(start_seq_position))
+                q_lens_list.append(extend_seq_len)
+                k_bases_list.append(k_cursor)
+                q_bases_list.append(q_cursor)
+                k_cursor += kv_len_i
+                q_cursor += extend_seq_len
+
+            if current_index_kv is None:
+                assert index_buffer is not None
+                assert block_tables is not None
+                descriptor_device = q_fp8.device
+                tai_batch_prepared = try_tai_prepare_cp_mqa_index_batch(
+                    index_buffer=index_buffer,
+                    block_tables=block_tables,
+                    batch_indices=torch.tensor(
+                        batch_idx_list, dtype=torch.int64, device=descriptor_device
+                    ),
+                    kv_lens=torch.tensor(
+                        kv_lens_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    q_starts=torch.tensor(
+                        q_starts_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    q_lens=torch.tensor(
+                        q_lens_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    k_bases=torch.tensor(
+                        k_bases_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    q_bases=torch.tensor(
+                        q_bases_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    total_kv_len=k_cursor,
+                    total_q_count=q_cursor,
+                    max_kv_len=max(kv_lens_list, default=0),
+                    max_q_len=max(q_lens_list, default=0),
+                    page_size=page_size,
+                    index_head_dim=forward_batch.token_to_kv_pool.index_head_dim,
+                )
+                if tai_batch_prepared is not None:
+                    k_fp8_u8, k_scale, ks, ke_offset = tai_batch_prepared
+                    k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
+                    kv_fp8 = (k_fp8, k_scale)
+                    actual_seq_q = torch.tensor(
+                        q_lens_list, dtype=torch.int32, device=q_fp8.device
+                    )
+                    ke = ks + ke_offset
+                    with self._with_real_sm_count():
+                        logits = deep_gemm.fp8_mqa_logits(
+                            q_fp8,
+                            kv_fp8,
+                            weights,
+                            ks,
+                            ke,
+                            clean_logits=False,
+                        )
+                    topk_result = metadata.topk_transform(
+                        logits,
+                        self.index_topk,
+                        ks=ks,
+                        cu_seqlens_q=actual_seq_q,
+                        ke_offset=ke_offset,
+                        batch_idx_list=batch_idx_list,
+                    )
+                    return topk_result
+            else:
+                assert current_req_offsets is not None
+                descriptor_device = q_fp8.device
+                current_bases_list = [
+                    int(current_req_offsets[batch_idx]) for batch_idx in batch_idx_list
+                ]
+                current_index_head_dim = getattr(
+                    forward_batch.token_to_kv_pool,
+                    "index_head_dim",
+                    int(current_index_kv[0].reshape(current_index_kv[0].shape[0], -1).shape[1]),
+                )
+                tai_current_prepared = try_tai_prepare_cp_mqa_current_index_batch(
+                    current_index_k=current_index_kv[0],
+                    current_index_scale=current_index_kv[1],
+                    current_bases=torch.tensor(
+                        current_bases_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    kv_lens=torch.tensor(
+                        kv_lens_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    q_starts=torch.tensor(
+                        q_starts_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    q_lens=torch.tensor(
+                        q_lens_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    k_bases=torch.tensor(
+                        k_bases_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    q_bases=torch.tensor(
+                        q_bases_list, dtype=torch.int32, device=descriptor_device
+                    ),
+                    total_kv_len=k_cursor,
+                    total_q_count=q_cursor,
+                    max_kv_len=max(kv_lens_list, default=0),
+                    max_q_len=max(q_lens_list, default=0),
+                    index_head_dim=current_index_head_dim,
+                )
+                if tai_current_prepared is not None:
+                    k_fp8_u8, k_scale, ks, ke_offset = tai_current_prepared
+                    k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn)
+                    kv_fp8 = (k_fp8, k_scale)
+                    actual_seq_q = torch.tensor(
+                        q_lens_list, dtype=torch.int32, device=q_fp8.device
+                    )
+                    ke = ks + ke_offset
+                    with self._with_real_sm_count():
+                        logits = deep_gemm.fp8_mqa_logits(
+                            q_fp8,
+                            kv_fp8,
+                            weights,
+                            ks,
+                            ke,
+                            clean_logits=False,
+                        )
+                    topk_result = metadata.topk_transform(
+                        logits,
+                        self.index_topk,
+                        ks=ks,
+                        cu_seqlens_q=actual_seq_q,
+                        ke_offset=ke_offset,
+                        batch_idx_list=batch_idx_list,
+                    )
+                    return topk_result
+
+            for (
+                batch_idx,
+                start_seq_position,
+                end_seq_position,
+                extend_seq_len,
+                _kv_len_i,
+                segment_k_base,
+                _segment_q_base,
+            ) in segment_records:
+                if current_index_kv is None:
+                    k_fp8 = index_buf_accessor.GetK.execute(
+                        forward_batch.token_to_kv_pool,
+                        index_buffer,
+                        seq_len=end_seq_position,
+                        page_indices=block_tables[batch_idx],
+                    )
+                    k_scale = index_buf_accessor.GetS.execute(
+                        forward_batch.token_to_kv_pool,
+                        index_buffer,
+                        seq_len=end_seq_position,
+                        page_indices=block_tables[batch_idx],
+                    )
+                else:
+                    if pre_chunk_offset != 0:
+                        raise RuntimeError(
+                            "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                            "reason=batch_cp_index_current_with_prefix "
+                            f"batch_idx={batch_idx} prefix={pre_chunk_offset}"
+                        )
+                    assert current_req_offsets is not None
+                    req_base = current_req_offsets[batch_idx]
+                    req_end = req_base + int(end_seq_position)
+                    if (
+                        req_end > int(current_index_kv[0].shape[0])
+                        or req_end > int(current_index_kv[1].shape[0])
+                    ):
+                        raise RuntimeError(
+                            "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                            "reason=batch_cp_index_current_rows_short "
+                            f"batch_idx={batch_idx} req_end={req_end} "
+                            f"k_rows={int(current_index_kv[0].shape[0])} "
+                            f"scale_rows={int(current_index_kv[1].shape[0])}"
+                        )
+                    k_fp8 = current_index_kv[0][req_base:req_end].contiguous()
+                    k_scale = current_index_kv[1][req_base:req_end].contiguous()
 
-                extend_seq_len = end_seq_position - start_seq_position
                 ks = torch.full(
-                    (extend_seq_len,), offset, dtype=torch.int32, device="cuda"
+                    (extend_seq_len,),
+                    segment_k_base,
+                    dtype=torch.int32,
+                    device=q_fp8.device,
                 )
                 k_fp8_list.append(k_fp8)
                 k_scale_list.append(k_scale)
@@ -1211,14 +1422,13 @@ class Indexer(MultiPlatformOp):
                     start_seq_position + 1,
                     end_seq_position + 1,
                     dtype=torch.int32,
-                    device="cuda",
+                    device=q_fp8.device,
                 )
                 ke_offset_list.append(ke_offset)
                 actual_seq_q = torch.tensor(
-                    [extend_seq_len], dtype=torch.int32, device="cuda"
+                    [extend_seq_len], dtype=torch.int32, device=q_fp8.device
                 )
                 actual_seq_q_list.append(actual_seq_q)
-                batch_idx_list.append(batch_idx)
 
             k_fp8 = torch.cat(k_fp8_list, dim=0).view(torch.float8_e4m3fn)
             k_scale = torch.cat(k_scale_list, dim=0).view(torch.float32).squeeze(-1)
@@ -1553,68 +1763,83 @@ class Indexer(MultiPlatformOp):
                 )
             )
 
-        outputs = []
         cursor = 0
+        output_cursor = 0
+        cp_index: List[Tuple[int, int, int]] = []
+        compact_q_chunks = []
+        compact_weight_chunks = []
+        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)
+        )
+        page_table_1 = None if current_only else metadata.get_page_table_1()
 
-        def call_segment(
+        def collect_segment(
             *,
             req_id: int,
             segment_len: int,
             kv_len: int,
-        ) -> torch.Tensor:
-            nonlocal cursor
+        ) -> None:
+            nonlocal cursor, output_cursor
             segment_len = int(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 torch.empty(
-                    (0, self.index_topk),
-                    dtype=torch.int32,
-                    device=q_fp8.device,
+                return
+
+            seq_len = int(forward_batch.seq_lens_cpu[req_id].item())
+            extend_seq_len = int(forward_batch.extend_seq_lens_cpu[req_id])
+            cp_kv_end = seq_len - extend_seq_len + kv_len
+            if current_only:
+                logical_kv_limit = seq_len
+            else:
+                if page_table_1 is None:
+                    raise RuntimeError(
+                        "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                        "reason=batch_gt1_missing_page_table_1"
+                    )
+                if page_table_1.dim() == 1:
+                    logical_kv_limit = min(seq_len, int(page_table_1.numel()))
+                else:
+                    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,
+                logical_kv_limit=logical_kv_limit,
+            )
+            if valid_q_count <= 0:
+                return
+
+            start_abs = cp_kv_end - segment_len
+            end_abs = start_abs + valid_q_count
+            pre_chunk_offset = seq_len - extend_seq_len
+            cp_index.append(
+                (
+                    req_id,
+                    start_abs - pre_chunk_offset,
+                    end_abs - pre_chunk_offset,
                 )
-            actual_seq_q_tensor = torch.tensor(
-                [segment_len],
-                dtype=torch.int32,
-                device=q_fp8.device,
-            )
-            actual_seq_q_cu_tensor = torch.tensor(
-                [0, segment_len],
-                dtype=torch.int32,
-                device=q_fp8.device,
-            )
-            return self._get_topk_ragged_with_cp(
-                forward_batch,
-                layer_id,
-                q_segment,
-                weights_segment,
-                metadata,
-                kv_len,
-                segment_len,
-                current_index_kv=current_index_kv_for_topk,
-                shared_index_buffer=shared_index_buffer,
-                shared_block_tables=shared_block_tables,
-                actual_seq_q_tensor=actual_seq_q_tensor,
-                actual_seq_q_cu_tensor=actual_seq_q_cu_tensor,
-                batch_idx=req_id,
             )
+            compact_q_chunks.append(q_segment[:valid_q_count])
+            compact_weight_chunks.append(weights_segment[:valid_q_count])
+            compact_output_spans.append((output_cursor, valid_q_count))
 
         for req_id in range(batch_size):
-            outputs.append(
-                call_segment(
-                    req_id=req_id,
-                    segment_len=request_actual_seq_q_prev[req_id],
-                    kv_len=request_kv_len_prev[req_id],
-                )
+            collect_segment(
+                req_id=req_id,
+                segment_len=request_actual_seq_q_prev[req_id],
+                kv_len=request_kv_len_prev[req_id],
             )
-            outputs.append(
-                call_segment(
-                    req_id=req_id,
-                    segment_len=request_actual_seq_q_next[req_id],
-                    kv_len=request_kv_len_next[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],
+                kv_len=request_kv_len_next[req_id],
             )
+            output_cursor += int(request_actual_seq_q_next[req_id])
 
         if cursor != int(q_fp8.shape[0]) or cursor != int(weights.shape[0]):
             raise RuntimeError(
@@ -1624,9 +1849,46 @@ class Indexer(MultiPlatformOp):
                 f"q_tokens={int(q_fp8.shape[0])} weights_tokens={int(weights.shape[0])}"
             )
 
-        if not outputs:
-            return torch.empty((0, self.index_topk), dtype=torch.int32, device=q_fp8.device)
-        return torch.cat(outputs, dim=0)
+        result = torch.full(
+            (output_cursor, self.index_topk),
+            -1,
+            dtype=torch.int32,
+            device=q_fp8.device,
+        )
+        if not compact_q_chunks:
+            return result
+
+        compact_q = torch.cat(compact_q_chunks, dim=0)
+        compact_weights = torch.cat(compact_weight_chunks, dim=0)
+        compact_rows = int(compact_q.shape[0])
+        compact_topk = self._get_topk_ragged_with_cp(
+            forward_batch,
+            layer_id,
+            compact_q,
+            compact_weights,
+            metadata,
+            0,
+            compact_rows,
+            cp_index=cp_index,
+            current_index_kv=current_index_kv_for_topk,
+            shared_index_buffer=shared_index_buffer,
+            shared_block_tables=shared_block_tables,
+            actual_seq_q_tensor=None,
+            actual_seq_q_cu_tensor=None,
+        )
+        if int(compact_topk.shape[0]) != compact_rows:
+            raise RuntimeError(
+                "[CP_SHARED_KV_FAIL_FAST][index_topk] "
+                "reason=batch_gt1_compact_topk_rows_mismatch "
+                f"expected={compact_rows} got={int(compact_topk.shape[0])}"
+            )
+        compact_cursor = 0
+        for output_start, valid_q_count in compact_output_spans:
+            result[output_start : output_start + valid_q_count] = compact_topk[
+                compact_cursor : compact_cursor + valid_q_count
+            ]
+            compact_cursor += valid_q_count
+        return result
 
     def forward_indexer(
         self,

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -1203,6 +1203,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             def get_page_table_64(self):
                 return logical_pages
 
+            def get_page_table_1(self):
+                return torch.empty((2, 1000), dtype=torch.int32)
+
+            def get_page_table_1(self):
+                return torch.empty((2, 1000), dtype=torch.int32)
+
+            def get_page_table_1(self):
+                return torch.empty((2, 1000), 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
@@ -1304,6 +1313,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             def get_page_table_64(self):
                 return logical_pages
 
+            def get_page_table_1(self):
+                return torch.empty((2, 1000), 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
@@ -1329,6 +1341,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                     "batch_idx": batch_idx,
                     "kv_len": kv_len,
                     "actual_seq_q": actual_seq_q,
+                    "cp_index": cp_index,
                     "q": q_fp8.flatten().tolist(),
                     "weights": weights.flatten().tolist(),
                     "actual_seq_q_tensor": actual_seq_q_tensor,
@@ -1338,17 +1351,20 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                     "current_index_kv": current_index_kv,
                 }
             )
-            return torch.full(
-                (actual_seq_q, 2),
-                len(topk_calls),
-                dtype=torch.int32,
-            )
+            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=2,
+            forward_mode=SimpleNamespace(
+                is_extend_without_speculative=lambda: True,
+            ),
+            extend_prefix_lens_cpu=[0, 0],
+            extend_seq_lens_cpu=[1000, 1000],
+            seq_lens_cpu=torch.tensor([1000, 1000], dtype=torch.int64),
             nsa_cp_metadata=NSAContextParallelMetadata(
                 batch_size=2,
                 kv_len_prev=100,
@@ -1357,8 +1373,8 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
                 actual_seq_q_next=1,
                 actual_seq_q_prev_cu_tensor=torch.tensor([0, 2], dtype=torch.int32),
                 actual_seq_q_next_cu_tensor=torch.tensor([0, 1], dtype=torch.int32),
-                request_kv_len_prev=[100, 300],
-                request_kv_len_next=[200, 400],
+                request_kv_len_prev=[2, 1],
+                request_kv_len_next=[3, 4],
                 request_actual_seq_q_prev=[2, 1],
                 request_actual_seq_q_next=[1, 3],
             ),
@@ -1378,32 +1394,26 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         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]["batch_idx"], 0)
+        self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
         self.assertEqual(
-            [
-                (
-                    call["batch_idx"],
-                    call["kv_len"],
-                    call["actual_seq_q"],
-                    call["q"],
-                    call["weights"],
-                    call["actual_seq_q_tensor"].tolist(),
-                    call["actual_seq_q_cu_tensor"].tolist(),
-                )
-                for call in topk_calls
-            ],
-            [
-                (0, 100, 2, [0.0, 1.0], [100.0, 101.0], [2], [0, 2]),
-                (0, 200, 1, [2.0], [102.0], [1], [0, 1]),
-                (1, 300, 1, [3.0], [103.0], [1], [0, 1]),
-                (1, 400, 3, [4.0, 5.0, 6.0], [104.0, 105.0, 106.0], [3], [0, 3]),
-            ],
+            topk_calls[0]["cp_index"],
+            [(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)],
         )
+        self.assertEqual(topk_calls[0]["q"], [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0])
+        self.assertEqual(
+            topk_calls[0]["weights"],
+            [100.0, 101.0, 102.0, 103.0, 104.0, 105.0, 106.0],
+        )
+        self.assertIsNone(topk_calls[0]["actual_seq_q_tensor"])
+        self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"])
         self.assertTrue(all(call["shared_index_buffer"] is materialized_index for call in topk_calls))
         self.assertTrue(all(call["shared_block_tables"] is dense_pages for call in topk_calls))
         self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls))
         self.assertEqual(
             result.tolist(),
-            [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]],
+            [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
         )
 
     def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self):
@@ -1413,7 +1423,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         indexer = object.__new__(Indexer)
         indexer.index_topk = 2
-        current_index_kv = (torch.tensor([1]), torch.tensor([2]))
+        current_index_kv = (torch.arange(7), torch.arange(7))
         logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
         materialized_index = torch.tensor([11], dtype=torch.int32)
         dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32)
@@ -1424,6 +1434,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             def get_page_table_64(self):
                 return logical_pages
 
+            def get_page_table_1(self):
+                return torch.empty((2, 512), dtype=torch.int32)
+
         def fake_materialize(
             forward_batch,
             layer_id,
@@ -1458,23 +1471,33 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             topk_calls.append(
                 {
                     "batch_idx": batch_idx,
+                    "actual_seq_q": actual_seq_q,
+                    "cp_index": cp_index,
+                    "q_rows": int(q_fp8.shape[0]),
                     "current_index_kv": current_index_kv,
                     "shared_index_buffer": shared_index_buffer,
                     "shared_block_tables": shared_block_tables,
                     "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
                 }
             )
-            return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
+            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=2,
+            forward_mode=SimpleNamespace(
+                is_extend_without_speculative=lambda: True,
+            ),
+            extend_prefix_lens_cpu=[64, 64],
+            extend_seq_lens_cpu=[3, 4],
+            seq_lens_cpu=torch.tensor([67, 68], dtype=torch.int64),
             nsa_cp_metadata=NSAContextParallelMetadata(
                 batch_size=2,
-                request_kv_len_prev=[100, 300],
-                request_kv_len_next=[200, 400],
+                request_kv_len_prev=[2, 1],
+                request_kv_len_next=[3, 4],
                 request_actual_seq_q_prev=[2, 1],
                 request_actual_seq_q_next=[1, 3],
             ),
@@ -1492,7 +1515,13 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertEqual(len(materialize_calls), 1)
         self.assertIs(materialize_calls[0]["logical_page_table"], logical_pages)
         self.assertIs(materialize_calls[0]["current_index_kv"], current_index_kv)
-        self.assertEqual(len(topk_calls), 4)
+        self.assertEqual(len(topk_calls), 1)
+        self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
+        self.assertEqual(
+            topk_calls[0]["cp_index"],
+            [(0, 0, 2), (0, 2, 3), (1, 0, 1), (1, 1, 4)],
+        )
+        self.assertEqual(topk_calls[0]["q_rows"], 7)
         self.assertTrue(
             all(call["current_index_kv"] is None for call in topk_calls)
         )
@@ -1502,14 +1531,11 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertTrue(
             all(call["shared_block_tables"] is dense_pages for call in topk_calls)
         )
-        self.assertEqual([call["batch_idx"] for call in topk_calls], [0, 0, 1, 1])
-        self.assertEqual(
-            [call["actual_seq_q_cu_tensor"].tolist() for call in topk_calls],
-            [[0, 2], [0, 1], [0, 1], [0, 3]],
-        )
+        self.assertEqual([call["batch_idx"] for call in topk_calls], [0])
+        self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"])
         self.assertEqual(
             result.tolist(),
-            [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]],
+            [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
         )
 
     def test_indexer_in_seq_cp_pair_batch_reuses_current_index_without_materialize(self):
@@ -1519,7 +1545,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
 
         indexer = object.__new__(Indexer)
         indexer.index_topk = 2
-        current_index_kv = (torch.tensor([1]), torch.tensor([2]))
+        current_index_kv = (torch.arange(7), torch.arange(7))
         topk_calls = []
 
         class Mode:
@@ -1552,13 +1578,17 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             topk_calls.append(
                 {
                     "batch_idx": batch_idx,
+                    "actual_seq_q": actual_seq_q,
+                    "cp_index": cp_index,
+                    "q_rows": int(q_fp8.shape[0]),
                     "current_index_kv": current_index_kv,
                     "shared_index_buffer": shared_index_buffer,
                     "shared_block_tables": shared_block_tables,
                     "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor,
                 }
             )
-            return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
+            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
@@ -1587,7 +1617,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
             current_index_kv=current_index_kv,
         )
 
-        self.assertEqual(len(topk_calls), 4)
+        self.assertEqual(len(topk_calls), 1)
         self.assertTrue(
             all(call["current_index_kv"] is current_index_kv for call in topk_calls)
         )
@@ -1597,16 +1627,305 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertTrue(
             all(call["shared_block_tables"] is None for call in topk_calls)
         )
-        self.assertEqual([call["batch_idx"] for call in topk_calls], [0, 0, 1, 1])
+        self.assertEqual(topk_calls[0]["batch_idx"], 0)
+        self.assertEqual(topk_calls[0]["actual_seq_q"], 7)
         self.assertEqual(
-            [call["actual_seq_q_cu_tensor"].tolist() for call in topk_calls],
-            [[0, 2], [0, 1], [0, 1], [0, 3]],
+            topk_calls[0]["cp_index"],
+            [(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
         )
+        self.assertEqual(topk_calls[0]["q_rows"], 7)
+        self.assertIsNone(topk_calls[0]["actual_seq_q_cu_tensor"])
         self.assertEqual(
             result.tolist(),
-            [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]],
+            [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]],
         )
 
+    def test_indexer_ragged_cp_index_current_batch_does_not_materialize(self):
+        import contextlib
+        import torch
+        from types import SimpleNamespace
+        from unittest.mock import patch
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer.index_topk = 2
+        indexer._with_real_sm_count = lambda: contextlib.nullcontext()
+
+        deep_gemm_calls = []
+
+        def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
+            deep_gemm_calls.append(
+                {
+                    "q_rows": int(q_fp8.shape[0]),
+                    "kv_rows": int(kv_fp8[0].shape[0]),
+                    "weights_rows": int(weights.shape[0]),
+                    "ks": ks.tolist(),
+                    "ke": ke.tolist(),
+                }
+            )
+            return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
+
+        class Metadata:
+            def get_page_table_64(self):
+                raise AssertionError("current cp_index path must not materialize index pages")
+
+            def topk_transform(self, logits, topk, **kwargs):
+                return (
+                    torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
+                    .view(-1, 1)
+                    .repeat(1, topk)
+                )
+
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=SimpleNamespace(page_size=64),
+            seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
+            extend_seq_lens_cpu=[3, 4],
+        )
+        q_fp8 = torch.empty((7, 1), dtype=torch.float32)
+        weights = torch.empty((7, 1, 1), dtype=torch.float32)
+        current_index_kv = (
+            torch.arange(7, dtype=torch.uint8).view(7, 1),
+            torch.arange(7, dtype=torch.float32).view(7, 1),
+        )
+
+        with patch.object(
+            nsa_indexer,
+            "deep_gemm",
+            SimpleNamespace(fp8_mqa_logits=fake_logits),
+        ):
+            result = Indexer._get_topk_ragged_with_cp(
+                indexer,
+                forward_batch,
+                layer_id=7,
+                q_fp8=q_fp8,
+                weights=weights,
+                metadata=Metadata(),
+                kv_len=0,
+                actual_seq_q=7,
+                cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
+                current_index_kv=current_index_kv,
+            )
+
+        self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]])
+        self.assertEqual(len(deep_gemm_calls), 1)
+        self.assertEqual(deep_gemm_calls[0]["q_rows"], 7)
+        self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7)
+        self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14)
+        self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
+        self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
+
+    def test_indexer_ragged_cp_index_shared_batch_uses_tai_prepare_once(self):
+        import contextlib
+        import torch
+        from types import SimpleNamespace
+        from unittest.mock import patch
+
+        from sglang.srt.layers.attention.nsa import index_buf_accessor, nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer.index_topk = 2
+        indexer._with_real_sm_count = lambda: contextlib.nullcontext()
+
+        prepare_calls = []
+        deep_gemm_calls = []
+
+        def fake_prepare(**kwargs):
+            prepare_calls.append(kwargs)
+            total_kv_len = int(kwargs["total_kv_len"])
+            total_q_count = int(kwargs["total_q_count"])
+            return (
+                torch.zeros((total_kv_len, 1), dtype=torch.uint8),
+                torch.zeros((total_kv_len,), dtype=torch.float32),
+                torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32),
+                torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32),
+            )
+
+        def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
+            deep_gemm_calls.append(
+                {
+                    "q_rows": int(q_fp8.shape[0]),
+                    "kv_rows": int(kv_fp8[0].shape[0]),
+                    "weights_rows": int(weights.shape[0]),
+                    "ks": ks.tolist(),
+                    "ke": ke.tolist(),
+                }
+            )
+            return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
+
+        class Metadata:
+            def topk_transform(self, logits, topk, **kwargs):
+                return (
+                    torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
+                    .view(-1, 1)
+                    .repeat(1, topk)
+                )
+
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
+            seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
+            extend_seq_lens_cpu=[3, 4],
+        )
+        q_fp8 = torch.empty((7, 1), dtype=torch.float32)
+        weights = torch.empty((7, 1, 1), dtype=torch.float32)
+        shared_index_buffer = torch.zeros((8, 264), dtype=torch.uint8)
+        shared_block_tables = torch.arange(8, dtype=torch.int64).view(2, 4)
+
+        with patch.object(
+            nsa_indexer,
+            "try_tai_prepare_cp_mqa_index_batch",
+            side_effect=fake_prepare,
+            create=True,
+        ), patch.object(
+            index_buf_accessor.GetK,
+            "execute",
+            side_effect=AssertionError("batched path must not call per-segment GetK"),
+        ), patch.object(
+            index_buf_accessor.GetS,
+            "execute",
+            side_effect=AssertionError("batched path must not call per-segment GetS"),
+        ), patch.object(
+            nsa_indexer,
+            "deep_gemm",
+            SimpleNamespace(fp8_mqa_logits=fake_logits),
+        ):
+            result = Indexer._get_topk_ragged_with_cp(
+                indexer,
+                forward_batch,
+                layer_id=7,
+                q_fp8=q_fp8,
+                weights=weights,
+                metadata=Metadata(),
+                kv_len=0,
+                actual_seq_q=7,
+                cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
+                shared_index_buffer=shared_index_buffer,
+                shared_block_tables=shared_block_tables,
+            )
+
+        self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]])
+        self.assertEqual(len(prepare_calls), 1)
+        call = prepare_calls[0]
+        self.assertIs(call["index_buffer"], shared_index_buffer)
+        self.assertIs(call["block_tables"], shared_block_tables)
+        self.assertEqual(call["batch_indices"].tolist(), [0, 0, 1, 1])
+        self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4])
+        self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1])
+        self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3])
+        self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10])
+        self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4])
+        self.assertEqual(call["total_kv_len"], 14)
+        self.assertEqual(call["total_q_count"], 7)
+        self.assertEqual(call["max_kv_len"], 4)
+        self.assertEqual(call["max_q_len"], 3)
+        self.assertEqual(len(deep_gemm_calls), 1)
+        self.assertEqual(deep_gemm_calls[0]["q_rows"], 7)
+        self.assertEqual(deep_gemm_calls[0]["weights_rows"], 7)
+        self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14)
+        self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
+        self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
+
+    def test_indexer_ragged_cp_index_current_batch_uses_tai_compact_once(self):
+        import contextlib
+        import torch
+        from types import SimpleNamespace
+        from unittest.mock import patch
+
+        from sglang.srt.layers.attention.nsa import nsa_indexer
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer.index_topk = 2
+        indexer._with_real_sm_count = lambda: contextlib.nullcontext()
+
+        prepare_calls = []
+        deep_gemm_calls = []
+
+        def fake_prepare(**kwargs):
+            prepare_calls.append(kwargs)
+            total_kv_len = int(kwargs["total_kv_len"])
+            return (
+                torch.zeros((total_kv_len, 1), dtype=torch.uint8),
+                torch.zeros((total_kv_len,), dtype=torch.float32),
+                torch.tensor([0, 0, 3, 6, 10, 10, 10], dtype=torch.int32),
+                torch.tensor([2, 3, 3, 4, 2, 3, 4], dtype=torch.int32),
+            )
+
+        def fake_logits(q_fp8, kv_fp8, weights, ks, ke, clean_logits=False):
+            deep_gemm_calls.append(
+                {
+                    "q_rows": int(q_fp8.shape[0]),
+                    "kv_rows": int(kv_fp8[0].shape[0]),
+                    "ks": ks.tolist(),
+                    "ke": ke.tolist(),
+                }
+            )
+            return torch.zeros((int(q_fp8.shape[0]), 8), dtype=torch.float32)
+
+        class Metadata:
+            def get_page_table_64(self):
+                raise AssertionError("current cp_index path must not materialize index pages")
+
+            def topk_transform(self, logits, topk, **kwargs):
+                return (
+                    torch.arange(1, int(logits.shape[0]) + 1, dtype=torch.int32)
+                    .view(-1, 1)
+                    .repeat(1, topk)
+                )
+
+        forward_batch = SimpleNamespace(
+            token_to_kv_pool=SimpleNamespace(page_size=64, index_head_dim=1),
+            seq_lens_cpu=torch.tensor([3, 4], dtype=torch.int64),
+            extend_seq_lens_cpu=[3, 4],
+        )
+        current_index_kv = (
+            torch.arange(7, dtype=torch.uint8).view(7, 1),
+            torch.arange(7, dtype=torch.float32).view(7, 1),
+        )
+
+        with patch.object(
+            nsa_indexer,
+            "try_tai_prepare_cp_mqa_current_index_batch",
+            side_effect=fake_prepare,
+            create=True,
+        ), patch.object(
+            nsa_indexer,
+            "deep_gemm",
+            SimpleNamespace(fp8_mqa_logits=fake_logits),
+        ):
+            result = Indexer._get_topk_ragged_with_cp(
+                indexer,
+                forward_batch,
+                layer_id=7,
+                q_fp8=torch.empty((7, 1), dtype=torch.float32),
+                weights=torch.empty((7, 1, 1), dtype=torch.float32),
+                metadata=Metadata(),
+                kv_len=0,
+                actual_seq_q=7,
+                cp_index=[(0, 1, 3), (0, 2, 3), (1, 3, 4), (1, 1, 4)],
+                current_index_kv=current_index_kv,
+            )
+
+        self.assertEqual(result.tolist(), [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]])
+        self.assertEqual(len(prepare_calls), 1)
+        call = prepare_calls[0]
+        self.assertIs(call["current_index_k"], current_index_kv[0])
+        self.assertIs(call["current_index_scale"], current_index_kv[1])
+        self.assertEqual(call["current_bases"].tolist(), [0, 0, 3, 3])
+        self.assertEqual(call["kv_lens"].tolist(), [3, 3, 4, 4])
+        self.assertEqual(call["q_starts"].tolist(), [1, 2, 3, 1])
+        self.assertEqual(call["q_lens"].tolist(), [2, 1, 1, 3])
+        self.assertEqual(call["k_bases"].tolist(), [0, 3, 6, 10])
+        self.assertEqual(call["q_bases"].tolist(), [0, 2, 3, 4])
+        self.assertEqual(call["total_kv_len"], 14)
+        self.assertEqual(call["total_q_count"], 7)
+        self.assertEqual(len(deep_gemm_calls), 1)
+        self.assertEqual(deep_gemm_calls[0]["kv_rows"], 14)
+        self.assertEqual(deep_gemm_calls[0]["ks"], [0, 0, 3, 6, 10, 10, 10])
+        self.assertEqual(deep_gemm_calls[0]["ke"], [2, 3, 6, 10, 12, 13, 14])
+
     def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
         import torch
 

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -1191,6 +1191,55 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
         self.assertEqual(current_mask.tolist(), [[False, True, True, False, False]])
         self.assertEqual(mixed_locs.tolist(), [[4, 12, 13, -1, -1]])
 
+    def test_fill_current_index_page_slots_uses_tai_kernel_when_available(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+
+        class FakeKernels:
+            calls = []
+
+            @staticmethod
+            def fill_current_index_page_slots(*args, **kwargs):
+                FakeKernels.calls.append((args, kwargs))
+                dense_page_buffer = args[0]
+                dense_page_buffer[2, 0] = 99
+                return dense_page_buffer
+
+        dense_page_buffer = torch.zeros((3, 32), dtype=torch.uint8)
+        current_k = torch.ones((2, 4), dtype=torch.uint8)
+        current_scale = torch.ones((2, 1), dtype=torch.float32)
+        current_locs = torch.tensor([8, 9], dtype=torch.int64)
+        page_inverse = torch.tensor([0, -1, 2], dtype=torch.int64)
+
+        with patch.object(
+            runtime,
+            "_tai_materialize_runtime_enabled",
+            return_value=True,
+        ), patch.object(
+            runtime,
+            "_load_tai_materialize_kernels",
+            return_value=FakeKernels,
+        ):
+            result = runtime.fill_current_index_page_slots(
+                dense_page_buffer=dense_page_buffer,
+                current_index_k=current_k,
+                current_index_scale=current_scale,
+                current_locs=current_locs,
+                page_inverse=page_inverse,
+                page_size=4,
+                index_head_dim=4,
+            )
+
+        self.assertIs(result, dense_page_buffer)
+        self.assertEqual(int(result[2, 0]), 99)
+        self.assertEqual(len(FakeKernels.calls), 1)
+        args, kwargs = FakeKernels.calls[0]
+        self.assertIs(args[0], dense_page_buffer)
+        self.assertTrue(torch.equal(args[1], current_k))
+        self.assertTrue(torch.equal(args[2], current_scale))
+        self.assertTrue(torch.equal(args[3], current_locs))
+        self.assertTrue(torch.equal(args[4], page_inverse))
+        self.assertEqual(kwargs, {"page_size": 4, "index_head_dim": 4})
+
     def test_tai_current_slot_fill_is_skipped_when_sparse_page_self_test_fails(self):
         from sglang.srt.environ import envs
         from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime