Batch CP shared-KV index work for bs>1 fast paths
The bs>1 path needs index top-k, shared-index prepare, current-index compact, and current-slot compose to consume flattened batch descriptors instead of falling back to per-request or per-segment Python/Torch work. This change wires SGLang to the new TAI batch prepare kernels, keeps fallback explicit, and records the remaining HiCache/load-backup gaps in the bs>1 workstream docs. Constraint: CP shared-KV bs>1 must reuse fast paths rather than adding slow batch-only fallbacks Constraint: No new collective operations were introduced Rejected: Leave current-only cp_index as Python slice/cat | it keeps per-segment overhead in the short-extend bs>1 case Rejected: Infer max segment lengths from CUDA descriptor tensors | .item() would add CPU synchronization on the hot path Confidence: medium Scope-risk: moderate Directive: Do not remove the explicit fallback warnings without verifying the corresponding TAI symbols are present in production Tested: local py_compile for touched SGLang files Tested: remote g0034 test_nsa_cp_utils.py passed, 53 tests Tested: remote g0034 test_fill_current_index_page_slots_uses_tai_kernel_when_available passed Not-tested: full ETE bs>1 traffic with HiCache load/backup and draft/EAGLE enabled
This commit is contained in:
laoyao0822
@@ -651,6 +651,15 @@ Phase 0 -> Phase 2
|
||||
5. **workspace pressure：** bs>1 materialize buffer 更大，必须测显存和临时 buffer，不只看 latency。
|
||||
6. **collective 顺序：** 所有 CP rank 必须按完全相同的 layer/request 顺序执行 collective。
|
||||
7. **fallback 可见性：** 开发阶段优先 fail-fast；生产 fallback 必须 warning 且限频。
|
||||
8. **W2 allocator CPU descriptor overhead：** 当前 W2 correctness path 仍在
|
||||
allocation 热路径中用 Python list/tensor 临时对象构造 batch owner descriptor：
|
||||
`extend_lens` / `extend_prefix_lens` list、`flat_page_compute_owners`
|
||||
list、`positions_by_owner` list，以及每个 owner lane 的 `position_tensor`。
|
||||
这仍复用了 owner free/release bucket 本地记账，比旧全量 scan/sort 路径轻，
|
||||
但在线上大量 200-2000 token 短 extend 叠 batch 时可能成为 CPU hot path。
|
||||
后续应让 W1 batch plan 直接携带 page-owner descriptor，W2 消费
|
||||
tensor/固定 buffer，并用 kernel 或固定 workspace 完成 selected-pages
|
||||
gather/scatter，避免每次 allocation 重建 Python descriptor。
|
||||
|
||||
---
|
||||
|
||||
|
||||
@@ -253,6 +253,44 @@ out_cache_loc: torch.Tensor # flattened request order
|
||||
- CP shared-KV bs>1 进入 allocator 时不再出现 `multi_batch` fallback。
|
||||
- W3 direct write 可以信任 `out_cache_loc` 的 owner-lane 合同。
|
||||
|
||||
### 当前实现状态与性能风险
|
||||
|
||||
PR #10 已把 W2 正确性路径接入：
|
||||
|
||||
- bs>1 会按 request 独立调用 in-seq page owner planner；
|
||||
- owner list 按 request order flatten；
|
||||
- `alloc_extend_compute_owner()` 删除 bs=1 hard guard；
|
||||
- allocator 复用已有 `alloc_extend_naive()` 生成 flattened `out_cache_loc`；
|
||||
- supported CP shared-KV case 不再走 `multi_batch` legacy fallback。
|
||||
|
||||
但当前实现仍有一段 Python/control-path descriptor overhead，主要包括：
|
||||
|
||||
1. `alloc_paged_token_slots_extend()` 每次 allocation 从 CPU tensor `.item()` 构造
|
||||
`extend_lens` / `extend_prefix_lens` Python lists。
|
||||
2. `build_batch_in_seq_page_compute_owners()` 用 Python loop 生成
|
||||
`flat_page_compute_owners: List[int]`。
|
||||
3. `_select_compute_owner_pages()` 再用 Python loop 构造
|
||||
`required_by_owner` / `positions_by_owner`。
|
||||
4. 每个 owner lane 会临时构造 `position_tensor`，再把 lane pages scatter 回
|
||||
`selected_pages`。
|
||||
|
||||
这比旧的全量 free-page scan / `torch.isin` / sort-merge 路径轻，因为 allocator
|
||||
仍使用 `_owner_free_pages` / `_owner_release_pages` 本地 bucket 记账；但它还不是最终
|
||||
descriptor 化 fast path。在线上大量 200-2000 token 短 extend 叠 batch 时，allocation
|
||||
频率高，这段 Python list/tensor descriptor 构造可能重新成为 CPU hot path。
|
||||
|
||||
后续优化方向：
|
||||
|
||||
```text
|
||||
W1 batch plan 直接携带 page-owner descriptor
|
||||
-> W2 allocator 消费 tensor/descriptor 而不是重新 Python 推导
|
||||
-> selected_pages gather/scatter 用固定 buffer 或 kernel 完成
|
||||
-> 避免每 tick/request 反复创建 Python lists 和临时 position tensors
|
||||
```
|
||||
|
||||
这个风险不影响当前 W2 correctness，但需要在 bs>1 ETE/perf 阶段用 micro benchmark
|
||||
和线上 trace 单独验证。
|
||||
|
||||
---
|
||||
|
||||
## 6. W3：CP split/rebuild + direct write fast path
|
||||
@@ -662,3 +700,48 @@ runtime / kernel 都消费 CPSharedKVBatchPlan descriptors
|
||||
- 设计 variable-length descriptor benchmark，覆盖 bf16/fp8、page_first_direct、random/owner-lane pages。
|
||||
|
||||
第一批完成后，再开始 W3/W4/W6/W7 draft 接入。
|
||||
|
||||
---
|
||||
|
||||
## 15. 2026-06-03 W5/W4 kernel 接入状态更新
|
||||
|
||||
### 已接入并验证
|
||||
|
||||
1. **NSA index top-k 的 batch cp_index 调用已从 per request/segment MQA top-k 收敛为单次调用。**
|
||||
- SGLang 路径：`python/sglang/srt/layers/attention/nsa/nsa_indexer.py::_get_topk_in_seq_cp_pair_batch`。
|
||||
- 语义：多个 request 的 prev/next CP segment 先 compact 成一个 `cp_index` descriptor，再调用一次 `_get_topk_ragged_with_cp`。
|
||||
- 已有测试：`test_indexer_in_seq_cp_pair_batch_*` 与 `test_indexer_ragged_cp_index_current_batch_does_not_materialize`。
|
||||
|
||||
2. **index partial/current compose 的 current-slot fill 已接入 TAI kernel。**
|
||||
- SGLang 路径：`cp_shared_kv_runtime.fill_current_index_page_slots`。
|
||||
- TAI kernel：`tai_kernel.nsa_prefill.cp_shared_kv_materialize.fill_current_index_page_slots`。
|
||||
- 输入是 flatten 后的 current rows，可一次处理 bs>1 的 current suffix；不再在 CUDA 上默认走 PyTorch advanced indexing。
|
||||
- fallback 仍为显式 warning：`[CP_SHARED_KV_FALLBACK][tai_materialize] reason=index_current_fill_*`。
|
||||
- 远端 CUDA smoke：通过 64/page slot compose 的 K/scale 字节布局检查。
|
||||
|
||||
3. **shared-index batch cp_index 的 K/S + range prepare 已接入 TAI batch descriptor kernel。**
|
||||
- SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_index_batch`。
|
||||
- TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_kv_and_range_batch`。
|
||||
- descriptor：`batch_indices, kv_lens, q_starts, q_lens, k_bases, q_bases, total_kv_len, total_q_count, max_kv_len, max_q_len`。
|
||||
- 关键约束：`max_kv_len/max_q_len` 由 Python planner 传入，避免在 wrapper 内对 CUDA tensor 做 `.max().item()` 同步。
|
||||
- 输出合同：最终仍是按 segment concat 的 dense K buffer / scale buffer；`ks` 为 segment 的 K base；`ke_offset` 为 segment-local causal end offset，调用端计算 `ke = ks + ke_offset`。
|
||||
- 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。
|
||||
|
||||
### 仍需继续处理的 kernel/runtime 缺口
|
||||
|
||||
1. **current-only cp_index 的 K/S compact copy 已接入 TAI batch kernel。**
|
||||
- SGLang wrapper：`cp_shared_kv_runtime.try_tai_prepare_cp_mqa_current_index_batch`。
|
||||
- TAI kernel：`tai_kernel.nsa_prefill.cp_index_mqa_prepare.prepare_cp_mqa_current_kv_and_range_batch`。
|
||||
- 语义：从 flatten current rows 按 segment `current_base/current_len` 拷到 concat K/S buffer，同时生成 `ks/ke_offset`。
|
||||
- 远端 CUDA smoke：K bytes、scale float、`ks`、`ke_offset` 与 CPU reference 一致。
|
||||
|
||||
2. **HiCache per-layer D2H backup 仍按 reservation/node 提交。**
|
||||
- `scheduler._prepare_hicache_write_backups_before_forward()` 仍 `for req in batch.reqs: prepare_write_backup_for_req(req)`。
|
||||
- `cache_controller.submit_write_cp_layer()` 每个 reservation 在每层调用 host pool backup。
|
||||
- 底层 TAI direct transfer kernel 接收 flatten indices，本身可以吃 batched descriptor；但 runtime 还没有把多个 request/reservation 合并成一个 layer-level descriptor。
|
||||
|
||||
3. **L2->L1 load/prefetch 需要继续确认实际 queue 合并粒度。**
|
||||
- `cache_controller.load_cp()` / `start_loading()` 已比 backup 更接近 batched queue，但仍需用 ETE/NVTX 确认 layer-level descriptor 是否足够大，是否存在 per-node launch。
|
||||
|
||||
4. **MLA current/partial-current 和 direct store 已有 flatten kernel 基础，但还需用 bs>1 ETE 验证没有 fallback hot path。**
|
||||
- 特别关注 FP8/BF16 两种 dtype、draft/EAGLE 路径、prefetch 关闭时 full/current reuse 是否仍启用。
|
||||
|
||||
Reference in New Issue
Block a user