Fail fast when CP compose would hide dense fallback

CP shared-KV bs>1 compose must not silently fall back to dense full-buffer collectives when CUDA TAI materialize is expected. The fallback masks both correctness contract drift and severe synchronization/communication regressions, especially while comparing the symm path with the older IPC path.\n\nThis keeps CPU/unit-test fallback available, but makes production CUDA+TAI runs raise an explicit compose_v2 fail-fast for token-KV and index dense fallback. It also records the symm-vs-IPC comparison contract so barrier and collective counts are evaluated alongside elapsed time.\n\nConstraint: Production cache-hit-heavy bs>1 paths must expose unexpected dense collectives instead of silently taking them.\nRejected: Cherry-pick the old IPC branch wholesale | it conflicts with the symm compose design and would mix two transport protocols before benchmarking.\nRejected: Allow dense fallback with warning only | warning can be missed and still corrupts performance conclusions.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not re-enable dense full fallback in CUDA+TAI compose paths without a benchmark proving it is intentional and a correctness test covering cache-hit bs>1.\nTested: python -m py_compile for cp_shared_kv_runtime.py and test_cp_shared_kv_runtime.py; git diff --check.\nNot-tested: Remote container pytest/ETE; local pytest is not reliable in this workspace because dependencies such as orjson are missing.
2026-06-12 23:57:40 +08:00
parent 9d65bdba95
commit 2387787ebc
3 changed files with 174 additions and 0 deletions
@@ -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 二轮精度不掉点。
@@ -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,
@@ -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