diff --git a/docs/advanced_features/nsa_prefill_cp_compute_padding_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_compute_padding_plan_zh.md new file mode 100644 index 000000000..fa332d4ea --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_compute_padding_plan_zh.md @@ -0,0 +1,869 @@ +# NSA Prefill CP Compute Padding 实现计划 + +> **For agentic workers:** REQUIRED SUB-SKILL: Use `superpowers:subagent-driven-development` 或 `superpowers:executing-plans` 逐任务实现。本文记录当前代码依据、目标合同、分阶段实现和验证点。 + +**Goal:** 用 compute padding 消除 tiny extend 在 CP shared KV in-seq split 下产生 zero-lane/mostly-zero-lane 的 hang 风险,同时不把 dummy token 写入 radix/HiCache/KV cache。 + +**Architecture:** cache/radix/HiCache 继续以真实 page 为最小单位;forward compute 额外补 dummy query row,让每个 CP rank 至少拿到一个 suffix page 的 compute work。metadata 需要同时表达 compute split 与 valid/cache split,所有写 cache、current reuse、backup/load 只看 valid split,attention/last-token 的本地排布看 compute split。 + +**Tech Stack:** Python/SGLang NSA prefill CP metadata, CP shared KV direct write/current reuse, HiCache, unit tests under `test/registered/unit`. + +--- + +## 1. 当前代码依据 + +### 1.1 tiny skip 目前只支持 bs=1 + +代码位置: + +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` + - `can_cp_split()` + +现状: + +- `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` 要求 `len(extend_seq_lens_cpu) == 1`。 +- bs>1 时,即使每个 request 都是 tiny extend,`can_cp_split()` 仍可能因为 flattened total token 足够而进入 CP split。 +- 这会重新制造之前 bs=1 修过的 zero-lane / mostly-zero-lane 分布。 + +### 1.2 现有 batched plan 是 valid-token split,不是 compute split + +代码位置: + +- `CPSharedKVBatchPlan` +- `build_batch_page_aligned_in_seq_split_plan()` +- `build_page_aligned_in_seq_split_list()` + +现状: + +- 每个 request 独立 page-rounded,但 `request_split_lists` 仍按真实 valid token 计数。 +- `extend=65,page=64,cp=8` 会得到类似 `[64, 1, 0, 0, 0, 0, 0, 0, ...]` 的 valid split。 +- 这对 cache page contract 是正确的,但对 distributed compute 是不稳定的:大多数 rank 没有 query。 + +### 1.3 不能直接把 total_len 传大来“补齐” + +`build_page_aligned_in_seq_split_list(total_len=512, extend_len=65, ...)` 当前会把 `padding_tokens = total_len - extend_len` 加到最后一个 segment,而不是把 8 个 page 分散到前 8 个 segment。 + +因此 compute padding 需要新的 split 构造逻辑,不能复用当前 `padding_tokens` 参数。 + +### 1.4 关键 consumer 目前默认 split rows == valid rows + +受影响路径: + +- `split_tensor_by_cp_batch_plan()`:要求 `tensor.shape[0] == sum(request_extend_lens)`。 +- `cp_split_and_rebuild_position()`:bs>1 目前仍用 scalar `metadata.split_list/zigzag_index`,没有 batch-aware。 +- `get_cp_shared_kv_local_out_cache_loc()`:按 valid `out_cache_loc` split,且 direct write 要求 local KV rows 与 local loc 数量一致。 +- `DeepseekV2Model.forward_core()`:进入模型后对 `hidden_states/positions` 做 CP split。 +- `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()`:要求 `k_nope/k_pe` rows 与 local loc rows 一致。 +- `nsa_indexer.py::_store_cp_shared_local_index_k_cache()`:要求 local index KV rows 与 local loc rows 一致。 +- `nsa_indexer.py::_get_topk_in_seq_cp_pair_batch()`:使用 `request_actual_seq_q_prev/next` 顺序消费 `q_fp8/weights`。 +- `_in_seq_collect_last_token_batch()`:用 `request_last_token_owner/local_offset/rank_local_offsets` 从本地 hidden 中取 compact last hidden。 + +结论:compute padding 后必须区分 “本地 compute rows” 与 “本地 valid/cache rows”。不能把 dummy rows 直接送到 direct write/current reuse。 + +--- + +## 2. 目标合同 + +### 2.1 cache padding 与 compute padding 分离 + +对每个 request: + +```text +valid_tokens = extend_len +valid_pages = ceil_div(valid_tokens, page_size) +valid_padded_tokens = valid_pages * page_size + +compute_pages = max(valid_pages, cp_size) +compute_tokens = compute_pages * page_size +compute_padding_tokens = compute_tokens - valid_tokens +``` + +约束: + +- radix/HiCache/device KV allocation 只使用 `valid_tokens/valid_pages`。 +- compute padding 只存在于 forward local tensor 排布。 +- dummy rows 不进入 out_cache_loc、direct write、backup、load、radix insert。 + +### 2.2 示例合同 + +`extend_len=65,page_size=64,cp_size=8`: + +```text +valid_pages = 2 +valid_padded_tokens = 128 +compute_pages = 8 +compute_tokens = 512 + +valid_split = [64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0] +compute_split = [64,64,64,64,64,64,64,64,0,0,0,0,0,0,0,0] +``` + +CP rank local compute rows: + +```text +rank0: segment0 + segment15 = 64 + 0 +rank1: segment1 + segment14 = 64 + 0 +... +rank7: segment7 + segment8 = 64 + 0 +``` + +真实最后一个 token: + +```text +token index = 64 +owner rank = 1 +local offset = 0 +``` + +### 2.3 bs>1 合同 + +每个 request 独立 compute padding,不把 batch 当成一条长序列。 + +示例: + +```text +extend_lens=[65, 1024], page_size=64, cp_size=8 + +req0: + valid_pages=2, compute_pages=8 +req1: + valid_pages=16, compute_pages=16 +``` + +每个 rank 的本地 tensor 是按 request 顺序拼接: + +```text +rank_local = req0.local_compute_rows + req1.local_compute_rows + ... +``` + +--- + +## 3. Metadata 设计 + +### 3.1 保留 valid/cache split,新增 compute split + +推荐字段: + +```python +request_valid_split_lists: List[List[int]] +request_valid_rank_local_tokens: List[int] +request_valid_rank_local_offsets: List[int] +request_valid_padded_pages: List[int] +request_valid_padded_tokens: List[int] +request_valid_padding_tokens: List[int] + +request_compute_split_lists: List[List[int]] +request_compute_rank_local_tokens: List[int] +request_compute_rank_local_offsets: List[int] +request_compute_padded_pages: List[int] +request_compute_padded_tokens: List[int] +request_compute_padding_tokens: List[int] + +compute_padding_enabled: bool +``` + +兼容策略: + +- 在 compute padding 模式下,`request_split_lists` 应明确作为 compute split 使用,因为它驱动 `cp_split_and_rebuild_*`、attention local q 长度、`max_rank_len`。 +- valid/cache consumer 必须改为显式使用 `request_valid_split_lists`。 +- bs=1 也应走同一套 batch plan 合同,避免 scalar path 继续隐藏 tiny 形状问题。 + +### 3.2 actual_seq_q/kv_len 语义 + +需要拆开: + +- `request_actual_seq_q_prev/next`:给 attention/indexer 的 compute query row 数,应该来自 compute split。 +- `request_valid_seq_q_prev/next`:给 direct write/current reuse/cache 写入的真实 row 数,来自 valid split。 +- `request_kv_len_prev/next`:attention 的 KV 可见长度仍要按真实逻辑位置计算,不应把 dummy rows 当成已经写入 cache 的 KV。 + +注意:这部分是实现的核心风险点。attention 可以计算 dummy query,但 dummy query 不能增加真实 KV cache length。 + +--- + +## 4. 分阶段实现计划 + +### P1:planner 纯 CPU 合同 + +**Files:** + +- Modify: `python/sglang/srt/layers/attention/nsa/utils.py` +- Test: `test/registered/unit/layers/test_nsa_cp_utils.py` + +步骤: + +1. 增加 compute padding split helper,例如: + +```python +def build_page_aligned_compute_padding_split_list( + *, + extend_len: int, + extend_prefix_len: int, + page_size: int, + cp_size: int, +) -> tuple[list[int], list[int], PageAlignedInSeqSplitInfo]: + ... +``` + +返回: + +- `valid_split_list` +- `compute_split_list` +- split info 内含 valid/compute padded pages/tokens + +2. 扩展 `CPSharedKVBatchPlan`: + +- 新增 valid split 字段。 +- 新增 compute split 字段。 +- `request_split_lists` 临时/兼容指向 compute split。 + +3. 更新 `build_batch_page_aligned_in_seq_split_plan()`: + +- 每个 request 独立计算 valid/compute split。 +- last-token owner 用 compute split + actual valid token count 计算。 +- rank-local offsets 用 compute local tokens 计算。 + +4. 单测: + +- `extend=65,page=64,cp=8`: + - valid split 是 `[64,1,0,...]` + - compute split 是前 8 段各 64。 + - last owner 是 rank1。 +- `extend=100,page=64,cp=8`: + - compute tokens 512,不是 1024。 +- `extend=1024,page=64,cp=8`: + - compute split 与 valid split 等价,不额外 padding。 +- `extend_lens=[65,1024]`: + - 每个 request 独立 padding,rank offsets 按 request 顺序累加。 + +### P2:CP split/rebuild 支持 compute rows + +**Files:** + +- Modify: `python/sglang/srt/layers/attention/nsa/utils.py` +- Test: `test/registered/unit/layers/test_nsa_cp_utils.py` + +步骤: + +1. 修改 `split_tensor_by_cp_batch_plan()`: + +- 输入是 valid flattened rows 时,按 request 先 pad 到 compute rows,再按 compute split 切分。 +- `mode="data"`:dummy rows 填 0。 +- `mode="1d"`:需要由 caller 指定 pad 值;input_ids 默认 0。 +- `mode="position"`:按 request 生成 dummy positions,建议用最后一个真实 position 继续递增,保证 RoPE 输入合法;dummy 输出后续不被收集/写 cache。 + +2. 修复 `cp_split_and_rebuild_position()`: + +- bs>1 不能继续 scalar split。 +- 统一调用 batch plan split helper。 + +3. 单测: + +- `hidden_states` 输入 65 rows,rank1 输出 64 rows,其中第 1 row 是真实 token 64,其余 dummy 为 0。 +- positions 65 rows 后 rank0-rank7 都有 64 positions。 +- bs>1 时 request 边界不被打乱。 + +### P3:narrow last-token output 基于 compute offset + +**Files:** + +- Modify: `python/sglang/srt/layers/attention/nsa/utils.py` +- Test: `test/registered/unit/layers/test_nsa_cp_utils.py` + +步骤: + +1. `_get_in_seq_last_token_owner_and_offset()` 输入 compute split,`actual_token_count` 仍是真实 extend_len。 +2. `_in_seq_collect_last_token_batch()` 使用 compute rank offsets。 +3. full rerange/logprob/capture-hidden 在 compute padding 模式下先 fail-fast: + +```text +[CP_SHARED_KV_FAIL_FAST][compute_padding_full_rerange_unsupported] +``` + +原因:full output 需要把 all-gather 后的 compute dummy rows trim 回 valid rows,不能复用当前 valid/full-rerange 假设。 + +### P4:valid-row selector,禁止 dummy rows 写 cache + +**Files:** + +- Modify: `python/sglang/srt/layers/attention/nsa/utils.py` +- Modify: `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` +- Modify: `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` +- Test: `test/registered/unit/layers/test_nsa_cp_utils.py` + +步骤: + +1. 增加 helper: + +```python +def select_cp_local_valid_rows_for_cache_write( + forward_batch, + local_tensor: torch.Tensor, +) -> torch.Tensor: + ... +``` + +行为: + +- 输入 local compute rows。 +- 按 `request_valid_split_lists` + `request_compute_split_lists` 为当前 rank 选出真实 valid rows。 +- 输出 rows 数必须等于 `get_cp_shared_kv_local_out_cache_loc()` 的 rows 数。 + +2. `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()` 前过滤 `k_nope/k_pe`。 +3. `nsa_indexer.py::_store_cp_shared_local_index_k_cache()` 前过滤 local index key。 +4. current reuse 不能再用 `key[:valid_current_rows]` 这种全局 prefix slice;需要同样用 valid selector 得到当前真实 rows。 + +单测: + +- rank1 的 `extend=65` local compute rows=64,但 valid rows=1。 +- direct write loc rows=1。 +- 只有第 1 个真实 row 被写入,dummy 63 rows 不写。 + +### P5:index/top-k 与 attention metadata + +**Files:** + +- Modify: `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` +- Modify: `python/sglang/srt/layers/attention/nsa/utils.py` +- Test: `test/registered/unit/layers/test_nsa_cp_utils.py` + +步骤: + +1. `request_actual_seq_q_prev/next` 使用 compute lengths,保证 attention/indexer 输入 rows 与本地 compute tensor 一致。 +2. 新增 valid q metadata,供 top-k compact 真实 query: + +```python +request_valid_seq_q_prev +request_valid_seq_q_next +request_valid_query_row_spans +``` + +3. `_get_topk_in_seq_cp_pair_batch()` 对 compute rows 输出同长度 result: + +- valid rows 正常 top-k。 +- dummy rows 填 `-1`。 +- output_cursor 按 compute rows 前进。 + +4. `current_index_kv` reuse 使用 valid rows,不使用 dummy rows。 + +风险: + +- 如果 attention kernel 不接受 dummy top-k `-1` rows,需要把 dummy top-k 映射到安全 page,而不是 `-1`。这一步必须用远端 CUDA/ETE 验证。 + +### P6:can_cp_split gate 与 bs>1 接入 + +**Files:** + +- Modify: `python/sglang/srt/layers/attention/nsa/utils.py` +- Test: `test/registered/unit/layers/test_nsa_cp_utils.py` + +步骤: + +1. 删除/替换 `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` 的 bs=1-only tiny skip 逻辑。 +2. `can_cp_split()` 对 CP shared KV 使用 per-request planner 判断: + +- prefix 必须 page-aligned,否则 fail-fast。 +- tiny request 不再 skip CP split,而是启用 compute padding。 +- 仍要保证 `use_nsa`、`context_parallel_extend`、`nsa_enable_prefill_cp` 等原条件。 + +3. 对非 shared-KV 路径保持原逻辑。 + +### P7:HiCache/backup/load 不扩容到 compute padding + +**Files:** + +- Verify/possibly modify: + - `python/sglang/srt/managers/schedule_batch.py` + - `python/sglang/srt/model_executor/forward_batch_info.py` + - `python/sglang/srt/mem_cache/hicache_controller.py` + - `python/sglang/srt/managers/cache_controller.py` + +检查点: + +- `alloc_for_extend()` 仍按 valid `extend_num_tokens` 分配。 +- `out_cache_loc` 仍只有 valid rows。 +- HiCache per-layer backup batching 只提交 valid rows。 +- load-back/prefetch 不因为 compute padding 请求更多 host/device slots。 + +### P8:验证与上线保护 + +**本地 CPU 单测:** + +```bash +cd /root/sglang-work/sglang-dev +PYTHONPATH=python python -m pytest -q \ + test/registered/unit/layers/test_nsa_cp_utils.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_layout.py +``` + +**远端 CUDA 验证:** + +```bash +scp -o ControlMaster=no -o ControlPath=none \ + python/sglang/srt/layers/attention/nsa/utils.py \ + python/sglang/srt/layers/attention/nsa/nsa_indexer.py \ + python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py \ + g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/layers/attention/nsa/ +``` + +实际同步时 `forward_mla.py` 路径要单独同步到: + +```bash +g0034:/mnt/beegfs/cjy/sglang-dev/python/sglang/srt/models/deepseek_common/attention_forward_methods/ +``` + +远端先跑 targeted unit,再由用户启动 ETE: + +```bash +ssh -o ControlMaster=no -o ControlPath=none g0034 \ + "docker exec sglang-glm5-dev-2 bash -lc 'cd /sgl-workspace/sglang-tai && PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py'" +``` + +ETE 重点观察: + +- repeated cache-hit `extend=65` 不再 hang。 +- accept len 不因为 dummy rows 掉到 1。 +- HiCache fallback/warning 不出现 dummy write。 +- output len 不为 0。 +- cache hit 后 direct write/current reuse 仍走 fast path。 + +--- + +## 5. 当前最关键风险 + +1. **attention 对 dummy query 的接受程度未验证。** 如果 top-k/attention kernel 不接受 dummy top-k `-1`,需要使用 safe page id 或 valid-query mask。 +2. **current reuse 当前大量代码按 prefix slice。** compute padding 后必须改成按 local valid span select,否则会把 dummy rows 或错位 rows 当成 current KV。 +3. **full rerange 暂不能直接复用。** compute padding 会让 all-gather 后 rows 包含 dummy,必须先 fail-fast 或实现 trim。 +4. **positions padding 需要远端验证。** dummy position 不写 cache,但会进 q/projection/attention,必须保证不会触发 kernel shape/causal 边界问题。 +5. **不能增加 cache capacity 需求。** 所有 allocation/backup/load 必须继续按 valid page,不按 compute page。 + +--- + +## 6. 推荐实施顺序 + +先做 P1-P3,锁住 metadata、split/rebuild 和 last-token narrow output;这部分可以主要靠 CPU/unit test。 + +然后做 P4-P5,处理 cache write/current reuse/top-k 的 valid-row 选择;这是 correctness 核心。 + +最后做 P6-P8,把 tiny skip gate 改为 compute padding 并做远端 ETE。 + +不要先改 `can_cp_split()` 放开 tiny bs>1,否则现有 consumer 仍会看到 zero/dummy 不一致,风险最高。 + +--- + +## 7. Implementation Ledger + +### P1 complete: planner exposes valid/compute split side-by-side + +Date: 2026-06-03 + +Implemented: + +- `CPSharedKVBatchPlan` now exposes explicit valid/cache fields: + - `request_valid_split_lists` + - `request_valid_padded_pages` + - `request_valid_padded_tokens` + - `request_valid_padding_tokens` + - `request_valid_rank_local_tokens` + - `request_valid_rank_local_offsets` + - `request_valid_actual_seq_q_prev/next` +- It also exposes compute-padding fields: + - `request_compute_split_lists` + - `request_compute_padded_pages` + - `request_compute_padded_tokens` + - `request_compute_padding_tokens` + - `request_compute_padded_token_offsets` + - `request_compute_rank_local_tokens` + - `request_compute_rank_local_offsets` + - `request_compute_actual_seq_q_prev/next` + - `compute_padding_enabled` +- `NSAContextParallelMetadata` mirrors the new fields from the batch plan. +- Planner computes last-token owner/local offset using compute split plus real + `extend_len`, so tiny suffixes can locate the real final token inside a + padded local compute segment. + +Important boundary: + +- P1 does **not** switch runtime split helpers to compute rows. +- `request_split_lists`, `request_padded_pages`, `request_rank_local_tokens`, + and `request_actual_seq_q_prev/next` intentionally keep the old valid-token + semantics until P2 updates consumers. +- This avoids breaking direct write/current reuse/top-k paths before valid-row + selection exists. + +Verification: + +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py` + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py` + - Result: `55 passed`. + +Next: + +- P2 should add compute-row split/rebuild support and then decide when + `request_split_lists` aliases should move from valid split to compute split. + +### P2 complete: split/rebuild can materialize local compute rows + +Date: 2026-06-03 + +Implemented: + +- `split_tensor_by_cp_batch_plan()` now uses `request_compute_split_lists` and + `request_compute_padded_tokens` when `compute_padding_enabled=True`. +- Its input contract remains valid flattened rows. The helper pads each + request independently before CP segment splitting: + - `mode="data"`: dummy rows are zeros. + - `mode="1d"`: dummy token/id rows are zeros. + - `mode="position"`: dummy positions continue from the last real position. +- `cp_split_and_rebuild_data()` and `cp_split_and_rebuild_1d()` now use the + batch plan path whenever `nsa_cp_metadata.batch_plan` exists, including + bs=1 plans used by compute-padding tests. +- `cp_split_and_rebuild_position()` is batch-plan aware and no longer forces + scalar `metadata.split_list` when a batch plan exists. +- Added explicit fail-fast when padded request rows do not match the selected + split total: + `[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_target_len_mismatch]`. + +Example verified: + +- `extend=65,page_size=64,cp_size=8,cp_rank=1` + - valid rows input: 65. + - local compute rows output: 64. + - first row is real token 64. + - remaining 63 rows are dummy zeros for data/1d. + - positions become `[40384, ..., 40447]`. + +Important boundary: + +- P2 still does not make dummy rows safe for cache writes/current reuse/top-k. +- `can_cp_split()` should not be changed to force tiny requests into CP split + until P4/P5 provide valid-row selection and dummy-safe top-k/attention + metadata. +- Existing metadata aliases (`request_split_lists`, `request_actual_seq_q_*`) + still preserve prior valid-token semantics unless a caller explicitly uses + the new compute fields through the split helper. + +Verification: + +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py` + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Result: `200 passed, 2 subtests passed`. + +### P3 complete: narrow last-token collect uses compute offsets + +Date: 2026-06-04 + +Implemented: + +- `_in_seq_collect_last_token()` now dispatches to the batch-plan path whenever + `nsa_cp_metadata.batch_plan` exists, including bs=1 compute-padding plans. +- `_in_seq_collect_last_token_batch()` now resolves last-token metadata from + either `NSAContextParallelMetadata` or its `batch_plan`. +- When `compute_padding_enabled=True`, last-token collection uses + `request_compute_rank_local_offsets` instead of valid/cache + `request_rank_local_offsets`. +- Existing valid-offset behavior remains for non-compute-padding metadata. + +Example verified: + +- Single request `extend=65,page_size=64,cp_size=8,cp_rank=1`: + - local compute rows: 64. + - real last token is local row 0. + - dummy row 1 is ignored. +- Batch request with two compute-padded suffixes: + - valid offsets would select the wrong second request row. + - compute offsets select the correct per-request last hidden. + +Important boundary: + +- P3 only fixes narrow output collection. +- Full rerange/logprob/capture-hidden still require a separate + compute-to-valid trim implementation before they can be enabled under + compute padding. +- Direct write/current reuse/top-k still need P4/P5 valid-row selection before + tiny CP split can be opened in `can_cp_split()`. + +Verification: + +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - Targeted P3 tests passed. + - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py test/registered/unit/layers/test_nsa_cp_utils.py` + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Result: `202 passed, 2 subtests passed`. + +### P4 complete: cache/current paths select valid rows, not dummy compute rows + +Date: 2026-06-04 + +Implemented: + +- `split_tensor_by_cp_batch_plan()` now has an explicit `split_kind`: + - `split_kind="compute"` keeps the P2 behavior and materializes compute + padding rows. + - `split_kind="valid"` uses valid/cache split metadata and never pads. +- `get_cp_shared_kv_local_out_cache_loc()` now uses the valid split whenever a + batch plan exists, including bs=1 compute-padding plans. This keeps + `out_cache_loc` rows equal to real cache rows, not local compute rows. +- Added `select_cp_local_valid_rows_for_cache_write()`: + - builds and caches per-forward-batch local row indices from + `request_compute_split_lists`, `request_valid_split_lists`, and + `request_zigzag_indices`; + - fails fast if a caller passes rows that are not the expected local compute + rows; + - returns only real valid rows for cache/current consumers. +- `forward_mla.py::_maybe_write_cp_shared_local_mla_kv()` filters `k_nope/k_pe` + before MLA direct write. +- `nsa_indexer.py::_store_cp_shared_local_index_k_cache()` filters local index + K before direct write. +- Index partial-current reuse no longer assumes `key[:valid_current_rows]` + under compute padding: + - current index preparation uses the valid-row selector and local valid + `out_cache_loc`; + - `_maybe_materialize_shared_index_buffer()` accepts local valid + `current_index_kv` rows and local valid locs when compute padding is + enabled. +- MLA partial-current reuse in `nsa_backend.py` similarly selects local valid + rows and local valid locs before page-slot compose. + +Example verified: + +- `extend=65,page_size=64,cp_size=8,cp_rank=1` + - local compute rows: 64. + - local valid/cache rows: 1. + - direct write writes only row 0 and loc page 2. + - dummy rows 1..63 are not passed to MLA/index direct write or current + compose. + +Important boundary: + +- P4 makes cache writes/current reuse dummy-safe. +- It does **not** make top-k/attention dummy rows fully safe yet; P5 still must + ensure index/top-k output has a valid compute-row shape while dummy rows do + not affect real outputs. +- It does **not** open `can_cp_split()` for tiny bs>1 yet; P6 remains blocked + on P5. + +Verification: + +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - RED observed first: + - missing selector import failed collection; + - direct write tests failed with `*_local_shape_mismatch`; + - index partial-current compose failed because old code expected global + valid rows. + - Targeted P4 tests passed: + - valid-row selector filters 64 compute rows to 1 valid row. + - local `out_cache_loc` uses valid split under compute padding. + - MLA/index direct write filters dummy rows. + - index partial-current compose accepts local valid rows. + - `python -m py_compile` passed for: + - `python/sglang/srt/layers/attention/nsa/utils.py` + - `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + - `python/sglang/srt/layers/attention/nsa_backend.py` + - `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` + - `test/registered/unit/layers/test_nsa_cp_utils.py` + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Result: `207 passed, 5 warnings, 2 subtests passed`. + +### P5 complete: top-k/index consumes compute rows but compacts valid queries + +Date: 2026-06-04 + +Implemented: + +- `CPSharedKVBatchPlan` / `NSAContextParallelMetadata` now expose explicit + top-k query-length aliases: + - `request_compute_seq_q_prev/next` + - `request_valid_seq_q_prev/next` + - `request_valid_query_row_spans` +- `request_actual_seq_q_prev/next` now uses compute lengths in the batch plan. + This supersedes the earlier P1/P2 boundary where actual-q aliases remained + valid-token based. Cache/page accounting aliases (`request_split_lists`, + `request_padded_pages`) still remain valid-token based for compatibility. +- `_get_topk_in_seq_cp_pair()` routes any forward batch carrying a + `batch_plan` into the batch-aware path, including bs=1 compute-padding plans. + This avoids silently using the scalar path with compute-padded query tensors. +- `_get_topk_in_seq_cp_pair_batch()` now: + - consumes local q/weight rows by compute segment length; + - compacts only valid rows into `_get_topk_ragged_with_cp()`; + - advances output offsets by compute rows; + - fills dummy rows with `-1`; + - computes `cp_index` from the valid tail, not from the padded compute tail. +- Existing bs>1 non-padding metadata remains compatible: if explicit valid-q + aliases are absent and compute padding is disabled, valid q falls back to + actual q. Compute-padding metadata stays fail-fast when valid lengths are + missing. + +Example verified: + +- `extend=65,page_size=64,cp_size=8,cp_rank=1` + - local compute rows: 64. + - local valid q rows: 1. + - `_get_topk_ragged_with_cp()` receives only 1 compact row with + `cp_index=[(0,64,65)]`. + - returned top-k rows keep compute shape `(64, topk)`; + row 0 is real and rows 1..63 are `-1`. + +Important boundary: + +- P5 makes top-k/index dummy-row safe at the Python metadata/dispatch layer. +- It does **not** open `can_cp_split()` for tiny bs>1 yet; P6 still owns the + runtime gate change. +- CUDA/ETE still needs to prove the downstream attention kernel accepts dummy + top-k rows filled with `-1`. If not, dummy rows must map to a safe page id + plus valid-query masking instead of `-1`. + +Verification: + +- Local: + - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py` +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - RED observed first: + - compute-padded bs=1 batch-plan top-k initially used the scalar path and + expected q rows to equal valid rows. + - old bs>1 tests exposed missing valid-q fallback for non-padding metadata. + - scalar current-index reuse exposed an accidentally inserted undefined + `compute_padding_enabled` reference. + - Targeted P5 tests: + - `test_batch_plan_exposes_compute_padding_without_inflating_valid_cache_extent` + - `test_indexer_in_seq_cp_pair_compute_padding_outputs_dummy_safe_rows` + - plus the previous four failing batch/scalar top-k reuse tests. + - Full unit slice: + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Result: `208 passed, 5 warnings, 2 subtests passed`. + +### P6 complete: tiny CP shared-KV requests enter compute-padding split + +Date: 2026-06-04 + +Implemented: + +- Removed the old effective tiny-suffix skip for CP shared KV: + `should_skip_cp_shared_kv_cp_split_for_short_page_extent()` is now a + compatibility hook that validates the page-aligned contract and returns + `False`. +- `can_cp_split()` now treats CP shared-KV in-seq split differently from the + generic token-balanced path: + - validates every request has matching extend/prefix metadata; + - requires `token_to_kv_pool.page_size`; + - fail-fasts on negative lengths and non-page-aligned prefixes; + - allows positive tiny suffixes even when valid token count is below + `cp_size`, because compute padding supplies one page of work per CP lane. +- `prepare_input_dp_with_cp_dsa()` now builds a `CPSharedKVBatchPlan` for + **all** CP shared-KV requests, including bs=1. This is required because bs=1 + tiny suffixes need the same valid/compute split contract as bs>1. +- `_build_batch_metadata_from_plan()` now separates scalar compatibility fields: + - communication/rerange fields (`split_list`, `split_list_tensor`, + `max_rank_len`, `per_rank_actual_token`, `reverse_split_len`, + `cp_reverse_index`) use compute split when compute padding is enabled; + - attention KV-length fields (`kv_len_prev/next`) still use valid split, so + dummy rows do not extend visible KV length; + - request-level cache fields remain valid-token based. +- NSA index partial-current preparation now selects valid current rows from + `local_key` under compute padding, not from the already all-gathered key. + +Example verified: + +- `extend=65,page_size=64,cp_size=8,cp_rank=1` + - `can_cp_split(...)=True`. + - `prepare_input_dp_with_cp_dsa()` returns metadata with `batch_plan`. + - scalar `metadata.split_list` is compute split + `[64,64,64,64,64,64,64,64,0,0,0,0,0,0,0,0]`. + - valid split remains `[64,1,0,...]`. + - `max_rank_len=[64]*8`, while `kv_len_prev/next` remain valid-length based. + +Important boundary: + +- P6 opens the CP split gate, but still does not prove CUDA/ETE attention + accepts dummy top-k `-1` rows. That remains a P8 runtime validation item. +- Full rerange/logprob/capture-hidden are still not enabled for compute + padding; narrow output remains the supported path. +- P7 still needs to verify HiCache allocation/backup/load never scales to + compute padding rows. + +Verification: + +- Local: + - `python -m py_compile python/sglang/srt/layers/attention/nsa/utils.py python/sglang/srt/layers/attention/nsa/nsa_indexer.py test/registered/unit/layers/test_nsa_cp_utils.py` +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - targeted P6 tests: + - tiny single request `can_cp_split=True`; + - tiny per-request bs>1 `can_cp_split=True`; + - non-page-aligned prefix fail-fast; + - bs=1 `prepare_input_dp_with_cp_dsa()` returns compute-padding + batch-plan metadata; + - P5 dummy-safe top-k remains green. + - Result: `7 passed, 5 warnings`. + - Added current-reuse regression: + - `test_indexer_current_reuse_compute_padding_selects_local_key_not_gathered_key` + constructs `local_key != gathered_key` and verifies compute-padding index + current reuse quantizes the local valid row, not the all-gathered key. + - Result: `1 passed, 5 warnings`. + - full unit slice: + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Result after regression test: `211 passed, 5 warnings, 2 subtests passed`. + +### P7 complete: HiCache allocation/backup/load stay on valid/page extent + +Date: 2026-06-04 + +Code audit: + +- Scheduler allocation remains valid-token based: + - `ScheduleBatch.prepare_for_extend()` sets `extend_num_tokens` from real + `fill_ids[prefix:]`, not compute-padded rows. + - `alloc_for_extend()` passes that same valid `extend_num_tokens` into + `alloc_paged_token_slots_extend()`. +- CP shared-KV L1 owner-lane allocation remains valid-page based: + - `build_in_seq_page_compute_owners()` returns one owner per real new page, + including a tail page, and does **not** expand to `cp_size` compute pages. + - For `extend=65,page_size=64,cp_size=8`, allocation owners are `[0, 1]`, + output cache loc length is `65`, and no legacy allocation fallback is used. +- CP HiCache write reservation pads only to the physical tail page: + - `HiCacheController.reserve_write_cp()` calls + `pad_token_locs_to_page_boundary()`, whose contract is page-tail padding + only and explicitly not CP-size padding. + - For a 65-token logical span starting on a page boundary: + `logical_len=65`, `padded_len=128`, `page_owners=[0,1]`; a rank owning + the tail page reserves exactly one physical page (`64` slots), not + `8*64` compute-padding slots. +- Per-layer backup uses reservation indices only: + - `_submit_write_cp_layer_states()` concatenates + `state.host_indices` and `state.physical_device_indices` created by + `reserve_write_cp()`, so grouped bs>1 per-layer backup inherits the same + valid/page extent. + - No per-layer backup descriptor is built from local compute rows. +- CP HiCache load restores padded physical pages but exposes valid length: + - `load_cp()` replays `page_owners` through `alloc_pages_with_owners()`. + - It queues H2D only for this rank's saved `owned_positions/host_indices`. + - It returns `visible_device_indices` trimmed to `valid_len`, so scheduler / + radix-visible cache hit length does not become the padded physical length. + +Important boundary: + +- P7 did not require production code changes; the existing L1 allocation, + HiCache reservation, per-layer backup, and load paths already separate + compute padding from page-tail physical padding. +- The protected contract is now covered by regression tests. Future changes + must not pass compute-padded local rows into HiCache reservation or load. +- P8 still needs runtime CUDA/ETE validation that dummy top-k rows and + compute-padded attention do not regress accept length or cause hangs. + +Verification: + +- Local: + - `python -m py_compile test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Local pytest is still blocked by missing local dependencies + (`transformers` during `sglang.test.test_utils` import). +- Remote container `g0034:/sgl-workspace/sglang-tai`: + - Added P7 regressions: + - `test_alloc_extend_compute_owner_uses_valid_pages_not_compute_padding_pages` + - `test_cp_hicache_write_reservation_uses_page_tail_not_compute_padding_extent` + - `test_cp_hicache_load_returns_valid_visible_len_while_loading_owned_page_tail` + - Targeted P7 result: `3 passed, 3 warnings`. + - Full unit slice: + - `PYTHONPATH=python python -m pytest -q test/registered/unit/layers/test_nsa_cp_utils.py test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py test/registered/unit/mem_cache/test_cp_shared_kv_layout.py` + - Result after P7 regressions: `214 passed, 5 warnings, 2 subtests passed`. diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 248a4b791..6f426868e 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -67,12 +67,14 @@ 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_batch_plan, get_cp_shared_kv_local_out_cache_loc, get_cp_shared_kv_local_physical_out_cache_loc, is_nsa_enable_prefill_cp, is_nsa_prefill_cp_in_seq_split, nsa_use_prefill_cp, raise_cp_shared_kv_direct_write_error, + select_cp_local_valid_rows_for_cache_write, split_in_seq_cp_local_pair, ) from sglang.srt.layers.communicator import ScatterMode @@ -358,29 +360,57 @@ class Indexer(MultiPlatformOp): f"logical_page_table_shape={tuple(logical_page_table.shape)} " f"page_size={page_size}" ) - current_locs = forward_batch.out_cache_loc - valid_current_rows = current_extend_kv_rows_for_reuse( - forward_batch, - current_index_kv[0], - current_index_kv[1], + batch_plan = get_cp_shared_kv_batch_plan(forward_batch) + compute_padding_current = batch_plan is not None and bool( + getattr(batch_plan, "compute_padding_enabled", False) ) - if valid_current_rows is None: - raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " - "CP shared KV index partial-current compose received " - "current_index_kv that does not satisfy current reuse " - "metadata. " - f"cp_rank={layout.cp_rank} layer_id={layer_id} " - f"prefix_lens={prefix_lens} extend_lens={extend_lens} " - f"current_k_shape={tuple(current_index_kv[0].shape)} " - f"current_scale_shape={tuple(current_index_kv[1].shape)} " - f"out_cache_loc_shape={tuple(current_locs.shape)}" + if compute_padding_current: + current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) + if current_locs is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " + "CP shared KV index partial-current compose requires local " + "valid out_cache_loc when compute padding is enabled. " + f"cp_rank={layout.cp_rank} layer_id={layer_id}" + ) + valid_current_rows = int(current_locs.numel()) + if ( + int(current_index_kv[0].shape[0]) != valid_current_rows + or int(current_index_kv[1].shape[0]) != valid_current_rows + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " + "CP shared KV index partial-current compose received " + "local current rows that do not match local valid locs. " + f"cp_rank={layout.cp_rank} layer_id={layer_id} " + f"current_k_shape={tuple(current_index_kv[0].shape)} " + f"current_scale_shape={tuple(current_index_kv[1].shape)} " + f"current_locs_shape={tuple(current_locs.shape)}" + ) + else: + current_locs = forward_batch.out_cache_loc + valid_current_rows = current_extend_kv_rows_for_reuse( + forward_batch, + current_index_kv[0], + current_index_kv[1], + ) + if valid_current_rows is None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_partial_current_sync] " + "CP shared KV index partial-current compose received " + "current_index_kv that does not satisfy current reuse " + "metadata. " + f"cp_rank={layout.cp_rank} layer_id={layer_id} " + f"prefix_lens={prefix_lens} extend_lens={extend_lens} " + f"current_k_shape={tuple(current_index_kv[0].shape)} " + f"current_scale_shape={tuple(current_index_kv[1].shape)} " + f"out_cache_loc_shape={tuple(current_locs.shape)}" + ) + current_locs = current_locs[:valid_current_rows] + current_index_kv = ( + current_index_kv[0][:valid_current_rows], + current_index_kv[1][:valid_current_rows], ) - current_locs = current_locs[:valid_current_rows] - current_index_kv = ( - current_index_kv[0][:valid_current_rows], - current_index_kv[1][:valid_current_rows], - ) prefix_slot_span = None if len(prefix_lens_cpu) == 1: prefix_pages = int(prefix_lens_cpu[0]) // page_size @@ -1614,7 +1644,10 @@ class Indexer(MultiPlatformOp): current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: assert forward_batch.nsa_cp_metadata is not None - if int(getattr(forward_batch.nsa_cp_metadata, "batch_size", 1) or 1) > 1: + if ( + int(getattr(forward_batch.nsa_cp_metadata, "batch_size", 1) or 1) > 1 + or get_cp_shared_kv_batch_plan(forward_batch) is not None + ): return self._get_topk_in_seq_cp_pair_batch( forward_batch, layer_id, @@ -1648,9 +1681,8 @@ class Indexer(MultiPlatformOp): shared_index_buffer = None shared_block_tables = None current_index_kv_for_topk = current_index_kv - if current_index_kv is not None and not is_current_only_extend_batch( - forward_batch - ): + current_only_batch = is_current_only_extend_batch(forward_batch) + if current_index_kv is not None and not current_only_batch: current_index_kv_for_topk = None shared_block_tables = metadata.get_page_table_64() shared_index_buffer, shared_block_tables = ( @@ -1714,35 +1746,84 @@ class Indexer(MultiPlatformOp): cp_metadata = forward_batch.nsa_cp_metadata assert cp_metadata is not None batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1) + batch_plan = get_cp_shared_kv_batch_plan(forward_batch) + compute_padding_enabled = bool( + getattr(cp_metadata, "compute_padding_enabled", False) + or bool(getattr(batch_plan, "compute_padding_enabled", False)) + ) + + def metadata_list(name: str, fallback_name: Optional[str] = None) -> List[int]: + values = getattr(cp_metadata, name, None) + if values is None and batch_plan is not None: + values = getattr(batch_plan, name, None) + if values is None and fallback_name is not None: + values = getattr(cp_metadata, fallback_name, None) + if values is None and batch_plan is not None: + values = getattr(batch_plan, fallback_name, None) + return list(values or []) request_kv_len_prev = list(getattr(cp_metadata, "request_kv_len_prev", []) or []) request_kv_len_next = list(getattr(cp_metadata, "request_kv_len_next", []) or []) - request_actual_seq_q_prev = list( - getattr(cp_metadata, "request_actual_seq_q_prev", []) or [] + if not request_kv_len_prev and batch_plan is not None: + request_kv_len_prev = list( + getattr(batch_plan, "request_kv_len_prev", []) or [] + ) + if not request_kv_len_next and batch_plan is not None: + request_kv_len_next = list( + getattr(batch_plan, "request_kv_len_next", []) or [] + ) + request_actual_seq_q_prev = metadata_list( + "request_compute_seq_q_prev" + if compute_padding_enabled + else "request_actual_seq_q_prev", + fallback_name="request_actual_seq_q_prev", ) - request_actual_seq_q_next = list( - getattr(cp_metadata, "request_actual_seq_q_next", []) or [] + request_actual_seq_q_next = metadata_list( + "request_compute_seq_q_next" + if compute_padding_enabled + else "request_actual_seq_q_next", + fallback_name="request_actual_seq_q_next", ) + request_valid_seq_q_prev = metadata_list( + "request_valid_seq_q_prev", + fallback_name="request_valid_actual_seq_q_prev", + ) + request_valid_seq_q_next = metadata_list( + "request_valid_seq_q_next", + fallback_name="request_valid_actual_seq_q_next", + ) + if not compute_padding_enabled: + # Older bs>1 metadata did not have explicit valid-q aliases because + # actual q length was also the valid q length. Keep that path + # compatible while compute-padding remains fail-fast if valid + # lengths are missing. + if not request_valid_seq_q_prev: + request_valid_seq_q_prev = request_actual_seq_q_prev + if not request_valid_seq_q_next: + request_valid_seq_q_next = request_actual_seq_q_next if not ( len(request_kv_len_prev) == batch_size and len(request_kv_len_next) == batch_size and len(request_actual_seq_q_prev) == batch_size and len(request_actual_seq_q_next) == batch_size + and len(request_valid_seq_q_prev) == batch_size + and len(request_valid_seq_q_next) == batch_size ): raise RuntimeError( "[CP_SHARED_KV_FAIL_FAST][index_topk] " "reason=batch_gt1_index_metadata_incomplete " f"batch_size={batch_size} layer_id={layer_id} " f"kv_prev={request_kv_len_prev} kv_next={request_kv_len_next} " - f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next}" + f"q_prev={request_actual_seq_q_prev} q_next={request_actual_seq_q_next} " + f"valid_q_prev={request_valid_seq_q_prev} " + f"valid_q_next={request_valid_seq_q_next}" ) shared_index_buffer = None shared_block_tables = None current_index_kv_for_topk = current_index_kv - if current_index_kv is not None and not is_current_only_extend_batch( - forward_batch - ): + current_only_batch = is_current_only_extend_batch(forward_batch) + if current_index_kv is not None and not current_only_batch: current_index_kv_for_topk = None shared_block_tables = metadata.get_page_table_64() shared_index_buffer, shared_block_tables = ( @@ -1771,7 +1852,7 @@ class Indexer(MultiPlatformOp): compact_output_spans: List[Tuple[int, int]] = [] current_only = ( current_index_kv_for_topk is not None - and is_current_only_extend_batch(forward_batch) + and current_only_batch ) page_table_1 = None if current_only else metadata.get_page_table_1() @@ -1779,16 +1860,27 @@ class Indexer(MultiPlatformOp): *, req_id: int, segment_len: int, + valid_segment_len: int, kv_len: int, ) -> None: nonlocal cursor, output_cursor segment_len = int(segment_len) + valid_segment_len = int(valid_segment_len) kv_len = int(kv_len) q_segment = q_fp8[cursor : cursor + segment_len] weights_segment = weights[cursor : cursor + segment_len] cursor += segment_len if segment_len == 0: return + if valid_segment_len < 0 or valid_segment_len > segment_len: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_gt1_index_valid_q_len_mismatch " + f"req_id={req_id} segment_len={segment_len} " + f"valid_segment_len={valid_segment_len}" + ) + if valid_segment_len == 0: + return seq_len = int(forward_batch.seq_lens_cpu[req_id].item()) extend_seq_len = int(forward_batch.extend_seq_lens_cpu[req_id]) @@ -1807,13 +1899,13 @@ class Indexer(MultiPlatformOp): logical_kv_limit = min(seq_len, int(page_table_1.shape[1])) valid_q_count = _compute_contiguous_valid_cp_query_count( cp_kv_end=cp_kv_end, - actual_seq_q=segment_len, + actual_seq_q=valid_segment_len, logical_kv_limit=logical_kv_limit, ) if valid_q_count <= 0: return - start_abs = cp_kv_end - segment_len + start_abs = cp_kv_end - valid_segment_len end_abs = start_abs + valid_q_count pre_chunk_offset = seq_len - extend_seq_len cp_index.append( @@ -1831,12 +1923,14 @@ class Indexer(MultiPlatformOp): collect_segment( req_id=req_id, segment_len=request_actual_seq_q_prev[req_id], + valid_segment_len=request_valid_seq_q_prev[req_id], kv_len=request_kv_len_prev[req_id], ) output_cursor += int(request_actual_seq_q_prev[req_id]) collect_segment( req_id=req_id, segment_len=request_actual_seq_q_next[req_id], + valid_segment_len=request_valid_seq_q_next[req_id], kv_len=request_kv_len_next[req_id], ) output_cursor += int(request_actual_seq_q_next[req_id]) @@ -2051,6 +2145,10 @@ class Indexer(MultiPlatformOp): local_out_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch) if local_out_loc is None: return False + local_key = select_cp_local_valid_rows_for_cache_write( + forward_batch, + local_key, + ) if local_key.shape[0] != local_out_loc.numel(): raise_cp_shared_kv_direct_write_error( "index_local_shape_mismatch", @@ -2239,10 +2337,41 @@ class Indexer(MultiPlatformOp): current_index_kv = None if self._can_reuse_current_index_kv(forward_batch): - valid_current_rows = current_extend_kv_rows_for_reuse(forward_batch, key) - if valid_current_rows is not None and key.shape[0] >= valid_current_rows: + batch_plan = get_cp_shared_kv_batch_plan(forward_batch) + compute_padding_current = batch_plan is not None and bool( + getattr(batch_plan, "compute_padding_enabled", False) + ) + if compute_padding_current: + current_key = select_cp_local_valid_rows_for_cache_write( + forward_batch, + local_key, + ) + current_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) + valid_current_rows = ( + int(current_locs.numel()) if current_locs is not None else None + ) + else: + valid_current_rows = current_extend_kv_rows_for_reuse( + forward_batch, key + ) + current_key = ( + key[:valid_current_rows] + if valid_current_rows is not None + else None + ) + current_locs = ( + forward_batch.out_cache_loc[:valid_current_rows] + if valid_current_rows is not None + else None + ) + if ( + valid_current_rows is not None + and current_key is not None + and current_locs is not None + and current_key.shape[0] == valid_current_rows + ): current_k_fp8, current_k_scale = act_quant( - key[:valid_current_rows].contiguous(), + current_key.contiguous(), self.block_size, self.scale_fmt, ) @@ -2258,9 +2387,7 @@ class Indexer(MultiPlatformOp): if forward_batch.cp_shared_kv_layout is not None else None, layer_id, - tensor_debug_summary( - forward_batch.out_cache_loc[:valid_current_rows] - ), + tensor_debug_summary(current_locs), tensor_debug_checksum(current_index_kv[0]), tensor_debug_checksum(current_index_kv[1]), ) diff --git a/python/sglang/srt/layers/attention/nsa/utils.py b/python/sglang/srt/layers/attention/nsa/utils.py index 56a3dbb98..d7186dbb2 100644 --- a/python/sglang/srt/layers/attention/nsa/utils.py +++ b/python/sglang/srt/layers/attention/nsa/utils.py @@ -320,6 +320,33 @@ class NSAContextParallelMetadata: flat_zigzag_index: List[int] = None flat_segment_request_ids: List[int] = None flat_segment_offsets: List[int] = None + compute_padding_enabled: bool = False + request_valid_split_lists: List[List[int]] = None + request_valid_segment_page_starts: List[List[int]] = None + request_valid_segment_page_ends: List[List[int]] = None + request_valid_padded_pages: List[int] = None + request_valid_padded_tokens: List[int] = None + request_valid_padding_tokens: List[int] = None + request_valid_rank_local_tokens: List[int] = None + request_valid_rank_local_offsets: List[int] = None + request_valid_actual_seq_q_prev: List[int] = None + request_valid_actual_seq_q_next: List[int] = None + request_valid_seq_q_prev: List[int] = None + request_valid_seq_q_next: List[int] = None + request_valid_query_row_spans: List[List[Tuple[int, int]]] = None + request_compute_split_lists: List[List[int]] = None + request_compute_segment_page_starts: List[List[int]] = None + request_compute_segment_page_ends: List[List[int]] = None + request_compute_padded_pages: List[int] = None + request_compute_padded_tokens: List[int] = None + request_compute_padding_tokens: List[int] = None + request_compute_padded_token_offsets: List[int] = None + request_compute_rank_local_tokens: List[int] = None + request_compute_rank_local_offsets: List[int] = None + request_compute_actual_seq_q_prev: List[int] = None + request_compute_actual_seq_q_next: List[int] = None + request_compute_seq_q_prev: List[int] = None + request_compute_seq_q_next: List[int] = None batch_plan: object = None @@ -355,6 +382,33 @@ class CPSharedKVBatchPlan: flat_zigzag_index: List[int] flat_segment_request_ids: List[int] flat_segment_offsets: List[int] + compute_padding_enabled: bool + request_valid_split_lists: List[List[int]] + request_valid_segment_page_starts: List[List[int]] + request_valid_segment_page_ends: List[List[int]] + request_valid_padded_pages: List[int] + request_valid_padded_tokens: List[int] + request_valid_padding_tokens: List[int] + request_valid_rank_local_tokens: List[int] + request_valid_rank_local_offsets: List[int] + request_valid_actual_seq_q_prev: List[int] + request_valid_actual_seq_q_next: List[int] + request_valid_seq_q_prev: List[int] + request_valid_seq_q_next: List[int] + request_valid_query_row_spans: List[List[Tuple[int, int]]] + request_compute_split_lists: List[List[int]] + request_compute_segment_page_starts: List[List[int]] + request_compute_segment_page_ends: List[List[int]] + request_compute_padded_pages: List[int] + request_compute_padded_tokens: List[int] + request_compute_padding_tokens: List[int] + request_compute_padded_token_offsets: List[int] + request_compute_rank_local_tokens: List[int] + request_compute_rank_local_offsets: List[int] + request_compute_actual_seq_q_prev: List[int] + request_compute_actual_seq_q_next: List[int] + request_compute_seq_q_prev: List[int] + request_compute_seq_q_next: List[int] def build_token_balanced_in_seq_split_list(total_len: int, cp_size: int) -> List[int]: @@ -378,6 +432,54 @@ def _prefix_offsets(lengths: List[int]) -> List[int]: return offsets +def _build_full_page_unit_split( + *, + page_units: int, + extend_prefix_len: int, + page_size: int, + cp_size: int, +) -> Tuple[List[int], List[int], List[int]]: + """Split full physical pages across in-seq CP segments. + + This is the compute-side counterpart of + `build_page_aligned_in_seq_split_list`: every assigned page contributes a + full `page_size` rows because dummy compute rows fill any valid-token tail. + """ + + if page_units < 0: + raise ValueError(f"page_units must be non-negative, got {page_units}") + if page_size <= 0: + raise ValueError(f"page_size must be positive, got {page_size}") + if cp_size <= 0: + raise ValueError(f"cp_size must be positive, got {cp_size}") + if extend_prefix_len < 0 or extend_prefix_len % page_size != 0: + raise ValueError( + "extend_prefix_len must be non-negative and page-aligned, " + f"got extend_prefix_len={extend_prefix_len} page_size={page_size}" + ) + + cp_segment_num = cp_size * 2 + base_units = page_units // cp_segment_num + remainder_units = page_units % cp_segment_num + unit_counts = [ + base_units + (1 if i < remainder_units else 0) + for i in range(cp_segment_num) + ] + + split_list: List[int] = [] + segment_page_starts: List[int] = [] + segment_page_ends: List[int] = [] + unit_cursor = 0 + base_page = extend_prefix_len // page_size + for unit_count in unit_counts: + segment_page_starts.append(base_page + unit_cursor) + unit_cursor += unit_count + segment_page_ends.append(base_page + unit_cursor) + split_list.append(unit_count * page_size) + + return split_list, segment_page_starts, segment_page_ends + + def build_batch_page_aligned_in_seq_split_plan( *, extend_lens: List[int], @@ -445,6 +547,29 @@ def build_batch_page_aligned_in_seq_split_plan( request_actual_seq_q_next: List[int] = [] request_last_token_owner: List[int] = [] request_last_token_local_offset: List[int] = [] + request_valid_split_lists: List[List[int]] = [] + request_valid_segment_page_starts: List[List[int]] = [] + request_valid_segment_page_ends: List[List[int]] = [] + request_valid_padded_pages: List[int] = [] + request_valid_padded_tokens: List[int] = [] + request_valid_padding_tokens: List[int] = [] + request_valid_rank_local_tokens: List[int] = [] + request_valid_actual_seq_q_prev: List[int] = [] + request_valid_actual_seq_q_next: List[int] = [] + request_valid_seq_q_prev: List[int] = [] + request_valid_seq_q_next: List[int] = [] + request_valid_query_row_spans: List[List[Tuple[int, int]]] = [] + request_compute_split_lists: List[List[int]] = [] + request_compute_segment_page_starts: List[List[int]] = [] + request_compute_segment_page_ends: List[List[int]] = [] + request_compute_padded_pages: List[int] = [] + request_compute_padded_tokens: List[int] = [] + request_compute_padding_tokens: List[int] = [] + request_compute_rank_local_tokens: List[int] = [] + request_compute_actual_seq_q_prev: List[int] = [] + request_compute_actual_seq_q_next: List[int] = [] + request_compute_seq_q_prev: List[int] = [] + request_compute_seq_q_next: List[int] = [] flat_split_list: List[int] = [] flat_zigzag_index: List[int] = [] flat_segment_request_ids: List[int] = [] @@ -468,8 +593,17 @@ def build_batch_page_aligned_in_seq_split_plan( f"extend_len={extend_len} page_size={page_size}" ) + compute_pages = max(split_info.extend_padded_pages, cp_size) + compute_split_list, compute_page_starts, compute_page_ends = ( + _build_full_page_unit_split( + page_units=compute_pages, + extend_prefix_len=prefix_len, + page_size=page_size, + cp_size=cp_size, + ) + ) owner, local_offset = _get_in_seq_last_token_owner_and_offset( - split_list=split_list, + split_list=compute_split_list, cp_size=cp_size, actual_token_count=extend_len, ) @@ -479,6 +613,9 @@ def build_batch_page_aligned_in_seq_split_plan( rank_local_tokens = ( split_list[cp_rank] + split_list[mirror_idx] ) + compute_rank_local_tokens = ( + compute_split_list[cp_rank] + compute_split_list[mirror_idx] + ) split_prefix_list = [0] + prefix_sum_list[:-1] request_split_infos.append(split_info) @@ -492,10 +629,38 @@ def build_batch_page_aligned_in_seq_split_plan( request_rank_local_tokens.append(rank_local_tokens) request_kv_len_prev.append(prefix_sum_list[cp_rank]) request_kv_len_next.append(prefix_sum_list[mirror_idx]) - request_actual_seq_q_prev.append(split_list[cp_rank]) - request_actual_seq_q_next.append(split_list[mirror_idx]) + request_actual_seq_q_prev.append(compute_split_list[cp_rank]) + request_actual_seq_q_next.append(compute_split_list[mirror_idx]) request_last_token_owner.append(owner) request_last_token_local_offset.append(local_offset) + request_valid_split_lists.append(split_list) + request_valid_segment_page_starts.append(split_info.segment_page_starts) + request_valid_segment_page_ends.append(split_info.segment_page_ends) + request_valid_padded_pages.append(split_info.extend_padded_pages) + request_valid_padded_tokens.append(split_info.extend_padded_tokens) + request_valid_padding_tokens.append(split_info.extend_padding_tokens) + request_valid_rank_local_tokens.append(rank_local_tokens) + request_valid_actual_seq_q_prev.append(split_list[cp_rank]) + request_valid_actual_seq_q_next.append(split_list[mirror_idx]) + request_valid_seq_q_prev.append(split_list[cp_rank]) + request_valid_seq_q_next.append(split_list[mirror_idx]) + request_valid_query_row_spans.append( + [ + (0, split_list[cp_rank]), + (compute_split_list[cp_rank], split_list[mirror_idx]), + ] + ) + request_compute_split_lists.append(compute_split_list) + request_compute_segment_page_starts.append(compute_page_starts) + request_compute_segment_page_ends.append(compute_page_ends) + request_compute_padded_pages.append(compute_pages) + request_compute_padded_tokens.append(compute_pages * page_size) + request_compute_padding_tokens.append(compute_pages * page_size - extend_len) + request_compute_rank_local_tokens.append(compute_rank_local_tokens) + request_compute_actual_seq_q_prev.append(compute_split_list[cp_rank]) + request_compute_actual_seq_q_next.append(compute_split_list[mirror_idx]) + request_compute_seq_q_prev.append(compute_split_list[cp_rank]) + request_compute_seq_q_next.append(compute_split_list[mirror_idx]) flat_split_list.extend(split_list) segment_base = req_id * cp_segment_num flat_zigzag_index.extend(segment_base + idx for idx in zigzag_index) @@ -533,6 +698,42 @@ def build_batch_page_aligned_in_seq_split_plan( flat_zigzag_index=flat_zigzag_index, flat_segment_request_ids=flat_segment_request_ids, flat_segment_offsets=flat_segment_offsets, + compute_padding_enabled=any( + compute_tokens != valid_tokens + for compute_tokens, valid_tokens in zip( + request_compute_padded_tokens, request_valid_padded_tokens + ) + ), + request_valid_split_lists=request_valid_split_lists, + request_valid_segment_page_starts=request_valid_segment_page_starts, + request_valid_segment_page_ends=request_valid_segment_page_ends, + request_valid_padded_pages=request_valid_padded_pages, + request_valid_padded_tokens=request_valid_padded_tokens, + request_valid_padding_tokens=request_valid_padding_tokens, + request_valid_rank_local_tokens=request_valid_rank_local_tokens, + request_valid_rank_local_offsets=_prefix_offsets(request_valid_rank_local_tokens), + request_valid_actual_seq_q_prev=request_valid_actual_seq_q_prev, + request_valid_actual_seq_q_next=request_valid_actual_seq_q_next, + request_valid_seq_q_prev=request_valid_seq_q_prev, + request_valid_seq_q_next=request_valid_seq_q_next, + request_valid_query_row_spans=request_valid_query_row_spans, + request_compute_split_lists=request_compute_split_lists, + request_compute_segment_page_starts=request_compute_segment_page_starts, + request_compute_segment_page_ends=request_compute_segment_page_ends, + request_compute_padded_pages=request_compute_padded_pages, + request_compute_padded_tokens=request_compute_padded_tokens, + request_compute_padding_tokens=request_compute_padding_tokens, + request_compute_padded_token_offsets=_prefix_offsets( + request_compute_padded_tokens + ), + request_compute_rank_local_tokens=request_compute_rank_local_tokens, + request_compute_rank_local_offsets=_prefix_offsets( + request_compute_rank_local_tokens + ), + request_compute_actual_seq_q_prev=request_compute_actual_seq_q_prev, + request_compute_actual_seq_q_next=request_compute_actual_seq_q_next, + request_compute_seq_q_prev=request_compute_seq_q_prev, + request_compute_seq_q_next=request_compute_seq_q_next, ) @@ -555,25 +756,43 @@ def split_tensor_by_cp_batch_plan( plan, *, mode: str = "data", + split_kind: str = "compute", ) -> torch.Tensor: """Split a flattened batch tensor by per-request in-seq CP plan. `mode` is kept explicit for future shape-specific kernels. The current CPU/Python planner path splits along dim0 for 1d, position, and data views. + `split_kind="compute"` materializes padded compute rows. Cache writes must + use `split_kind="valid"` so dummy compute rows never receive cache locs. """ if mode not in ("1d", "data", "position"): raise ValueError(f"unsupported CP batch split mode={mode!r}") + if split_kind not in ("compute", "valid"): + raise ValueError(f"unsupported CP batch split_kind={split_kind!r}") request_extend_lens = getattr(plan, "request_extend_lens", None) - request_split_lists = getattr(plan, "request_split_lists", None) + compute_padding_enabled = bool(getattr(plan, "compute_padding_enabled", False)) + if split_kind == "valid": + request_split_lists = getattr( + plan, "request_valid_split_lists", None + ) or getattr(plan, "request_split_lists", None) + request_target_lens = request_extend_lens + elif compute_padding_enabled: + request_split_lists = getattr(plan, "request_compute_split_lists", None) + request_target_lens = getattr(plan, "request_compute_padded_tokens", None) + else: + request_split_lists = getattr(plan, "request_split_lists", None) + request_target_lens = request_extend_lens request_zigzag_indices = getattr(plan, "request_zigzag_indices", None) batch_size = int(getattr(plan, "batch_size", 1) or 1) if ( request_extend_lens is None + or request_target_lens is None or request_split_lists is None or request_zigzag_indices is None or len(request_extend_lens) != batch_size + or len(request_target_lens) != batch_size or len(request_split_lists) != batch_size or len(request_zigzag_indices) != batch_size ): @@ -591,9 +810,28 @@ def split_tensor_by_cp_batch_plan( local_chunks = [] request_tensors = torch.split(tensor, [int(x) for x in request_extend_lens], dim=0) - for req_tensor, split_list, zigzag_index in zip( - request_tensors, request_split_lists, request_zigzag_indices + for req_id, (req_tensor, target_len, split_list, zigzag_index) in enumerate( + zip( + request_tensors, + request_target_lens, + request_split_lists, + request_zigzag_indices, + ) ): + req_tensor = _pad_cp_request_tensor_for_split( + req_tensor, + target_len=int(target_len), + mode=mode, + req_id=req_id, + ) + split_total = sum(int(x) for x in split_list) + if split_total != int(req_tensor.shape[0]): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][batch_gt1_split_target_len_mismatch] " + "request split rows must equal padded request rows. " + f"req_id={req_id} split_total={split_total} " + f"request_rows={int(req_tensor.shape[0])} mode={mode}" + ) req_segments = list(torch.split(req_tensor, [int(x) for x in split_list], dim=0)) local_chunks.extend(req_segments[int(index)] for index in zigzag_index) @@ -602,6 +840,150 @@ def split_tensor_by_cp_batch_plan( return torch.cat(local_chunks, dim=0).view(-1, *tensor.shape[1:]) +def _get_cp_local_valid_row_indices_cache(forward_batch, plan, device: torch.device): + cached = getattr(forward_batch, "cp_local_valid_row_indices_for_cache_write", None) + cached_expected_rows = getattr( + forward_batch, "cp_local_valid_compute_rows_for_cache_write", None + ) + if cached is not None and cached_expected_rows is not None: + if cached.device == device: + return cached, int(cached_expected_rows) + + batch_size = int(getattr(plan, "batch_size", 1) or 1) + request_compute_split_lists = getattr(plan, "request_compute_split_lists", None) + request_valid_split_lists = getattr( + plan, "request_valid_split_lists", None + ) or getattr(plan, "request_split_lists", None) + request_zigzag_indices = getattr(plan, "request_zigzag_indices", None) + if ( + request_compute_split_lists is None + or request_valid_split_lists is None + or request_zigzag_indices is None + or len(request_compute_split_lists) != batch_size + or len(request_valid_split_lists) != batch_size + or len(request_zigzag_indices) != batch_size + ): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] " + "CP shared-KV cache writes require valid and compute split metadata " + "when compute padding is enabled." + ) + + chunks = [] + local_cursor = 0 + for req_id in range(batch_size): + compute_split = request_compute_split_lists[req_id] + valid_split = request_valid_split_lists[req_id] + if len(compute_split) != len(valid_split): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] " + "valid and compute split metadata disagree. " + f"req_id={req_id} compute_segments={len(compute_split)} " + f"valid_segments={len(valid_split)}" + ) + for segment_id in request_zigzag_indices[req_id]: + segment_id = int(segment_id) + compute_len = int(compute_split[segment_id]) + valid_len = int(valid_split[segment_id]) + if valid_len > compute_len: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_metadata] " + "valid segment is longer than compute segment. " + f"req_id={req_id} segment_id={segment_id} " + f"valid_len={valid_len} compute_len={compute_len}" + ) + if valid_len > 0: + chunks.append( + torch.arange( + local_cursor, + local_cursor + valid_len, + device=device, + dtype=torch.long, + ) + ) + local_cursor += compute_len + + if chunks: + indices = torch.cat(chunks, dim=0) + else: + indices = torch.empty((0,), device=device, dtype=torch.long) + forward_batch.cp_local_valid_row_indices_for_cache_write = indices + forward_batch.cp_local_valid_compute_rows_for_cache_write = local_cursor + return indices, local_cursor + + +def select_cp_local_valid_rows_for_cache_write( + forward_batch, + local_tensor: torch.Tensor, +) -> torch.Tensor: + """Drop compute-padding rows before writing CP shared KV into persistent cache.""" + + plan = get_cp_shared_kv_batch_plan(forward_batch) + if plan is None or not bool(getattr(plan, "compute_padding_enabled", False)): + return local_tensor + + indices, expected_compute_rows = _get_cp_local_valid_row_indices_cache( + forward_batch, + plan, + local_tensor.device, + ) + local_rows = int(local_tensor.shape[0]) + if local_rows != expected_compute_rows: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cache_write_valid_selector_shape] " + "CP shared-KV cache write tensor must contain local compute rows " + "before valid-row selection. " + f"local_rows={local_rows} expected_compute_rows={expected_compute_rows} " + f"valid_rows={int(indices.numel())}" + ) + if indices.numel() == local_rows: + return local_tensor + if indices.numel() == 0: + return local_tensor.new_empty((0, *local_tensor.shape[1:])) + return local_tensor.index_select(0, indices) + + +def _pad_cp_request_tensor_for_split( + tensor: torch.Tensor, + *, + target_len: int, + mode: str, + req_id: int, +) -> torch.Tensor: + current_len = int(tensor.shape[0]) + if target_len < current_len: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_len_mismatch] " + "target split length is shorter than request valid rows. " + f"req_id={req_id} target_len={target_len} current_len={current_len}" + ) + pad_len = target_len - current_len + if pad_len == 0: + return tensor + if current_len <= 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][batch_gt1_compute_padding_empty_request] " + f"cannot compute-pad an empty request. req_id={req_id}" + ) + + if mode == "position": + if tensor.dim() != 1: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][batch_gt1_position_padding_rank] " + "position compute padding expects a 1D position tensor. " + f"req_id={req_id} shape={tuple(tensor.shape)}" + ) + start = tensor[-1] + 1 + padding = torch.arange( + pad_len, + device=tensor.device, + dtype=tensor.dtype, + ) + start + else: + padding = tensor.new_zeros((pad_len, *tensor.shape[1:])) + return torch.cat([tensor, padding], dim=0) + + def build_flat_page_owner_plan(plan) -> List[int]: from sglang.srt.mem_cache.cp_shared_kv_compute_owner import ( build_in_seq_page_compute_owners, @@ -761,11 +1143,16 @@ def _build_in_seq_split_for_forward_batch( def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan): + communication_split_lists = ( + plan.request_compute_split_lists + if plan.compute_padding_enabled + else plan.request_split_lists + ) per_rank_actual_token = [] for rank in range(plan.cp_size): rank_tokens = 0 mirror = plan.cp_size * 2 - rank - 1 - for split_list in plan.request_split_lists: + for split_list in communication_split_lists: rank_tokens += split_list[rank] + split_list[mirror] per_rank_actual_token.append(rank_tokens) max_rank_token = max(per_rank_actual_token) if per_rank_actual_token else 0 @@ -773,31 +1160,58 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan): # Scalar fields remain populated for compatibility, but scalar-only # consumers must not use them when batch_size > 1. - first_split = plan.request_split_lists[0] if plan.request_split_lists else [] + first_split = communication_split_lists[0] if communication_split_lists else [] + first_valid_split = ( + plan.request_valid_split_lists[0] + if plan.request_valid_split_lists + else plan.request_split_lists[0] + if plan.request_split_lists + else [] + ) first_info = plan.request_split_infos[0] if plan.request_split_infos else None first_zigzag = plan.request_zigzag_indices[0] if plan.request_zigzag_indices else [] - first_prefix_sum = list(accumulate(first_split)) - first_kv_len_prev = first_prefix_sum[plan.cp_rank] if first_prefix_sum else 0 + first_valid_prefix_sum = list(accumulate(first_valid_split)) + first_kv_len_prev = ( + first_valid_prefix_sum[plan.cp_rank] if first_valid_prefix_sum else 0 + ) first_mirror = plan.cp_size * 2 - plan.cp_rank - 1 - first_kv_len_next = first_prefix_sum[first_mirror] if first_prefix_sum else 0 + first_kv_len_next = ( + first_valid_prefix_sum[first_mirror] if first_valid_prefix_sum else 0 + ) first_actual_seq_q_prev = first_split[plan.cp_rank] if first_split else 0 first_actual_seq_q_next = first_split[first_mirror] if first_split else 0 + flat_communication_split_list = [ + token_count + for split_list in communication_split_lists + for token_count in split_list + ] + first_reverse_split_len = [ + element + for i in range(plan.cp_size) + for element in (first_split[i], first_split[plan.cp_size * 2 - i - 1]) + ] + first_cp_reverse_index = ( + list(range(0, plan.cp_size * 2, 2)) + + list(range(plan.cp_size * 2 - 1, 0, -2)) + if first_split + else [] + ) return NSAContextParallelMetadata( split_list=first_split, split_list_tensor=torch.tensor( - plan.flat_split_list, device="cuda", dtype=torch.int32 + flat_communication_split_list, device="cuda", dtype=torch.int32 ), split_prefix_tensor=torch.tensor( - [0] + list(accumulate(plan.flat_split_list))[:-1], + [0] + list(accumulate(flat_communication_split_list))[:-1], device="cuda", dtype=torch.int32, ), max_rank_len=max_rank_len, zigzag_index=first_zigzag, per_rank_actual_token=per_rank_actual_token, - reverse_split_len=None, - cp_reverse_index=None, + reverse_split_len=first_reverse_split_len, + cp_reverse_index=first_cp_reverse_index, kv_len_prev=first_kv_len_prev, kv_len_next=first_kv_len_next, actual_seq_q_prev=first_actual_seq_q_prev, @@ -892,6 +1306,33 @@ def _build_batch_metadata_from_plan(plan: CPSharedKVBatchPlan): flat_zigzag_index=plan.flat_zigzag_index, flat_segment_request_ids=plan.flat_segment_request_ids, flat_segment_offsets=plan.flat_segment_offsets, + compute_padding_enabled=plan.compute_padding_enabled, + request_valid_split_lists=plan.request_valid_split_lists, + request_valid_segment_page_starts=plan.request_valid_segment_page_starts, + request_valid_segment_page_ends=plan.request_valid_segment_page_ends, + request_valid_padded_pages=plan.request_valid_padded_pages, + request_valid_padded_tokens=plan.request_valid_padded_tokens, + request_valid_padding_tokens=plan.request_valid_padding_tokens, + request_valid_rank_local_tokens=plan.request_valid_rank_local_tokens, + request_valid_rank_local_offsets=plan.request_valid_rank_local_offsets, + request_valid_actual_seq_q_prev=plan.request_valid_actual_seq_q_prev, + request_valid_actual_seq_q_next=plan.request_valid_actual_seq_q_next, + request_valid_seq_q_prev=plan.request_valid_seq_q_prev, + request_valid_seq_q_next=plan.request_valid_seq_q_next, + request_valid_query_row_spans=plan.request_valid_query_row_spans, + request_compute_split_lists=plan.request_compute_split_lists, + request_compute_segment_page_starts=plan.request_compute_segment_page_starts, + request_compute_segment_page_ends=plan.request_compute_segment_page_ends, + request_compute_padded_pages=plan.request_compute_padded_pages, + request_compute_padded_tokens=plan.request_compute_padded_tokens, + request_compute_padding_tokens=plan.request_compute_padding_tokens, + request_compute_padded_token_offsets=plan.request_compute_padded_token_offsets, + request_compute_rank_local_tokens=plan.request_compute_rank_local_tokens, + request_compute_rank_local_offsets=plan.request_compute_rank_local_offsets, + request_compute_actual_seq_q_prev=plan.request_compute_actual_seq_q_prev, + request_compute_actual_seq_q_next=plan.request_compute_actual_seq_q_next, + request_compute_seq_q_prev=plan.request_compute_seq_q_prev, + request_compute_seq_q_next=plan.request_compute_seq_q_next, batch_plan=plan, ) @@ -914,20 +1355,26 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent( forward_batch: "ForwardBatch", cp_size: int, ) -> bool: - """Avoid in-seq CP split when a shared-KV suffix has too few pages. + """Compatibility hook for the old tiny-suffix skip gate. - CP shared KV is page-owned. A cache-hit suffix with fewer physical pages - than CP lanes creates mostly-zero in-seq segments; the distributed NSA path - has repeatedly shown hangs on that shape. Keep the page cache contract by - running those tiny suffixes without NSA in-seq CP split instead of falling - back to token-balanced page-splitting. + Compute padding now handles suffixes with fewer physical pages than CP + lanes, so this function validates the page-aligned shared-KV contract and + no longer requests a skip. """ - if ( - forward_batch is None - or not getattr(forward_batch, "uses_cp_shared_kv", False) - or cp_size <= 1 - ): + if not _is_cp_shared_kv_forward_batch(forward_batch) or cp_size <= 1: + return False + _validate_cp_shared_kv_cp_split_plan_inputs(forward_batch, cp_size) + return False + + +def _validate_cp_shared_kv_cp_split_plan_inputs( + forward_batch: "ForwardBatch", + cp_size: int, +) -> bool: + if not _is_cp_shared_kv_forward_batch(forward_batch): + return False + if cp_size <= 1: return False extend_seq_lens_cpu = getattr(forward_batch, "extend_seq_lens_cpu", None) @@ -936,32 +1383,55 @@ def should_skip_cp_shared_kv_cp_split_for_short_page_extent( page_size = int(getattr(token_to_kv_pool, "page_size", 0) or 0) if ( extend_seq_lens_cpu is None - or len(extend_seq_lens_cpu) != 1 or extend_prefix_lens_cpu is None - or len(extend_prefix_lens_cpu) != 1 or page_size <= 0 ): - return False - - extend_len = int(extend_seq_lens_cpu[0]) - if extend_len <= 0: - return False - - prefix_len = int(extend_prefix_lens_cpu[0]) - if prefix_len < 0: - return False - if prefix_len % page_size != 0: raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] " - "CP shared KV NSA in-seq split requires a page-aligned prefix. " - "The radix/HiCache match path should floor cache hits to the " - "previous page boundary before CP split planning. " - f"prefix_len={prefix_len} extend_len={extend_len} " + "[CP_SHARED_KV_FAIL_FAST][cp_split_missing_page_plan_inputs] " + "CP shared KV NSA in-seq split requires extend lengths, prefix " + "lengths, and token_to_kv_pool.page_size before planning. " + f"extend_seq_lens_cpu={extend_seq_lens_cpu} " + f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} " f"page_size={page_size} cp_size={cp_size}" ) + if len(extend_seq_lens_cpu) != len(extend_prefix_lens_cpu): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cp_split_length_mismatch] " + "CP shared KV NSA in-seq split requires one prefix length per " + "extend length. " + f"extend_seq_lens_cpu={extend_seq_lens_cpu} " + f"extend_prefix_lens_cpu={extend_prefix_lens_cpu} " + f"cp_size={cp_size}" + ) + if len(extend_seq_lens_cpu) == 0: + return False - padded_pages = ceil_div(extend_len, page_size) - return padded_pages < cp_size + has_extend = False + for req_id, (extend_len_raw, prefix_len_raw) in enumerate( + zip(extend_seq_lens_cpu, extend_prefix_lens_cpu) + ): + extend_len = int(extend_len_raw) + prefix_len = int(prefix_len_raw) + if extend_len < 0 or prefix_len < 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cp_split_negative_len] " + "CP shared KV NSA in-seq split received a negative length. " + f"req_id={req_id} prefix_len={prefix_len} " + f"extend_len={extend_len} page_size={page_size} " + f"cp_size={cp_size}" + ) + if prefix_len % page_size != 0: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][cp_split_non_page_aligned_prefix] " + "CP shared KV NSA in-seq split requires a page-aligned prefix. " + "The radix/HiCache match path should floor cache hits to the " + "previous page boundary before CP split planning. " + f"req_id={req_id} prefix_len={prefix_len} " + f"extend_len={extend_len} page_size={page_size} " + f"cp_size={cp_size}" + ) + has_extend = has_extend or extend_len > 0 + return has_extend def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): @@ -977,18 +1447,28 @@ def can_cp_split(seq_len: int, cp_size: int, use_nsa: bool, forward_batch): # 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 - if should_skip_cp_shared_kv_cp_split_for_short_page_extent( - forward_batch, cp_size - ): - return False - cur_cp_seq_len = seq_len // (cp_size * 2) + if _is_cp_shared_kv_forward_batch(forward_batch): + cur_cp_seq_len = ( + 1 + if _validate_cp_shared_kv_cp_split_plan_inputs( + forward_batch, cp_size + ) + else 0 + ) + else: + cur_cp_seq_len = seq_len // (cp_size * 2) + extend_token_count = sum(forward_batch.extend_seq_lens_cpu) + if _is_cp_shared_kv_forward_batch(forward_batch): + min_extend_token_count = 1 + else: + min_extend_token_count = cp_size if ( cur_cp_seq_len != 0 and cp_size > 1 and use_nsa and forward_batch.forward_mode.is_context_parallel_extend() and is_nsa_enable_prefill_cp() - and sum(forward_batch.extend_seq_lens_cpu) >= cp_size + and extend_token_count >= min_extend_token_count ): return True else: @@ -1005,7 +1485,7 @@ def cp_split_and_rebuild_data(forward_batch, input_: torch.Tensor): return nsa_cp_round_robin_split_data(input_) metadata = forward_batch.nsa_cp_metadata - if getattr(metadata, "batch_size", 1) > 1: + if get_cp_shared_kv_batch_plan(forward_batch) is not None: return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_) input_list = list( @@ -1027,7 +1507,7 @@ def cp_split_and_rebuild_1d(forward_batch, input_: torch.Tensor): return nsa_cp_round_robin_split_data(input_) metadata = forward_batch.nsa_cp_metadata - if getattr(metadata, "batch_size", 1) > 1: + if get_cp_shared_kv_batch_plan(forward_batch) is not None: return _cp_split_and_rebuild_batch_in_seq(forward_batch, input_) input_list = list( @@ -1135,7 +1615,7 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"): "nsa_prefill_cp_mode is not in-seq-split", ) batch_plan = get_cp_shared_kv_batch_plan(forward_batch) - if batch_plan is not None and int(getattr(batch_plan, "batch_size", 1) or 1) > 1: + if batch_plan is not None: split_tokens = sum(int(x) for x in getattr(batch_plan, "request_extend_lens", [])) mismatch_reason = "batch_split_out_cache_len_mismatch" else: @@ -1150,10 +1630,18 @@ def get_cp_shared_kv_local_out_cache_loc(forward_batch: "ForwardBatch"): out_cache_tokens, ) - local_out_cache_loc = cp_split_and_rebuild_1d( - forward_batch, - out_cache_loc.contiguous(), - ) + if batch_plan is not None: + local_out_cache_loc = split_tensor_by_cp_batch_plan( + out_cache_loc.contiguous(), + batch_plan, + mode="1d", + split_kind="valid", + ) + else: + 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 @@ -1231,6 +1719,13 @@ def cp_split_and_rebuild_position(forward_batch, positions: torch.Tensor): ) return nsa_cp_round_robin_split_data(positions) + if get_cp_shared_kv_batch_plan(forward_batch) is not None: + return split_tensor_by_cp_batch_plan( + positions, + get_cp_shared_kv_batch_plan(forward_batch), + mode="position", + ) + position_id_list = list( torch.split(positions, forward_batch.nsa_cp_metadata.split_list, dim=-1) ) @@ -1855,11 +2350,10 @@ def prepare_input_dp_with_cp_dsa( and getattr(forward_batch, "uses_cp_shared_kv", False) and getattr(forward_batch, "extend_seq_lens_cpu", None) is not None and getattr(forward_batch, "extend_prefix_lens_cpu", None) is not None - and len(forward_batch.extend_seq_lens_cpu) > 1 ): if page_size is None: raise RuntimeError( - "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_page_size] " + "[CP_SHARED_KV_FAIL_FAST][cp_shared_missing_page_size] " "CP shared-KV batch planning requires token_to_kv_pool.page_size" ) batch_plan = build_batch_page_aligned_in_seq_split_plan( @@ -2098,7 +2592,7 @@ def _in_seq_collect_last_token( cp_rank = get_attention_cp_rank() bs = len(forward_batch.extend_seq_lens_cpu) metadata = getattr(forward_batch, "nsa_cp_metadata", None) - if bs > 1: + if get_cp_shared_kv_batch_plan(forward_batch) is not None or bs > 1: return _in_seq_collect_last_token_batch(hidden_states, metadata, cp_size, cp_rank, bs) owner = 0 @@ -2134,9 +2628,25 @@ def _in_seq_collect_last_token_batch( "[CP_SHARED_KV_FAIL_FAST][batch_gt1_missing_metadata] " "CP in-seq bs>1 last-token collection requires batch metadata." ) - owners = getattr(metadata, "request_last_token_owner", None) - local_offsets = getattr(metadata, "request_last_token_local_offset", None) - rank_offsets = getattr(metadata, "request_rank_local_offsets", None) + plan = getattr(metadata, "batch_plan", None) + owners = _get_cp_last_token_metadata_list( + metadata, plan, "request_last_token_owner" + ) + local_offsets = _get_cp_last_token_metadata_list( + metadata, plan, "request_last_token_local_offset" + ) + compute_padding_enabled = bool( + getattr(metadata, "compute_padding_enabled", False) + or bool(getattr(plan, "compute_padding_enabled", False)) + ) + if compute_padding_enabled: + rank_offsets = _get_cp_last_token_metadata_list( + metadata, plan, "request_compute_rank_local_offsets" + ) + else: + rank_offsets = _get_cp_last_token_metadata_list( + metadata, plan, "request_rank_local_offsets" + ) if ( owners is None or local_offsets is None @@ -2181,3 +2691,16 @@ def _in_seq_collect_last_token_batch( dtype=torch.long, ) return gathered.index_select(0, gather_indices) + + +def _get_cp_last_token_metadata_list( + metadata: NSAContextParallelMetadata, + plan, + field_name: str, +): + value = getattr(metadata, field_name, None) + if value is not None: + return value + if plan is not None: + return getattr(plan, field_name, None) + return None diff --git a/python/sglang/srt/layers/attention/nsa_backend.py b/python/sglang/srt/layers/attention/nsa_backend.py index fac68f791..e809d041e 100644 --- a/python/sglang/srt/layers/attention/nsa_backend.py +++ b/python/sglang/srt/layers/attention/nsa_backend.py @@ -51,11 +51,14 @@ from sglang.srt.layers.attention.nsa.transform_index import ( from sglang.srt.layers.attention.nsa.utils import ( can_nsa_prefill_cp_round_robin_split, compute_nsa_seqlens, + get_cp_shared_kv_batch_plan, + get_cp_shared_kv_local_out_cache_loc, is_nsa_enable_prefill_cp, nsa_cp_round_robin_split_data, nsa_cp_round_robin_split_q_seqs, nsa_use_prefill_cp, pad_nsa_cache_seqlens, + select_cp_local_valid_rows_for_cache_write, ) from sglang.srt.layers.attention.utils import ( concat_mla_absorb_q_general, @@ -1807,12 +1810,41 @@ class NativeSparseAttnBackend( mla_prefetcher = getattr( forward_batch, "cp_shared_kv_mla_prefetcher", None ) - current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse( - forward_batch, - k, - k_rope, + batch_plan = get_cp_shared_kv_batch_plan(forward_batch) + compute_padding_current = batch_plan is not None and bool( + getattr(batch_plan, "compute_padding_enabled", False) ) - can_reuse_current_kv = current_kv_rows_for_reuse is not None + current_locs_for_reuse = None + current_k_nope = None + current_k_rope = None + if compute_padding_current: + assert k is not None and k_rope is not None + current_k_nope = select_cp_local_valid_rows_for_cache_write( + forward_batch, k + ) + current_k_rope = select_cp_local_valid_rows_for_cache_write( + forward_batch, k_rope + ) + current_locs_for_reuse = get_cp_shared_kv_local_out_cache_loc( + forward_batch + ) + current_kv_rows_for_reuse = ( + int(current_locs_for_reuse.numel()) + if current_locs_for_reuse is not None + else None + ) + can_reuse_current_kv = ( + current_kv_rows_for_reuse is not None + and int(current_k_nope.shape[0]) == current_kv_rows_for_reuse + and int(current_k_rope.shape[0]) == current_kv_rows_for_reuse + ) + else: + current_kv_rows_for_reuse = current_extend_kv_rows_for_reuse( + forward_batch, + k, + k_rope, + ) + can_reuse_current_kv = current_kv_rows_for_reuse is not None if cp_shared_kv_mla_prefetch_log_enabled(): if cp_shared_kv_mla_prefetch_should_log_layer(layer.layer_id): prefix_lens_cpu = getattr( @@ -1850,8 +1882,15 @@ class NativeSparseAttnBackend( assert k is not None and k_rope is not None assert current_kv_rows_for_reuse is not None valid_current_rows = int(current_kv_rows_for_reuse) - current_k_nope = k[:valid_current_rows] - current_k_rope = k_rope[:valid_current_rows] + if not compute_padding_current: + current_k_nope = k[:valid_current_rows] + current_k_rope = k_rope[:valid_current_rows] + current_locs_for_reuse = forward_batch.out_cache_loc[ + :valid_current_rows + ] + assert current_k_nope is not None + assert current_k_rope is not None + assert current_locs_for_reuse is not None if is_packed_fp8_mla_kv_cache(kv_cache): current_kv_cache = pack_current_mla_kv_for_reuse( current_k_nope, @@ -1860,9 +1899,6 @@ class NativeSparseAttnBackend( ) else: current_kv_cache = _cat([current_k_nope, current_k_rope], dim=-1) - current_locs_for_reuse = forward_batch.out_cache_loc[ - :valid_current_rows - ] logical_page_table_1 = page_table_1 current_remap_page_size, current_remap_logical_page_capacity = ( current_loc_remap_fast_path_args(forward_batch) 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 6198523a3..a18c1b765 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 @@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.utils import ( log_cp_draft_shared_kv_debug, nsa_use_prefill_cp, raise_cp_shared_kv_direct_write_error, + select_cp_local_valid_rows_for_cache_write, ) from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( should_reuse_current_extend_kv, @@ -557,6 +558,8 @@ class DeepseekMLAForwardMixin: local_out_cache_loc = get_cp_shared_kv_local_out_cache_loc(forward_batch) if local_out_cache_loc is None: return False + k_nope = select_cp_local_valid_rows_for_cache_write(forward_batch, k_nope) + k_pe = select_cp_local_valid_rows_for_cache_write(forward_batch, k_pe) if ( k_nope.shape[0] != local_out_cache_loc.numel() or k_pe.shape[0] != local_out_cache_loc.numel() diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index 0dddcb75c..f284ac516 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.utils import ( build_batch_page_aligned_in_seq_split_plan, build_page_aligned_cache_extent, _get_in_seq_last_token_owner_and_offset, + _build_batch_metadata_from_plan, build_page_aligned_in_seq_split_list, build_token_balanced_in_seq_split_list, can_cp_split, @@ -19,6 +20,7 @@ from sglang.srt.layers.attention.nsa.utils import ( cp_collect_last_token_hidden, cp_split_and_rebuild_1d, cp_split_and_rebuild_data, + cp_split_and_rebuild_position, _torch_batch_in_seq_all_gather_rerange, get_cp_shared_kv_batch_plan, get_cp_shared_kv_local_out_cache_loc, @@ -26,6 +28,7 @@ from sglang.srt.layers.attention.nsa.utils import ( get_cp_local_embedding_padded_token_count, pad_cp_local_input_ids_for_embedding, prepare_input_dp_with_cp_dsa, + select_cp_local_valid_rows_for_cache_write, split_tensor_by_cp_batch_plan, split_in_seq_cp_local_pair, ) @@ -242,7 +245,7 @@ class TestNSAInSeqCPUtils(unittest.TestCase): split_list, extend_prefix_len=54464, extend_len=256, page_size=64 ) - def test_can_cp_split_skips_cp_when_radix_hit_suffix_has_too_few_pages(self): + def test_can_cp_split_uses_compute_padding_for_short_radix_hit_suffix(self): class Mode: def is_context_parallel_extend(self): return True @@ -265,9 +268,11 @@ class TestNSAInSeqCPUtils(unittest.TestCase): return_value=True, ), ): - self.assertFalse(can_cp_split(128, 8, True, forward_batch)) + self.assertTrue(can_cp_split(128, 8, True, forward_batch)) - def test_can_cp_split_skips_cp_when_page_units_do_not_cover_all_lanes(self): + def test_can_cp_split_uses_compute_padding_when_page_units_do_not_cover_all_lanes( + self, + ): class Mode: def is_context_parallel_extend(self): return True @@ -290,9 +295,9 @@ class TestNSAInSeqCPUtils(unittest.TestCase): return_value=True, ), ): - self.assertFalse(can_cp_split(256, 8, True, forward_batch)) + self.assertTrue(can_cp_split(256, 8, True, forward_batch)) - def test_can_cp_split_skips_current_only_when_page_units_do_not_cover_all_lanes( + def test_can_cp_split_uses_compute_padding_for_current_only_one_page_suffix( self, ): class Mode: @@ -317,7 +322,34 @@ class TestNSAInSeqCPUtils(unittest.TestCase): return_value=True, ), ): - self.assertFalse(can_cp_split(64, 8, True, forward_batch)) + self.assertTrue(can_cp_split(64, 8, True, forward_batch)) + + def test_can_cp_split_uses_compute_padding_per_request_for_batched_tiny_suffix( + self, + ): + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + extend_seq_lens_cpu=[65, 64], + extend_prefix_lens_cpu=[54464, 8192], + 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(129, 8, True, forward_batch)) def test_can_cp_split_fails_on_non_page_aligned_cp_shared_prefix(self): class Mode: @@ -445,6 +477,70 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(plan.flat_segment_request_ids, [0, 0, 0, 0, 1, 1, 1, 1]) self.assertEqual(plan.flat_segment_offsets, [0, 4, 4, 4, 0, 4, 8, 9]) + def test_batch_plan_exposes_compute_padding_without_inflating_valid_cache_extent( + self, + ): + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[40320], + page_size=64, + cp_size=8, + cp_rank=1, + ) + + self.assertTrue(plan.compute_padding_enabled) + self.assertEqual( + plan.request_valid_split_lists, + [[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], + ) + self.assertEqual( + plan.request_compute_split_lists, + [[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]], + ) + self.assertEqual(plan.request_valid_padded_pages, [2]) + self.assertEqual(plan.request_valid_padded_tokens, [128]) + self.assertEqual(plan.request_compute_padded_pages, [8]) + self.assertEqual(plan.request_compute_padded_tokens, [512]) + self.assertEqual(plan.request_compute_padding_tokens, [447]) + self.assertEqual(plan.request_compute_rank_local_tokens, [64]) + self.assertEqual(plan.request_compute_rank_local_offsets, [0]) + self.assertEqual(plan.request_valid_rank_local_tokens, [1]) + self.assertEqual(plan.request_valid_rank_local_offsets, [0]) + self.assertEqual(plan.request_last_token_owner, [1]) + self.assertEqual(plan.request_last_token_local_offset, [0]) + + # Compatibility aliases for cache/page accounting stay valid-token + # based. Query-length metadata is split separately below: attention and + # top-k consume compute rows, cache/current paths consume valid rows. + self.assertEqual(plan.request_split_lists, plan.request_valid_split_lists) + self.assertEqual(plan.request_padded_pages, plan.request_valid_padded_pages) + self.assertEqual(plan.request_actual_seq_q_prev, [64]) + self.assertEqual(plan.request_actual_seq_q_next, [0]) + self.assertEqual(plan.request_valid_seq_q_prev, [1]) + self.assertEqual(plan.request_valid_seq_q_next, [0]) + self.assertEqual(plan.request_compute_seq_q_prev, [64]) + self.assertEqual(plan.request_compute_seq_q_next, [0]) + + def test_batch_plan_compute_padding_is_per_request_not_batch_total(self): + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65, 1024], + prefix_lens=[40320, 8192], + page_size=64, + cp_size=8, + cp_rank=0, + ) + + self.assertTrue(plan.compute_padding_enabled) + self.assertEqual(plan.request_valid_padded_pages, [2, 16]) + self.assertEqual(plan.request_compute_padded_pages, [8, 16]) + self.assertEqual(plan.request_compute_padded_tokens, [512, 1024]) + self.assertEqual(plan.request_compute_padding_tokens, [447, 0]) + self.assertEqual(plan.request_compute_rank_local_tokens, [64, 128]) + self.assertEqual(plan.request_compute_rank_local_offsets, [0, 64]) + self.assertEqual(plan.request_valid_rank_local_tokens, [64, 128]) + self.assertEqual(plan.request_valid_rank_local_offsets, [0, 64]) + self.assertEqual(plan.request_last_token_owner, [1, 0]) + def test_batch_plan_stable_helpers_split_and_build_page_owner_plan(self): import torch @@ -461,12 +557,15 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(build_flat_page_owner_plan(plan), [0, 0, 1, 1]) local_1d = split_tensor_by_cp_batch_plan(torch.arange(13), plan, mode="1d") - self.assertEqual(local_1d.tolist(), [8, 9, 10, 11, 12]) + self.assertEqual(local_1d.tolist(), [0, 0, 0, 0, 8, 9, 10, 11, 12, 0, 0, 0]) local_data = split_tensor_by_cp_batch_plan( torch.arange(13 * 2).view(13, 2), plan, mode="data" ) - self.assertEqual(local_data[:, 0].tolist(), list(range(16, 26, 2))) + self.assertEqual( + local_data[:, 0].tolist(), + [0, 0, 0, 0, 16, 18, 20, 22, 24, 0, 0, 0], + ) def test_collect_last_token_hidden_uses_batch_owner_metadata(self): import torch @@ -506,6 +605,91 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(collected.tolist(), [[13.0], [99.0]]) + def test_collect_last_token_hidden_uses_compute_padding_for_single_request(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[40320], + page_size=64, + cp_size=8, + cp_rank=1, + ) + hidden_states = torch.zeros((64, 1), dtype=torch.float32) + hidden_states[0] = 123.0 + hidden_states[1] = 999.0 + forward_batch = SimpleNamespace( + extend_seq_lens_cpu=[65], + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + ), + ) + + def fake_all_gather(output, local_last): + self.assertEqual(local_last.tolist(), [[123.0]]) + output.zero_() + output[1] = local_last[0] + + 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.get_attention_cp_rank", + return_value=1, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", + side_effect=fake_all_gather, + ), + ): + collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 8) + + self.assertEqual(collected.tolist(), [[123.0]]) + + def test_collect_last_token_hidden_uses_compute_rank_offsets_for_batch(self): + import torch + + hidden_states = torch.zeros((8, 1), dtype=torch.float32) + hidden_states[0] = 10.0 + hidden_states[1] = 99.0 + hidden_states[4] = 20.0 + forward_batch = SimpleNamespace( + extend_seq_lens_cpu=[5, 5], + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_last_token_owner=[1, 1], + request_last_token_local_offset=[0, 0], + request_rank_local_offsets=[0, 1], + request_compute_rank_local_offsets=[0, 4], + compute_padding_enabled=True, + ), + ) + + def fake_all_gather(output, local_last): + self.assertEqual(local_last.tolist(), [[10.0], [20.0]]) + output.copy_(torch.tensor([[0.0], [0.0], [10.0], [20.0]])) + + 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.get_attention_cp_rank", + return_value=1, + ), + patch( + "sglang.srt.layers.attention.nsa.utils.attn_cp_all_gather_into_tensor", + side_effect=fake_all_gather, + ), + ): + collected = cp_collect_last_token_hidden(hidden_states, forward_batch, 2) + + self.assertEqual(collected.tolist(), [[10.0], [20.0]]) + def test_collect_last_token_hidden_fails_fast_without_batch_owner_metadata(self): import torch @@ -577,6 +761,51 @@ class TestNSAInSeqCPUtils(unittest.TestCase): page_size=64, ) + def test_cp_shared_kv_prepare_uses_batch_plan_for_bs1_compute_padding(self): + class Mode: + def is_context_parallel_extend(self): + return True + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + extend_seq_lens_cpu=[65], + extend_prefix_lens_cpu=[0], + 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, + ): + metadata = prepare_input_dp_with_cp_dsa( + 65, + cp_rank=1, + cp_size=8, + seqs_len=[65], + forward_batch=forward_batch, + page_size=64, + ) + + self.assertIsNotNone(metadata.batch_plan) + self.assertEqual(metadata.batch_size, 1) + self.assertTrue(metadata.compute_padding_enabled) + self.assertEqual( + metadata.request_valid_split_lists, + [[64, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], + ) + self.assertEqual( + metadata.request_compute_split_lists, + [[64, 64, 64, 64, 64, 64, 64, 64, 0, 0, 0, 0, 0, 0, 0, 0]], + ) + self.assertEqual(metadata.split_list, metadata.request_compute_split_lists[0]) + self.assertEqual(metadata.max_rank_len, [64] * 8) + self.assertEqual(metadata.per_rank_actual_token, [64] * 8) + self.assertEqual(metadata.actual_seq_q_prev, 64) + self.assertEqual(metadata.actual_seq_q_next, 0) + self.assertEqual(metadata.request_valid_seq_q_prev, [1]) + self.assertEqual(metadata.request_valid_seq_q_next, [0]) + def test_cp_shared_kv_all_gather_rejects_round_robin_mode(self): import torch @@ -792,6 +1021,34 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(local[:, 0].tolist(), list(range(0, 16, 2))) + def test_cp_split_and_rebuild_data_uses_compute_padding_rows(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[40320], + page_size=64, + cp_size=8, + cp_rank=1, + ) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + ) + ) + tensor = torch.arange(65 * 2, dtype=torch.float32).view(65, 2) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local = cp_split_and_rebuild_data(forward_batch, tensor) + + self.assertEqual(local.shape, (64, 2)) + self.assertEqual(local[0].tolist(), [128.0, 129.0]) + self.assertTrue(torch.equal(local[1:], torch.zeros((63, 2)))) + def test_cp_split_and_rebuild_1d_keeps_batch_request_boundaries(self): import torch @@ -812,6 +1069,85 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(local.tolist(), [8, 9, 10, 11, 12]) + def test_cp_split_and_rebuild_1d_uses_zero_compute_padding_rows(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[40320], + page_size=64, + cp_size=8, + cp_rank=1, + ) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + ) + ) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local = cp_split_and_rebuild_1d(forward_batch, torch.arange(65)) + + self.assertEqual(local.shape, (64,)) + self.assertEqual(local[0].item(), 64) + self.assertEqual(local[1:].tolist(), [0] * 63) + + def test_select_cp_local_valid_rows_filters_compute_padding_rows(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=64, + cp_size=8, + cp_rank=1, + ) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + ) + ) + local_compute_rows = torch.full((64, 2), -1.0) + local_compute_rows[0] = torch.tensor([50.0, 51.0]) + + selected = select_cp_local_valid_rows_for_cache_write( + forward_batch, local_compute_rows + ) + + self.assertEqual(selected.tolist(), [[50.0, 51.0]]) + + def test_cp_split_and_rebuild_position_is_batch_aware_and_compute_padded(self): + import torch + + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[40320], + page_size=64, + cp_size=8, + cp_rank=1, + ) + forward_batch = SimpleNamespace( + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + ) + ) + positions = torch.arange(40320, 40385, dtype=torch.int32) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local = cp_split_and_rebuild_position(forward_batch, positions) + + self.assertEqual(local.shape, (64,)) + self.assertEqual(local.tolist(), list(range(40384, 40448))) + def test_cp_local_embedding_pad_len_uses_metadata_max_rank_len(self): from types import SimpleNamespace @@ -952,6 +1288,48 @@ class TestNSAInSeqCPUtils(unittest.TestCase): list(range(2 * page_size, 3 * page_size)) + [4 * page_size], ) + def test_local_out_cache_loc_uses_valid_rows_under_compute_padding(self): + import torch + + page_size = 64 + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=page_size, + cp_size=8, + cp_rank=1, + ) + out_cache_loc = torch.cat( + [ + torch.arange(page_size, 2 * page_size), + torch.tensor([2 * page_size]), + ] + ).to(torch.int64) + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, + cp_size=8, + cp_rank=1, + ), + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + page_aligned=True, + page_size=page_size, + extend_prefix_len=0, + ), + out_cache_loc=out_cache_loc, + ) + + with patch( + "sglang.srt.layers.attention.nsa.utils.is_nsa_prefill_cp_round_robin_split", + return_value=False, + ): + local_locs = get_cp_shared_kv_local_out_cache_loc(forward_batch) + + self.assertEqual(local_locs.tolist(), [2 * page_size]) + def test_batch_local_physical_out_cache_loc_reuses_layer_invariant_plan(self): import torch from types import SimpleNamespace @@ -1135,6 +1513,67 @@ class TestNSAInSeqCPUtils(unittest.TestCase): act_quant=None, ) + def test_indexer_direct_write_filters_compute_padding_rows(self): + import torch + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + page_size = 64 + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=page_size, + cp_size=8, + cp_rank=1, + ) + out_cache_loc = torch.cat( + [ + torch.arange(page_size, 2 * page_size), + torch.tensor([2 * page_size]), + ] + ).to(torch.int64) + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, + cp_size=8, + cp_rank=1, + ), + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + page_aligned=True, + page_size=page_size, + extend_prefix_len=0, + ), + out_cache_loc=out_cache_loc, + token_to_kv_pool=SimpleNamespace(page_size=page_size), + ) + indexer = object.__new__(Indexer) + indexer.nsa_enable_prefill_cp = True + calls = [] + + def fake_store_index_k_cache(**kwargs): + calls.append(kwargs) + + indexer._store_index_k_cache = fake_store_index_k_cache + local_key = torch.full((64, 2), -1.0) + local_key[0] = torch.tensor([9.0, 10.0]) + + with patch.object(nsa_indexer, "nsa_use_prefill_cp", return_value=True): + stored = Indexer._store_cp_shared_local_index_k_cache( + indexer, + forward_batch, + layer_id=0, + local_key=local_key, + act_quant=None, + ) + + self.assertTrue(stored) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["key"].tolist(), [[9.0, 10.0]]) + def test_mla_direct_write_fails_fast_on_local_shape_mismatch(self): import torch @@ -1161,6 +1600,292 @@ class TestNSAInSeqCPUtils(unittest.TestCase): k_pe=torch.empty((2, 8)), ) + def test_mla_direct_write_filters_compute_padding_rows(self): + import torch + + from sglang.srt.models.deepseek_common.attention_forward_methods import ( + forward_mla, + ) + + page_size = 64 + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=page_size, + cp_size=8, + cp_rank=1, + ) + out_cache_loc = torch.cat( + [ + torch.arange(page_size, 2 * page_size), + torch.tensor([2 * page_size]), + ] + ).to(torch.int64) + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, + cp_size=8, + cp_rank=1, + ), + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + page_aligned=True, + page_size=page_size, + extend_prefix_len=0, + ), + out_cache_loc=out_cache_loc, + token_to_kv_pool=SimpleNamespace(page_size=page_size), + ) + mla = SimpleNamespace(attn_mqa=SimpleNamespace(layer_id=0)) + calls = [] + + def fake_tai_store(**kwargs): + calls.append(kwargs) + return True + + k_nope = torch.full((64, 2), -1.0) + k_nope[0] = torch.tensor([1.0, 2.0]) + k_pe = torch.full((64, 2), -1.0) + k_pe[0] = torch.tensor([3.0, 4.0]) + + with patch.object(forward_mla, "try_tai_fused_mla_store", fake_tai_store): + stored = ( + forward_mla.DeepseekMLAForwardMixin._maybe_write_cp_shared_local_mla_kv( + mla, + forward_batch, + k_nope=k_nope, + k_pe=k_pe, + ) + ) + + self.assertTrue(stored) + self.assertEqual(len(calls), 1) + self.assertEqual(calls[0]["k_nope"].tolist(), [[1.0, 2.0]]) + self.assertEqual(calls[0]["k_rope"].tolist(), [[3.0, 4.0]]) + self.assertEqual(calls[0]["logical_locs"].tolist(), [2 * page_size]) + + def test_index_partial_current_compose_accepts_local_valid_compute_padding_rows( + self, + ): + import torch + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + page_size = 64 + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=page_size, + cp_size=8, + cp_rank=1, + ) + + class FakePool: + page_size = 64 + index_head_dim = 2 + + def get_index_k_with_scale_buffer(self, layer_id): + return torch.zeros((4, 3), dtype=torch.float32) + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, + cp_size=8, + cp_rank=1, + ), + cp_shared_kv_index_prefetcher=None, + token_to_kv_pool=FakePool(), + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=1, + batch_plan=plan, + page_aligned=True, + page_size=page_size, + extend_prefix_len=0, + ), + out_cache_loc=torch.cat( + [ + torch.arange(page_size, 2 * page_size), + torch.tensor([2 * page_size]), + ] + ).to(torch.int64), + extend_prefix_lens_cpu=[page_size], + extend_seq_lens_cpu=[65], + ) + logical_page_table = torch.tensor([[1, 2, 3]], dtype=torch.int32) + current_index_kv = ( + torch.tensor([[7.0, 8.0]], dtype=torch.float32), + torch.tensor([[0.5]], dtype=torch.float32), + ) + materialize_calls = [] + expected_buffer = torch.ones((3, 3), dtype=torch.float32) + expected_pages = torch.tensor([[0, 1, 2]], dtype=torch.int32) + indexer = object.__new__(Indexer) + + def fake_materialize(**kwargs): + materialize_calls.append(kwargs) + return expected_buffer, expected_pages + + with patch.object( + nsa_indexer, + "get_or_build_shared_paged_buffer_slot_remap", + return_value=torch.tensor([0, 1, 2], dtype=torch.int64), + ), patch.object( + nsa_indexer, + "materialize_prefix_and_reuse_current_index_page_slots", + side_effect=fake_materialize, + ): + dense_buffer, dense_pages = indexer._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id=0, + logical_page_table=logical_page_table, + current_index_kv=current_index_kv, + ) + + self.assertIs(dense_buffer, expected_buffer) + self.assertIs(dense_pages, expected_pages) + self.assertEqual(len(materialize_calls), 1) + self.assertIs(materialize_calls[0]["current_index_k"], current_index_kv[0]) + self.assertEqual(materialize_calls[0]["current_locs"].tolist(), [2 * page_size]) + + def test_indexer_current_reuse_compute_padding_selects_local_key_not_gathered_key( + self, + ): + import torch + import types + + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + page_size = 64 + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=page_size, + cp_size=8, + cp_rank=1, + ) + metadata_obj = _build_batch_metadata_from_plan(plan) + + class Mode: + def is_extend_without_speculative(self): + return True + + def is_decode_or_idle(self): + return False + + def is_target_verify(self): + return False + + def is_draft_extend(self, include_v2=False): + return False + + def is_context_parallel_extend(self): + return True + + class AttnBackend: + def get_indexer_metadata(self, layer_id, forward_batch): + return object() + + forward_batch = SimpleNamespace( + uses_cp_shared_kv=True, + cp_shared_kv_layout=CpSharedKVLayout( + page_size=page_size, + cp_size=8, + cp_rank=1, + ), + token_to_kv_pool=SimpleNamespace(page_size=page_size), + nsa_cp_metadata=metadata_obj, + out_cache_loc=torch.cat( + [ + torch.arange(page_size, 2 * page_size), + torch.tensor([2 * page_size]), + ] + ).to(torch.int64), + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[65], + seq_lens_cpu=torch.tensor([65], dtype=torch.int64), + forward_mode=Mode(), + attn_backend=AttnBackend(), + hisparse_coordinator=None, + ) + + local_key = torch.full((64, 2), -1.0, dtype=torch.float32) + local_key[0] = torch.tensor([11.0, 12.0]) + gathered_key = torch.full((64, 2), 99.0, dtype=torch.float32) + gathered_key[0] = torch.tensor([101.0, 102.0]) + query = torch.zeros((64, 2), dtype=torch.float32) + act_quant_inputs = [] + topk_current_index_kv = [] + + def fake_act_quant(tensor, block_size, scale_fmt): + act_quant_inputs.append(tensor.detach().clone()) + return tensor.detach().clone(), torch.ones( + (int(tensor.shape[0]), 1), dtype=torch.float32 + ) + + fake_triton_kernel = types.ModuleType( + "sglang.srt.layers.attention.nsa.triton_kernel" + ) + fake_triton_kernel.act_quant = fake_act_quant + + indexer = object.__new__(Indexer) + indexer.alt_stream = None + indexer.nsa_enable_prefill_cp = True + indexer.index_topk = 2 + indexer.block_size = 64 + indexer.scale_fmt = None + indexer._get_q_k_bf16 = ( + lambda *args, **kwargs: (query, gathered_key, local_key) + ) + indexer._store_cp_shared_local_index_k_cache = lambda **kwargs: True + indexer._can_reuse_current_index_kv = lambda forward_batch: True + indexer._get_logits_head_gate = ( + lambda x_for_gate, q_scale: torch.zeros((64, 1), dtype=torch.float32) + ) + + def fake_topk(*args, **kwargs): + topk_current_index_kv.append(kwargs["current_index_kv"]) + return torch.zeros((64, 2), dtype=torch.int32) + + indexer._get_topk_in_seq_cp_pair = fake_topk + + with ( + patch.dict( + sys.modules, + { + "sglang.srt.layers.attention.nsa.triton_kernel": fake_triton_kernel + }, + ), + patch.object(nsa_indexer, "_is_cuda", True), + patch.object(nsa_indexer, "_is_hip", False), + patch.object(nsa_indexer, "_is_npu", False), + patch.object( + nsa_indexer, + "is_nsa_prefill_cp_in_seq_split", + return_value=True, + ), + ): + result = Indexer.forward_cuda( + indexer, + x=torch.zeros((64, 2), dtype=torch.float32), + q_lora=torch.zeros((64, 2), dtype=torch.float32), + positions=torch.arange(64, dtype=torch.int64), + forward_batch=forward_batch, + layer_id=0, + return_indices=True, + ) + + self.assertEqual(result.shape, (64, 2)) + self.assertEqual(len(act_quant_inputs), 2) + self.assertEqual(act_quant_inputs[1].tolist(), [[11.0, 12.0]]) + self.assertNotEqual(act_quant_inputs[1].tolist(), [[101.0, 102.0]]) + self.assertEqual(len(topk_current_index_kv), 1) + self.assertEqual(topk_current_index_kv[0][0].tolist(), [[11.0, 12.0]]) + def test_indexer_direct_write_does_not_log_missing_metadata_for_non_cp_batch(self): import torch from types import SimpleNamespace @@ -1416,6 +2141,113 @@ class TestNSAInSeqCPUtils(unittest.TestCase): [[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6], [7, 7]], ) + def test_indexer_in_seq_cp_pair_compute_padding_outputs_dummy_safe_rows(self): + import torch + + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + page_size = 64 + plan = build_batch_page_aligned_in_seq_split_plan( + extend_lens=[65], + prefix_lens=[0], + page_size=page_size, + cp_size=8, + cp_rank=1, + ) + metadata_obj = _build_batch_metadata_from_plan(plan) + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + logical_pages = torch.tensor([[1, 2]], dtype=torch.int32) + materialized_index = torch.tensor([11], dtype=torch.int32) + dense_pages = torch.tensor([[1, 2]], dtype=torch.int32) + materialize_calls = [] + topk_calls = [] + + class Metadata: + def get_page_table_64(self): + return logical_pages + + def get_page_table_1(self): + return torch.empty((1, 65), dtype=torch.int32) + + def fake_materialize(forward_batch, layer_id, logical_page_table): + materialize_calls.append((layer_id, logical_page_table)) + return materialized_index, dense_pages + + 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, + actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, + batch_idx=0, + ): + topk_calls.append( + { + "actual_seq_q": actual_seq_q, + "cp_index": cp_index, + "q": q_fp8.flatten().tolist(), + "weights": weights.flatten().tolist(), + "shared_index_buffer": shared_index_buffer, + "shared_block_tables": shared_block_tables, + } + ) + rows = int(q_fp8.shape[0]) + return ( + torch.arange(1, rows + 1, dtype=torch.int32) + .view(rows, 1) + .repeat(1, 2) + ) + + indexer._maybe_materialize_shared_index_buffer = fake_materialize + indexer._get_topk_ragged_with_cp = fake_get_topk + + forward_batch = SimpleNamespace( + batch_size=1, + forward_mode=SimpleNamespace( + is_extend_without_speculative=lambda: True, + ), + extend_prefix_lens_cpu=[0], + extend_seq_lens_cpu=[65], + seq_lens_cpu=torch.tensor([65], dtype=torch.int64), + nsa_cp_metadata=metadata_obj, + ) + q_fp8 = torch.arange(64, dtype=torch.float32).view(64, 1) + weights = (torch.arange(64, dtype=torch.float32) + 100).view(64, 1) + + 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), 1) + self.assertEqual(topk_calls[0]["actual_seq_q"], 1) + self.assertEqual(topk_calls[0]["cp_index"], [(0, 64, 65)]) + self.assertEqual(topk_calls[0]["q"], [0.0]) + self.assertEqual(topk_calls[0]["weights"], [100.0]) + self.assertIs(topk_calls[0]["shared_index_buffer"], materialized_index) + self.assertIs(topk_calls[0]["shared_block_tables"], dense_pages) + self.assertEqual(result.shape, (64, 2)) + self.assertEqual(result[0].tolist(), [1, 1]) + self.assertTrue( + torch.equal(result[1:], torch.full((63, 2), -1, dtype=torch.int32)) + ) + def test_indexer_in_seq_cp_pair_batch_materializes_partial_current_index_reuse_once(self): import torch 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 c5acd9c50..8b0e419a5 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 @@ -297,6 +297,173 @@ class TestCPSharedPagedAllocator(CustomTestCase): self.assertEqual(owners, [0, 1]) + def test_alloc_extend_compute_owner_uses_valid_pages_not_compute_padding_pages( + self, + ): + from types import SimpleNamespace + + from sglang.srt.mem_cache import common + + page_size = 64 + + class FakeAllocator: + def __init__(self): + self.page_size = page_size + self.cp_size = 8 + self.owner_calls = [] + self.extend_num_tokens = [] + + def alloc_extend_compute_owner( + self, + _prefix_lens, + _prefix_lens_cpu, + _seq_lens, + _seq_lens_cpu, + _last_loc, + extend_num_tokens, + page_compute_owners, + ): + self.extend_num_tokens.append(int(extend_num_tokens)) + self.owner_calls.append(list(page_compute_owners)) + return torch.arange( + 1024, 1024 + int(extend_num_tokens), dtype=torch.int64 + ) + + def alloc_extend(self, *_args, **_kwargs): + raise AssertionError("legacy allocation should not be used") + + class FakeTreeCache: + def __init__(self): + self.token_to_kv_pool_allocator = FakeAllocator() + + def is_chunk_cache(self): + return False + + def evict(self, *_args, **_kwargs): + raise AssertionError("eviction should not be needed") + + server_args = SimpleNamespace( + enable_nsa_prefill_cp_shared_kv=True, + enable_nsa_prefill_context_parallel=True, + nsa_prefill_cp_mode="in-seq-split", + ) + tree_cache = FakeTreeCache() + + with patch.object(common, "get_global_server_args", return_value=server_args): + out_cache_loc = common.alloc_paged_token_slots_extend( + tree_cache=tree_cache, + prefix_lens=torch.tensor([0], dtype=torch.int64), + prefix_lens_cpu=torch.tensor([0], dtype=torch.int64), + seq_lens=torch.tensor([65], dtype=torch.int64), + seq_lens_cpu=torch.tensor([65], dtype=torch.int64), + last_loc=torch.tensor([-1], dtype=torch.int64), + extend_num_tokens=65, + ) + + self.assertEqual(out_cache_loc.numel(), 65) + self.assertEqual(tree_cache.token_to_kv_pool_allocator.extend_num_tokens, [65]) + self.assertEqual(tree_cache.token_to_kv_pool_allocator.owner_calls, [[0, 1]]) + + def test_cp_hicache_write_reservation_uses_page_tail_not_compute_padding_extent( + self, + ): + from sglang.srt.managers.cache_controller import HiCacheController + + page_size = 64 + + class FakeHostPool: + def __init__(self): + self.alloc_sizes = [] + + def alloc_contiguous_preferred(self, need_size): + self.alloc_sizes.append(int(need_size)) + return torch.arange(1024, 1024 + int(need_size), dtype=torch.int64) + + def alloc(self, need_size): + return self.alloc_contiguous_preferred(need_size) + + def free(self, _indices): + raise AssertionError("reservation should not roll back") + + controller = HiCacheController.__new__(HiCacheController) + controller.page_size = page_size + controller.cp_shared_kv_layout = CpSharedKVLayout( + page_size=page_size, cp_size=8, cp_rank=1 + ) + controller.mem_pool_host = FakeHostPool() + controller.draft_mem_pool_host = None + controller.draft_mem_pool_device = None + + reservation = controller.reserve_write_cp( + torch.arange(page_size, page_size + 65, dtype=torch.int64), + node_id=123, + ) + + self.assertEqual(controller.mem_pool_host.alloc_sizes, [page_size]) + self.assertEqual(reservation.metadata.logical_len, 65) + self.assertEqual(reservation.metadata.padded_len, page_size * 2) + self.assertEqual(reservation.metadata.page_owners.tolist(), [0, 1]) + self.assertEqual(reservation.host_indices.numel(), page_size) + self.assertEqual(reservation.physical_device_indices.numel(), page_size) + self.assertEqual( + reservation.metadata.owned_positions.tolist(), + list(range(64, 128)), + ) + + def test_cp_hicache_load_returns_valid_visible_len_while_loading_owned_page_tail( + self, + ): + from types import SimpleNamespace + + from sglang.srt.managers.cache_controller import HiCacheController + from sglang.srt.mem_cache.hiradix_cache import CpHiCacheNodeMetadata + + page_size = 64 + + class FakeDeviceAllocator: + def __init__(self): + self.owner_calls = [] + self.freed = [] + + def alloc_pages_with_owners(self, page_owners): + self.owner_calls.append(list(page_owners)) + return torch.arange(page_size, page_size * 3, dtype=torch.int64) + + def free(self, indices): + self.freed.append(indices.clone()) + + controller = HiCacheController.__new__(HiCacheController) + controller.page_size = page_size + controller.cp_shared_kv_layout = CpSharedKVLayout( + page_size=page_size, cp_size=8, cp_rank=1 + ) + controller.mem_pool_device_allocator = FakeDeviceAllocator() + controller.load_queue = [] + controller.draft_load_queue = [] + controller.draft_mem_pool_host = None + controller.draft_mem_pool_device = None + + metadata = CpHiCacheNodeMetadata( + logical_len=65, + padded_len=page_size * 2, + owned_positions=torch.arange(page_size, page_size * 2, dtype=torch.int64), + host_indices=torch.arange(1024, 1024 + page_size, dtype=torch.int64), + page_owners=torch.tensor([0, 1], dtype=torch.int8), + page_size=page_size, + ) + node = SimpleNamespace(cp_hicache=metadata, host_len=65, id=321) + + visible_device_indices = controller.load_cp([node], node_id=321) + + self.assertEqual(controller.mem_pool_device_allocator.owner_calls, [[0, 1]]) + self.assertEqual(controller.mem_pool_device_allocator.freed, []) + self.assertEqual(visible_device_indices.numel(), 65) + self.assertEqual(visible_device_indices.tolist(), list(range(64, 129))) + self.assertEqual(len(controller.load_queue), 1) + load_op = controller.load_queue[0] + self.assertEqual(load_op.host_indices.tolist(), list(range(1024, 1024 + 64))) + self.assertEqual(load_op.device_indices.tolist(), list(range(64, 128))) + def test_compute_owner_page_assignment_allows_radix_hit_suffix_with_one_page_per_rank( self, ):