Keep CP current IPC staging proportional to touched pages
Cache-hit bs>1 current reuse can create very large dense attention buffers while touching only a small set of current pages. The previous SGLang runtime asked tai-kernel for a staging buffer sized like the full dense tensor, which caused CUDA OOM before the current IPC fast path could run. Switch token and index current IPC helpers to descriptor-compact staging: publish the dense destination pages into compact staging slots and materialize peers from compact source page ids back to the original dense destination pages. Document the failure mode and the compact-staging contract so the dense-sized contract is not reintroduced. Constraint: CUDA + SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 must fail fast instead of silently falling back to current-slot all_reduce Rejected: Let staging allocation failure fall back to all_reduce | hides the bug and restores the expensive collective path Rejected: Size staging by the full dense tensor | reproduces the 965MB staging OOM on long-prefix cache-hit batches Confidence: high Scope-risk: moderate Directive: Current IPC helper source ids are compact staging ids; destination ids remain dense slot pages Tested: Remote cjy-glm5-new PYTHONPATH=python:/mnt/beegfs/cjy/tai-kernel/python python -m pytest -q test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py -> 144 passed, 2 subtests passed Tested: Local py_compile cp_shared_kv_runtime.py Not-tested: Full ETE service restart with production traffic after this commit (cherry picked from commit 906ecbe5d4f08b73242e98e2b628e26516d5b04a)
This commit is contained in:
@@ -269,3 +269,60 @@ PYTHONPATH=python torchrun --standalone --nproc_per_node=8 \
|
||||
| 5 | 200k | 65k | 7.387 ms | 7.366 ms | 基本持平 |
|
||||
|
||||
BF16 下收益不明显,原因是两阶段 current publish 的额外 copy 被放大;当前线上 GLM5 使用 fp8 KV cache,因此优先级仍然是把 fp8 bs>1 fast path 接稳。后续若要兼顾 BF16,需要 fused current fill+publish,减少 current staging 的额外 HBM copy。
|
||||
|
||||
## 2026-06-12 current IPC staging OOM 修复
|
||||
|
||||
### 现象
|
||||
|
||||
远端 cache-hit bs>1 跑到 `materialize_prefix_and_reuse_current_kv_page_slots` 时触发 fail-fast:
|
||||
|
||||
```text
|
||||
[CP_SHARED_KV_FAIL_FAST][tai_ipc_materialize]
|
||||
reason=token_current_ipc_unavailable
|
||||
dense_shape=(1470528, 1, 656)
|
||||
```
|
||||
|
||||
前置 warning 指向真实原因:
|
||||
|
||||
```text
|
||||
[CP_SHARED_KV_FALLBACK][tai_ipc_materialize]
|
||||
reason=current_staging_setup_failed
|
||||
required_nbytes=964666368
|
||||
error=CUDA error: out of memory
|
||||
```
|
||||
|
||||
### 根因
|
||||
|
||||
旧 current IPC staging 合同是:
|
||||
|
||||
```text
|
||||
staging[dense_page_id] = dense_current[dense_page_id]
|
||||
peer read staging[dense_page_id]
|
||||
```
|
||||
|
||||
这要求 staging 至少和整个 attention-visible dense buffer 一样大。cache-hit bs>1 下 dense buffer 包含长 prefix slot + current slot,即使 current spans 只有几百页,也会额外申请接近 1GB 的 staging。显存紧张时 staging 分配失败,随后 fail-fast;如果放开 fallback,则会退回 current-slot all_reduce,重新引入同步和性能问题。
|
||||
|
||||
### 修复合同
|
||||
|
||||
current IPC 改为 compact staging:
|
||||
|
||||
```text
|
||||
staging[compact_i] = dense_current[dense_page_id_i]
|
||||
peer read staging[compact_i] -> dense_current[dense_page_id_i]
|
||||
```
|
||||
|
||||
其中 `compact_i` 只覆盖本次 current spans 的 page 数。以上述报错为例,实际 spans 约 676 pages,fp8 page bytes 为 `64 * 656`,staging 从约 965MB 降到约 28MB。
|
||||
|
||||
### 约束
|
||||
|
||||
1. compact staging 只用于 current 临时数据;persistent prefix/L1 cache IPC 仍直接从 owner page buffer 读取。
|
||||
2. zero/invalid slot 仍保留 `owner=-1/src=-1` 语义,materialize 端负责 zero-fill;compact source index 只对 valid slot 生效。
|
||||
3. 不允许回到静默 all_reduce fallback。CUDA + TAI materialize 开启时,compact IPC 不可用仍应 fail-fast。
|
||||
4. 后续 fused current fill+publish 仍有价值:compact staging 解决容量问题,但 current 仍有一次额外 HBM copy。
|
||||
|
||||
### 验证
|
||||
|
||||
- 新增 TAI CUDA 单测:dense 10 pages、staging 3 pages,只发布 `[1, 4, 9]` 到 compact staging,验证不需要 dense-sized staging。
|
||||
- SGLang runtime 单测覆盖 token/index current helper:publish 使用 dense destination pages,peer materialize 使用 compact source pages。
|
||||
- `benchmark_cp_shared_kv_ipc_gather.py` 已同步改成 compact current staging 合同;quick smoke:
|
||||
`bs=2 cached=4096 extend=1024 fp8/uint8` 下 dense all_reduce p50 0.459ms,IPC compose p50 0.310ms。
|
||||
|
||||
Reference in New Issue
Block a user