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，不影响正确性。