From bafb55044bdbc6a502944e40207db2ddccef1fcf Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Fri, 12 Jun 2026 03:49:10 +0800 Subject: [PATCH] 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) --- ...d_kv_ipc_collective_replacement_plan_zh.md | 57 ++++++++++++++++ .../attention/nsa/cp_shared_kv_runtime.py | 68 ++++++++++++++----- .../mem_cache/test_cp_shared_kv_runtime.py | 12 ++-- 3 files changed, 115 insertions(+), 22 deletions(-) diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md index 288ed36ea..4306da81b 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_ipc_collective_replacement_plan_zh.md @@ -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。 diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 1eeb7797b..0ccaa6514 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -460,6 +460,7 @@ def _load_tai_ipc_kernels(): "materialize_cuda_ipc_peer_pages_slot_indices", "materialize_cuda_ipc_peer_pages_slot_indices_wait_ready", "publish_cuda_ipc_slot_pages_and_mark_ready", + "publish_cuda_ipc_slot_pages_compact_and_mark_ready", ) missing = [name for name in required if not hasattr(tai_ipc, name)] if missing: @@ -2534,15 +2535,17 @@ def _build_current_staging_ipc_descriptors( layout: CpSharedKVLayout, spans: list[tuple[int, int]], device: torch.device, + slot_indices: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: flat_slot_logical_pages = _contiguous_for_tai( slot_logical_pages.reshape(-1).to(device=device) ) - slot_indices = _slot_spans_to_cuda_slot_indices( - spans, - total_slots=int(flat_slot_logical_pages.numel()), - device=device, - ) + if slot_indices is None: + slot_indices = _slot_spans_to_cuda_slot_indices( + spans, + total_slots=int(flat_slot_logical_pages.numel()), + device=device, + ) if slot_indices.numel() == 0: empty = torch.empty((0,), dtype=torch.long, device=device) return empty, empty, empty @@ -2559,6 +2562,21 @@ def _build_current_staging_ipc_descriptors( return owner_ranks.contiguous(), src_page_indices, dense_page_indices +def _build_compact_current_staging_src_page_indices( + src_page_indices: torch.Tensor, +) -> torch.Tensor: + compact_src = torch.arange( + int(src_page_indices.numel()), + dtype=torch.long, + device=src_page_indices.device, + ) + return torch.where( + src_page_indices >= 0, + compact_src, + torch.full_like(compact_src, -1), + ).contiguous() + + def _try_tai_ipc_materialize_token_kv_page_slot_spans_into( *, kv_cache: torch.Tensor, @@ -2646,9 +2664,15 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into( if not spans: return True page_nbytes = _token_kv_page_nbytes(dense_kv_cache, page_size) - required_nbytes = int(dense_kv_cache.shape[0]) * _page_nbytes_from_page_tensor( - dense_kv_cache + slot_indices = _slot_spans_to_cuda_slot_indices( + spans, + total_slots=int(slot_logical_pages.reshape(-1).numel()), + device=dense_kv_cache.device, ) + if slot_indices.numel() == 0: + return True + dense_page_indices = (slot_indices + 1).to(torch.long).contiguous() + required_nbytes = int(dense_page_indices.numel()) * int(page_nbytes) staging_state = _get_or_create_tai_ipc_current_staging( kind="token", dense_tensor=dense_kv_cache, @@ -2665,13 +2689,15 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into( layout=layout, spans=spans, device=dense_kv_cache.device, + slot_indices=slot_indices, ) ) - if dense_page_indices.numel() == 0: - return True + compact_src_page_indices = _build_compact_current_staging_src_page_indices( + src_page_indices + ) state.ready_seq += 1 ready_seq = int(state.ready_seq) - kernels.publish_cuda_ipc_slot_pages_and_mark_ready( + kernels.publish_cuda_ipc_slot_pages_compact_and_mark_ready( dense_kv_cache, state.staging, dense_page_indices, @@ -2684,7 +2710,7 @@ def _try_tai_ipc_materialize_current_token_kv_page_slot_spans_into( state.ready_peer_ptrs, dense_kv_cache, owner_ranks, - src_page_indices, + compact_src_page_indices, dense_page_indices, ready_seq=ready_seq, page_nbytes=page_nbytes, @@ -2866,7 +2892,15 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into( if not spans: return True page_nbytes = _page_nbytes_from_page_tensor(dense_page_buffer) - required_nbytes = int(dense_page_buffer.shape[0]) * page_nbytes + slot_indices = _slot_spans_to_cuda_slot_indices( + spans, + total_slots=int(slot_logical_pages.reshape(-1).numel()), + device=dense_page_buffer.device, + ) + if slot_indices.numel() == 0: + return True + dense_page_indices = (slot_indices + 1).to(torch.long).contiguous() + required_nbytes = int(dense_page_indices.numel()) * int(page_nbytes) staging_state = _get_or_create_tai_ipc_current_staging( kind="paged", dense_tensor=dense_page_buffer, @@ -2883,13 +2917,15 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into( layout=layout, spans=spans, device=dense_page_buffer.device, + slot_indices=slot_indices, ) ) - if dense_page_indices.numel() == 0: - return True + compact_src_page_indices = _build_compact_current_staging_src_page_indices( + src_page_indices + ) state.ready_seq += 1 ready_seq = int(state.ready_seq) - kernels.publish_cuda_ipc_slot_pages_and_mark_ready( + kernels.publish_cuda_ipc_slot_pages_compact_and_mark_ready( dense_page_buffer, state.staging, dense_page_indices, @@ -2902,7 +2938,7 @@ def _try_tai_ipc_materialize_current_paged_buffer_page_slot_spans_into( state.ready_peer_ptrs, dense_page_buffer, owner_ranks, - src_page_indices, + compact_src_page_indices, dense_page_indices, ready_seq=ready_seq, page_nbytes=page_nbytes, diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index f09110f0b..4400375b2 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -2874,7 +2874,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(mixed_locs.tolist(), [[4, 5]]) self.assertTrue(torch.equal(mixed_kv[4:6], current_kv)) - def test_current_token_ipc_helper_uses_dense_slot_pages_for_staging(self): + def test_current_token_ipc_helper_uses_compact_pages_for_staging(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -2891,7 +2891,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): calls = [] class FakeKernels: - def publish_cuda_ipc_slot_pages_and_mark_ready(self, *args, **kwargs): + def publish_cuda_ipc_slot_pages_compact_and_mark_ready(self, *args, **kwargs): calls.append(("publish", args, kwargs)) def materialize_cuda_ipc_peer_pages_slot_indices_wait_ready( @@ -2922,7 +2922,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): gather_name, gather_args, gather_kwargs = calls[1] self.assertEqual(gather_name, "gather") self.assertTrue(torch.equal(gather_args[3], torch.tensor([0, 1]))) - self.assertTrue(torch.equal(gather_args[4], torch.tensor([1, 2]))) + self.assertTrue(torch.equal(gather_args[4], torch.tensor([0, 1]))) self.assertTrue(torch.equal(gather_args[5], torch.tensor([1, 2]))) self.assertEqual(gather_kwargs["ready_seq"], 1) @@ -3463,7 +3463,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(dense_page_buffer[1, 0:4], current_k[0])) self.assertTrue(torch.equal(dense_page_buffer[1, 4:8], current_k[1])) - def test_current_index_ipc_helper_uses_dense_slot_pages_for_staging(self): + def test_current_index_ipc_helper_uses_compact_pages_for_staging(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -3480,7 +3480,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): calls = [] class FakeKernels: - def publish_cuda_ipc_slot_pages_and_mark_ready(self, *args, **kwargs): + def publish_cuda_ipc_slot_pages_compact_and_mark_ready(self, *args, **kwargs): calls.append(("publish", args, kwargs)) def materialize_cuda_ipc_peer_pages_slot_indices_wait_ready( @@ -3505,7 +3505,7 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(calls[0][1][2], torch.tensor([2, 3]))) self.assertEqual(calls[0][2]["ready_seq"], 7) self.assertTrue(torch.equal(calls[1][1][3], torch.tensor([1, -1]))) - self.assertTrue(torch.equal(calls[1][1][4], torch.tensor([2, -1]))) + self.assertTrue(torch.equal(calls[1][1][4], torch.tensor([0, -1]))) self.assertTrue(torch.equal(calls[1][1][5], torch.tensor([2, 3]))) self.assertEqual(calls[1][2]["ready_seq"], 7)