Protect CP shared-KV cache-hit correctness under batched FP8 reuse

Cache-hit GSM8K regressions only appeared after the second pass reused request-specific suffix pages, so this change adds fail-fast transfer validation, masks stale rectangular page-table tails, and extends CUDA/unit coverage across FP8 CP shared-KV write, load, top-k, and materialization paths. The temporary ledger records eliminated hypotheses to prevent re-debugging the same L2 and persistent-cache paths.\n\nConstraint: CP shared KV stores physical pages but scheduler-visible semantics must remain valid-token/page-bounded.\nConstraint: bs>1 FP8 prefill must preserve existing CP shared-KV fast paths without silent fallback.\nRejected: Blame raw HiCache L2 load without tests | L2 KV and index backup/load/materialize roundtrips pass on remote CUDA.\nRejected: Disable current/partial reuse broadly | hides the cache-hit contract regression and costs performance.\nConfidence: medium\nScope-risk: moderate\nDirective: Do not weaken CP shared-KV fail-fast or rectangular-tail masking without rerunning second-pass cache-hit accuracy tests.\nTested: remote CUDA pytest for fused FP8 MLA store, fused persistent index store, L2-loaded FP8 KV materialize, L2-loaded index materialize, ragged top-k offset, TAI batched index MQA prepare.\nTested: local py_compile for touched test files and git diff --check.\nNot-tested: full second-pass GSM8K accuracy after these diagnostic tests; root cause remains under investigation.
2026-06-07 13:26:49 +08:00
parent b17976b60d
commit f75ffff8d9
10 changed files with 2645 additions and 1 deletions
@@ -4,8 +4,12 @@ import unittest
 from types import SimpleNamespace
 from unittest.mock import MagicMock, patch

+import torch
+import torch.distributed as dist
+
 from sglang.srt.disaggregation.base import KVPoll
 from sglang.srt.disaggregation.prefill import PrefillBootstrapQueue
+from sglang.srt.disaggregation.utils import poll_and_all_reduce_attn_cp_tp_group
 from sglang.srt.disaggregation.utils import ReqToMetadataIdxAllocator
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
@@ -25,6 +29,9 @@ class FakeSender:
        if self.should_fail:
            raise RuntimeError("boom")

+    def poll(self):
+        return KVPoll.Success
+

 class TestPrefillBootstrapQueue(CustomTestCase):
    def _make_req(self, rid, bootstrap_room, origin_input_ids, sender):
@@ -124,6 +131,70 @@ class TestPrefillBootstrapQueue(CustomTestCase):
        self.assertEqual(skipped.disagg_kv_sender.init_calls, [])
        self.assertEqual(checked.disagg_kv_sender.init_calls, [(3, 0)])

+    def test_poll_consensus_debug_fails_before_shape_mismatch_hang(self):
+        reduce_calls = []
+
+        def fake_all_reduce(tensor, op, group):
+            reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group))
+            # The new debug guard uses an int64 [queue_len, queue_hash] scalar
+            # vector before building the uint8 per-request poll tensor.  Simulate
+            # another rank having a different queue length and assert that we
+            # fail fast before reaching the old variable-length uint8 all_reduce.
+            if tensor.dtype == torch.int64 and int(tensor.numel()) == 2:
+                if op == dist.ReduceOp.MIN:
+                    tensor[0] = 1
+                elif op == dist.ReduceOp.MAX:
+                    tensor[0] = 2
+                return
+            raise AssertionError(
+                "poll tensor all_reduce should not run after queue mismatch"
+            )
+
+        with (
+            patch(
+                "sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get",
+                return_value=True,
+            ),
+            patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce),
+        ):
+            with self.assertRaisesRegex(RuntimeError, "poll_queue.*inflight"):
+                poll_and_all_reduce_attn_cp_tp_group(
+                    [FakeSender(), FakeSender()],
+                    MagicMock(name="cp_group"),
+                    MagicMock(name="tp_group"),
+                    debug_label="inflight",
+                    debug_ids=["rid-a", "rid-b"],
+                )
+
+        self.assertTrue(reduce_calls)
+
+    def test_poll_consensus_debug_disabled_preserves_old_collective_shape(self):
+        reduce_calls = []
+
+        def fake_all_reduce(tensor, op, group):
+            reduce_calls.append((tensor.dtype, int(tensor.numel()), op, group))
+
+        with (
+            patch(
+                "sglang.srt.disaggregation.utils.envs.SGLANG_CP_SHARED_KV_BS_GT1_DEBUG.get",
+                return_value=False,
+            ),
+            patch("sglang.srt.disaggregation.utils.dist.all_reduce", fake_all_reduce),
+        ):
+            polls = poll_and_all_reduce_attn_cp_tp_group(
+                [FakeSender(), FakeSender()],
+                MagicMock(name="cp_group"),
+                MagicMock(name="tp_group"),
+                debug_label="inflight",
+                debug_ids=["rid-a", "rid-b"],
+            )
+
+        self.assertEqual(polls, [KVPoll.Success, KVPoll.Success])
+        self.assertEqual(
+            [(dtype, size) for dtype, size, _op, _group in reduce_calls],
+            [(torch.uint8, 2), (torch.uint8, 2)],
+        )
+

 if __name__ == "__main__":
    unittest.main()