diff --git a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py b/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py index 58480570d..ac9a451b4 100644 --- a/test/registered/radix_cache/unified_radix_tree/test_unified_radix_cache_kl_hybrid_bitexact.py +++ b/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.