Add bit-exact guard for extra_buffer_lazy (#35030)

test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py

@@ -19,8 +19,13 @@ Measured avg_kl_div:
            test_prefill_cache_hit             6.22e-06    4.40e-06        0.0
            test_decode_cache_hit                   0.0         0.0        0.0
            multiturn branching                     0.0    2.01e-07        0.0
+           lazy test_prefill_cache_hit        5.56e-06         n/a        0.0
   #29792   multiturn branching       9.43e-06/1.16e-05    5.14e-04        0.0
 
+The lazy row was measured in the same session as the 1.18e-05 the non-lazy
+`test_prefill_cache_hit` read on that GPU, so read the pair as "both fire",
+not as one being weaker.
+
 `test_prefill_cache_hit` is the only case that fires on both, so read the rest as
 extra coverage rather than as the guard for one fix. CI runs `1-gpu-large`, which
 is SM90.
@@ -64,7 +69,7 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=600, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=790, 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")
@@ -92,7 +97,7 @@ def _random_suffixes(n: int, length: int, seed: int) -> list[list[int]]:
     return [[rng.randint(1, 30000) for _ in range(length)] for _ in range(n)]
 
 
-def _base_args() -> list[str]:
+def _base_args(mamba_strategy: str = "extra_buffer") -> list[str]:
     return [
         "--trust-remote-code",
         "--attention-backend",
@@ -100,7 +105,7 @@ def _base_args() -> list[str]:
         "--page-size",
         str(PAGE_SIZE),
         "--mamba-radix-cache-strategy",
-        "extra_buffer",
+        mamba_strategy,
         "--swa-full-tokens-ratio",
         "0.1",
         "--mamba-full-memory-ratio",
@@ -175,6 +180,40 @@ class TestUnifiedHybridBitExact(CustomTestCase):
         self._run(assert_decode_cache_hit)
 
 
+class TestUnifiedHybridLazyBitExact(TestUnifiedHybridBitExact):
+    """Same exactness bar on the lazy extra-buffer strategy.
+
+    Lazy is a different code path, not the same code under a flag: a request
+    holds one ping-pong slot instead of two and the second is allocated only on
+    the forward that crosses a track boundary, then freed again afterwards. The
+    checkpoint therefore lands in a slot chosen per forward rather than one the
+    request owns for its lifetime, and picking the wrong one restores another
+    request's conv window.
+
+    Admitted the same way the class it inherits from was: reverting #34184 takes
+    test_prefill_cache_hit here to 5.56e-06, so this is a guard for that class of
+    bug on the lazy path rather than coverage of a path nothing checks.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        cls.model = _MODEL_PATH
+        cls.base_url = DEFAULT_URL_FOR_TEST
+        other_args = _base_args("extra_buffer_lazy") + [
+            "--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"},
+        )
+
+
 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.