diff --git a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md b/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md index d65d45f13..3b310109c 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase4_page_aligned_split_plan.md +++ b/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 diff --git a/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md b/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md index 48b4cd64f..412e60412 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase5_compute_owner_kv_layout_plan.md +++ b/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: diff --git a/docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md b/docs/advanced_features/nsa_prefill_cp_phase6_index_materialize_reuse_plan.md new file mode 100644 index 000000000..e2b155974 --- /dev/null +++ b/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 +``` diff --git a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md b/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md new file mode 100644 index 000000000..7100da517 --- /dev/null +++ b/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,不影响正确性。 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index d1dbdab33..805ac6967 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 2c41b412c..d868529d2 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/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[ diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 803ebd8a4..54e5cb889 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/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, diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 5b18b1fb3..1fb5a6038 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/allocator.py b/python/sglang/srt/mem_cache/allocator.py index c7d7a1a9d..6d62a21fc 100644 --- a/python/sglang/srt/mem_cache/allocator.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 5a759ed11..732d63909 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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 = ( diff --git a/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py b/python/sglang/srt/mem_cache/cp_shared_kv_compute_owner.py new file mode 100644 index 000000000..8df8e6b6c --- /dev/null +++ b/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 diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index 0ae3e6614..0ef8700d4 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/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 diff --git a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py b/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py index c5bf33336..42c650ded 100644 --- a/python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py +++ b/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, diff --git a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py b/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py index fe1b3c966..9ecca8d07 100644 --- a/python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py +++ b/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: diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index be84e5b41..ce8237cfc 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py b/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py index b4f00efe0..5b569ae74 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +++ b/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() diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index acf69c604..13005505c 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/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()