chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 14 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
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_implementation_plan_zh.md
+5 -5
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。
测试重点：
docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_w3_w4_plan_zh.md
+30 -30
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