Add bit-exact unified radix cache KL test for hybrid SWA + mamba (#34607)

2026-08-13 02:15:36 +08:00
parent 773faf992d
commit 3974b00359
1 changed files with 251 additions and 0 deletions
@@ -0,0 +1,251 @@
+"""Bit-exact KL guards for the unified radix cache on a hybrid SWA + mamba model.
+
+The other KL tests in this directory gate on a loose threshold because their
+models cannot score a token identically twice: Qwen3-Next's chunkwise prefill
+scan and its decode recurrence are different algorithms and land an ulp apart, so
+a tight floor there would fail on float noise. The shrunken Inkling checkpoint
+reproduces every logprob exactly under deterministic inference, which turns the
+same comparison into an exact one -- any nonzero KL is a state-reuse bug. It also
+fits on one GPU, so these run per-commit rather than on a 4-GPU stage.
+
+Each class below reproduces a specific merged regression when its fix is
+reverted; the measured pre-fix divergence is recorded in the class docstring so a
+later threshold change has to argue with a number.
+
+These classes do not use UnifiedRadixTreeTestMixin: it bundles gsm8k and mmlu,
+which an undertrained checkpoint cannot gate on, and each class here runs the
+harness its regression was actually reproduced with.
+
+The imported `test_`-prefixed helpers are aliased so pytest does not collect them
+as tests.
+"""
+
+import os
+import random
+import unittest
+
+from sglang.srt.utils import kill_process_tree
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kl_multiturn_utils import (
+    make_mamba_decode_assert,
+)
+from sglang.test.kl_multiturn_utils import (
+    test_input_output_logprobs_match_decode_cache_hit_helper as assert_multiturn_decode_cache_hit,
+)
+from sglang.test.kl_test_utils import (
+    get_input_ids,
+)
+from sglang.test.kl_test_utils import (
+    test_input_output_logprobs_match_decode_cache_hit_helper as assert_decode_cache_hit,
+)
+from sglang.test.kl_test_utils import (
+    test_input_output_logprobs_match_helper as assert_logprobs_match,
+)
+from sglang.test.kl_test_utils import (
+    test_input_output_logprobs_match_prefill_cache_hit_helper as assert_prefill_cache_hit,
+)
+from sglang.test.test_utils import (
+    DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+    DEFAULT_URL_FOR_TEST,
+    CustomTestCase,
+    popen_launch_server,
+)
+
+register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large")
+
+_MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inkling")
+_MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test")
+
+# Both classes measure exactly 0 in their fixed state -- every logprob matches bit
+# for bit. The floor only keeps a stray ulp from failing the run; a state-reuse
+# bug lands orders of magnitude above it. It cannot be 0.0: the comparison is a
+# strict `<`, so an exact 0 would fail its own threshold.
+KL_DIV_THRESHOLD = 1e-9
+
+# Equal to the page size below. Out-of-window SWA slots are freed a page at a
+# time, so only a checkpoint sitting on a page boundary still has a full window of
+# SWA data below it -- at the default 256 half the sequence lengths land off that
+# boundary and lose their decode prefix entirely.
+TRACK_INTERVAL = 128
+PAGE_SIZE = 128
+
+# Past the 512-token sliding window, so decode carries the window through the
+# handover from prompt tokens to generated ones.
+MAX_NEW_TOKENS = 1024
+
+
+def _random_suffixes(n: int, length: int, seed: int) -> list[list[int]]:
+    rng = random.Random(seed)
+    return [[rng.randint(1, 30000) for _ in range(length)] for _ in range(n)]
+
+
+def _base_args() -> list[str]:
+    return [
+        "--trust-remote-code",
+        "--attention-backend",
+        "fa4",
+        "--page-size",
+        str(PAGE_SIZE),
+        "--mamba-radix-cache-strategy",
+        "extra_buffer",
+        "--swa-full-tokens-ratio",
+        "0.1",
+        "--mamba-full-memory-ratio",
+        "0.1",
+        # 0.85 was carried over from the 4-GPU B200 test and OOMs an 80 GB card:
+        # the static pool leaves ~19 GB for the prefill graphs, the fa4 workspace
+        # and the chunked-prefill activations, which is what this config needs.
+        "--mem-fraction-static",
+        "0.6",
+        "--mamba-track-interval",
+        str(TRACK_INTERVAL),
+        "--enable-deterministic-inference",
+    ]
+
+
+class TestUnifiedHybridBitExact(CustomTestCase):
+    """Prefill and decode must score a token identically once every kernel on the
+    path is batch-invariant, so any drift is a stale conv/mamba checkpoint or a
+    prefix the cache restored wrong.
+
+    Guards #34184 (stale track rows corrupting conv checkpoints under the prefill
+    graph). Reverting that fix here measures avg_kl_div 5.58e-07 on
+    test_logprobs_match and 6.22e-06 on test_prefill_cache_hit, against 0.0 with
+    it in place. test_decode_cache_hit is 0.0 either way -- it guards decode-region
+    state reuse in general, not that regression.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = _MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = _base_args() + [
+            # Pinned, not incidental: the prefill graph derives its fixed
+            # request-slot count from this (chunked_prefill_size // 512), and those
+            # slots are exactly what #34184 left stale. Lowering it shrinks the
+            # sentinel tail and the guard stops firing while still passing.
+            "--chunked-prefill-size",
+            "16384",
+        ]
+        if _MODEL_REVISION:
+            other_args += ["--revision", _MODEL_REVISION]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if getattr(cls, "process", None) is not None:
+            kill_process_tree(cls.process.pid)
+
+    def _run(self, helper):
+        helper(
+            self.base_url,
+            {self.model: {"kl_div": KL_DIV_THRESHOLD}},
+            self.model,
+            max_samples=32,
+            max_new_tokens=MAX_NEW_TOKENS,
+            trust_remote_code=True,
+        )
+
+    def test_logprobs_match(self):
+        self._run(assert_logprobs_match)
+
+    def test_prefill_cache_hit(self):
+        self._run(assert_prefill_cache_hit)
+
+    def test_decode_cache_hit(self):
+        self._run(assert_decode_cache_hit)
+
+
+class TestUnifiedHybridHiCacheBitExact(CustomTestCase):
+    """Same exactness bar with the host tier in the loop, over interleaved
+    branches so hits land at many prefix lengths rather than one aligned one.
+
+    Guards #29792 (decode track save picking its slot from the producer-side
+    pointer). Without that fix this measures avg_kl_div 9.43e-06 and 1.16e-05 over
+    two rounds, with 3 of 9 samples dirty and the rest exactly 0; with it in place
+    both rounds are 0.0.
+
+    Runs the multi-turn branching harness because the single-turn helpers above
+    cannot produce a non-aligned hit length, which this regression needs.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = _MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = _base_args() + [
+            "--enable-hierarchical-cache",
+            "--hicache-ratio",
+            "4",
+            "--hicache-write-policy",
+            "write_through",
+            "--hicache-io-backend",
+            "direct",
+            # The mamba host pool only supports page_first and page_first_direct.
+            "--hicache-mem-layout",
+            "page_first_direct",
+            # Tight pools and a small budget so decode crosses a track boundary and
+            # the host tier is actually exercised instead of everything staying
+            # resident on device.
+            "--chunked-prefill-size",
+            "2048",
+            "--max-total-tokens",
+            "65536",
+            "--max-mamba-cache-size",
+            "500",
+            "--max-running-requests",
+            "4",
+        ]
+        if _MODEL_REVISION:
+            other_args += ["--revision", _MODEL_REVISION]
+        cls.process = popen_launch_server(
+            cls.model,
+            cls.base_url,
+            timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
+            other_args=other_args,
+            env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+        )
+        cls.input_ids = get_input_ids(
+            tokenizer_path=cls.model, num_samples=9, trust_remote_code=True
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if getattr(cls, "process", None) is not None:
+            kill_process_tree(cls.process.pid)
+
+    def test_multiturn_decode_cache_hit_branching(self):
+        groups, branches = 3, 3
+        n = groups * branches
+        first_turn = []
+        for g in range(groups):
+            base = self.input_ids[g][:512]
+            for _ in range(branches):
+                first_turn.append(list(base))
+
+        assert_multiturn_decode_cache_hit(
+            self.base_url,
+            self.model,
+            KL_DIV_THRESHOLD,
+            first_turn,
+            turn_suffixes=[
+                _random_suffixes(n, 512, seed=300),
+                _random_suffixes(n, 256, seed=400),
+            ],
+            # Not the default exact equality: a mamba checkpoint lands on a track
+            # boundary, so the reusable prefix is floor-aligned to the interval.
+            assert_decode_cached_tokens=make_mamba_decode_assert(TRACK_INTERVAL),
+            branches_per_group=branches,
+            max_new_tokens=512,
+            sampling_temperature=0,
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()