[UnifiedTree]: Sync mamba int8 checkpoint (#30626)

test/registered/radix_cache/test_int8_mamba_checkpoint_e2e.py

@@ -21,16 +21,24 @@ Usage:
     python3 -m unittest test_int8_mamba_checkpoint_e2e
 """
 
+import time
 import unittest
 from types import SimpleNamespace
 from urllib.parse import urlparse
 
+from sglang.srt.utils import kill_process_tree
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.kits.kl_divergence_kit import KLDivergenceMixin
-from sglang.test.server_fixtures.default_fixture import DefaultServerBase
-from sglang.test.test_utils import DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST
+from sglang.test.server_fixtures.default_fixture import (
+    DefaultServerBase,
+    openai_api_env,
+)
+from sglang.test.test_utils import (
+    DEFAULT_HYBRID_MAMBA_MODEL_NAME_FOR_TEST,
+    popen_launch_server,
+)
 
-register_cuda_ci(est_time=400, stage="extra-b", runner_config="4-gpu-h100")
+register_cuda_ci(est_time=800, stage="extra-b", runner_config="4-gpu-h100")
 
 
 class TestInt8MambaCheckpointE2E(KLDivergenceMixin, DefaultServerBase):
@@ -89,5 +97,27 @@ class TestInt8MambaCheckpointE2E(KLDivergenceMixin, DefaultServerBase):
         self.assertGreaterEqual(metrics["accuracy"], self.gsm8k_threshold)
 
 
+class TestUnifiedRadixTreeInt8MambaCheckpointE2E(TestInt8MambaCheckpointE2E):
+    """Run the same int8 mamba checkpoint checks with UnifiedRadixTree forced on."""
+
+    @classmethod
+    def setUpClass(cls):
+        assert cls.model is not None, "Please set cls.model in subclass"
+
+        with openai_api_env(cls.api_key):
+            cls.process = popen_launch_server(
+                cls.model,
+                cls.base_url,
+                timeout=cls.timeout,
+                other_args=cls.other_args,
+                env={"SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
+            )
+
+    @classmethod
+    def tearDownClass(cls):
+        kill_process_tree(cls.process.pid, wait_timeout=60)
+        time.sleep(2)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -91,6 +91,7 @@ class CacheConfig:
     mamba_head_dim: int = 16
     mamba_state_size: int = 16
     mamba_conv_kernel: int = 4
+    enable_int8_mamba_checkpoint: bool = False
 
     # Model / pool
     kv_size: int = 256
@@ -214,7 +215,11 @@ class TestUnifiedTreeNodeGetPrefixHashValues(CustomTestCase):
 
 def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
     """Create (tree, allocator, req_to_token_pool) from a CacheConfig."""
-    server_args = ServerArgs(model_path="dummy", page_size=cfg.page_size)
+    server_args = ServerArgs(
+        model_path="dummy",
+        page_size=cfg.page_size,
+        enable_int8_mamba_checkpoint=cfg.enable_int8_mamba_checkpoint,
+    )
     # MambaRadixCache reads mamba_cache_chunk_size, whose property otherwise
     # loads the HF config for self.model_path — impossible for the dummy model.
     # Mirror the property's default for a dummy HF config: FLA_CHUNK_SIZE.
@@ -4037,6 +4042,73 @@ class UnifiedLRUListBoundedRefreshTest(CustomTestCase):
         self.assertEqual(self._lru_order(lru), before)
 
 
+class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
+    cfg = CacheConfig(
+        components=(ComponentType.FULL, ComponentType.MAMBA),
+        enable_mamba_extra_buffer=True,
+        enable_int8_mamba_checkpoint=True,
+        mamba_cache_size=8,
+        kv_size=64,
+        max_context_len=64,
+    )
+
+    def _make_req(self, req_to_token_pool, tokens):
+        req = Req(
+            rid="int8-mamba",
+            origin_input_text="",
+            origin_input_ids=array("q", tokens),
+            sampling_params=SamplingParams(temperature=0, max_new_tokens=1),
+        )
+        req_to_token_pool.alloc([req])
+        req.output_ids = array("q")
+        req.kv_committed_len = len(tokens)
+        req.kv_allocated_len = len(tokens)
+        req.cache_protected_len = 0
+        req.swa_uuid_for_lock = None
+        req.extra_key = None
+        req.mamba_last_track_seqlen = len(tokens)
+        return req
+
+    def _cache_finished(self, cache, allocator, req_to_token_pool, tokens):
+        req = self._make_req(req_to_token_pool, tokens)
+        kv_indices = allocator.alloc(len(tokens))
+        self.assertIsNotNone(kv_indices)
+        req_to_token_pool.write((req.req_pool_idx, slice(0, len(tokens))), kv_indices)
+        req.last_node = cache.root_node
+
+        cache.cache_finished_req(req, is_insert=True)
+
+    def test_finished_req_stores_radix_mamba_state_in_int8_pool(self):
+        cache, allocator, req_to_token_pool = build_fixture(self.cfg)
+        ckpt_pool = req_to_token_pool.mamba_ckpt_pool
+        self.assertIsNotNone(ckpt_pool)
+        active_initial = req_to_token_pool.mamba_allocator.available_size()
+        ckpt_initial = ckpt_pool.available_size()
+        tokens = [1, 2, 3, 4]
+
+        self._cache_finished(cache, allocator, req_to_token_pool, tokens)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), active_initial
+        )
+        self.assertEqual(ckpt_pool.available_size(), ckpt_initial - 1)
+        self.assertEqual(cache.mamba_evictable_size(), 1)
+
+        self._cache_finished(cache, allocator, req_to_token_pool, tokens)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), active_initial
+        )
+        self.assertEqual(ckpt_pool.available_size(), ckpt_initial - 1)
+        self.assertEqual(cache.mamba_evictable_size(), 1)
+
+        result = cache.evict(EvictParams(mamba_num=1))
+        self.assertEqual(result.mamba_num_evicted, 1)
+        self.assertEqual(
+            req_to_token_pool.mamba_allocator.available_size(), active_initial
+        )
+        self.assertEqual(ckpt_pool.available_size(), ckpt_initial)
+        self.assertEqual(cache.mamba_evictable_size(), 0)
+
+
 _CONFIGS: list[CacheConfig] = [
     CacheConfig(page_size=1, components=(ComponentType.FULL,)),
     CacheConfig(page_size=4, components=(ComponentType.FULL,)),