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)