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
@@ -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