diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md index a3012e440..8dc58a9ce 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md @@ -9,7 +9,7 @@ **当前约束：** 先保证 target model 同步路径正确，再恢复 current/partial-current reuse，再处理 EAGLE/draft，最后打开 MLA/index L1 prefetch。不能通过删除 `batch_size != 1` guard 来“强行开启”。 - 并行派工版计划：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_parallel_workstreams_zh.md`。 -- W3/W4 target sync 细化文档：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md`。 +- W3/W4-1 target sync 细化文档：`docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md`。 --- @@ -199,7 +199,7 @@ target sync correctness - `python/sglang/srt/layers/attention/nsa_backend.py:1940-2020` - MLA partial-current compose 同样是单 prefix 合同。 -结论：Phase 5 需要做 batched current suffix slicing 和 remap。Phase 4 可先只做 sync correctness。 +结论：Phase 5 需要做 batched current suffix slicing 和 remap。W4-1 可先只做 target index/top-k sync correctness。 ### C8. L1 shared-KV prefetch 是单请求结构 @@ -220,7 +220,7 @@ target sync correctness - `_get_topk_in_seq_cp_pair()` 使用 scalar `kv_len_prev/next` 和 scalar `actual_seq_q_prev/next`； - 源码里记录 multi-batch 尚未支持。 -结论：Phase 4 必须改成 request-local prev/next slices。 +结论：W4-1 必须改成 request-local prev/next slices。 ### C10. EAGLE / draft local path 对单请求 metadata 敏感 @@ -449,7 +449,7 @@ cp_size=8 - target MLA KV 和 index direct-write 可以使用 batched local loc。 -### Phase 4：target index/top-k 同步正确性 +### W4-1：target index/top-k 同步正确性 目标：不依赖 async prefetch，先让 target index/top-k 在 bs>1 下正确。 @@ -466,7 +466,7 @@ cp_size=8 2. 每个 request 生成自己的 prev/next top-k pair。 3. 按 flattened local query order 重新拼接 top-k result。 4. index materialize 先走同步路径：构造 per-request prefix page range，再 flatten。 -5. Phase 4 不打开 bs>1 index prefetch。 +5. W4-1 不打开 bs>1 index prefetch。 测试重点： diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md index 990a76697..fa1b6d5c0 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md @@ -1,20 +1,20 @@ -# NSA Prefill CP Shared-KV bs>1 W3/W4 实现文档 +# NSA Prefill CP Shared-KV bs>1 W3/W4-1 实现文档 > 日期：2026-06-03 > 分支：`cjy-cp-refactor` > 当前基线：`e4cf8d18b` -> 范围：W3 `local out_cache_loc + direct write`，W4 `target index/top-k sync correctness`。 +> 范围：W3 `local out_cache_loc + direct write`，W4-1 `target index/top-k sync correctness`。 -> **命名说明：** 本文的 W4 指顺序实现计划中的 **Phase 4 target index/top-k sync correctness**。并行派工文档里的 W4 是 current/partial-current reuse；那部分在本文中仍视为后续阶段，不在本轮实现范围内。 +> **命名说明：** 本文的 **W4-1** 指顺序实现计划中的 **target index/top-k sync correctness**。并行派工文档/原定 W4 是 **current/partial-current reuse**；那部分在本文中仍视为后续阶段，不在本轮实现范围内。 ## 0. 目标和非目标 -目标是在 **不等待 W2 allocator 最终完成** 的前提下，先把 target model 的 W3/W4 runtime consumer 做成 batch-aware： +目标是在 **不等待 W2 allocator 最终完成** 的前提下，先把 target model 的 W3/W4-1 runtime consumer 做成 batch-aware： 1. bs>1 的 CP shared-KV direct write 使用每个 request 独立的 page-aligned split，再按 request order 拼接本 rank local rows。 2. MLA KV direct write 和 index KV direct write 共享同一个 ForwardBatch 级 page/local-loc plan。 3. target index/top-k sync path 使用 per-request metadata，不能把 batch flatten 成一条长序列。 -4. W4 第一版只做 sync correctness，不启用 bs>1 current reuse、partial-current reuse、L1 prefetch、draft/EAGLE。 +4. W4-1 第一版只做 sync correctness，不启用 bs>1 current reuse、partial-current reuse、L1 prefetch、draft/EAGLE。 非目标： @@ -36,7 +36,7 @@ CP shared-KV 下 page 相关规划是 **request/batch 级别**，不是 layer 只有 KV bytes、index K/scale、materialized dense buffer 内容、topk/logits 变化 ``` -因此 W3/W4 的实现不能每层重新推导 page plan。应该在 `ForwardBatch` 上缓存： +因此 W3/W4-1 的实现不能每层重新推导 page plan。应该在 `ForwardBatch` 上缓存： - local logical `out_cache_loc`； - local physical `out_cache_loc`； @@ -50,7 +50,7 @@ CP shared-KV 下 page 相关规划是 **request/batch 级别**，不是 layer - `ForwardBatch.cp_shared_kv_paged_slot_remap_key/remap` 已在 `forward_batch_info.py:432-433`。 - `get_or_build_shared_paged_buffer_slot_remap()` 已在 `cp_shared_kv_runtime.py:2486-2520` 按 key 缓存 paged remap。 -W3/W4 的实现原则： +W3/W4-1 的实现原则： ```text prepare/batch plan 阶段确定 page 和 segment。 @@ -164,7 +164,7 @@ CP shared-KV 合同内 + local loc unavailable / owner mismatch / shape mismatch 可以返回 False 走原路径；这不是 CP shared-KV fallback ``` -### C7. 当前 W4 `_get_topk_in_seq_cp_pair()` 是单请求 +### C7. 当前 W4-1 `_get_topk_in_seq_cp_pair()` 是单请求 `_get_topk_in_seq_cp_pair()` 在 `nsa_indexer.py:1371-1455`。 @@ -202,13 +202,13 @@ forward_batch.extend_seq_lens_cpu[0] block_tables[0] ``` -W4 如果想复用这个函数处理 bs>1 的每个 request，就必须新增 `batch_idx: int = 0` 参数，并用 `block_tables[batch_idx]`、`seq_lens_cpu[batch_idx]`、`extend_seq_lens_cpu[batch_idx]`。 +W4-1 如果想复用这个函数处理 bs>1 的每个 request，就必须新增 `batch_idx: int = 0` 参数，并用 `block_tables[batch_idx]`、`seq_lens_cpu[batch_idx]`、`extend_seq_lens_cpu[batch_idx]`。 -### C9. 不应使用现有 `cp_index` 分支作为 W4 第一版 +### C9. 不应使用现有 `cp_index` 分支作为 W4-1 第一版 `_get_topk_ragged_with_cp()` 的 `cp_index` 分支在 `nsa_indexer.py:1155-1222`，源码注释明确写着 `TODO Multi-batch support has accuracy issues`。 -W4 第一版为了 correctness 应该避开该分支，采用 per-request/per-segment 同步调用： +W4-1 第一版为了 correctness 应该避开该分支，采用 per-request/per-segment 同步调用： ```text for req_id: @@ -218,7 +218,7 @@ for req_id: 后续性能优化可以再把多个 segment 合并成 batched top-k descriptor。 -### C10. W4 current/partial-current reuse 仍是单请求合同 +### C10. 原定 W4 current/partial-current reuse 仍是单请求合同 `_maybe_materialize_shared_index_buffer()` 在 `nsa_indexer.py:309-517`。 @@ -231,13 +231,13 @@ positive page-aligned prefix 见 `nsa_indexer.py:341-355`。 -因此 W4 第一版遇到 bs>1 + `current_index_kv is not None` 必须 fail-fast： +因此 W4-1 第一版遇到 bs>1 + `current_index_kv is not None` 必须 fail-fast： ```text [CP_SHARED_KV_FAIL_FAST][batch_gt1_index_current_reuse_unsupported] ``` -不要在 W4 里临时拼 current reuse，否则会和 W5 的 partial/current reuse 工作混在一起。 +不要在 W4-1 里临时拼 current reuse，否则会和 W5 的 partial/current reuse 工作混在一起。 ### C11. `_build_batch_metadata_from_plan()` 的 batch cu tensor 不能直接传给单 segment top-k @@ -250,7 +250,7 @@ request_actual_seq_q_next_cu_tensor = [0] + cumsum(request_actual_seq_q_next) 这是 batch-level cumulative tensor，不是单个 request segment 的 `[0, segment_len]`。 -W4 的 per-request/per-segment top-k 调用应该传： +W4-1 的 per-request/per-segment top-k 调用应该传： ```text actual_seq_q_tensor = tensor([segment_len]) @@ -411,11 +411,11 @@ mock `_store_index_k_cache()`，在 bs>1 local loc 不可用时确保不会调 构造 bs>1 local loc tokens=N，但 `k_nope/k_pe` tokens != N，断言 fail-fast。 -## 4. W4 设计：target index/top-k sync correctness +## 4. W4-1 设计：target index/top-k sync correctness ### 4.1 输入合同 -W4 依赖： +W4-1 依赖： - W1 batch metadata； - W3 local q/weights 已按 request boundary split； @@ -440,7 +440,7 @@ concat( ### 4.3 实现步骤 -#### W4-S1：给 `_get_topk_in_seq_cp_pair()` 增加 batch dispatch +#### W4-1-S1：给 `_get_topk_in_seq_cp_pair()` 增加 batch dispatch ```python if getattr(metadata, "batch_size", 1) > 1: @@ -449,7 +449,7 @@ if getattr(metadata, "batch_size", 1) > 1: 保留原 scalar path。 -#### W4-S2：新增 `_get_topk_in_seq_cp_pair_batch()` +#### W4-1-S2：新增 `_get_topk_in_seq_cp_pair_batch()` 伪代码： @@ -490,7 +490,7 @@ zero segment： segment_len == 0 -> 返回 shape (0, index_topk) 的 empty tensor，不调用底层 MQA/topk ``` -#### W4-S3：泛化 `_get_topk_ragged_with_cp(..., batch_idx=0)` +#### W4-1-S3：泛化 `_get_topk_ragged_with_cp(..., batch_idx=0)` 新增参数： @@ -506,7 +506,7 @@ batch_idx: int = 0 scalar path 默认 `batch_idx=0`，行为不变。 -#### W4-S4：单 segment cu tensor +#### W4-1-S4：单 segment cu tensor 每个 segment 调用 `_get_topk_ragged_with_cp()` 时传： @@ -517,11 +517,11 @@ actual_seq_q_cu_tensor = torch.tensor([0, segment_len], device=q_fp8.device, dty 不要传 `request_actual_seq_q_prev_cu_tensor` 的 batch cumulative view。 -#### W4-S5：不要使用 `cp_index` 分支 +#### W4-1-S5：不要使用 `cp_index` 分支 -现有 `cp_index` 分支虽然看起来能表达 multi-batch，但源码已标注 accuracy issue。W4 第一版只追 correctness，明确不使用。 +现有 `cp_index` 分支虽然看起来能表达 multi-batch，但源码已标注 accuracy issue。W4-1 第一版只追 correctness，明确不使用。 -### 4.4 W4 测试 +### 4.4 W4-1 测试 #### T1：bs>1 per-request prev/next 调用顺序 @@ -570,13 +570,13 @@ mock `_maybe_materialize_shared_index_buffer()`，断言 bs>1 top-k 只调用一 ## 5. 执行顺序 -建议按以下顺序实施，避免把 W3/W4 的 bug 混在一起： +建议按以下顺序实施，避免把 W3/W4-1 的 bug 混在一起： 1. W3 tests：batch local loc / owner mismatch / physical loc cache。 2. W3 implementation：修 `get_cp_shared_kv_local_out_cache_loc()` 的 bs>1 expected length。 3. W3 direct-write guard：index/MLA bs>1 local loc unavailable 不 silent fallback。 -4. W4 tests：batch top-k per-request order / materialize once / zero segment / current reuse fail-fast。 -5. W4 implementation：batch dispatch + `_get_topk_ragged_with_cp(batch_idx=...)`。 +4. W4-1 tests：batch top-k per-request order / materialize once / zero segment / current reuse fail-fast。 +5. W4-1 implementation：batch dispatch + `_get_topk_ragged_with_cp(batch_idx=...)`。 6. 只跑 unit，不跑 ETE；W2 allocator 未合入前，ETE 仍可能被 allocator blocker 阻塞。 ## 6. 验证命令 @@ -622,15 +622,15 @@ out_cache_loc.numel() == sum(request_extend_lens) ### R2. per-segment top-k 会增加调用次数 -W4 第一版 bs=N 会最多调用 `2N` 次 `_get_topk_ragged_with_cp()`。这是 correctness-first 设计。后续性能优化可以引入 batched segment descriptor 或修复 `cp_index` 分支，但不能在第一版混入。 +W4-1 第一版 bs=N 会最多调用 `2N` 次 `_get_topk_ragged_with_cp()`。这是 correctness-first 设计。后续性能优化可以引入 batched segment descriptor 或修复 `cp_index` 分支，但不能在第一版混入。 ### R3. current reuse 不支持会影响 cache-hit 性能 -W4 只保证 sync correctness。cache-hit 的性能收益需要 W5 current/partial-current reuse 恢复后再评估。 +W4-1 只保证 sync correctness。cache-hit 的性能收益需要 W5 current/partial-current reuse 恢复后再评估。 ### R4. fallback 策略需要保持醒目 -W3/W4 支持范围内不能 silent fallback。尤其 index direct-write 当前 fallback 很隐蔽，必须在 bs>1 CP shared-KV 下收窄。 +W3/W4-1 支持范围内不能 silent fallback。尤其 index direct-write 当前 fallback 很隐蔽，必须在 bs>1 CP shared-KV 下收窄。 ### R5. page descriptor 必须 layer-invariant diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 50721e7ca..acc3b4586 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -1099,6 +1099,7 @@ class Indexer(MultiPlatformOp): shared_block_tables: Optional[torch.Tensor] = None, actual_seq_q_tensor: Optional[torch.Tensor] = None, actual_seq_q_cu_tensor: Optional[torch.Tensor] = None, + batch_idx: int = 0, ) -> torch.Tensor: if TYPE_CHECKING: assert isinstance(forward_batch.token_to_kv_pool, NSATokenToKVPool) @@ -1221,14 +1222,16 @@ class Indexer(MultiPlatformOp): batch_idx_list=batch_idx_list, ) else: + seq_len = int(forward_batch.seq_lens_cpu[batch_idx].item()) + extend_seq_len = int(forward_batch.extend_seq_lens_cpu[batch_idx]) cp_kv_end = ( - forward_batch.seq_lens_cpu[0].item() - - forward_batch.extend_seq_lens_cpu[0] + seq_len + - extend_seq_len + kv_len ) page_table_1 = metadata.get_page_table_1() logical_kv_limit = min( - int(forward_batch.seq_lens_cpu[0].item()), + seq_len, int(page_table_1.shape[1]), ) valid_q_count = _compute_contiguous_valid_cp_query_count( @@ -1253,7 +1256,7 @@ class Indexer(MultiPlatformOp): assert index_buffer is not None tai_prepared = try_tai_prepare_cp_mqa_index( index_buffer=index_buffer, - page_indices=block_tables[0], + page_indices=block_tables[batch_idx], kv_len=kv_len, valid_q_count=valid_q_count, ke_start=ke_start, @@ -1283,13 +1286,13 @@ class Indexer(MultiPlatformOp): forward_batch.token_to_kv_pool, index_buffer, seq_len=kv_len, - page_indices=block_tables[0], + page_indices=block_tables[batch_idx], ) k_scale = index_buf_accessor.GetS.execute( forward_batch.token_to_kv_pool, index_buffer, seq_len=kv_len, - page_indices=block_tables[0], + page_indices=block_tables[batch_idx], ) k_fp8 = k_fp8.view(torch.float8_e4m3fn) @@ -1378,6 +1381,16 @@ 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: + return self._get_topk_in_seq_cp_pair_batch( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + current_index_kv=current_index_kv, + ) + kv_len_prev = forward_batch.nsa_cp_metadata.kv_len_prev kv_len_next = forward_batch.nsa_cp_metadata.kv_len_next actual_seq_q_prev = forward_batch.nsa_cp_metadata.actual_seq_q_prev @@ -1456,6 +1469,131 @@ class Indexer(MultiPlatformOp): ) return torch.cat([topk_result_prev, topk_result_next], dim=0) + def _get_topk_in_seq_cp_pair_batch( + self, + forward_batch: ForwardBatch, + layer_id: int, + q_fp8: torch.Tensor, + weights: torch.Tensor, + metadata: BaseIndexerMetadata, + current_index_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + ) -> torch.Tensor: + cp_metadata = forward_batch.nsa_cp_metadata + assert cp_metadata is not None + batch_size = int(getattr(cp_metadata, "batch_size", 1) or 1) + if current_index_kv is not None: + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_gt1_index_current_reuse_unsupported " + f"batch_size={batch_size} layer_id={layer_id}" + ) + + 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 [] + ) + request_actual_seq_q_next = list( + getattr(cp_metadata, "request_actual_seq_q_next", []) or [] + ) + 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 + ): + 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}" + ) + + shared_block_tables = metadata.get_page_table_64() + shared_index_buffer, shared_block_tables = ( + self._maybe_materialize_shared_index_buffer( + forward_batch, + layer_id, + shared_block_tables, + ) + ) + + outputs = [] + cursor = 0 + + def call_segment( + *, + req_id: int, + segment_len: int, + kv_len: int, + ) -> torch.Tensor: + nonlocal cursor + segment_len = int(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 torch.empty( + (0, self.index_topk), + dtype=torch.int32, + device=q_fp8.device, + ) + actual_seq_q_tensor = torch.tensor( + [segment_len], + dtype=torch.int32, + device=q_fp8.device, + ) + actual_seq_q_cu_tensor = torch.tensor( + [0, segment_len], + dtype=torch.int32, + device=q_fp8.device, + ) + return self._get_topk_ragged_with_cp( + forward_batch, + layer_id, + q_segment, + weights_segment, + metadata, + kv_len, + segment_len, + current_index_kv=None, + shared_index_buffer=shared_index_buffer, + shared_block_tables=shared_block_tables, + actual_seq_q_tensor=actual_seq_q_tensor, + actual_seq_q_cu_tensor=actual_seq_q_cu_tensor, + batch_idx=req_id, + ) + + for req_id in range(batch_size): + outputs.append( + call_segment( + req_id=req_id, + segment_len=request_actual_seq_q_prev[req_id], + kv_len=request_kv_len_prev[req_id], + ) + ) + outputs.append( + call_segment( + req_id=req_id, + segment_len=request_actual_seq_q_next[req_id], + kv_len=request_kv_len_next[req_id], + ) + ) + + if cursor != int(q_fp8.shape[0]) or cursor != int(weights.shape[0]): + raise RuntimeError( + "[CP_SHARED_KV_FAIL_FAST][index_topk] " + "reason=batch_gt1_index_q_length_mismatch " + f"batch_size={batch_size} layer_id={layer_id} cursor={cursor} " + f"q_tokens={int(q_fp8.shape[0])} weights_tokens={int(weights.shape[0])}" + ) + + if not outputs: + return torch.empty((0, self.index_topk), dtype=torch.int32, device=q_fp8.device) + return torch.cat(outputs, dim=0) + def forward_indexer( self, q_fp8: torch.Tensor, diff --git a/test/registered/unit/layers/test_nsa_cp_utils.py b/test/registered/unit/layers/test_nsa_cp_utils.py index e7e2c5e6d..db2496f73 100644 --- a/test/registered/unit/layers/test_nsa_cp_utils.py +++ b/test/registered/unit/layers/test_nsa_cp_utils.py @@ -1090,6 +1090,156 @@ class TestNSAInSeqCPUtils(unittest.TestCase): self.assertEqual(topk_calls[1]["actual_seq_q_cu_tensor"].tolist(), [0, 2]) self.assertEqual(result.tolist(), [[1, 1], [1, 1], [1, 1], [2, 2], [2, 2]]) + def test_indexer_in_seq_cp_pair_batch_preserves_request_segment_order(self): + import torch + + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer.index_topk = 2 + logical_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) + materialized_index = torch.tensor([11], dtype=torch.int32) + dense_pages = torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.int32) + materialize_calls = [] + topk_calls = [] + + class Metadata: + def get_page_table_64(self): + return logical_pages + + def fake_materialize(forward_batch, layer_id, logical_page_table): + materialize_calls.append((layer_id, logical_page_table)) + return materialized_index, dense_pages + + def fake_get_topk( + forward_batch, + layer_id, + q_fp8, + weights, + metadata, + kv_len, + actual_seq_q, + cp_index=None, + current_index_kv=None, + shared_index_buffer=None, + shared_block_tables=None, + actual_seq_q_tensor=None, + actual_seq_q_cu_tensor=None, + batch_idx=0, + ): + topk_calls.append( + { + "batch_idx": batch_idx, + "kv_len": kv_len, + "actual_seq_q": actual_seq_q, + "q": q_fp8.flatten().tolist(), + "weights": weights.flatten().tolist(), + "actual_seq_q_tensor": actual_seq_q_tensor, + "actual_seq_q_cu_tensor": actual_seq_q_cu_tensor, + "shared_index_buffer": shared_index_buffer, + "shared_block_tables": shared_block_tables, + "current_index_kv": current_index_kv, + } + ) + return torch.full( + (actual_seq_q, 2), + len(topk_calls), + dtype=torch.int32, + ) + + indexer._maybe_materialize_shared_index_buffer = fake_materialize + indexer._get_topk_ragged_with_cp = fake_get_topk + + forward_batch = SimpleNamespace( + batch_size=2, + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + kv_len_prev=100, + kv_len_next=200, + actual_seq_q_prev=2, + actual_seq_q_next=1, + actual_seq_q_prev_cu_tensor=torch.tensor([0, 2], dtype=torch.int32), + actual_seq_q_next_cu_tensor=torch.tensor([0, 1], dtype=torch.int32), + request_kv_len_prev=[100, 300], + request_kv_len_next=[200, 400], + request_actual_seq_q_prev=[2, 1], + request_actual_seq_q_next=[1, 3], + ), + ) + q_fp8 = torch.arange(7, dtype=torch.float32).view(7, 1) + weights = (torch.arange(7, dtype=torch.float32) + 100).view(7, 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( + [ + ( + call["batch_idx"], + call["kv_len"], + call["actual_seq_q"], + call["q"], + call["weights"], + call["actual_seq_q_tensor"].tolist(), + call["actual_seq_q_cu_tensor"].tolist(), + ) + for call in topk_calls + ], + [ + (0, 100, 2, [0.0, 1.0], [100.0, 101.0], [2], [0, 2]), + (0, 200, 1, [2.0], [102.0], [1], [0, 1]), + (1, 300, 1, [3.0], [103.0], [1], [0, 1]), + (1, 400, 3, [4.0, 5.0, 6.0], [104.0, 105.0, 106.0], [3], [0, 3]), + ], + ) + self.assertTrue(all(call["shared_index_buffer"] is materialized_index for call in topk_calls)) + self.assertTrue(all(call["shared_block_tables"] is dense_pages for call in topk_calls)) + self.assertTrue(all(call["current_index_kv"] is None for call in topk_calls)) + self.assertEqual( + result.tolist(), + [[1, 1], [1, 1], [2, 2], [3, 3], [4, 4], [4, 4], [4, 4]], + ) + + def test_indexer_in_seq_cp_pair_batch_rejects_current_index_reuse(self): + import torch + + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + forward_batch = SimpleNamespace( + batch_size=2, + nsa_cp_metadata=NSAContextParallelMetadata( + batch_size=2, + request_kv_len_prev=[100, 300], + request_kv_len_next=[200, 400], + request_actual_seq_q_prev=[2, 1], + request_actual_seq_q_next=[1, 3], + ), + ) + + with self.assertRaisesRegex( + RuntimeError, + "CP_SHARED_KV_FAIL_FAST.*batch_gt1_index_current_reuse_unsupported", + ): + Indexer._get_topk_in_seq_cp_pair( + indexer, + forward_batch, + layer_id=7, + q_fp8=torch.empty(7, 1), + weights=torch.empty(7, 1), + metadata=SimpleNamespace(), + current_index_kv=(torch.empty(1), torch.empty(1)), + ) + def test_indexer_in_seq_cp_pair_skips_materialize_when_current_index_reused(self): import torch