Fix mamba checkpoint depth under dcp (#34808)

python/sglang/srt/managers/schedule_batch.py

@@ -8,6 +8,7 @@ from sglang.srt.runtime_context import (
     get_serving,
     get_spec,
     mamba_cache_chunk_size,
+    mamba_checkpoint_grid,
     mamba_extra_buffer_enabled,
     mamba_extra_buffer_lazy_enabled,
 )
@@ -2630,6 +2631,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         req: Req,
     ) -> _MambaRadixCacheV2TrackEntry:
         chunk_size = mamba_cache_chunk_size()
+        # The donated depth has to be a radix node boundary. Read the tree's own
+        # page rather than re-deriving how DCP widens it; the kernel still
+        # snapshots on the chunk_size grid.
+        checkpoint_grid = mamba_checkpoint_grid(self.tree_cache.page_size)
 
         def _force_track_h(i: int) -> int:
             assert i % chunk_size == 0
@@ -2642,7 +2647,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             # to force the math calculation to retrieve the correct mamba state from h.
             return i + 1
 
-        mask = req.extend_range.length >= chunk_size
+        mask = req.extend_range.length >= checkpoint_grid
         track_index = req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item()
         mamba_track_seqlen = -1
         if mask:
@@ -2659,13 +2664,13 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             # mamba radix cache to track which seqlen this mamba state should store at.
             mamba_track_seqlen_aligned = (
                 len(req.prefix_indices)
-                + (req.extend_range.length // chunk_size) * chunk_size
+                + (req.extend_range.length // checkpoint_grid) * checkpoint_grid
             )
 
             # mamba_track_fla_chunk_aligned is the aligned seqlen based on chunk_size
-            # If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which can be true when
+            # If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which is true when
-            # page_size > chunk_size, we need to force the math calculation to retrieve the correct mamba state from h
+            # checkpoint_grid is coarser than chunk_size, we need to force the math calculation to
-            # by _force_track_h()
+            # retrieve the correct mamba state from h by _force_track_h()
             mamba_track_fla_chunk_aligned = (
                 len(req.prefix_indices)
                 + (req.extend_range.length // chunk_size) * chunk_size

python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py

@@ -38,6 +38,7 @@ from sglang.srt.mem_cache.unified_cache.components.tree_component import (
 from sglang.srt.runtime_context import (
     get_exec,
     mamba_cache_chunk_size,
+    mamba_checkpoint_grid,
 )
 
 if TYPE_CHECKING:
@@ -69,6 +70,9 @@ class MambaComponent(TreeComponent):
             ), f"MambaComponent requires page_size=1 when mamba_extra_buffer is disabled, got {params.page_size}"
         super().__init__(cache, params)
         self.mamba_cache_chunk_size = mamba_cache_chunk_size()
+        # params.page_size is the tree page the allocator actually uses, already
+        # widened by dcp_size, so it is the one grid a checkpoint depth can land on.
+        self.mamba_checkpoint_grid = mamba_checkpoint_grid(params.page_size)
         self.mamba_max_states_per_path = get_exec().mamba.mamba_max_states_per_path
         # HiCache state
         self._mamba_pool_host = None  # set to host mamba pool when HiCache enabled
@@ -164,8 +168,8 @@ class MambaComponent(TreeComponent):
         # persistence of a new branching state is currently write-through only;
         # write-back eviction may discard the device-only state.
         aligned_seqlen = (
-            result.full_kv_hit_length // self.mamba_cache_chunk_size
+            result.full_kv_hit_length // self.mamba_checkpoint_grid
-        ) * self.mamba_cache_chunk_size
+        ) * self.mamba_checkpoint_grid
         branching_seqlen = (
             aligned_seqlen if aligned_seqlen > mamba_boundary_len else None
         )

python/sglang/srt/runtime_context.py

@@ -48,6 +48,7 @@ from __future__ import annotations
 
 import dataclasses
 import functools
+import math
 import os
 import sys
 from contextlib import contextmanager
@@ -1435,6 +1436,14 @@ def mamba_cache_chunk_size() -> int:
     return get_server_args().mamba_cache_chunk_size
 
 
+def mamba_checkpoint_grid(tree_page: int) -> int:
+    """The granularity a donated mamba checkpoint's depth must land on so the
+    radix tree can name it. Pass the page the tree actually allocates on: DCP
+    widens it past ``mamba_cache_chunk_size``, and deriving that here would be a
+    second copy of a predicate that already lives in the cache builder."""
+    return math.lcm(mamba_cache_chunk_size(), tree_page)
+
+
 def max_speculative_num_draft_tokens() -> int | None:
     """The largest draft-token count speculative decoding may use.
 

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):
+def build_fixture(
-    """Create (tree, allocator, req_to_token_pool) from a CacheConfig."""
+    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),