From c20aceeb88d65cd4815551d8e2a75f7a1988ed72 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Sat, 15 Aug 2026 00:34:00 +0800 Subject: [PATCH] Fix mamba checkpoint depth under dcp (#34808) --- python/sglang/srt/managers/schedule_batch.py | 15 ++-- .../components/mamba_component.py | 8 +- python/sglang/srt/runtime_context.py | 9 +++ .../managers/test_mamba_checkpoint_depth.py | 74 +++++++++++++++++++ .../test_unified_radix_cache_unittest.py | 59 ++++++++++++++- 5 files changed, 154 insertions(+), 11 deletions(-) create mode 100644 test/registered/unit/managers/test_mamba_checkpoint_depth.py diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 3f7465f4f..62a39ba74 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 - # page_size > chunk_size, we need to force the math calculation to retrieve the correct mamba state from h - # by _force_track_h() + # If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which is true when + # checkpoint_grid is coarser than chunk_size, we need to force the math calculation to + # 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 diff --git a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py b/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py index 11773c07f..fe5347ac3 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/mamba_component.py +++ b/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 - ) * self.mamba_cache_chunk_size + result.full_kv_hit_length // self.mamba_checkpoint_grid + ) * self.mamba_checkpoint_grid branching_seqlen = ( aligned_seqlen if aligned_seqlen > mamba_boundary_len else None ) diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index 95173d9ff..b572581bd 100644 --- a/python/sglang/srt/runtime_context.py +++ b/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. diff --git a/test/registered/unit/managers/test_mamba_checkpoint_depth.py b/test/registered/unit/managers/test_mamba_checkpoint_depth.py new file mode 100644 index 000000000..beb937949 --- /dev/null +++ b/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() diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index f8bb9866d..f9e7c6cee 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/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),