# 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` 下降。

### P0 fused MLA persistent store 接入

在 materialize 之前，shared KV direct-write 的 MLA persistent KV 写入还有一段
独立开销：

```text
logical_locs_to_physical(...)
quantize_k_cache_separate(...)
set_mla_kv_buffer_triton(...)
```

对应的 tai-kernel fast path 已通过 SGLang 实验开关接入：

```text
SGLANG_CP_SHARED_KV_FUSED_MLA_STORE=0/1
```

默认关闭。开启后只影响 prefill CP shared-KV 的 MLA local persistent write：

```text
forward_mla.py
  _maybe_write_cp_shared_local_mla_kv(...)
    -> try_tai_fused_mla_store(...)
       -> tai_kernel.nsa_prefill.fused_store_mla_kv(
              k_nope, k_rope, raw_kv_buffer, logical_locs,
              page_size, cp_size)
    -> fallback logical_locs_to_physical + token_to_kv_pool.set_mla_kv_buffer
```

接入约束：

1. `fused_store_mla_kv` 根据 raw KV buffer dtype 分发：
   - `uint8 [capacity, 1, 656]`：packed NSA FP8 路径，
     fused logical->physical + FP8 quant + scale/rope store；
   - `bf16/fp16 [capacity, 1, 576]`：non-FP8 direct-store 路径，
     fused logical->physical + MLA KV direct write。
2. tai-kernel import 失败、shape/dtype 不支持、kernel runtime 失败时自动 fallback。
3. `SGLANG_DEBUG_CP_SHARED_KV=1` 时强制保留原 PyTorch path，便于继续暴露
   shared KV correctness 问题。
4. fused kernel 消费 shared logical token locs；调用前仍要求 direct-write gate
   已确认 `local_out_cache_loc` 属于当前 CP owner。
5. 该路径不替代后续 attention 前的 shared KV materialize，只减少 persistent
   write 阶段的 remap/quant/store kernel 碎片。

### P0 fused NSA index MQA prepare 接入

在 shared NSA index K/scale materialize 之后、`deep_gemm.fp8_mqa_logits(...)`
之前，CP in-seq-split pair path 原来还会执行：

```text
GetK(index_buffer, block_tables[0], kv_len)
GetS(index_buffer, block_tables[0], kv_len)
torch.arange(ke_start, ke_start + valid_q_count)
torch.zeros_like(ke)
```

对应的 tai-kernel fast path 已通过实验开关接入：

```text
SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE=0/1
```

默认关闭。开启后只影响 `_get_topk_ragged_with_cp(...)` 的单序列
`cp_index is None` 且 `current_index_kv is None` 路径：

```text
nsa_indexer.py
  _get_topk_ragged_with_cp(...)
    -> try_tai_prepare_cp_mqa_index(...)
       -> tai_kernel.nsa_prefill.prepare_cp_mqa_kv_and_range(
              index_buffer, block_tables[0], kv_len,
              valid_q_count, ke_start, page_size=64, index_head_dim=128)
    -> fallback GetK + GetS + torch arange/zeros
```

接入约束：

1. 仅支持 CUDA、`page_size=64`、`index_head_dim=128`、`index_buffer.dtype=uint8`。
2. 不覆盖 multi-batch `cp_index is not None` path。
3. 不覆盖 `current_index_kv` reuse path。
4. 不融合 `fp8_mqa_logits` 本身，也不改变 distributed top-k 语义。

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