Reduce CP shared KV materialize and direct-write overhead

Shared KV now relies on page-aligned CP metadata and compute-owner page allocation so persistent MLA KV and NSA index shards can be written by the rank that computed them. The compatibility read path keeps the dense full-view contract for existing topk and attention kernels, but removes duplicated prev/next index materialize, adds optional tai materialize integration, and tightens tests/docs around the fallback boundaries.

Constraint: Decode remains non-CP while prefill CP owns the shared-KV changes

Constraint: Existing attention/topk kernels still expect dense full-view KV/index inputs

Rejected: Change attention kernels to read owner-sharded KV directly | larger semantic change reserved for later phases

Rejected: Merge index K/scale storage with MLA KV storage | would couple topk and attention cache lifecycles before materialize overhead is isolated

Confidence: medium

Scope-risk: broad

Directive: Do not remove fallback logging or debug-gated assertions without reproducing long-context chunked/radix-hit paths

Tested: git diff --check --cached

Not-tested: Local pytest/runtime server verification not run in this commit step per current workflow constraints

docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md

@@ -234,14 +234,21 @@ Phase 4 MVP 建议：
   fallback 到旧 token-average split
 ```
 
-当前实现采用保守 gate：
+当前实现采用保守 gate，但对 radix-hit suffix 放宽：
 
 ```text
-如果 num_units < 2 * cp_size:
+如果 extend_prefix_len == 0 且 num_units < 2 * cp_size:
   fallback 到旧 token-average split
+
+如果 extend_prefix_len > 0 且 prefix page-aligned:
+  允许 cp_size <= num_units < 2 * cp_size，未覆盖的 segment 长度为 0
+
+如果 extend_prefix_len > 0 且 prefix page-aligned 且 num_units < cp_size:
+  不启用 CP split；保留 replicated compute（所有 rank 计算 short suffix），
+  但 compute-owner allocator 仍按 page owner 分配 logical page。
 ```
 
-原因是 CP 本身不适合短序列；真实长上下文场景下 page unit 数通常远大于 `2 * cp_size`，先保证每个 zigzag segment 至少拿到一个完整 page unit，可以避免 zero-token segment 给通信、attention kernel 和后续 compute-owner layout 带来额外风险。后续如果要让短序列直接不走 CP，应单独收紧 `can_cp_split(...)` 的启用阈值，而不是混入 Phase 4 的 page-aligned split 逻辑。
+原因是 CP 本身不适合无 cache 的短序列；但长 agent 上下文里 radix cache hit 很频繁，命中后 current suffix 可能只有少量 page。如果继续因为 suffix 太短 fallback，会导致 compute-owner allocation/write 在高频 radix-hit 路径失效。放宽后的约束仍然是：prefix 必须 page-aligned，实际 suffix page 仍不被切开。`cp_size <= num_units < 2 * cp_size` 时仍走 page-aligned CP split，每个 CP rank 至少拿到一个 page unit；`num_units < cp_size` 时避免构造 zero-token CP rank，改走 replicated compute，由已有 shared-KV write filter 只写本 rank 拥有的 page。
 
 ---
 
@@ -461,7 +468,8 @@ test/registered/unit/attention/test_nsa_cp_page_aligned_split.py
 5. `extend_prefix_len=1`
    - MVP fallback 到旧 split，或显式返回 `page_aligned=False`。
 6. too-short case
-   - `num_units < 2 * cp_size` fallback 或 `page_aligned=False`。
+   - cache-miss: `num_units < 2 * cp_size` fallback 或 `page_aligned=False`。
+   - radix-hit 且 prefix page-aligned: `cp_size <= num_units < 2 * cp_size` 允许；`num_units < cp_size` 不构造 CP split，改走 replicated compute，但 compute-owner allocation 仍可用。
 
 ### 7.2 Invariant tests
 

docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md

@@ -490,6 +490,11 @@ Phase 5 MVP 建议：
 
 radix prefix 命中会让 current extend 从已有 logical pages 之后开始。只要 `extend_prefix_len` page-aligned，new pages 可以按 compute owner lane 分配。
 
+Phase 4/5 对 radix-hit short suffix 放宽 too-short gate：
+
+- 当 `extend_prefix_len > 0` 且 prefix page-aligned，并且 current suffix page 数至少为 `cp_size` 时，即使 page 数小于 `2 * cp_size`，仍生成 page-aligned split / compute-owner page owner list；未覆盖的第二段 zigzag segment 长度为 0。
+- 当 suffix page 数小于 `cp_size` 时，仍允许 compute-owner page allocation，但不启用 CP split。此时沿用 SGLang 原有 replicated compute 行为：所有 rank 计算 short suffix，shared-KV 的 MLA/index write filter 只保留本 rank owner page 的写入。这样避免 zero-token CP rank 的通信/kernel 边界问题，同时避免 radix-hit 高频短 suffix 回退到 legacy allocation。
+
 如果命中到 partial page，MVP fallback。
 
 ### 7.3 Page owner lane free/evict
@@ -606,9 +611,63 @@ Phase 5 完成时应满足：
 8. 非 page-aligned / unsupported case 有明确 fallback。
 ```
 
+## 10. 当前实现切入点
+
+第一版实现按 **allocation-aware modulo owner** 落地：
+
+```text
+mem_cache/cp_shared_kv_compute_owner.py
+  根据 Phase 4 page-aligned in-seq split 规则生成 current page -> compute owner。
+
+mem_cache/allocator.py
+  CPSharedPagedTokenToKVPoolAllocator.alloc_extend_compute_owner(...)
+  从对应 modulo owner lane 选择 logical page。
+
+mem_cache/common.py
+  在 shared KV + in-seq-split + 单请求 page-aligned 场景调用 compute-owner allocation；
+  lane 不足或不满足 gate 时 fallback 到旧 allocator。
+
+layers/attention/nsa/utils.py
+  cp_split_and_rebuild_1d(...)
+  get_cp_shared_kv_local_out_cache_loc(...)
+  生成并缓存本 rank local out_cache_loc，且只有 owner 校验通过才启用 direct write。
+
+forward_mla.py / nsa_indexer.py
+  MLA KV 与 NSA index K/scale 在 all-gather 前用 local KV/key + local physical loc 直接写 persistent pool；
+  attention/topk 计算路径仍保留原有 all-gather。
+  NSA index direct-write 只在 nsa_use_prefill_cp(...) 为 true 时尝试，避免 warmup/短 batch/
+  decode 等没有 nsa_cp_metadata 的非 CP 阶段刷 missing_metadata。
+```
+
+当前仍保留 fallback：
+
+```text
+- 非 page_aligned batch；
+- split_list 与 out_cache_loc 长度不一致（例如 padding 场景）；
+- local logical loc 不属于当前 cp_rank；
+- local KV/index key token 数与 local loc 数不一致；
+- compute-owner lane 分配失败。
+```
+
+compute-owner lane 分配失败通常不是总 free page 不足，而是 modulo owner lane
+不均衡：例如历史 legacy allocation / radix eviction 释放了足够总页数，但某个
+`(logical_page - 1) % cp_size == r` lane 的 free page 不够。当前实现会在
+compute-owner allocation 前按 owner lane deficit 主动触发 radix cache eviction，
+尽量释放对应 lane 的 logical page；只有多次 eviction 后仍不足才 fallback，并在
+fallback log 中打印 `required_by_owner / available_by_owner / deficit_by_owner`。
+
+这些 fallback 会通过 logger 每次触发都提示，不能按 reason 去重。PD warmup
+可能先触发和真实请求相同的 fallback reason；如果只提示一次，会隐藏后续真实
+请求仍在 fallback 的问题：
+
+```text
+CP shared KV compute-owner allocation fallback (...)
+CP shared KV direct-write fallback (...)
+```
+
 ---
 
-## 10. 后续 Phase 候选
+## 11. 后续 Phase 候选
 
 Phase 5 后，如果仍然慢，下一步应集中在 runtime compute path：
 

docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md

@@ -0,0 +1,93 @@
+# 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
+```

docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md

@@ -0,0 +1,553 @@
+# NSA Prefill CP Phase 7: Triton materialize kernels in tai-kernel
+
+Phase 7 的目标是在不改变 Phase 2-6 shared KV 语义的前提下，把 CP shared KV read compatibility 路径里的 materialize 本地 remap/copy 从多段 PyTorch tensor op 改成少量 Triton kernel。kernel 源码计划放在独立包 `tai-kernel` 中，SGLang 通过可选 import 和环境变量接入，保留现有 PyTorch fallback。
+
+## 背景
+
+Phase 2-5 将 persistent KV/index cache 从“每个 CP rank 都保存完整逻辑 KV”改成“每个 CP rank 只保存自己 owner 的 shard”。Phase 6 已经把 in-seq-split 的 prev/next NSA index materialize 从两次合并成一次。
+
+当前 read path 为了兼容现有 NSA topk 和 attention kernel，仍会在每层 attention 前把 shard 形式恢复成 dense full-view：
+
+```text
+owner-sharded physical cache on each CP rank
+  -> local materialize: owned pages copied, non-owned pages zero-filled
+  -> CP all_reduce(sum)
+  -> dense full-view cache consumed by existing topk / attention kernels
+```
+
+profiling 显示真实 all-reduce 不是唯一瓶颈，`materialize` 内部的 remap/local copy 也很重，主要原因是当前实现由多段 PyTorch op 拼接完成，包含 large allocation、zero fill、advanced indexing、`torch.where`、scatter/gather 和多次 kernel launch。
+
+Phase 7 只优化这一层 compatibility materialize 的本地计算；不改变 persistent layout、topk 语义、attention kernel、PD transfer 或 radix cache 语义。
+
+## 当前数据流
+
+### 1. NSA index K/scale materialize
+
+调用链：
+
+```text
+nsa_indexer.py::_maybe_materialize_shared_index_buffer(...)
+  -> cp_shared_kv_runtime.py::materialize_shared_paged_buffer(...)
+```
+
+输入：
+
+```text
+page_buffer    = token_to_kv_pool.get_index_k_with_scale_buffer(layer_id)
+logical_pages  = metadata.real_page_table / metadata.get_page_table_64()
+layout         = CpSharedKVLayout(page_size, cp_size, cp_rank)
+```
+
+当前流程：
+
+```text
+build_slot_page_remap(logical_pages)
+  -> slot_logical_pages = logical_pages.flatten()
+  -> dense_pages: positive logical page -> flat_slot + 1, 0/-1 sentinel 保留
+
+materialize_local_paged_buffer_page_slots(page_buffer, slot_logical_pages, layout)
+  -> owner = (logical_page - 1) % cp_size
+  -> physical_page = (logical_page - 1) // cp_size + 1
+  -> owner == cp_rank 时 copy page_buffer[physical_page]
+  -> 非 owner 或 invalid page 写 zero
+
+_all_reduce_materialized_buffer(dense_page_buffer)
+  -> 所有 rank 得到 dense full-view index buffer
+```
+
+输出：
+
+```text
+dense index buffer
+dense_pages / remapped block table
+```
+
+### 2. MLA KV materialize
+
+调用链：
+
+```text
+nsa_backend.py::forward_extend(...)
+  -> cp_shared_kv_runtime.py::materialize_shared_token_kv_buffer(...)
+```
+
+常见 paged path 输入：
+
+```text
+kv_cache             = persistent MLA KV cache, physical owner-sharded layout
+logical_locs         = page_table_1 after topk transform
+remap_logical_locs   = metadata.page_table_1
+remap_logical_pages  = metadata.real_page_table
+layout               = CpSharedKVLayout(page_size, cp_size, cp_rank)
+```
+
+当前流程：
+
+```text
+build_slot_page_remap(remap_logical_pages)
+  -> materialized_logical_pages / slot logical page table
+
+build_slot_page_inverse(materialized_logical_pages, logical_page_capacity)
+  -> logical_page -> dense slot page
+
+remap_logical_locs_to_slot_dense_locs(logical_locs, page_inverse, page_size)
+  -> logical token loc -> dense token loc
+
+materialize_local_token_kv_page_slots(kv_cache, materialized_logical_pages, layout, page_size)
+  -> owner page copied into dense page slot
+  -> non-owner / invalid page zero-filled
+
+_all_reduce_materialized_buffer(dense_kv_cache)
+  -> 所有 rank 得到 dense full-view MLA KV cache
+```
+
+输出：
+
+```text
+dense kv_cache
+dense_locs / remapped page_table_1
+```
+
+## 当前瓶颈
+
+当前代码路径的问题不是公式复杂，而是执行方式低效：
+
+1. `new_zeros(...)` 对完整 dense buffer 做大块初始化。
+2. `page_buffer[safe_physical_pages]` / `kv_cache[src_tokens]` 触发 advanced indexing，并产生大临时 tensor。
+3. `torch.where(...)` 对完整 dense output 再走一遍，用于 owner/non-owner 选择。
+4. `copy_(...)` 再写一次 dense buffer。
+5. remap 由 `arange/div/remainder/where/scatter/indexing` 多个 PyTorch kernel 拼接，kernel launch 和临时 tensor 都偏多。
+6. 每个 CP rank 都处理完整 dense view，但其中大部分 page 对当前 rank 只是 zero。
+
+因此 Phase 7 的优化重点是：
+
+```text
+用 Triton kernel 一次性完成 slot remap + owner 判断 + physical page 计算 + copy/zero 写出。
+```
+
+CP all-reduce 先保持不变，因为它涉及 distributed group 和 NCCL/torch.distributed 语义，不适合在这一阶段塞进 tai-kernel。
+
+## Phase 7 范围
+
+### In scope
+
+- 在 `tai-kernel` 新增 NSA prefill CP shared KV materialize Triton kernel。
+- SGLang 增加可选接入：环境变量开启时优先调用 tai-kernel，失败或 unsupported shape 回退 PyTorch path。
+- P7A：page materialize copy kernel。
+- P7B：logical loc remap kernel。
+- 单测、benchmark、runtime fallback 保护。
+
+### Out of scope
+
+- 不改 CP all-reduce。
+- 不改 persistent KV/index layout。
+- 不改 NSA topk 语义。
+- 不改 attention kernel 直接读 shared layout。
+- 不合并 index cache 与 MLA KV cache 的存储格式。
+- 不改变 PD transfer 协议。
+- 不改变 radix cache eviction/ownership 策略。
+
+## P7A: page materialize copy kernel
+
+P7A 替换当前的 local page copy/zero path。
+
+### 目标
+
+把当前多段 PyTorch：
+
+```text
+owned_mask = owner(logical_pages) == cp_rank
+physical_pages = logical_pages_to_physical(logical_pages)
+gathered = src[safe_physical_pages]
+dst[1:] = where(owned_mask, gathered, zero)
+```
+
+改成单个 Triton kernel：
+
+```text
+for slot in logical_pages:
+  lp = logical_pages[slot]
+  dense_page = slot + 1 if lp > 0 else lp
+
+  if lp > 0 and ((lp - 1) % cp_size) == cp_rank:
+      physical_page = (lp - 1) // cp_size + 1
+      copy src[physical_page] -> dst[slot + 1]
+  else:
+      write zero -> dst[slot + 1]
+```
+
+### 通用接口草案
+
+`tai-kernel` Python wrapper：
+
+```python
+def materialize_shared_pages(
+    src_pages: torch.Tensor,
+    logical_pages: torch.Tensor,
+    *,
+    page_nbytes: int,
+    cp_rank: int,
+    cp_size: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Materialize owner-sharded pages into a dense slot view.
+
+    Args:
+        src_pages: uint8 view shaped [physical_num_pages, page_nbytes].
+        logical_pages: int tensor with arbitrary shape. Positive values are logical pages;
+            0/-1 are sentinels.
+        page_nbytes: bytes per page in src/dst view.
+        cp_rank/cp_size: CP owner mapping parameters.
+
+    Returns:
+        dense_pages: same shape as logical_pages. Positive logical pages become slot_id + 1;
+            0/-1 sentinel values are preserved.
+        dense_page_buffer: uint8 tensor shaped [num_slots + 1, page_nbytes]. Page 0 is zero.
+    """
+```
+
+Triton launch shape：
+
+```text
+grid = (num_slots, ceil_div(page_nbytes, BLOCK_BYTES))
+```
+
+每个 program 处理一个 page slot 的一个 byte block：
+
+```text
+slot_id       = tl.program_id(0)
+byte_block_id = tl.program_id(1)
+byte_offsets  = byte_block_id * BLOCK_BYTES + tl.arange(0, BLOCK_BYTES)
+```
+
+### 用于 NSA index
+
+输入 view：
+
+```text
+index_buffer: [physical_num_pages, page_size * index_head_dim + page_size * scale_nbytes]
+```
+
+GLM/NSA 常见参数：
+
+```text
+page_size = 64
+index_head_dim = 128
+scale_nbytes = 4
+page_nbytes = 64 * 128 + 64 * 4 = 8448 bytes
+```
+
+替代函数：
+
+```text
+materialize_local_paged_buffer_page_slots(...)
+```
+
+### 用于 MLA KV
+
+输入 view：
+
+```text
+kv_cache: [physical_num_tokens, kv_dim]
+```
+
+先 view 成 page bytes：
+
+```text
+src_pages = kv_cache.view(uint8).reshape(physical_num_pages, page_nbytes)
+```
+
+输出再 view 回：
+
+```text
+dense_kv_cache = dense_page_buffer.view(original_dtype).reshape(
+    (num_slots + 1) * page_size,
+    *kv_cache.shape[1:],
+)
+```
+
+替代函数：
+
+```text
+materialize_local_token_kv_page_slots(...)
+```
+
+### P7A correctness contract
+
+P7A 必须保持：
+
+1. positive logical page 的 dense page id 等于 `flat_slot + 1`。
+2. page 0 dummy 全 zero。
+3. `0/-1` sentinel 在 `dense_pages` 中保留。
+4. 非 owner slot 写 zero，不能留 uninitialized bytes。
+5. owner slot 字节级等价于 PyTorch reference。
+6. 所有 CP rank 经过 sum all-reduce 后结果等价于原 dense full-view。
+
+## P7B: logical loc remap kernel
+
+P7B 替换 MLA KV materialize 中 logical token loc 到 dense token loc 的 remap。
+
+### 目标
+
+当前 PyTorch path：
+
+```text
+page_inverse = build_slot_page_inverse(materialized_logical_pages, logical_page_capacity)
+dense_locs = remap_logical_locs_to_slot_dense_locs(logical_locs, page_inverse, page_size)
+```
+
+语义：
+
+```text
+logical_loc  = logical_page * page_size + offset
+dense_page   = page_inverse[logical_page]
+dense_loc    = dense_page * page_size + offset
+```
+
+P7B 用 Triton 实现两个 kernel：
+
+```text
+1. build_page_inverse_kernel
+   remap_logical_pages slots -> logical_page_capacity inverse table
+
+2. remap_logical_locs_kernel
+   logical_locs -> dense_locs using page_inverse
+```
+
+### 为什么先保留 page_inverse
+
+更激进的做法是对 `remap_logical_pages` 做 per-loc search 或 compact hash table，但当前目标是低风险替代 PyTorch reference。full `page_inverse` 虽然有额外内存，但语义最接近当前实现，也最容易做 exact equality test。
+
+后续如果 P7B 仍然成为瓶颈，再考虑：
+
+- compact hash inverse；
+- 与 topk transform 融合，直接输出 dense loc；
+- attention kernel 直接消费 logical/shared layout。
+
+这些不放入 Phase 7。
+
+### P7B correctness contract
+
+P7B 必须保持：
+
+1. `logical_locs < 0` 输出 `-1`。
+2. page 0 映射到 dense page 0。
+3. 不在 `remap_logical_pages` 中的 logical page 在 debug 模式下仍应暴露错误；非 debug 模式保持现有容错行为。
+4. 输出 dtype/shape 与输入 `logical_locs` 一致。
+5. 对 paged topk path 的 `page_table_1` remap 与 PyTorch reference 完全一致。
+
+## SGLang 接入计划
+
+新增 env：
+
+```text
+SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=0/1
+```
+
+默认关闭。开启后走 tai-kernel fast path：
+
+```text
+cp_shared_kv_runtime.py
+  materialize_shared_paged_buffer(...)
+    -> try tai materialize_shared_pages for index buffer
+    -> fallback PyTorch reference
+
+  materialize_shared_token_kv_buffer(...)
+    -> try tai page materialize for MLA KV
+    -> try tai loc remap for dense_locs
+    -> fallback PyTorch reference
+```
+
+接入原则：
+
+1. tai-kernel import 失败时不影响启动。
+2. unsupported dtype/shape 时回退 PyTorch path。
+3. debug env `SGLANG_DEBUG_CP_SHARED_KV=1` 时保留现有 validate/checksum 能力。
+4. fallback 需要 log reason，但避免在 hot path 无限制刷日志；可复用当前 shared KV fallback logger 风格。
+
+### 当前接入状态
+
+已在 SGLang 侧加入实验开关：
+
+```text
+SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1
+```
+
+默认仍关闭。开启后：
+
+- `materialize_shared_paged_buffer(...)` 优先调用
+  `tai_kernel.nsa_prefill.cp_shared_kv_materialize.materialize_shared_pages(...)`
+  处理 NSA index K/scale page buffer；
+- `materialize_shared_token_kv_buffer(..., remap_logical_pages=...)` 优先调用
+  tai-kernel 的 `build_slot_page_inverse(...)`、
+  `remap_logical_locs_to_slot_dense_locs(...)` 和
+  `materialize_shared_token_kv_pages(...)` 处理 MLA KV page materialize 与
+  dense loc remap；
+- tai token fast path 直接使用 `remap_logical_pages.reshape(-1)` 作为 slot
+  logical pages，避免在 fast path 前额外执行 PyTorch
+  `build_slot_page_remap(...)` 的 `clone/arange/where` 开销；只有 tai
+  fallback 时才构建 PyTorch slot remap；
+- `SGLANG_DEBUG_CP_SHARED_KV=1` 时强制保留原 PyTorch path，避免绕过现有
+  debug assert/checksum；
+- tai-kernel import 失败、CPU tensor、非 contiguous / unsupported shape 等异常
+  会回退 PyTorch reference path，并按 reason 限流 warning。
+
+当前接入仍只替换 local materialize/remap，CP all-reduce 不变。
+
+## tai-kernel 文件计划
+
+新增：
+
+```text
+tai-kernel/python/tai_kernel/nsa_prefill/__init__.py
+tai-kernel/python/tai_kernel/nsa_prefill/cp_shared_kv_materialize.py
+```
+
+可选新增 benchmark/test：
+
+```text
+tai-kernel/tests/nsa_prefill/test_cp_shared_kv_materialize.py
+tai-kernel/benchmark/nsa_prefill/benchmark_cp_shared_kv_materialize.py
+```
+
+当前 `tai-kernel` 包主要已有 quantization extension；`nsa_prefill` 目录只有 pycache 残留，没有可维护源码。因此 Phase 7 会正式建立 `tai_kernel.nsa_prefill` Python/Triton module。
+
+## Test plan
+
+### 1. tai-kernel unit tests
+
+构造 PyTorch reference，对每个 CP rank 单独 materialize，然后模拟 all-reduce：
+
+```text
+sum(local_dense_buffer_per_rank) == reference_dense_full_view
+```
+
+覆盖：
+
+- `cp_size = 1/2/8`
+- `cp_rank = 0..cp_size-1`
+- `page_size = 64`
+- random logical pages
+- logical pages 包含 `0` 和 `-1`
+- duplicated logical pages
+- owner-empty rank
+- index buffer dtype uint8
+- MLA KV buffer dtype bfloat16 / float16 / uint8 view
+
+### 2. SGLang unit tests
+
+在 SGLang 侧增加 env-on/off 对比：
+
+```text
+SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=0 -> PyTorch reference
+SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 -> tai fast path
+```
+
+验证：
+
+```text
+dense_pages equal
+dense_locs equal
+dense_page_buffer byte-level equal before all_reduce simulation
+dense_kv_cache byte-level equal before all_reduce simulation
+```
+
+### 3. Runtime validation
+
+使用现有 GLM5 prefill CP 启动命令，分别测试：
+
+```text
+SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=0 baseline
+SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 fast path
+```
+
+请求类型：
+
+1. 短 prompt warmup。
+2. 长 prompt cache miss。
+3. 长 prompt radix hit + short suffix。
+4. 多请求并发。
+
+检查：
+
+- 输出内容正常；
+- 无 materialize invalid loc/page；
+- fallback reason 可解释；
+- profiler 中 `materialize.token.local_copy` / `materialize.paged.local_copy` 下降。
+
+## Benchmark plan
+
+新增 benchmark 只测本地 materialize，不包含 CP all-reduce：
+
+```text
+python benchmark/nsa_prefill/benchmark_cp_shared_kv_materialize.py
+```
+
+case matrix：
+
+```text
+cp_size: 8
+page_size: 64
+num_slots/pages: 512, 1024, 2048, 4096
+index page bytes: 8448
+MLA KV page bytes: derived from kv_cache shape
+logical page pattern:
+  - contiguous owner-balanced
+  - random owner-balanced
+  - duplicated pages
+  - pages with 0/-1 sentinels
+```
+
+输出指标：
+
+```text
+PyTorch reference ms
+Triton P7A ms
+Triton P7B ms
+speedup
+allocated bytes if measurable
+```
+
+## Rollout plan
+
+1. 实现 tai-kernel P7A page materialize kernel。
+2. 在 tai-kernel 内用 reference 单测保证 byte-level correctness。
+3. SGLang 接入 index materialize fast path，默认 env off。
+4. SGLang 接入 MLA KV page copy fast path，默认 env off。
+5. 实现 P7B loc remap kernel。
+6. 增加 SGLang env-on/off 对比测试。
+7. 在 g0034/g0035/g0036 跑 prefill/decode/router runtime 验证。
+8. profiler 对比 materialize 本地开销。
+9. 若稳定，再考虑把 env 默认值改为 on；否则保持实验开关。
+
+## Risks and mitigations
+
+### Risk 1: Triton kernel 对 dtype/view 处理错误
+
+Mitigation：P7A 统一用 `uint8` view 做 byte copy，避免 dtype-specific copy kernel。输出再 view 回原 dtype。
+
+### Risk 2: 输出未完整初始化
+
+Mitigation：不用 `empty` 后依赖部分写入；kernel 必须覆盖 page 0 和所有 slot page 的所有 byte。测试中加入 owner-empty rank 和 sentinel pages。
+
+### Risk 3: dense_locs remap 边界错误导致 attention 读错 KV
+
+Mitigation：P7B 独立 exact equality test；debug 模式保留现有 invalid page/loc assert；runtime 先 env-gated。
+
+### Risk 4: 多 batch page table 语义不一致
+
+Mitigation：P7A 对 flattened slot table 天然 batch-agnostic。P7B 先按当前 `build_slot_page_inverse` 语义实现 global inverse，不引入 batch-specific search。
+
+### Risk 5: all-reduce 仍是瓶颈，P7A/P7B 收益有限
+
+Mitigation：benchmark 分离 local materialize 与 all-reduce；Phase 7 只承诺降低 local materialize。通信量压缩或 layout-aware attention 留给后续 phase。
+
+## Acceptance criteria
+
+Phase 7 完成标准：
+
+1. tai-kernel 提供 P7A/P7B Triton wrapper，SGLang 可选启用。
+2. `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=0` 行为保持现有 PyTorch path。
+3. `SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1` 时，unit test 与 PyTorch reference byte-level equal。
+4. GLM5 prefill CP + decode + router 长 prompt 请求输出正常。
+5. profiler 中 local materialize remap/copy 时间相比 baseline 下降。
+6. 若 tai-kernel 不存在或 kernel 不支持当前 shape，服务可自动 fallback，不影响正确性。

python/sglang/srt/environ.py

@@ -204,6 +204,7 @@ class Envs:
     SGLANG_DEBUG_MEMORY_POOL = EnvBool(False)
     SGLANG_DEBUG_CP_SHARED_KV = EnvBool(False)
     SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False)
+    SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False)
     SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False)
     SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False)
     SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1)

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

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+from functools import lru_cache
 
 import torch
 
@@ -11,6 +12,7 @@ from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 logger = logging.getLogger(__name__)
 
 _DEBUG_LOG_COUNTS: dict[str, int] = {}
+_TAI_MATERIALIZE_FALLBACK_LOG_COUNTS: dict[str, int] = {}
 
 
 def cp_shared_kv_debug_enabled() -> bool:
@@ -21,6 +23,125 @@ def cp_shared_kv_current_reuse_enabled() -> bool:
     return envs.SGLANG_CP_SHARED_KV_CURRENT_REUSE.get()
 
 
+def cp_shared_kv_tai_materialize_enabled() -> bool:
+    return envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get()
+
+
+@lru_cache(maxsize=1)
+def _load_tai_materialize_kernels():
+    try:
+        from tai_kernel.nsa_prefill import cp_shared_kv_materialize
+
+        return cp_shared_kv_materialize
+    except Exception as exc:
+        _log_tai_materialize_fallback(
+            "import_failed",
+            "CP shared KV tai materialize kernels are unavailable; "
+            "falling back to torch materialize. error=%s",
+            exc,
+            limit=1,
+        )
+        return None
+
+
+def _tai_materialize_runtime_enabled() -> bool:
+    # Keep the debug path on the existing torch implementation.  The debug path
+    # intentionally preserves tensor summaries and value assertions used for
+    # diagnosing shared-KV correctness; the Triton kernels optimize the normal
+    # runtime path only.
+    return cp_shared_kv_tai_materialize_enabled() and not cp_shared_kv_debug_enabled()
+
+
+def _log_tai_materialize_fallback(
+    key: str,
+    message: str,
+    *args,
+    limit: int = 8,
+) -> None:
+    count = _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS.get(key, 0)
+    if count >= limit:
+        return
+    _TAI_MATERIALIZE_FALLBACK_LOG_COUNTS[key] = count + 1
+    logger.warning(message, *args)
+
+
+def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor:
+    return tensor if tensor.is_contiguous() else tensor.contiguous()
+
+
+def _try_tai_materialize_shared_pages(
+    page_buffer: torch.Tensor,
+    logical_pages: torch.Tensor,
+    layout: CpSharedKVLayout,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    if not _tai_materialize_runtime_enabled():
+        return None
+
+    kernels = _load_tai_materialize_kernels()
+    if kernels is None:
+        return None
+
+    try:
+        return kernels.materialize_shared_pages(
+            page_buffer,
+            _contiguous_for_tai(logical_pages),
+            cp_rank=layout.cp_rank,
+            cp_size=layout.cp_size,
+        )
+    except Exception as exc:
+        _log_tai_materialize_fallback(
+            "paged_failed",
+            "CP shared KV tai paged materialize failed; falling back to torch "
+            "materialize. error=%s",
+            exc,
+        )
+        return None
+
+
+def _try_tai_materialize_token_kv_pages_and_locs(
+    kv_cache: torch.Tensor,
+    logical_locs: torch.Tensor,
+    slot_logical_pages: torch.Tensor,
+    logical_page_capacity: int,
+    layout: CpSharedKVLayout,
+    page_size: int,
+) -> tuple[torch.Tensor, torch.Tensor] | None:
+    if not _tai_materialize_runtime_enabled():
+        return None
+
+    kernels = _load_tai_materialize_kernels()
+    if kernels is None:
+        return None
+
+    try:
+        tai_slot_logical_pages = _contiguous_for_tai(slot_logical_pages.reshape(-1))
+        page_inverse = kernels.build_slot_page_inverse(
+            tai_slot_logical_pages,
+            logical_page_capacity,
+        )
+        dense_locs = kernels.remap_logical_locs_to_slot_dense_locs(
+            _contiguous_for_tai(logical_locs),
+            page_inverse,
+            page_size=page_size,
+        )
+        dense_kv_cache = kernels.materialize_shared_token_kv_pages(
+            kv_cache,
+            tai_slot_logical_pages,
+            page_size=page_size,
+            cp_rank=layout.cp_rank,
+            cp_size=layout.cp_size,
+        )
+        return dense_kv_cache, dense_locs
+    except Exception as exc:
+        _log_tai_materialize_fallback(
+            "token_failed",
+            "CP shared KV tai token materialize failed; falling back to torch "
+            "materialize. error=%s",
+            exc,
+        )
+        return None
+
+
 def is_current_only_extend_batch(forward_batch) -> bool:
     """Return whether an extend batch has no cached/history tokens.
 
@@ -737,6 +858,7 @@ def materialize_shared_token_kv_buffer(
             physical_token_capacity=kv_cache.shape[0],
         )
 
+    dense_kv_cache = None
     if remap_logical_pages is None:
         remap_pages_from_locs = logical_pages_from_locs(remap_logical_locs, page_size)
         materialized_logical_pages, _ = build_dense_page_remap(remap_pages_from_locs)
@@ -757,21 +879,36 @@ def materialize_shared_token_kv_buffer(
             layout=layout,
             physical_page_capacity=kv_cache.shape[0] // page_size,
         )
-        materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
         logical_page_capacity = _logical_page_capacity_from_physical_page_capacity(
             kv_cache.shape[0] // page_size,
             layout,
         )
-        page_inverse = build_slot_page_inverse(
-            materialized_logical_pages,
-            logical_page_capacity=logical_page_capacity,
-        )
-        dense_locs = remap_logical_locs_to_slot_dense_locs(
-            logical_locs,
-            page_inverse=page_inverse,
-            page_size=page_size,
-        )
-        use_slot_materialize = True
+        tai_result = None
+        if _tai_materialize_runtime_enabled():
+            materialized_logical_pages = remap_logical_pages.reshape(-1)
+            tai_result = _try_tai_materialize_token_kv_pages_and_locs(
+                kv_cache=kv_cache,
+                logical_locs=logical_locs,
+                slot_logical_pages=materialized_logical_pages,
+                logical_page_capacity=logical_page_capacity,
+                layout=layout,
+                page_size=page_size,
+            )
+        if tai_result is None:
+            materialized_logical_pages, _ = build_slot_page_remap(remap_logical_pages)
+            page_inverse = build_slot_page_inverse(
+                materialized_logical_pages,
+                logical_page_capacity=logical_page_capacity,
+            )
+            dense_locs = remap_logical_locs_to_slot_dense_locs(
+                logical_locs,
+                page_inverse=page_inverse,
+                page_size=page_size,
+            )
+            use_slot_materialize = True
+        else:
+            dense_kv_cache, dense_locs = tai_result
+            use_slot_materialize = False
 
     if use_slot_materialize:
         dense_kv_cache = materialize_local_token_kv_page_slots(
@@ -780,7 +917,7 @@ def materialize_shared_token_kv_buffer(
             layout=layout,
             page_size=page_size,
         )
-    else:
+    elif dense_kv_cache is None:
         dense_kv_cache = materialize_local_token_kv_pages(
             kv_cache=kv_cache,
             unique_logical_pages=materialized_logical_pages,
@@ -830,6 +967,7 @@ def materialize_shared_token_kv_buffer(
         )
     return dense_kv_cache, dense_locs
 
+
 def materialize_shared_paged_buffer(
     page_buffer: torch.Tensor,
     logical_pages: torch.Tensor,
@@ -845,12 +983,21 @@ def materialize_shared_paged_buffer(
         layout=layout,
         physical_page_capacity=page_buffer.shape[0],
     )
-    materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
-    dense_page_buffer = materialize_local_paged_buffer_page_slots(
+    tai_result = _try_tai_materialize_shared_pages(
         page_buffer=page_buffer,
-        slot_logical_pages=materialized_logical_pages,
+        logical_pages=logical_pages,
         layout=layout,
     )
+    if tai_result is None:
+        materialized_logical_pages, dense_pages = build_slot_page_remap(logical_pages)
+        dense_page_buffer = materialize_local_paged_buffer_page_slots(
+            page_buffer=page_buffer,
+            slot_logical_pages=materialized_logical_pages,
+            layout=layout,
+        )
+    else:
+        dense_page_buffer, dense_pages = tai_result
+        materialized_logical_pages = logical_pages.reshape(-1)
 
     if cp_shared_kv_debug_enabled():
         owned_pages = materialized_logical_pages[

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

@@ -52,8 +52,11 @@ 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_local_out_cache_loc,
     is_nsa_enable_prefill_cp,
     is_nsa_prefill_cp_in_seq_split,
+    log_cp_shared_kv_direct_write_fallback,
+    nsa_use_prefill_cp,
     split_in_seq_cp_local_pair,
 )
 from sglang.srt.layers.communicator import ScatterMode
@@ -442,6 +445,7 @@ class Indexer(MultiPlatformOp):
             query = rotate_activation(query)
             key = rotate_activation(key)
 
+        local_key = key
         # allgather+rerrange
         if forward_batch.nsa_cp_metadata is not None and self.nsa_enable_prefill_cp:
             key = cp_all_gather_rerange_output(
@@ -450,7 +454,7 @@ class Indexer(MultiPlatformOp):
                 forward_batch,
                 torch.cuda.current_stream(),
             )
-        return query, key
+        return query, key, local_key
 
     def _get_k_bf16(
         self,
@@ -839,6 +843,8 @@ class Indexer(MultiPlatformOp):
         actual_seq_q: int,
         cp_index: List[Tuple[int, int, int]] = None,
         current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        shared_index_buffer: Optional[torch.Tensor] = None,
+        shared_block_tables: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
         if TYPE_CHECKING:
             assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool)
@@ -855,15 +861,23 @@ class Indexer(MultiPlatformOp):
         actual_seq_q_list = []
         batch_idx_list = []
 
-        block_tables = metadata.get_page_table_64()
         if current_index_kv is not None and cp_index is not None:
             current_index_kv = None
         if current_index_kv is None:
-            index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
-                forward_batch,
-                layer_id,
-                block_tables,
-            )
+            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:
+                    raise RuntimeError(
+                        "shared index buffer and block tables must be provided together"
+                    )
+                index_buffer = shared_index_buffer
+                block_tables = shared_block_tables
+            else:
+                block_tables = metadata.get_page_table_64()
+                index_buffer, block_tables = self._maybe_materialize_shared_index_buffer(
+                    forward_batch,
+                    layer_id,
+                    block_tables,
+                )
         else:
             index_buffer = None
             if cp_shared_kv_debug_enabled():
@@ -1032,6 +1046,76 @@ class Indexer(MultiPlatformOp):
 
         return topk_result
 
+    def _get_topk_in_seq_cp_pair(
+        self,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        q_fp8: torch.Tensor,
+        weights: torch.Tensor,
+        metadata: BaseIndexerMetadata,
+        current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+    ) -> torch.Tensor:
+        assert forward_batch.nsa_cp_metadata is not None
+        kv_len_prev = forward_batch.nsa_cp_metadata.kv_len_prev
+        kv_len_next = forward_batch.nsa_cp_metadata.kv_len_next
+        actual_seq_q_prev = forward_batch.nsa_cp_metadata.actual_seq_q_prev
+        actual_seq_q_next = forward_batch.nsa_cp_metadata.actual_seq_q_next
+
+        # TODO support mutil-batch
+        # cp_batch_seq_index_prev = forward_batch.nsa_cp_metadata["cp_batch_seq_index_prev"]
+        # cp_batch_seq_index_next = forward_batch.nsa_cp_metadata["cp_batch_seq_index_next"]
+        q_fp8_prev, q_fp8_next = split_in_seq_cp_local_pair(
+            q_fp8,
+            actual_seq_q_prev,
+            actual_seq_q_next,
+            name="q_fp8",
+        )
+        weights_prev, weights_next = split_in_seq_cp_local_pair(
+            weights,
+            actual_seq_q_prev,
+            actual_seq_q_next,
+            name="weights",
+        )
+
+        shared_index_buffer = None
+        shared_block_tables = None
+        if current_index_kv is None:
+            shared_block_tables = metadata.get_page_table_64()
+            shared_index_buffer, shared_block_tables = (
+                self._maybe_materialize_shared_index_buffer(
+                    forward_batch,
+                    layer_id,
+                    shared_block_tables,
+                )
+            )
+
+        topk_result_prev = self._get_topk_ragged_with_cp(
+            forward_batch,
+            layer_id,
+            q_fp8_prev,
+            weights_prev,
+            metadata,
+            kv_len_prev,
+            actual_seq_q_prev,
+            current_index_kv=current_index_kv,
+            shared_index_buffer=shared_index_buffer,
+            shared_block_tables=shared_block_tables,
+        )
+
+        topk_result_next = self._get_topk_ragged_with_cp(
+            forward_batch,
+            layer_id,
+            q_fp8_next,
+            weights_next,
+            metadata,
+            kv_len_next,
+            actual_seq_q_next,
+            current_index_kv=current_index_kv,
+            shared_index_buffer=shared_index_buffer,
+            shared_block_tables=shared_block_tables,
+        )
+        return torch.cat([topk_result_prev, topk_result_next], dim=0)
+
     def forward_indexer(
         self,
         q_fp8: torch.Tensor,
@@ -1129,6 +1213,7 @@ class Indexer(MultiPlatformOp):
         key: torch.Tensor,
         *,
         act_quant=None,  # fallback only
+        out_loc_override: Optional[torch.Tensor] = None,
     ) -> None:
         """
         Store NSA indexer K cache for current step.
@@ -1136,7 +1221,10 @@ class Indexer(MultiPlatformOp):
         Preferred: fused_store_index_k_cache(key, cache, out_cache_loc, page_size)
         Fallback : act_quant(key) + token_to_kv_pool.set_index_k_scale_buffer(...)
         """
-        out_loc, key = self._filter_shared_index_write(forward_batch, key)
+        if out_loc_override is None:
+            out_loc, key = self._filter_shared_index_write(forward_batch, key)
+        else:
+            out_loc = out_loc_override
         if out_loc.numel() == 0:
             return
         if not out_loc.is_contiguous():
@@ -1175,6 +1263,48 @@ class Indexer(MultiPlatformOp):
             index_k_scale=k_scale,
         )
 
+    def _store_cp_shared_local_index_k_cache(
+        self,
+        forward_batch: ForwardBatch,
+        layer_id: int,
+        local_key: torch.Tensor,
+        *,
+        act_quant,
+    ) -> bool:
+        if not nsa_use_prefill_cp(forward_batch, self.nsa_enable_prefill_cp):
+            return False
+
+        local_out_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+        if local_out_loc is None:
+            return False
+        if local_key.shape[0] != local_out_loc.numel():
+            log_cp_shared_kv_direct_write_fallback(
+                "index_local_shape_mismatch",
+                "NSA index local key token count does not match local out_cache_loc: "
+                "local_key=%s local_out_cache_loc=%s layer_id=%s",
+                local_key.shape[0],
+                local_out_loc.numel(),
+                layer_id,
+            )
+            return False
+        if local_out_loc.numel() == 0:
+            return True
+
+        assert forward_batch.cp_shared_kv_layout is not None
+        physical_out_loc = (
+            forward_batch.cp_shared_kv_layout.logical_locs_to_physical(
+                local_out_loc
+            ).contiguous()
+        )
+        self._store_index_k_cache(
+            forward_batch=forward_batch,
+            layer_id=layer_id,
+            key=local_key,
+            act_quant=act_quant,
+            out_loc_override=physical_out_loc,
+        )
+        return True
+
     def forward_cuda(
         self,
         x: torch.Tensor,
@@ -1236,21 +1366,27 @@ class Indexer(MultiPlatformOp):
             current_stream = torch.cuda.current_stream()
             self.alt_stream.wait_stream(current_stream)
             weights = self._project_and_scale_head_gates(x)
-            query, key = self._get_q_k_bf16(
+            query, key, local_key = self._get_q_k_bf16(
                 q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
             )
             q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
             with torch.cuda.stream(self.alt_stream):
-                self._store_index_k_cache(
+                if not self._store_cp_shared_local_index_k_cache(
                     forward_batch=forward_batch,
                     layer_id=layer_id,
-                    key=key,
+                    local_key=local_key,
                     act_quant=act_quant,
-                )
+                ):
+                    self._store_index_k_cache(
+                        forward_batch=forward_batch,
+                        layer_id=layer_id,
+                        key=key,
+                        act_quant=act_quant,
+                    )
             current_stream.wait_stream(self.alt_stream)
             weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
         else:
-            query, key = self._get_q_k_bf16(
+            query, key, local_key = self._get_q_k_bf16(
                 q_lora, x, positions, enable_dual_stream, forward_batch=forward_batch
             )
 
@@ -1260,21 +1396,33 @@ class Indexer(MultiPlatformOp):
 
                 q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
                 with torch.cuda.stream(self.alt_stream):
+                    if not self._store_cp_shared_local_index_k_cache(
+                        forward_batch=forward_batch,
+                        layer_id=layer_id,
+                        local_key=local_key,
+                        act_quant=act_quant,
+                    ):
+                        self._store_index_k_cache(
+                            forward_batch=forward_batch,
+                            layer_id=layer_id,
+                            key=key,
+                            act_quant=act_quant,
+                        )
+                current_stream.wait_stream(self.alt_stream)
+            else:
+                q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
+                if not self._store_cp_shared_local_index_k_cache(
+                    forward_batch=forward_batch,
+                    layer_id=layer_id,
+                    local_key=local_key,
+                    act_quant=act_quant,
+                ):
                     self._store_index_k_cache(
                         forward_batch=forward_batch,
                         layer_id=layer_id,
                         key=key,
                         act_quant=act_quant,
                     )
-                current_stream.wait_stream(self.alt_stream)
-            else:
-                q_fp8, q_scale = act_quant(query, self.block_size, self.scale_fmt)
-                self._store_index_k_cache(
-                    forward_batch=forward_batch,
-                    layer_id=layer_id,
-                    key=key,
-                    act_quant=act_quant,
-                )
 
             # `_get_logits_head_gate` expects a Tensor. For tuple activations, dequantize
             # to a float tensor here (callsite), keeping `_get_logits_head_gate` backend-agnostic.
@@ -1366,49 +1514,14 @@ class Indexer(MultiPlatformOp):
                     forward_batch.nsa_cp_metadata is not None
                     and is_nsa_prefill_cp_in_seq_split()
                 ):
-                    kv_len_prev = forward_batch.nsa_cp_metadata.kv_len_prev
-                    kv_len_next = forward_batch.nsa_cp_metadata.kv_len_next
-                    actual_seq_q_prev = forward_batch.nsa_cp_metadata.actual_seq_q_prev
-                    actual_seq_q_next = forward_batch.nsa_cp_metadata.actual_seq_q_next
-
-                    # TODO support mutil-batch
-                    # cp_batch_seq_index_prev = forward_batch.nsa_cp_metadata["cp_batch_seq_index_prev"]
-                    # cp_batch_seq_index_next = forward_batch.nsa_cp_metadata["cp_batch_seq_index_next"]
-                    # TODO prev, next, combined into a single call
-                    q_fp8_prev, q_fp8_next = split_in_seq_cp_local_pair(
+                    return self._get_topk_in_seq_cp_pair(
+                        forward_batch,
+                        layer_id,
                         q_fp8,
-                        actual_seq_q_prev,
-                        actual_seq_q_next,
-                        name="q_fp8",
-                    )
-                    weights_prev, weights_next = split_in_seq_cp_local_pair(
                         weights,
-                        actual_seq_q_prev,
-                        actual_seq_q_next,
-                        name="weights",
-                    )
-                    topk_result_prev = self._get_topk_ragged_with_cp(
-                        forward_batch,
-                        layer_id,
-                        q_fp8_prev,
-                        weights_prev,
                         metadata,
-                        kv_len_prev,
-                        actual_seq_q_prev,
                         current_index_kv=current_index_kv,
                     )
-
-                    topk_result_next = self._get_topk_ragged_with_cp(
-                        forward_batch,
-                        layer_id,
-                        q_fp8_next,
-                        weights_next,
-                        metadata,
-                        kv_len_next,
-                        actual_seq_q_next,
-                        current_index_kv=current_index_kv,
-                    )
-                    return torch.cat([topk_result_prev, topk_result_next], dim=0)
                 else:
                     topk_result = self._get_topk_ragged(
                         enable_dual_stream,

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

@@ -1,4 +1,5 @@
 # temp NSA debugging environ
+import logging
 from dataclasses import dataclass
 from itertools import accumulate
 from typing import TYPE_CHECKING, List, Tuple, Union
@@ -26,6 +27,25 @@ from sglang.srt.utils.common import ceil_align, ceil_div
 if TYPE_CHECKING:
     from sglang.srt.model_executor.forward_batch_info import ForwardBatch
 
+logger = logging.getLogger(__name__)
+
+
+def log_cp_shared_kv_direct_write_fallback(
+    reason: str,
+    message: str,
+    *args,
+) -> None:
+    """Log every direct-write fallback event.
+
+    Warmup can hit the same fallback reason as a later real request, so
+    de-duplicating by reason hides correctness/performance issues after startup.
+    """
+    logger.info(
+        "CP shared KV direct-write fallback (%s): " + message,
+        reason,
+        *args,
+    )
+
 
 def compute_nsa_seqlens(original_seq_lens, nsa_index_topk: int):
     return original_seq_lens.clamp(max=nsa_index_topk)
@@ -204,9 +224,12 @@ def build_page_aligned_in_seq_split_list(
 ) -> Tuple[List[int], PageAlignedInSeqSplitInfo]:
     """Build an in-seq split list whose real-token boundaries do not cut pages.
 
-    Phase 4 deliberately uses a conservative gate: at least `2 * cp_size` page
-    units are required so every zigzag segment has at least one page unit.  When
-    the gate does not hold, this helper falls back to the existing token-balanced
+    Phase 4 deliberately uses a conservative gate for cache-miss chunks: at
+    least `2 * cp_size` page units are required so every zigzag segment has at
+    least one page unit.  For radix-hit suffixes with a page-aligned prefix, the
+    gate is relaxed to `cp_size` page units so every CP rank still receives at
+    least one local page while second zigzag segments may be empty.  When the
+    gate does not hold, this helper falls back to the existing token-balanced
     split and marks the result as not page-aligned.
     """
 
@@ -230,7 +253,9 @@ def build_page_aligned_in_seq_split_list(
     tail_tokens = extend_len % page_size
     num_page_units = full_pages + (1 if tail_tokens > 0 else 0)
     cp_segment_num = cp_size * 2
-    if num_page_units < cp_segment_num:
+    if num_page_units < cp_size or (
+        num_page_units < cp_segment_num and extend_prefix_len == 0
+    ):
         return fallback_split, fallback_info
 
     base_units = num_page_units // cp_segment_num
@@ -303,6 +328,58 @@ def _build_in_seq_split_for_forward_batch(
     )
 
 
+def should_use_replicated_compute_for_short_radix_hit(
+    forward_batch: "ForwardBatch",
+    cp_size: int,
+) -> bool:
+    """Return whether a short radix-hit suffix should avoid CP splitting.
+
+    With CP shared KV, radix-hit suffixes can be page-aligned but shorter than
+    one page per CP rank.  A page-aligned CP split would give some ranks zero
+    local tokens, which is unsafe for parts of the current CP collective/kernel
+    path.  Instead, keep the original non-CP behavior: every rank computes the
+    short suffix, while shared-KV write filters persist only pages owned by the
+    local rank.
+    """
+
+    if (
+        forward_batch is None
+        or cp_size <= 0
+        or not getattr(forward_batch, "uses_cp_shared_kv", False)
+    ):
+        return False
+
+    extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None)
+    extend_prefix_lens_cpu = getattr(forward_batch, "extend_prefix_lens_cpu", None)
+    if (
+        extend_seq_lens_cpu is None
+        or extend_prefix_lens_cpu is None
+        or len(extend_seq_lens_cpu) != 1
+        or len(extend_prefix_lens_cpu) != 1
+    ):
+        return False
+
+    token_to_kv_pool = getattr(forward_batch, "token_to_kv_pool", None)
+    page_size = getattr(token_to_kv_pool, "page_size", None)
+    if page_size is None:
+        return False
+    page_size = int(page_size)
+    if page_size <= 1:
+        return False
+
+    extend_len = int(extend_seq_lens_cpu[0])
+    extend_prefix_len = int(extend_prefix_lens_cpu[0])
+    if (
+        extend_len <= 0
+        or extend_prefix_len <= 0
+        or extend_prefix_len % page_size != 0
+    ):
+        return False
+
+    num_page_units = ceil_div(extend_len, page_size)
+    return 0 < num_page_units < cp_size
+
+
 def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
     if is_nsa_prefill_cp_round_robin_split():
         cur_cp_seq_len = seq_len // cp_size
@@ -313,6 +390,8 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch):
         # TODO current just support prefill batch=1 and len(input_ids) > self.cp_size * 2
         # Note: (self.cp_size * 2) To achieve load balancing for seq computation,
         # 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
         cur_cp_seq_len = seq_len // (cp_size * 2)
     if (
         cur_cp_seq_len != 0
@@ -344,6 +423,110 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor):
     return result
 
 
+def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor):
+    try:
+        round_robin_split = is_nsa_prefill_cp_round_robin_split()
+    except ValueError:
+        round_robin_split = False
+    if round_robin_split:
+        return nsa_cp_round_robin_split_data(input_)
+
+    input_list = list(
+        torch.split(input_, forward_batch.nsa_cp_metadata.split_list, dim=0)
+    )
+    return torch.cat(
+        [input_list[i] for i in forward_batch.nsa_cp_metadata.zigzag_index], dim=0
+    ).view(-1)
+
+
+def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"):
+    """Return this CP rank's local logical out_cache_loc for direct writes.
+
+    `None` means the batch should keep using the compatibility path.  This path
+    is intentionally conservative: it only enables direct writes after Phase 4
+    page-aligned split and after the logical page ids prove they are owned by
+    this CP rank.
+    """
+
+    cached = getattr(forward_batch, "cp_local_out_cache_loc", None)
+    if cached is not None:
+        return cached
+
+    if not getattr(forward_batch, "uses_cp_shared_kv", False):
+        return None
+    metadata = getattr(forward_batch, "nsa_cp_metadata", None)
+    layout = getattr(forward_batch, "cp_shared_kv_layout", None)
+    out_cache_loc = getattr(forward_batch, "out_cache_loc", None)
+    if metadata is None:
+        log_cp_shared_kv_direct_write_fallback(
+            "missing_metadata",
+            "nsa_cp_metadata is missing",
+        )
+        return None
+    if layout is None:
+        log_cp_shared_kv_direct_write_fallback(
+            "missing_layout",
+            "cp_shared_kv_layout is missing",
+        )
+        return None
+    if out_cache_loc is None:
+        log_cp_shared_kv_direct_write_fallback(
+            "missing_out_cache_loc",
+            "out_cache_loc is missing",
+        )
+        return None
+    if not getattr(metadata, "page_aligned", False):
+        log_cp_shared_kv_direct_write_fallback(
+            "not_page_aligned",
+            "metadata is not page-aligned: page_size=%s extend_prefix_len=%s",
+            getattr(metadata, "page_size", None),
+            getattr(metadata, "extend_prefix_len", None),
+        )
+        return None
+    try:
+        in_seq_split = is_nsa_prefill_cp_in_seq_split()
+    except ValueError:
+        in_seq_split = True
+    if not in_seq_split:
+        log_cp_shared_kv_direct_write_fallback(
+            "not_in_seq_split",
+            "nsa_prefill_cp_mode is not in-seq-split",
+        )
+        return None
+    split_tokens = sum(int(x) for x in metadata.split_list)
+    out_cache_tokens = int(out_cache_loc.numel())
+    if split_tokens != out_cache_tokens:
+        log_cp_shared_kv_direct_write_fallback(
+            "split_out_cache_len_mismatch",
+            "split_list tokens=%s out_cache_loc tokens=%s",
+            split_tokens,
+            out_cache_tokens,
+        )
+        return None
+
+    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
+
+    valid_locs = local_out_cache_loc[local_out_cache_loc > 0]
+    if valid_locs.numel() > 0 and not torch.all(layout.owned_by_this_rank(valid_locs)):
+        log_cp_shared_kv_direct_write_fallback(
+            "local_loc_owner_mismatch",
+            "local out_cache_loc contains pages not owned by this rank: cp_rank=%s cp_size=%s page_size=%s",
+            layout.cp_rank,
+            layout.cp_size,
+            layout.page_size,
+        )
+        return None
+
+    forward_batch.cp_local_out_cache_loc = local_out_cache_loc
+    return local_out_cache_loc
+
+
 def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor):
     if is_nsa_prefill_cp_round_robin_split():
         cp_size = get_attention_cp_size()

python/sglang/srt/mem_cache/allocator.py

@@ -20,7 +20,7 @@ Page-aligned memory pool.
 """
 
 import abc
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, List, Optional
 
 import torch
 import triton
@@ -555,3 +555,130 @@ class CPSharedPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
         self.physical_size = physical_size
         self.cp_size = cp_size
         self.cp_rank = cp_rank
+
+    def compute_owner_lane_stats(
+        self,
+        page_compute_owners: List[int],
+    ) -> tuple[List[int], List[int], List[int]]:
+        required = [0 for _ in range(self.cp_size)]
+        for owner in page_compute_owners:
+            if owner < 0 or owner >= self.cp_size:
+                raise ValueError(
+                    f"compute owner must be in [0, {self.cp_size}), got {owner}"
+                )
+            required[owner] += 1
+
+        free_pages = self.free_pages
+        if len(self.release_pages) > 0:
+            free_pages = torch.cat((free_pages, self.release_pages))
+        available = [
+            int(
+                (
+                    torch.remainder(free_pages - 1, self.cp_size) == owner
+                ).sum().item()
+            )
+            for owner in range(self.cp_size)
+        ]
+        deficits = [
+            max(0, required_count - available_count)
+            for required_count, available_count in zip(required, available)
+        ]
+        return required, available, deficits
+
+    def _select_compute_owner_pages(
+        self,
+        page_compute_owners: List[int],
+    ) -> Optional[torch.Tensor]:
+        selected_pages = []
+        lane_offsets = [0 for _ in range(self.cp_size)]
+        lane_pages = [
+            self.free_pages[
+                torch.remainder(self.free_pages - 1, self.cp_size) == owner
+            ]
+            for owner in range(self.cp_size)
+        ]
+
+        for owner in page_compute_owners:
+            if owner < 0 or owner >= self.cp_size:
+                raise ValueError(
+                    f"compute owner must be in [0, {self.cp_size}), got {owner}"
+                )
+            lane_offset = lane_offsets[owner]
+            if lane_offset >= lane_pages[owner].numel():
+                return None
+            selected_pages.append(lane_pages[owner][lane_offset])
+            lane_offsets[owner] = lane_offset + 1
+
+        if not selected_pages:
+            return torch.empty((0,), dtype=torch.int64, device=self.device)
+        return torch.stack(selected_pages).to(torch.int64)
+
+    def alloc_extend_compute_owner(
+        self,
+        prefix_lens: torch.Tensor,
+        prefix_lens_cpu: torch.Tensor,
+        seq_lens: torch.Tensor,
+        seq_lens_cpu: torch.Tensor,
+        last_loc: torch.Tensor,
+        extend_num_tokens: int,
+        page_compute_owners: List[int],
+    ):
+        """Allocate extend KV locs so logical page owner matches CP compute rank.
+
+        The returned logical `out_cache_loc` is still full-order and identical on
+        every CP rank.  Only the chosen logical page ids change: each newly
+        allocated request page comes from the modulo-owner lane that will compute
+        and directly persist that page.
+        """
+
+        if len(prefix_lens_cpu) != 1 or len(seq_lens_cpu) != 1:
+            raise ValueError("compute-owner allocation supports batch size 1 only")
+
+        num_new_pages = get_num_new_pages(
+            seq_lens=seq_lens_cpu,
+            page_size=self.page_size,
+            prefix_lens=prefix_lens_cpu,
+        )
+        if num_new_pages != len(page_compute_owners):
+            raise ValueError(
+                "compute-owner page count mismatch: "
+                f"{num_new_pages=} page_compute_owners={len(page_compute_owners)}"
+            )
+
+        if self.need_sort and num_new_pages > len(self.free_pages):
+            self.merge_and_sort_free()
+
+        selected_pages = self._select_compute_owner_pages(page_compute_owners)
+        if selected_pages is None and self.need_sort and len(self.release_pages) > 0:
+            self.merge_and_sort_free()
+            selected_pages = self._select_compute_owner_pages(page_compute_owners)
+        if selected_pages is None:
+            return None
+
+        out_indices = torch.empty(
+            (extend_num_tokens,), dtype=torch.int64, device=self.device
+        )
+        alloc_extend_naive(
+            prefix_lens_cpu,
+            seq_lens_cpu,
+            last_loc,
+            selected_pages,
+            out_indices,
+            self.page_size,
+            self.device,
+        )
+
+        selected_mask = torch.isin(self.free_pages, selected_pages)
+        self.free_pages = self.free_pages[~selected_mask]
+
+        if self.debug_mode:
+            assert len(torch.unique(out_indices)) == len(out_indices)
+            selected_owners = torch.remainder(selected_pages - 1, self.cp_size)
+            expected_owners = torch.tensor(
+                page_compute_owners,
+                dtype=selected_owners.dtype,
+                device=selected_owners.device,
+            )
+            assert torch.equal(selected_owners, expected_owners)
+
+        return out_indices

python/sglang/srt/mem_cache/common.py

@@ -8,6 +8,10 @@ import triton
 import triton.language as tl
 
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, EvictParams
+from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+    build_in_seq_page_compute_owners,
+    get_in_seq_page_compute_owner_unavailable_reason,
+)
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.server_args import get_global_server_args
@@ -24,6 +28,18 @@ MAMBA_STATE_PER_REQ_NO_CACHE = 1
 logger = logging.getLogger(__name__)
 
 
+def _log_cp_shared_kv_alloc_fallback(
+    reason: str,
+    message: str,
+    *args,
+) -> None:
+    logger.info(
+        "CP shared KV compute-owner allocation fallback (%s): " + message,
+        reason,
+        *args,
+    )
+
+
 @triton.jit
 def write_req_to_token_pool_triton(
     req_to_token_ptr,  # [max_batch, max_context_len]
@@ -252,6 +268,47 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
             tree_cache.evict(EvictParams(num_tokens=num_tokens))
 
 
+def _evict_for_compute_owner_lanes(
+    *,
+    tree_cache: BasePrefixCache | None,
+    allocator,
+    page_compute_owners: list[int],
+) -> None:
+    if tree_cache is None or tree_cache.is_chunk_cache():
+        return
+
+    compute_owner_lane_stats = getattr(allocator, "compute_owner_lane_stats", None)
+    if compute_owner_lane_stats is None:
+        return
+
+    max_attempts = max(2, min(8, int(getattr(allocator, "cp_size", 1))))
+    for _ in range(max_attempts):
+        _required, _available, deficits = compute_owner_lane_stats(page_compute_owners)
+        deficit_pages = sum(deficits)
+        if deficit_pages <= 0:
+            return
+
+        try:
+            evictable_size = tree_cache.evictable_size()
+        except Exception:
+            evictable_size = allocator.page_size
+        if isinstance(evictable_size, tuple):
+            evictable_size = evictable_size[0]
+        if evictable_size <= 0:
+            return
+
+        evict_tokens = max(
+            allocator.page_size,
+            deficit_pages * allocator.page_size * int(getattr(allocator, "cp_size", 1)),
+        )
+        before_available = allocator.available_size()
+        evict_result = tree_cache.evict(EvictParams(num_tokens=evict_tokens))
+        after_available = allocator.available_size()
+        evicted_tokens = getattr(evict_result, "num_tokens_evicted", 0)
+        if after_available <= before_available and evicted_tokens <= 0:
+            return
+
+
 def alloc_paged_token_slots_extend(
     tree_cache: BasePrefixCache,
     prefix_lens: torch.Tensor,
@@ -267,18 +324,114 @@ def alloc_paged_token_slots_extend(
     num_tokens = extend_num_tokens + len(seq_lens_cpu) * allocator.page_size
     evict_from_tree_cache(tree_cache, num_tokens)
 
+    alloc_extend_compute_owner = getattr(
+        allocator, "alloc_extend_compute_owner", None
+    )
+    page_compute_owners = None
+    compute_owner_unavailable_reason = None
+    if alloc_extend_compute_owner is not None and len(prefix_lens_cpu) == 1:
+        try:
+            server_args = get_global_server_args()
+        except ValueError:
+            server_args = None
+        if (
+            server_args is not None
+            and server_args.enable_nsa_prefill_cp_shared_kv
+            and server_args.enable_nsa_prefill_context_parallel
+            and server_args.nsa_prefill_cp_mode == "in-seq-split"
+        ):
+            extend_len = int(seq_lens_cpu[0].item() - prefix_lens_cpu[0].item())
+            page_compute_owners = build_in_seq_page_compute_owners(
+                extend_len=extend_len,
+                extend_prefix_len=int(prefix_lens_cpu[0].item()),
+                page_size=int(allocator.page_size),
+                cp_size=int(allocator.cp_size),
+            )
+            if page_compute_owners is None:
+                compute_owner_unavailable_reason = (
+                    get_in_seq_page_compute_owner_unavailable_reason(
+                        extend_len=extend_len,
+                        extend_prefix_len=int(prefix_lens_cpu[0].item()),
+                        page_size=int(allocator.page_size),
+                        cp_size=int(allocator.cp_size),
+                    )
+                    or "unknown"
+                )
+        else:
+            compute_owner_unavailable_reason = "server_args_not_enabled"
+    elif alloc_extend_compute_owner is not None:
+        compute_owner_unavailable_reason = (
+            "multi_batch" if len(prefix_lens_cpu) != 1 else "unknown"
+        )
+
+    if page_compute_owners is not None:
+        _evict_for_compute_owner_lanes(
+            tree_cache=tree_cache,
+            allocator=allocator,
+            page_compute_owners=page_compute_owners,
+        )
+
     state = None
     if backup_state:
         state = allocator.backup_state()
 
-    out_cache_loc = allocator.alloc_extend(
-        prefix_lens,
-        prefix_lens_cpu,
-        seq_lens,
-        seq_lens_cpu,
-        last_loc,
-        extend_num_tokens,
-    )
+    if page_compute_owners is not None:
+        out_cache_loc = alloc_extend_compute_owner(
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+            extend_num_tokens,
+            page_compute_owners,
+        )
+        if out_cache_loc is None:
+            required = available = deficits = None
+            compute_owner_lane_stats = getattr(
+                allocator, "compute_owner_lane_stats", None
+            )
+            if compute_owner_lane_stats is not None:
+                required, available, deficits = compute_owner_lane_stats(
+                    page_compute_owners
+                )
+            _log_cp_shared_kv_alloc_fallback(
+                "owner_lane_exhausted",
+                "failed to allocate pages from compute-owner lanes; "
+                "falling back to legacy page allocation. extend_num_tokens=%s page_size=%s "
+                "required_by_owner=%s available_by_owner=%s deficit_by_owner=%s",
+                extend_num_tokens,
+                allocator.page_size,
+                required,
+                available,
+                deficits,
+            )
+            out_cache_loc = allocator.alloc_extend(
+                prefix_lens,
+                prefix_lens_cpu,
+                seq_lens,
+                seq_lens_cpu,
+                last_loc,
+                extend_num_tokens,
+            )
+    else:
+        if alloc_extend_compute_owner is not None:
+            _log_cp_shared_kv_alloc_fallback(
+                compute_owner_unavailable_reason or "compute_owner_not_available",
+                "page-aligned compute-owner page assignment is unavailable; "
+                "falling back to legacy page allocation. batch_size=%s extend_num_tokens=%s page_size=%s reason=%s",
+                len(prefix_lens_cpu),
+                extend_num_tokens,
+                allocator.page_size,
+                compute_owner_unavailable_reason or "compute_owner_not_available",
+            )
+        out_cache_loc = allocator.alloc_extend(
+            prefix_lens,
+            prefix_lens_cpu,
+            seq_lens,
+            seq_lens_cpu,
+            last_loc,
+            extend_num_tokens,
+        )
 
     if out_cache_loc is None:
         error_msg = (

python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+from typing import List, Optional
+
+
+def get_in_seq_page_compute_owner_unavailable_reason(
+    *,
+    extend_len: int,
+    extend_prefix_len: int,
+    page_size: int,
+    cp_size: int,
+) -> Optional[str]:
+    if cp_size <= 0:
+        raise ValueError(f"cp_size must be positive, got {cp_size}")
+    if extend_len < 0:
+        raise ValueError(f"extend_len must be non-negative, got {extend_len}")
+    if page_size <= 1:
+        return "page_size_le_one"
+    if extend_len <= 0:
+        return "empty_extend"
+    if extend_prefix_len % page_size != 0:
+        return "prefix_not_page_aligned"
+
+    full_pages = extend_len // page_size
+    tail_tokens = extend_len % page_size
+    num_page_units = full_pages + (1 if tail_tokens > 0 else 0)
+    if num_page_units < cp_size * 2 and extend_prefix_len == 0:
+        return "too_short_for_page_aligned"
+
+    return None
+
+
+def build_in_seq_page_compute_owners(
+    *,
+    extend_len: int,
+    extend_prefix_len: int,
+    page_size: int,
+    cp_size: int,
+) -> Optional[List[int]]:
+    """Return compute-owner CP rank for each newly allocated current page.
+
+    This mirrors the Phase 4 page-aligned `in-seq-split` segmentation for
+    normal CP chunks, but it only returns page-unit owners for the real extend
+    chunk.  Short radix-hit suffixes with fewer pages than CP ranks are also
+    allowed: runtime keeps replicated compute for those chunks and the shared
+    KV write filters persist only locally owned pages.  `None` means the batch
+    must stay on the legacy allocation/write path.
+    """
+
+    if cp_size <= 0:
+        raise ValueError(f"cp_size must be positive, got {cp_size}")
+    if extend_len < 0:
+        raise ValueError(f"extend_len must be non-negative, got {extend_len}")
+    if (
+        get_in_seq_page_compute_owner_unavailable_reason(
+            extend_len=extend_len,
+            extend_prefix_len=extend_prefix_len,
+            page_size=page_size,
+            cp_size=cp_size,
+        )
+        is not None
+    ):
+        return None
+
+    full_pages = extend_len // page_size
+    tail_tokens = extend_len % page_size
+    num_page_units = full_pages + (1 if tail_tokens > 0 else 0)
+    cp_segment_num = cp_size * 2
+
+    base_units = num_page_units // cp_segment_num
+    remainder_units = num_page_units % cp_segment_num
+    owners: List[int] = []
+    for segment_idx in range(cp_segment_num):
+        unit_count = base_units + (1 if segment_idx < remainder_units else 0)
+        if segment_idx < cp_size:
+            owner = segment_idx
+        else:
+            owner = cp_segment_num - segment_idx - 1
+        owners.extend([owner] * unit_count)
+
+    return owners

python/sglang/srt/model_executor/forward_batch_info.py

@@ -422,6 +422,8 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
     nsa_cp_metadata: Optional[NSAContextParallelMetadata] = None
     uses_cp_shared_kv: bool = False
     cp_shared_kv_layout: Optional[CpSharedKVLayout] = None
+    cp_local_out_cache_loc: Optional[torch.Tensor] = None
+    cp_shared_mla_direct_write_done: bool = False
 
     # For hidden states before normal
     return_hidden_states_before_norm: bool = False

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

@@ -860,7 +860,7 @@ class ModelRunnerKVCacheMixin:
 
         if self.server_args.enable_nsa_prefill_cp_shared_kv:
             logger.info(
-                "CP shared KV enabled. physical_tokens_per_rank=%s, logical_tokens=%s, cp_size=%s, shard_policy=page_interleaved",
+                "CP shared KV enabled. physical_tokens_per_rank=%s, logical_tokens=%s, cp_size=%s, shard_policy=compute_owner_page_aligned_when_available",
                 self.physical_max_total_num_tokens,
                 self.max_total_num_tokens,
                 self.server_args.attn_cp_size,

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

@@ -6,7 +6,11 @@ import torch
 
 from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.layers import deep_gemm_wrapper
-from sglang.srt.layers.attention.nsa.utils import nsa_use_prefill_cp
+from sglang.srt.layers.attention.nsa.utils import (
+    get_cp_shared_kv_local_out_cache_loc,
+    log_cp_shared_kv_direct_write_fallback,
+    nsa_use_prefill_cp,
+)
 from sglang.srt.layers.communicator import get_attn_tp_context
 from sglang.srt.layers.quantization.fp8_kernel import (
     fp8_dtype,
@@ -299,7 +303,16 @@ class DeepseekMLAForwardMixin:
         ):
             q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
 
+        shared_mla_direct_write_done = False
         if nsa_use_prefill_cp(forward_batch):
+            shared_mla_direct_write_done = self._maybe_write_cp_shared_local_mla_kv(
+                forward_batch,
+                k_nope,
+                k_pe,
+            )
+            forward_batch.cp_shared_mla_direct_write_done = (
+                shared_mla_direct_write_done
+            )
             # support allgather+rerrange
             k_nope, k_pe = self.rebuild_cp_kv_cache(
                 latent_cache, forward_batch, k_nope, k_pe
@@ -329,7 +342,9 @@ class DeepseekMLAForwardMixin:
         topk_indices,
         llama_4_scaling,
     ):
-        save_kv_cache = True
+        save_kv_cache = not getattr(
+            forward_batch, "cp_shared_mla_direct_write_done", False
+        )
 
         if self.current_attention_backend in FORWARD_ABSORB_CORE_ATTENTION_BACKENDS:
             extra_args = {}
@@ -345,6 +360,7 @@ class DeepseekMLAForwardMixin:
                 k_nope,
                 k_nope,
                 forward_batch,
+                save_kv_cache=save_kv_cache,
                 q_rope=q_pe,
                 k_rope=k_pe,
                 **extra_args,
@@ -512,6 +528,46 @@ class DeepseekMLAForwardMixin:
 
         return output
 
+    def _maybe_write_cp_shared_local_mla_kv(
+        self: DeepseekV2AttentionMLA,
+        forward_batch: ForwardBatch,
+        k_nope: torch.Tensor,
+        k_pe: torch.Tensor,
+    ) -> bool:
+        local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+        if local_out_cache_loc is None:
+            return False
+        if (
+            k_nope.shape[0] != local_out_cache_loc.numel()
+            or k_pe.shape[0] != local_out_cache_loc.numel()
+        ):
+            log_cp_shared_kv_direct_write_fallback(
+                "mla_local_shape_mismatch",
+                "MLA local KV token count does not match local out_cache_loc: "
+                "k_nope=%s k_pe=%s local_out_cache_loc=%s layer_id=%s",
+                k_nope.shape[0],
+                k_pe.shape[0],
+                local_out_cache_loc.numel(),
+                self.attn_mqa.layer_id,
+            )
+            return False
+        if local_out_cache_loc.numel() == 0:
+            return True
+
+        assert forward_batch.cp_shared_kv_layout is not None
+        physical_out_cache_loc = (
+            forward_batch.cp_shared_kv_layout.logical_locs_to_physical(
+                local_out_cache_loc
+            ).contiguous()
+        )
+        forward_batch.token_to_kv_pool.set_mla_kv_buffer(
+            self.attn_mqa,
+            physical_out_cache_loc,
+            k_nope,
+            k_pe,
+        )
+        return True
+
     def _fuse_rope_for_trtllm_mla(
         self: DeepseekV2AttentionMLA, forward_batch: ForwardBatch
     ) -> bool:

test/registered/unit/layers/test_nsa_cp_utils.py

@@ -1,11 +1,18 @@
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
 from sglang.srt.layers.attention.nsa.utils import (
+    NSAContextParallelMetadata,
     _get_in_seq_last_token_owner_and_offset,
     build_page_aligned_in_seq_split_list,
     build_token_balanced_in_seq_split_list,
+    can_cp_split,
+    cp_split_and_rebuild_1d,
+    get_cp_shared_kv_local_out_cache_loc,
     split_in_seq_cp_local_pair,
 )
+from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=1, suite="stage-a-test-cpu")
@@ -93,6 +100,85 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         self.assertFalse(info.page_aligned)
         self.assertEqual(split_list, build_token_balanced_in_seq_split_list(512, 8))
 
+    def test_page_aligned_split_allows_radix_hit_suffix_with_one_page_per_rank(self):
+        split_list, info = build_page_aligned_in_seq_split_list(
+            total_len=512,
+            extend_len=512,
+            extend_prefix_len=54464,
+            page_size=64,
+            cp_size=8,
+        )
+
+        self.assertTrue(info.page_aligned)
+        self.assertEqual(sum(split_list), 512)
+        self.assertEqual(split_list[:8], [64] * 8)
+        self.assertEqual(split_list[8:], [0] * 8)
+        self.assert_page_aligned_boundaries(
+            split_list, extend_prefix_len=54464, extend_len=512, page_size=64
+        )
+
+    def test_page_aligned_split_falls_back_when_radix_hit_suffix_has_zero_rank(self):
+        split_list, info = build_page_aligned_in_seq_split_list(
+            total_len=256,
+            extend_len=256,
+            extend_prefix_len=54464,
+            page_size=64,
+            cp_size=8,
+        )
+
+        self.assertFalse(info.page_aligned)
+        self.assertEqual(split_list, build_token_balanced_in_seq_split_list(256, 8))
+
+    def test_can_cp_split_uses_replicated_compute_for_short_radix_hit_suffix(self):
+        class Mode:
+            def is_context_parallel_extend(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            extend_seq_lens_cpu=[256],
+            extend_prefix_lens_cpu=[54464],
+            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.assertFalse(can_cp_split(256, 8, True, forward_batch))
+
+    def test_can_cp_split_keeps_cp_for_radix_hit_suffix_with_one_page_per_rank(self):
+        class Mode:
+            def is_context_parallel_extend(self):
+                return True
+
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            extend_seq_lens_cpu=[512],
+            extend_prefix_lens_cpu=[54464],
+            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(512, 8, True, forward_batch))
+
     def test_page_aligned_split_adds_padding_tokens_to_last_segment(self):
         split_list, info = build_page_aligned_in_seq_split_list(
             total_len=1040,
@@ -152,6 +238,309 @@ class TestNSAInSeqCPUtils(unittest.TestCase):
         with self.assertRaisesRegex(RuntimeError, "local in-seq CP length mismatch"):
             split_in_seq_cp_local_pair(torch.arange(9), 5, 5, name="q_fp8")
 
+    def test_cp_split_and_rebuild_1d_matches_in_seq_zigzag_order(self):
+        import torch
+        from types import SimpleNamespace
+
+        forward_batch = SimpleNamespace(
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[2, 2, 2, 2, 2, 2, 2, 2],
+                zigzag_index=[1, 6],
+            )
+        )
+
+        local_locs = cp_split_and_rebuild_1d(forward_batch, torch.arange(16))
+
+        self.assertEqual(local_locs.tolist(), [2, 3, 12, 13])
+
+    def test_local_out_cache_loc_requires_compute_owner_pages(self):
+        import torch
+        from types import SimpleNamespace
+
+        page_size = 4
+        # Segment order for cp_size=4, cp_rank=1 is segment 1 then 6.
+        # The logical page ids below deliberately encode the same owners through
+        # (logical_page - 1) % cp_size:
+        #   segment 1 -> page 2 owner 1
+        #   segment 6 -> page 6 owner 1
+        segment_pages = [1, 2, 3, 4, 8, 7, 6, 5]
+        out_cache_loc = torch.cat(
+            [
+                torch.arange(page * page_size, (page + 1) * page_size)
+                for page in segment_pages
+            ]
+        )
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+        )
+
+        local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch)
+
+        self.assertIsNotNone(local_locs)
+        self.assertEqual(
+            local_locs.tolist(),
+            list(range(2 * page_size, 3 * page_size))
+            + list(range(6 * page_size, 7 * page_size)),
+        )
+
+    def test_local_out_cache_loc_falls_back_when_owner_mismatch(self):
+        import torch
+        from types import SimpleNamespace
+
+        page_size = 4
+        out_cache_loc = torch.arange(page_size * 8, page_size * 16)
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=True,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=out_cache_loc,
+        )
+
+        self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
+
+    def test_local_out_cache_loc_logs_every_fallback_event(self):
+        import torch
+        from types import SimpleNamespace
+
+        from sglang.srt.layers.attention.nsa import utils as nsa_utils
+
+        page_size = 4
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            cp_shared_kv_layout=CpSharedKVLayout(
+                page_size=page_size,
+                cp_size=4,
+                cp_rank=1,
+            ),
+            nsa_cp_metadata=NSAContextParallelMetadata(
+                split_list=[page_size] * 8,
+                zigzag_index=[1, 6],
+                page_aligned=False,
+                page_size=page_size,
+                extend_prefix_len=0,
+            ),
+            out_cache_loc=torch.arange(page_size * 8, page_size * 16),
+        )
+
+        with self.assertLogs(
+            "sglang.srt.layers.attention.nsa.utils", level="INFO"
+        ) as cm:
+            self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
+            self.assertIsNone(get_cp_shared_kv_local_out_cache_loc(forward_batch))
+
+        self.assertEqual(len(cm.output), 2)
+        self.assertIn("CP shared KV direct-write fallback", cm.output[0])
+        self.assertIn("metadata is not page-aligned", cm.output[0])
+        self.assertIn("CP shared KV direct-write fallback", cm.output[1])
+        self.assertIn("metadata is not page-aligned", cm.output[1])
+
+    def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self):
+        import torch
+        from types import SimpleNamespace
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        indexer.nsa_enable_prefill_cp = True
+        forward_batch = SimpleNamespace(
+            uses_cp_shared_kv=True,
+            nsa_cp_metadata=None,
+        )
+
+        with self.assertNoLogs(
+            "sglang.srt.layers.attention.nsa.utils", level="INFO"
+        ):
+            stored = Indexer._store_cp_shared_local_index_k_cache(
+                indexer,
+                forward_batch,
+                layer_id=0,
+                local_key=torch.empty(0),
+                act_quant=None,
+            )
+
+        self.assertFalse(stored)
+
+    def test_indexer_in_seq_cp_pair_materializes_index_once_for_prev_next(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        logical_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)
+        materialized_index = torch.tensor([11], dtype=torch.int32)
+        dense_pages = torch.tensor([[1, 2, 3, 4]], dtype=torch.int32)
+        materialize_calls = []
+        topk_calls = []
+
+        class Metadata:
+            def get_page_table_64(self):
+                return logical_pages
+
+        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,
+        ):
+            topk_calls.append(
+                {
+                    "kv_len": kv_len,
+                    "actual_seq_q": actual_seq_q,
+                    "shared_index_buffer": shared_index_buffer,
+                    "shared_block_tables": shared_block_tables,
+                    "current_index_kv": current_index_kv,
+                }
+            )
+            return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
+
+        indexer._maybe_materialize_shared_index_buffer = fake_materialize
+        indexer._get_topk_ragged_with_cp = fake_get_topk
+
+        forward_batch = type(
+            "ForwardBatchStub",
+            (),
+            {
+                "nsa_cp_metadata": NSAContextParallelMetadata(
+                    kv_len_prev=5,
+                    kv_len_next=9,
+                    actual_seq_q_prev=3,
+                    actual_seq_q_next=2,
+                )
+            },
+        )()
+        q_fp8 = torch.arange(5 * 4, dtype=torch.float32).view(5, 4)
+        weights = torch.arange(5 * 2, dtype=torch.float32).view(5, 2)
+
+        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), 2)
+        self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index)
+        self.assertIs(topk_calls[1]["shared_index_buffer"], materialized_index)
+        self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages)
+        self.assertIs(topk_calls[1]["shared_block_tables"], dense_pages)
+        self.assertIsNone(topk_calls[0]["current_index_kv"])
+        self.assertEqual(topk_calls[0]["kv_len"], 5)
+        self.assertEqual(topk_calls[1]["kv_len"], 9)
+        self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
+
+    def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self):
+        import torch
+
+        from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
+
+        indexer = object.__new__(Indexer)
+        current_index_kv = (torch.tensor([1]), torch.tensor([2]))
+        materialize_calls = []
+        topk_calls = []
+
+        class Metadata:
+            def get_page_table_64(self):
+                raise AssertionError("current index reuse should not read page table")
+
+        def fake_materialize(forward_batch, layer_id, logical_page_table):
+            materialize_calls.append((layer_id, logical_page_table))
+            raise AssertionError("current index reuse should not materialize")
+
+        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,
+        ):
+            topk_calls.append(
+                {
+                    "current_index_kv": current_index_kv,
+                    "shared_index_buffer": shared_index_buffer,
+                    "shared_block_tables": shared_block_tables,
+                }
+            )
+            return torch.full((actual_seq_q, 2), len(topk_calls), dtype=torch.int32)
+
+        indexer._maybe_materialize_shared_index_buffer = fake_materialize
+        indexer._get_topk_ragged_with_cp = fake_get_topk
+
+        forward_batch = type(
+            "ForwardBatchStub",
+            (),
+            {
+                "nsa_cp_metadata": NSAContextParallelMetadata(
+                    kv_len_prev=5,
+                    kv_len_next=9,
+                    actual_seq_q_prev=3,
+                    actual_seq_q_next=2,
+                )
+            },
+        )()
+
+        result = Indexer._get_topk_in_seq_cp_pair(
+            indexer,
+            forward_batch,
+            layer_id=7,
+            q_fp8=torch.empty(5, 4),
+            weights=torch.empty(5, 2),
+            metadata=Metadata(),
+            current_index_kv=current_index_kv,
+        )
+
+        self.assertEqual(materialize_calls, [])
+        self.assertEqual(len(topk_calls), 2)
+        self.assertIs(topk_calls[0]["current_index_kv"], current_index_kv)
+        self.assertIs(topk_calls[1]["current_index_kv"], current_index_kv)
+        self.assertIsNone(topk_calls[0]["shared_index_buffer"])
+        self.assertIsNone(topk_calls[1]["shared_block_tables"])
+        self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]])
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_cp_shared_kv_layout.py

@@ -52,6 +52,27 @@ class TestCpSharedKVLayout(unittest.TestCase):
 
 
 class TestCPSharedPagedAllocator(unittest.TestCase):
+    def test_compute_owner_alloc_fallback_logs_every_event(self):
+        from sglang.srt.mem_cache import common
+
+        with self.assertLogs("sglang.srt.mem_cache.common", level="INFO") as cm:
+            common._log_cp_shared_kv_alloc_fallback(
+                "too_short_for_page_aligned",
+                "falling back for test event %s",
+                1,
+            )
+            common._log_cp_shared_kv_alloc_fallback(
+                "too_short_for_page_aligned",
+                "falling back for test event %s",
+                2,
+            )
+
+        self.assertEqual(len(cm.output), 2)
+        self.assertIn("too_short_for_page_aligned", cm.output[0])
+        self.assertIn("test event 1", cm.output[0])
+        self.assertIn("too_short_for_page_aligned", cm.output[1])
+        self.assertIn("test event 2", cm.output[1])
+
     def test_shared_allocator_exposes_logical_capacity(self):
         from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
 
@@ -73,6 +94,229 @@ class TestCPSharedPagedAllocator(unittest.TestCase):
         allocator.free(locs)
         self.assertEqual(allocator.available_size(), 64 * 8)
 
+    def test_compute_owner_page_assignment_matches_in_seq_zigzag(self):
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            build_in_seq_page_compute_owners,
+        )
+
+        owners = build_in_seq_page_compute_owners(
+            extend_len=64 * 16,
+            extend_prefix_len=0,
+            page_size=64,
+            cp_size=4,
+        )
+
+        self.assertEqual(
+            owners,
+            [0, 0, 1, 1, 2, 2, 3, 3, 3, 3, 2, 2, 1, 1, 0, 0],
+        )
+
+    def test_compute_owner_page_assignment_falls_back_for_short_extend(self):
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            build_in_seq_page_compute_owners,
+            get_in_seq_page_compute_owner_unavailable_reason,
+        )
+
+        owners = build_in_seq_page_compute_owners(
+            extend_len=64 * 7,
+            extend_prefix_len=0,
+            page_size=64,
+            cp_size=4,
+        )
+
+        self.assertIsNone(owners)
+        self.assertEqual(
+            get_in_seq_page_compute_owner_unavailable_reason(
+                extend_len=64 * 7,
+                extend_prefix_len=0,
+                page_size=64,
+                cp_size=4,
+            ),
+            "too_short_for_page_aligned",
+        )
+
+    def test_compute_owner_page_assignment_allows_radix_hit_suffix_with_one_page_per_rank(
+        self,
+    ):
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            build_in_seq_page_compute_owners,
+            get_in_seq_page_compute_owner_unavailable_reason,
+        )
+
+        owners = build_in_seq_page_compute_owners(
+            extend_len=64 * 4,
+            extend_prefix_len=54464,
+            page_size=64,
+            cp_size=4,
+        )
+
+        self.assertEqual(owners, [0, 1, 2, 3])
+        self.assertIsNone(
+            get_in_seq_page_compute_owner_unavailable_reason(
+                extend_len=64 * 4,
+                extend_prefix_len=54464,
+                page_size=64,
+                cp_size=4,
+            )
+        )
+
+    def test_compute_owner_page_assignment_allows_short_radix_hit_suffix_with_replicated_compute(
+        self,
+    ):
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            build_in_seq_page_compute_owners,
+            get_in_seq_page_compute_owner_unavailable_reason,
+        )
+
+        owners = build_in_seq_page_compute_owners(
+            extend_len=64 * 3,
+            extend_prefix_len=54464,
+            page_size=64,
+            cp_size=4,
+        )
+
+        self.assertEqual(owners, [0, 1, 2])
+        self.assertIsNone(
+            get_in_seq_page_compute_owner_unavailable_reason(
+                extend_len=64 * 3,
+                extend_prefix_len=54464,
+                page_size=64,
+                cp_size=4,
+            )
+        )
+
+    def test_compute_owner_page_assignment_reports_prefix_misalignment(self):
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            get_in_seq_page_compute_owner_unavailable_reason,
+        )
+
+        self.assertEqual(
+            get_in_seq_page_compute_owner_unavailable_reason(
+                extend_len=64 * 16,
+                extend_prefix_len=1,
+                page_size=64,
+                cp_size=4,
+            ),
+            "prefix_not_page_aligned",
+        )
+
+    def test_shared_allocator_can_allocate_pages_from_compute_owner_lanes(self):
+        from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
+        from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
+            build_in_seq_page_compute_owners,
+        )
+
+        page_size = 64
+        cp_size = 4
+        owners = build_in_seq_page_compute_owners(
+            extend_len=page_size * 16,
+            extend_prefix_len=0,
+            page_size=page_size,
+            cp_size=cp_size,
+        )
+        allocator = CPSharedPagedTokenToKVPoolAllocator(
+            logical_size=page_size * 32,
+            physical_size=page_size * 8,
+            page_size=page_size,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+            cp_size=cp_size,
+            cp_rank=0,
+        )
+
+        locs = allocator.alloc_extend_compute_owner(
+            prefix_lens=torch.tensor([0], dtype=torch.int64),
+            prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+            seq_lens=torch.tensor([page_size * 16], dtype=torch.int64),
+            seq_lens_cpu=torch.tensor([page_size * 16], dtype=torch.int64),
+            last_loc=torch.tensor([-1], dtype=torch.int64),
+            extend_num_tokens=page_size * 16,
+            page_compute_owners=owners,
+        )
+
+        self.assertIsNotNone(locs)
+        logical_pages = locs.view(-1, page_size)[:, 0] // page_size
+        self.assertEqual(
+            ((logical_pages - 1) % cp_size).tolist(),
+            owners,
+        )
+        self.assertEqual(allocator.available_size(), page_size * 16)
+
+    def test_compute_owner_lane_eviction_recovers_exhausted_owner_lane(self):
+        from sglang.srt.mem_cache.allocator import CPSharedPagedTokenToKVPoolAllocator
+        from sglang.srt.mem_cache.base_prefix_cache import EvictResult
+        from sglang.srt.mem_cache.common import _evict_for_compute_owner_lanes
+
+        page_size = 64
+        cp_size = 4
+        allocator = CPSharedPagedTokenToKVPoolAllocator(
+            logical_size=page_size * 16,
+            physical_size=page_size * 4,
+            page_size=page_size,
+            dtype=torch.bfloat16,
+            device="cpu",
+            kvcache=None,
+            need_sort=False,
+            cp_size=cp_size,
+            cp_rank=0,
+        )
+        lane0_locs = allocator.alloc_extend_compute_owner(
+            prefix_lens=torch.tensor([0], dtype=torch.int64),
+            prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+            seq_lens=torch.tensor([page_size * 4], dtype=torch.int64),
+            seq_lens_cpu=torch.tensor([page_size * 4], dtype=torch.int64),
+            last_loc=torch.tensor([-1], dtype=torch.int64),
+            extend_num_tokens=page_size * 4,
+            page_compute_owners=[0, 0, 0, 0],
+        )
+        self.assertIsNone(
+            allocator.alloc_extend_compute_owner(
+                prefix_lens=torch.tensor([0], dtype=torch.int64),
+                prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+                seq_lens=torch.tensor([page_size * 4], dtype=torch.int64),
+                seq_lens_cpu=torch.tensor([page_size * 4], dtype=torch.int64),
+                last_loc=torch.tensor([-1], dtype=torch.int64),
+                extend_num_tokens=page_size * 4,
+                page_compute_owners=[0, 0, 0, 0],
+            )
+        )
+
+        class FakeTreeCache:
+            def __init__(self):
+                self.evict_calls = []
+
+            def is_chunk_cache(self):
+                return False
+
+            def evictable_size(self):
+                return lane0_locs.numel()
+
+            def evict(self, params):
+                self.evict_calls.append(params.num_tokens)
+                allocator.free(lane0_locs)
+                return EvictResult(num_tokens_evicted=lane0_locs.numel())
+
+        tree_cache = FakeTreeCache()
+        _evict_for_compute_owner_lanes(
+            tree_cache=tree_cache,
+            allocator=allocator,
+            page_compute_owners=[0, 0, 0, 0],
+        )
+
+        self.assertGreaterEqual(len(tree_cache.evict_calls), 1)
+        locs = allocator.alloc_extend_compute_owner(
+            prefix_lens=torch.tensor([0], dtype=torch.int64),
+            prefix_lens_cpu=torch.tensor([0], dtype=torch.int64),
+            seq_lens=torch.tensor([page_size * 4], dtype=torch.int64),
+            seq_lens_cpu=torch.tensor([page_size * 4], dtype=torch.int64),
+            last_loc=torch.tensor([-1], dtype=torch.int64),
+            extend_num_tokens=page_size * 4,
+            page_compute_owners=[0, 0, 0, 0],
+        )
+        self.assertIsNotNone(locs)
+
 
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py

@@ -612,5 +612,224 @@ class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
         self.assertEqual(key_to_write.shape[0], 2)
 
 
+class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase):
+    def test_paged_materialize_uses_tai_kernel_when_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakeTaiKernels:
+            def __init__(self):
+                self.calls = []
+
+            def materialize_shared_pages(
+                self,
+                page_buffer,
+                logical_pages,
+                *,
+                cp_rank,
+                cp_size,
+            ):
+                self.calls.append((page_buffer, logical_pages, cp_rank, cp_size))
+                return (
+                    torch.full((logical_pages.numel() + 1, 3), 7, dtype=torch.uint8),
+                    torch.tensor([1, 0, 3], dtype=logical_pages.dtype),
+                )
+
+        fake_tai = FakeTaiKernels()
+        layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
+        page_buffer = torch.arange(0, 5 * 3, dtype=torch.uint8).view(5, 3)
+        logical_pages = torch.tensor([1, 0, 3], dtype=torch.int64)
+
+        with patch(
+            "sglang.srt.environ.envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get",
+            return_value=True,
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=False
+        ), patch.object(
+            runtime, "_load_tai_materialize_kernels", return_value=fake_tai
+        ), patch.object(
+            runtime, "_all_reduce_materialized_buffer", lambda x, _: x
+        ):
+            dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
+                page_buffer=page_buffer,
+                logical_pages=logical_pages,
+                layout=layout,
+            )
+
+        self.assertEqual(len(fake_tai.calls), 1)
+        self.assertIs(fake_tai.calls[0][0], page_buffer)
+        self.assertIs(fake_tai.calls[0][1], logical_pages)
+        self.assertEqual(fake_tai.calls[0][2:], (1, 2))
+        self.assertEqual(dense_pages.tolist(), [1, 0, 3])
+        self.assertEqual(int(dense_page_buffer.sum().item()), 7 * 4 * 3)
+
+    def test_token_materialize_uses_tai_kernel_for_slot_remap_when_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakeTaiKernels:
+            def __init__(self):
+                self.page_inverse_calls = []
+                self.remap_calls = []
+                self.token_calls = []
+
+            def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
+                self.page_inverse_calls.append((slot_logical_pages, logical_page_capacity))
+                return torch.tensor([0, 1, 2, -1, 3], dtype=torch.long)
+
+            def remap_logical_locs_to_slot_dense_locs(
+                self,
+                logical_locs,
+                page_inverse,
+                *,
+                page_size,
+            ):
+                self.remap_calls.append((logical_locs, page_inverse, page_size))
+                return torch.tensor([4, 8, -1], dtype=logical_locs.dtype)
+
+            def materialize_shared_token_kv_pages(
+                self,
+                kv_cache,
+                slot_logical_pages,
+                *,
+                page_size,
+                cp_rank,
+                cp_size,
+            ):
+                self.token_calls.append(
+                    (kv_cache, slot_logical_pages, page_size, cp_rank, cp_size)
+                )
+                return torch.full((16, 1, 1), 3.0, dtype=kv_cache.dtype)
+
+        fake_tai = FakeTaiKernels()
+        layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
+        kv_cache = torch.arange(0, 24, dtype=torch.float32).view(24, 1, 1)
+        logical_locs = torch.tensor([4, 8, -1], dtype=torch.int64)
+        remap_logical_pages = torch.tensor([[1, 2, 4]], dtype=torch.int64)
+
+        with patch(
+            "sglang.srt.environ.envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get",
+            return_value=True,
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=False
+        ), patch.object(
+            runtime, "_load_tai_materialize_kernels", return_value=fake_tai
+        ), patch.object(
+            runtime, "_all_reduce_materialized_buffer", lambda x, _: x
+        ):
+            dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
+                kv_cache=kv_cache,
+                logical_locs=logical_locs,
+                remap_logical_pages=remap_logical_pages,
+                layout=layout,
+                page_size=4,
+            )
+
+        self.assertEqual(len(fake_tai.page_inverse_calls), 1)
+        self.assertTrue(
+            torch.equal(
+                fake_tai.page_inverse_calls[0][0],
+                remap_logical_pages.reshape(-1),
+            )
+        )
+        self.assertEqual(len(fake_tai.remap_calls), 1)
+        self.assertEqual(len(fake_tai.token_calls), 1)
+        self.assertIs(fake_tai.token_calls[0][0], kv_cache)
+        self.assertTrue(
+            torch.equal(fake_tai.token_calls[0][1], remap_logical_pages.reshape(-1))
+        )
+        self.assertEqual(fake_tai.token_calls[0][2:], (4, 1, 2))
+        self.assertEqual(dense_locs.tolist(), [4, 8, -1])
+        self.assertEqual(float(dense_kv.sum().item()), 48.0)
+
+    def test_token_tai_path_skips_torch_slot_page_remap(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        class FakeTaiKernels:
+            def build_slot_page_inverse(self, slot_logical_pages, logical_page_capacity):
+                return torch.tensor([0, 1, 2, -1, 3], dtype=torch.long)
+
+            def remap_logical_locs_to_slot_dense_locs(
+                self,
+                logical_locs,
+                page_inverse,
+                *,
+                page_size,
+            ):
+                return torch.tensor([4, 8, -1], dtype=logical_locs.dtype)
+
+            def materialize_shared_token_kv_pages(
+                self,
+                kv_cache,
+                slot_logical_pages,
+                *,
+                page_size,
+                cp_rank,
+                cp_size,
+            ):
+                return torch.full((16, 1, 1), 3.0, dtype=kv_cache.dtype)
+
+        layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
+        kv_cache = torch.arange(0, 24, dtype=torch.float32).view(24, 1, 1)
+        logical_locs = torch.tensor([4, 8, -1], dtype=torch.int64)
+        remap_logical_pages = torch.tensor([[1, 2, 4]], dtype=torch.int64)
+
+        with patch(
+            "sglang.srt.environ.envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get",
+            return_value=True,
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=False
+        ), patch.object(
+            runtime, "_load_tai_materialize_kernels", return_value=FakeTaiKernels()
+        ), patch.object(
+            runtime,
+            "build_slot_page_remap",
+            side_effect=AssertionError("tai token path must not run torch remap"),
+        ), patch.object(
+            runtime, "_all_reduce_materialized_buffer", lambda x, _: x
+        ):
+            dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
+                kv_cache=kv_cache,
+                logical_locs=logical_locs,
+                remap_logical_pages=remap_logical_pages,
+                layout=layout,
+                page_size=4,
+            )
+
+        self.assertEqual(dense_locs.tolist(), [4, 8, -1])
+        self.assertEqual(float(dense_kv.sum().item()), 48.0)
+
+    def test_tai_materialize_is_not_used_when_debug_enabled(self):
+        from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
+        from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
+
+        layout = CpSharedKVLayout(page_size=4, cp_size=1, cp_rank=0)
+        page_buffer = torch.arange(0, 4 * 3, dtype=torch.uint8).view(4, 3)
+        logical_pages = torch.tensor([1, 2], dtype=torch.int64)
+
+        with patch(
+            "sglang.srt.environ.envs.SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE.get",
+            return_value=True,
+        ), patch.object(
+            runtime, "cp_shared_kv_debug_enabled", return_value=True
+        ), patch.object(
+            runtime,
+            "_load_tai_materialize_kernels",
+            side_effect=AssertionError("tai path must stay off in debug mode"),
+        ), patch.object(
+            runtime, "_all_reduce_materialized_buffer", lambda x, _: x
+        ):
+            dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
+                page_buffer=page_buffer,
+                logical_pages=logical_pages,
+                layout=layout,
+            )
+
+        self.assertEqual(dense_pages.tolist(), [1, 2])
+        self.assertTrue(torch.equal(dense_page_buffer[1], page_buffer[1]))
+        self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[2]))
+
+
 if __name__ == "__main__":
     unittest.main()