[Fix] Land the decode mamba checkpoint depth on the tree page under DCP (#35412)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Ke Bao <ispobaoke@gmail.com>

test/registered/unit/managers/test_batch_result_processor_mamba_boundary.py

@@ -17,6 +17,11 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=1, suite="base-a-test-cpu")
 
 
+# The decode checkpoint grid is lcm(mamba_cache_chunk_size, tree page,
+# interval); keeping all three equal leaves it at the interval under test.
+TRACK_INTERVAL = 4
+
+
 def _make_batch() -> tuple[Req, ScheduleBatch]:
     sampling_params = SamplingParams(max_new_tokens=32)
     sampling_params.normalize(None)
@@ -31,6 +36,7 @@ def _make_batch() -> tuple[Req, ScheduleBatch]:
     req.kv_committed_len = 2
 
     batch = ScheduleBatch(reqs=[req])
+    batch.tree_cache = SimpleNamespace(page_size=TRACK_INTERVAL)
     batch.device = "cpu"
     batch.model_config = SimpleNamespace(is_encoder_decoder=False)
     batch.enable_overlap = True
@@ -57,7 +63,7 @@ def _make_processor() -> SchedulerBatchResultProcessor:
         server_args=SimpleNamespace(),
         model_config=SimpleNamespace(think_end_ids=None),
         token_to_kv_pool_allocator=MagicMock(),
-        tree_cache=None,
+        tree_cache=SimpleNamespace(page_size=TRACK_INTERVAL),
         hisparse_coordinator=None,
         req_to_token_pool=None,
         decode_offload_manager=None,
@@ -154,7 +160,8 @@ class TestMambaBoundaryMaskReuse(unittest.TestCase):
                     # defaults.
                     get_context().override_server_args(
                         mamba_radix_cache_strategy="extra_buffer",
-                        mamba_track_interval=4,
+                        mamba_track_interval=TRACK_INTERVAL,
+                        _mamba_cache_chunk_size=TRACK_INTERVAL,
                     ),
                     patch(
                         "sglang.srt.managers.schedule_batch.alloc_for_decode",

test/registered/unit/managers/test_mamba_checkpoint_depth.py

@@ -14,6 +14,7 @@ from unittest.mock import MagicMock
 import torch
 
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
+from sglang.srt.runtime_context import get_context, mamba_track_grid
 from sglang.srt.sampling.sampling_params import SamplingParams
 from sglang.srt.server_args import (
     ServerArgs,
@@ -21,7 +22,7 @@ from sglang.srt.server_args import (
 )
 from sglang.test.ci.ci_register import register_cpu_ci
 
-register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
 
 CHUNK = 64
 
@@ -70,5 +71,24 @@ class TestMambaCheckpointDepth(unittest.TestCase):
         self.assertEqual(depth, 20416)
 
 
+class TestMambaTrackGrid(unittest.TestCase):
+    def _grid(self, *, interval: int, tree_page: int, chunk: int = CHUNK) -> int:
+        with get_context().override_server_args(
+            mamba_track_interval=interval, _mamba_cache_chunk_size=chunk
+        ):
+            return mamba_track_grid(tree_page)
+
+    def test_widened_tree_page_rounds_the_interval_up(self):
+        self.assertEqual(self._grid(interval=256, tree_page=512), 512)
+
+    def test_interval_already_on_the_tree_page_is_untouched(self):
+        self.assertEqual(self._grid(interval=256, tree_page=256), 256)
+        self.assertEqual(self._grid(interval=256, tree_page=128), 256)
+        self.assertEqual(self._grid(interval=256, tree_page=64), 256)
+
+    def test_grid_stays_on_the_chunk_size(self):
+        self.assertEqual(self._grid(interval=192, tree_page=64, chunk=128) % 128, 0)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/spec/test_ngram_mamba_verify_update.py

@@ -198,8 +198,8 @@ class TestNgramMambaVerifyUpdate(CustomTestCase):
             "sglang.srt.speculative.spec_utils.mambaish_config",
             return_value={"some": "config"},
         ), patch(
-            "sglang.srt.speculative.spec_utils.get_exec",
-            return_value=MagicMock(mamba=MagicMock(mamba_track_interval=256)),
+            "sglang.srt.speculative.spec_utils.mamba_track_grid",
+            return_value=256,
         ):
             commit_mamba_states_after_verify(
                 target_worker,