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:
laoyao0822
2026-06-12 03:49:10 +08:00
parent 142e7a5a64
commit bafb55044b
3 changed files with 115 additions and 22 deletions

View File

@@ -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 pagesfp8 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-fillcompact 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 helperpublish 使用 dense destination pagespeer 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.459msIPC compose p50 0.310ms。