chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Preserve request boundaries in CP shared-KV index top-k

Browse Source 
W4-1 needs target index/top-k sync correctness before current/partial-current reuse can be restored. Batch-size>1 in-seq CP produces local q/weights in request-segment order, so top-k must consume req0_prev, req0_next, req1_prev, req1_next rather than treating the flattened batch as one scalar prev/next pair.

The implementation adds a batch dispatch for _get_topk_in_seq_cp_pair, reuses one synchronous shared-index materialization per layer, and calls _get_topk_ragged_with_cp per request segment with an explicit batch_idx. The scalar bs=1 path remains unchanged.

Constraint: This is W4-1 target index/top-k sync correctness; original W4 current/partial-current reuse remains a separate follow-up.

Constraint: Phase W4-1 must not enable bs>1 index prefetch, current reuse, partial-current reuse, or the cp_index multi-batch branch.

Rejected: Use cp_index branch for multi-batch | source marks that path as having accuracy issues.

Rejected: Pad batch requests to max length | wastes compute and violates packed/ragged batch contract.

Confidence: high

Scope-risk: moderate

Directive: Keep bs>1 target top-k ordered by request segment unless a later fused descriptor proves identical ordering and correctness.

Tested: Remote g0034 py_compile for nsa_indexer.py

Tested: Remote g0034 PYTHONPATH=python pytest test/registered/unit/layers/test_nsa_cp_utils.py -> 45 passed

Tested: Remote g0034 PYTHONPATH=python pytest test_nsa_cp_utils.py test_cp_shared_kv_layout.py test_cp_shared_kv_runtime.py -> 172 passed, 2 subtests passed

Not-tested: Full ETE bs>1 serving run with live traffic
This commit is contained in:
laoyao0822
2026-06-03 02:36:31 +08:00
parent f8b4f1915e
commit a7472c415f
4 changed files with 329 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md
View File

@@ -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。
测试重点：

60
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md
View File

@@ -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

150
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -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,

150
test/registered/unit/layers/test_nsa_cp_utils.py
View File

@@ -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