Fix KDA prefix caching under mamba extra_buffer and enable it for kimi_linear (#31474)

test/registered/models_e2e/test_kimi_linear_models.py

@@ -3,7 +3,11 @@ from types import SimpleNamespace
 
 from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.kits.eval_accuracy_kit import GSM8KMixin
+from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
+from sglang.test.kits.prefix_cache_branching_kit import PrefixCacheBranchingMixin
 from sglang.test.run_eval import run_eval
+from sglang.test.server_fixtures.default_fixture import DefaultServerBase
 from sglang.test.test_utils import (
     DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
     DEFAULT_URL_FOR_TEST,
@@ -11,13 +15,15 @@ from sglang.test.test_utils import (
     popen_launch_server,
 )
 
-register_cuda_ci(est_time=178, stage="base-b", runner_config="2-gpu-large")
+register_cuda_ci(est_time=600, stage="base-b", runner_config="2-gpu-large")
+
+KIMI_LINEAR_MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
 
 
 class TestKimiLinear(CustomTestCase):
     @classmethod
     def setUpClass(cls):
-        cls.model = "moonshotai/Kimi-Linear-48B-A3B-Instruct"
+        cls.model = KIMI_LINEAR_MODEL
         cls.base_url = DEFAULT_URL_FOR_TEST
         cls.process = popen_launch_server(
             cls.model,
@@ -45,5 +51,32 @@ class TestKimiLinear(CustomTestCase):
         self.assertGreater(metrics["score"], 0.88)
 
 
+class TestKimiLinearExtraBuffer(
+    GSM8KMixin, KLDivergenceMixin, PrefixCacheBranchingMixin, DefaultServerBase
+):
+    """Regression guard: KDA never wrote mamba track snapshots, so states
+    donated to the radix cache under mamba_radix_cache_strategy=extra_buffer
+    were garbage and prefix-cache hits restored wrong KDA state (GSM8K
+    0.150 pre-fix vs 0.895 post-fix). Pre-fix, launching KimiLinear with
+    extra_buffer also fails the arch allowlist assert."""
+
+    model = KIMI_LINEAR_MODEL
+    cache_chunk_size = 64
+    gsm8k_score_threshold = 0.88
+    kl_div_thres = 0.02
+    kl_div_trust_remote_code = True
+    other_args = [
+        "--trust-remote-code",
+        "--tp-size",
+        "2",
+        "--chunked-prefill-size",
+        "2048",
+        "--mamba-radix-cache-strategy",
+        "extra_buffer",
+        "--mamba-track-interval",
+        "2",
+    ]
+
+
 if __name__ == "__main__":
     unittest.main()