diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md index 67cdf844a..c093b398e 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md @@ -734,3 +734,53 @@ P7 GSM8K/replay/Nsight 验证 ``` 不要先做 P6 再修 P2/P3。0SM kernel 只解决 transport 代价，不解决 bs>1 prefix/current slot 语义；如果语义仍是 scalar prefix，kernel 越快只会越快地产生错误。 + +## 2026-06-12 追加：symm compose 与旧 IPC kernel 的对比口径 + +`symm-syh` 分支的 compose/symm kernel 可能比旧 CUDA IPC current-staging kernel 单次更快，但不能只比较 kernel elapsed time。需要把同步点作为一等指标，否则可能出现“单 kernel 快，但每层/每 buffer 多一次同步，ETE 更慢”的误判。 + +对比 benchmark / Nsight trace 必须至少拆分以下事件： + +1. **IPC capability agreement** + - `_agreed_tai_ipc_peer_ptrs()` 里的 group agreement / all-reduce。 + - 应确认是每个 pool tensor 一次，还是每个 forward/layer 重复触发。 + +2. **symm barrier** + - `cp_symm_barrier()` 调用次数、耗时、等待方差。 + - 需要按 token KV / index buffer 分开统计。 + - 当前 tai-kernel 实现是 1 个 block / 1 个 warp 的 CUDA spin barrier， + 不是 copy-engine/0SM 路径；它占用很少 SM，但会在当前 stream 上形成 + 明确同步点。判断 symm 是否优于旧 IPC 时，必须把这个 barrier 的次数和 + rank 间等待方差算进去。 + +3. **symm mega gather** + - `materialize_cuda_ipc_peer_pages_slot_dense()` 在 symm combined ptr table 上的耗时。 + - 统计 prefix pages、current pages、request 数、dense pages。 + +4. **compact current reduce** + - symm 未开启或不可用时的 `_reduce_current_pages_compact()`。 + - 这是 compact current collective，不是 dense full fallback；但仍会同步/占用通信资源，需要单独计数。 + +5. **dense full fallback reduce** + - `_all_reduce_materialized_buffer(... v2_full ...)`。 + - 生产 CUDA + TAI materialize 开启时不允许静默发生；应 fail-fast 暴露 `CP_SHARED_KV_FAIL_FAST][compose_v2]`。 + +6. **CPU descriptor / plan 成本** + - `get_or_build_compose_plan()` cache hit/miss。 + - per-forward 是否复用 descriptor；不要把一次性 build 成本误算到每层 steady state。 + +建议 benchmark 矩阵： + +- dtype：bf16 / fp8_e4m3； +- batch size：1, 2, 5, 10； +- extend：1k, 2k, 10k, 40k, 65k； +- cached/prefix：100k, 200k, 300k； +- case：cache-hit partial-current、current-only、multi-request shared prefix； +- mode：legacy dense fallback（只作为基线，不允许生产静默）、old CUDA IPC current staging、symm staging、symm prefetch。 + +结论标准：保留默认路径必须同时满足： + +- 无 dense full fallback； +- 同步点数量不高于旧 IPC 路径，或同步耗时被更少 kernel/更高带宽抵消； +- ETE replay 在 cache-hit-heavy 短 extend 场景提升，而不是只在 micro benchmark 提升； +- GSM8K cache-hit 二轮精度不掉点。 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 db7c39c5a..8db9457ad 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 @@ -832,6 +832,27 @@ def _log_tai_index_mqa_prepare_fallback( ) +def _should_fail_fast_compose_v2_dense_fallback(dense_tensor: torch.Tensor) -> bool: + """Whether compose_v2 may fall back to dense full-buffer collectives. + + CPU/unit-test paths can still use the simple fallback. In production CUDA + runs with TAI materialize enabled, a prefix-IPC miss would silently turn a + bs>1 cache-hit compose into a dense all_reduce over the whole buffer. That + hides both correctness-contract drift and severe performance regressions, + so fail fast instead. + """ + + return bool(dense_tensor.is_cuda and cp_shared_kv_tai_materialize_enabled()) + + +def _raise_compose_v2_dense_fallback_required(reason: str, **details: Any) -> None: + detail_str = " ".join(f"{key}={value}" for key, value in details.items()) + message = f"[CP_SHARED_KV_FAIL_FAST][compose_v2] reason={reason}" + if detail_str: + message = f"{message} {detail_str}" + raise RuntimeError(message) + + def _tai_materialize_runtime_disabled_reason() -> str | None: if not cp_shared_kv_tai_materialize_enabled(): return "env_disabled" @@ -4838,6 +4859,19 @@ def _compose_token_kv_partial_current_v2( ) else: dense_kv_cache = kv_cache.new_zeros((dense_rows, *kv_cache.shape[1:])) + if prefix_spans and _should_fail_fast_compose_v2_dense_fallback( + dense_kv_cache + ): + _raise_compose_v2_dense_fallback_required( + "token_kv_dense_fallback", + cp_rank=layout.cp_rank, + cp_size=layout.cp_size, + layer_id=layer_id, + prefix_spans=prefix_spans, + current_spans=current_spans, + dense_shape=tuple(dense_kv_cache.shape), + kv_dtype=kv_cache.dtype, + ) for prefix_start_slot, prefix_end_slot in prefix_spans: materialize_local_token_kv_page_slots_into( kv_cache=kv_cache, @@ -5294,6 +5328,19 @@ def _compose_index_partial_current_v2( dense_page_buffer = page_buffer.new_zeros( (dense_num_pages, *page_buffer.shape[1:]) ) + if prefix_spans and _should_fail_fast_compose_v2_dense_fallback( + dense_page_buffer + ): + _raise_compose_v2_dense_fallback_required( + "index_dense_fallback", + cp_rank=layout.cp_rank, + cp_size=layout.cp_size, + layer_id=layer_id, + prefix_spans=prefix_spans, + current_spans=current_spans, + dense_shape=tuple(dense_page_buffer.shape), + index_dtype=page_buffer.dtype, + ) for prefix_start_slot, prefix_end_slot in prefix_spans: materialize_local_paged_buffer_page_slots_into( page_buffer=page_buffer, 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 260828567..bf23ad631 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 @@ -2256,6 +2256,44 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertEqual(mixed_kv[8].item(), 14) self.assertTrue(torch.equal(mixed_kv[12:14], current_kv)) + def test_materialize_prefix_current_token_kv_compose_v2_fails_on_dense_fallback(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 + + page_size = 4 + layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0) + kv_cache = torch.zeros((32, 1, 1), dtype=torch.float32) + logical_locs = torch.tensor([[4, 8, 20, 21]], dtype=torch.int64) + current_locs = torch.tensor([20, 21], dtype=torch.int64) + current_kv = torch.arange(100, 102, dtype=torch.float32).view(2, 1, 1) + slot_remap = runtime.build_shared_token_kv_slot_remap( + kv_cache=kv_cache, + logical_locs=logical_locs, + remap_logical_pages=torch.tensor([[1, 2, 5]], dtype=torch.int64), + layout=layout, + page_size=page_size, + ) + + with patch.object( + runtime, "_get_or_open_tai_ipc_peer_ptrs", return_value=None + ), patch.object( + runtime, "_should_fail_fast_compose_v2_dense_fallback", return_value=True + ), self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[compose_v2\].*token_kv_dense_fallback", + ): + runtime.materialize_prefix_and_reuse_current_kv_page_slots( + kv_cache=kv_cache, + logical_locs=logical_locs, + current_kv_cache=current_kv, + current_locs=current_locs, + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + prefix_pages=2, + layer_id=3, + ) + def test_mla_prefetch_consume_prefix_with_current_skips_suffix_materialize(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout @@ -2778,6 +2816,45 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): self.assertTrue(torch.equal(dense_page_buffer[2, 0:4], current_k[0])) self.assertTrue(torch.equal(dense_page_buffer[2, 4:8], current_k[1])) + def test_materialize_prefix_current_index_compose_v2_fails_on_dense_fallback(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 + + page_size = 4 + index_head_dim = 4 + scale_bytes = 4 + page_bytes = page_size * index_head_dim + page_size * scale_bytes + layout = CpSharedKVLayout(page_size=page_size, cp_size=2, cp_rank=0) + page_buffer = torch.zeros((8, page_bytes), dtype=torch.uint8) + slot_remap = runtime.build_shared_paged_buffer_slot_remap( + page_buffer, + torch.tensor([[1, 2]], dtype=torch.int64), + layout, + ) + current_k = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8]], dtype=torch.uint8) + current_scale = torch.tensor([[1.25], [2.5]], dtype=torch.float32) + + with patch.object( + runtime, "_get_or_open_tai_ipc_peer_ptrs", return_value=None + ), patch.object( + runtime, "_should_fail_fast_compose_v2_dense_fallback", return_value=True + ), self.assertRaisesRegex( + RuntimeError, + r"\[CP_SHARED_KV_FAIL_FAST\]\[compose_v2\].*index_dense_fallback", + ): + runtime.materialize_prefix_and_reuse_current_index_page_slots( + page_buffer=page_buffer, + current_index_k=current_k, + current_index_scale=current_scale, + current_locs=torch.tensor([8, 9], dtype=torch.int64), + slot_remap=slot_remap, + layout=layout, + page_size=page_size, + index_head_dim=index_head_dim, + prefix_pages=1, + layer_id=4, + ) + def test_index_prefetch_partial_current_compose_fills_current_page_slots(self): from sglang.srt.environ import envs from sglang.srt.layers.attention.nsa import cp_shared_kv_prefetch as prefetch