Fix mamba checkpoint depth under dcp (#34808)

test/registered/unit/managers/test_mamba_checkpoint_depth.py

@@ -0,0 +1,74 @@
+"""The donated mamba checkpoint depth must land on the tree page.
+
+DCP widens the tree page past the mamba chunk grid. A checkpoint picked on the
+finer grid names a depth no radix node can carry, so it gets attached to the
+preceding node and a later request resumes from a state that already covers
+tokens past that node.
+"""
+
+import unittest
+from array import array
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import torch
+
+from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.sampling.sampling_params import SamplingParams
+from sglang.srt.server_args import (
+    ServerArgs,
+    set_global_server_args_for_scheduler,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+CHUNK = 64
+
+
+def _track_seqlen(*, tree_page: int, prefix_len: int, extend_len: int) -> int:
+    """Run one extend through the tracker and report the donated depth."""
+    server_args = ServerArgs(model_path="dummy", page_size=CHUNK)
+    # The property would otherwise load the HF config for the dummy model.
+    server_args._mamba_cache_chunk_size = CHUNK
+    set_global_server_args_for_scheduler(server_args)
+
+    sampling_params = SamplingParams(max_new_tokens=1)
+    sampling_params.normalize(None)
+    req = Req(
+        rid="req",
+        origin_input_text="",
+        origin_input_ids=array("q", [1] * (prefix_len + extend_len)),
+        sampling_params=sampling_params,
+        vocab_size=128,
+    )
+    req.prefix_indices = torch.arange(prefix_len, dtype=torch.int64)
+    req.set_extend_range(prefix_len, prefix_len + extend_len)
+    req.mamba_ping_pong_track_buffer = torch.tensor([0, 1], dtype=torch.int64)
+    req.mamba_next_track_idx = 0
+    req.mamba_branching_seqlen = None
+
+    batch = ScheduleBatch(reqs=[req])
+    batch.tree_cache = SimpleNamespace(page_size=tree_page)
+    batch.req_to_token_pool = MagicMock()
+    batch.req_to_token_pool.get_mamba_ping_pong_other_idx.return_value = 1
+
+    batch._mamba_radix_cache_v2_req_prepare_for_extend(req)
+    return req.mamba_last_track_seqlen
+
+
+class TestMambaCheckpointDepth(unittest.TestCase):
+    def test_widened_tree_page_moves_the_donated_depth_onto_it(self):
+        # 4066 tokens past a 16384 prefix: the chunk grid would stop at 20416,
+        # which a 256-token page cannot name.
+        depth = _track_seqlen(tree_page=256, prefix_len=16384, extend_len=4066)
+        self.assertEqual(depth % 256, 0)
+        self.assertEqual(depth, 20224)
+
+    def test_unwidened_tree_page_keeps_the_chunk_grid(self):
+        depth = _track_seqlen(tree_page=CHUNK, prefix_len=16384, extend_len=4066)
+        self.assertEqual(depth, 20416)
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -261,8 +261,20 @@ def _write_backup(cache, node, write_back: bool = False) -> int:
     )
 
 
-def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
-    """Create (tree, allocator, req_to_token_pool) from a CacheConfig."""
+def build_fixture(
+    cfg: CacheConfig,
+    *,
+    enable_kv_cache_events: bool = False,
+    tree_page_size: Optional[int] = None,
+    mamba_cache_chunk_size: Optional[int] = None,
+):
+    """Create (tree, allocator, req_to_token_pool) from a CacheConfig.
+
+    ``tree_page_size`` stands in for DCP, which widens the tree page past the
+    ``page_size`` the rest of the config still sees. It only reaches values
+    derived from the tree page: the allocator keeps ``cfg.page_size``, whereas
+    DCP sets the two equal, so do not read insert or match behaviour off it.
+    """
     server_args = ServerArgs(
         model_path="dummy",
         page_size=cfg.page_size,
@@ -271,7 +283,11 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
     # 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.
-    server_args._mamba_cache_chunk_size = max(FLA_CHUNK_SIZE, cfg.page_size)
+    server_args._mamba_cache_chunk_size = (
+        max(FLA_CHUNK_SIZE, cfg.page_size)
+        if mamba_cache_chunk_size is None
+        else mamba_cache_chunk_size
+    )
     set_global_server_args_for_scheduler(server_args)
     device = get_device()
 
@@ -372,7 +388,7 @@ def build_fixture(cfg: CacheConfig, *, enable_kv_cache_events: bool = False):
     cache_init_params = CacheInitParams(
         req_to_token_pool=req_to_token_pool,
         token_to_kv_pool_allocator=allocator,
-        page_size=cfg.page_size,
+        page_size=cfg.page_size if tree_page_size is None else tree_page_size,
         disable=False,
         sliding_window_size=cfg.sliding_window_size,
         tree_components=cfg.components,
@@ -5161,6 +5177,41 @@ class TestUnifiedMambaLRUMatchRefresh(CustomTestCase):
         self.assertGreater(order.index(a1), order.index(a2))
 
 
+class TestMambaCheckpointGrid(CustomTestCase):
+    """A donated mamba checkpoint is only reusable at a depth the tree can name.
+
+    ``tree_page_size`` simulates DCP: it widens the page the tree allocates on
+    while ``page_size`` and the mamba chunk grid stay where they are, which is
+    exactly the split that lets a checkpoint land between two node boundaries.
+    """
+
+    cfg = CacheConfig(
+        page_size=64,
+        components=(ComponentType.FULL, ComponentType.MAMBA),
+        enable_mamba_extra_buffer=True,
+        kv_size=1024,
+        max_context_len=1024,
+    )
+
+    def _grid(self, cache):
+        component = next(
+            c
+            for c in cache._components_tuple
+            if c.component_type is ComponentType.MAMBA
+        )
+        return component.mamba_checkpoint_grid
+
+    def test_grid_follows_the_widened_tree_page(self):
+        cache, _, _ = build_fixture(
+            self.cfg, tree_page_size=256, mamba_cache_chunk_size=64
+        )
+        self.assertEqual(self._grid(cache), 256)
+
+    def test_grid_is_the_chunk_size_without_widening(self):
+        cache, _, _ = build_fixture(self.cfg, mamba_cache_chunk_size=64)
+        self.assertEqual(self._grid(cache), 64)
+
+
 class TestUnifiedRadixCacheInt8MambaCheckpoint(CustomTestCase):
     cfg = CacheConfig(
         components=(ComponentType.FULL, ComponentType.MAMBA),