From 07d996e6762fe8738925ab70d4a0f1456216918d Mon Sep 17 00:00:00 2001 From: kurt Date: Fri, 11 Sep 2026 15:24:59 +0000 Subject: [PATCH 1/2] [Mamba] Pick the radix checkpoint depth on the absolute grid when the prefix is off the tree page With DCP the allocator and radix-tree page widen (e.g. 64 -> 512) while PrefillAdder still aligns chunk budgets to the scheduler page, so a prefill chunk can end at 64. `cache_unfinished_req` inserts nothing for that chunk (no Mamba checkpoint yet, key shorter than a page) and leaves `prefix_indices` at 64. The next extend then measures the checkpoint depth relative to the prefix: 64 + floor(540 / 512) * 512 = 576, a position no 512-token page can name. On finish the Mamba component truncates the insert to 576 while the radix key is page-floored to 512, and the two freed kv-row tails [512, 576) and [576, 605) share a page (`_page_disjoint` assertion); without the assertion the tree would hold a 576-token Mamba state on a 512-token node. Pick the depth on the absolute checkpoint grid instead and donate only where an h snapshot exists ((depth - prefix) % cache_chunk_size == 0). Extends that cross no page boundary, or whose prefix is off the kernel chunk grid, donate nothing. Unit test: one table-driven case; three of its four rows fail before this change. --- python/sglang/srt/managers/schedule_batch.py | 17 ++++++---- .../managers/test_mamba_checkpoint_depth.py | 34 ++++++++++++++++--- 2 files changed, 41 insertions(+), 10 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index efde185a4de9..374e67ebe5ed 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2856,7 +2856,16 @@ def _force_track_h(i: int) -> int: # to force the math calculation to retrieve the correct mamba state from h. return i + 1 - mask = req.extend_range.length >= checkpoint_grid + # Pick the depth on the absolute checkpoint grid: a chunk boundary can leave + # the prefix off the (DCP-widened) tree page, and a prefix-relative depth then + # names a position no page can hold. Donate only where an h snapshot exists. + prefix_len = len(req.prefix_indices) + seq_end = prefix_len + req.extend_range.length + mamba_track_seqlen_aligned = (seq_end // checkpoint_grid) * checkpoint_grid + mask = ( + mamba_track_seqlen_aligned > prefix_len + and (mamba_track_seqlen_aligned - prefix_len) % cache_chunk_size == 0 + ) track_index = req.kv.mamba_ping_pong_track_buffer[ req.kv.mamba_next_track_idx ].item() @@ -2869,14 +2878,10 @@ def _force_track_h(i: int) -> int: # otherwise retrieved from h (i.e. unaligned). # We need to pass the non-aligned seqlen to the calculation. Even though # we pass in mamba_track_seqlen, the actual tracked seqlen is mamba_last_track_seqlen. - mamba_track_seqlen = len(req.prefix_indices) + req.extend_range.length + mamba_track_seqlen = seq_end # mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to # 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 // checkpoint_grid) * checkpoint_grid - ) # A coarser checkpoint grid may not be a model-state boundary, so # force retrieval from the intermediate h state in that case. diff --git a/test/registered/unit/managers/test_mamba_checkpoint_depth.py b/test/registered/unit/managers/test_mamba_checkpoint_depth.py index a08843cf6f6d..dd4e62c03bb5 100644 --- a/test/registered/unit/managers/test_mamba_checkpoint_depth.py +++ b/test/registered/unit/managers/test_mamba_checkpoint_depth.py @@ -3,7 +3,10 @@ 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. +tokens past that node. The same happens when the depth is measured from a +prefix that a chunked prefill left on the scheduler page but off the widened +tree page: 64 + 512 = 576 is no 512-page boundary, and the finished request +then frees two kv-row segments that share the page holding 576. """ import unittest @@ -27,8 +30,9 @@ 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.""" +def _track_seqlen(*, tree_page: int, prefix_len: int, extend_len: int) -> int | None: + """Run one extend through the tracker and report the donated depth, or + None when the extend donates no checkpoint.""" 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 @@ -57,7 +61,10 @@ def _track_seqlen(*, tree_page: int, prefix_len: int, extend_len: int) -> int: 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) + entry = batch._mamba_radix_cache_v2_req_prepare_for_extend(req) + if not entry.track_mask: + assert req.kv.mamba_last_track_seqlen is None + return None return req.kv.mamba_last_track_seqlen @@ -73,6 +80,25 @@ 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) + def test_depth_is_picked_on_the_absolute_grid(self): + # (tree_page, prefix_len, extend_len) -> donated depth, or None when the + # extend crosses no page or the prefix is off the kernel chunk grid. + for tree_page, prefix_len, extend_len, expected in ( + (512, 64, 540, 512), + (512, 448, 128, 512), + (512, 64, 400, None), + (512, 37, 540, None), + ): + with self.subTest(prefix_len=prefix_len, extend_len=extend_len): + self.assertEqual( + _track_seqlen( + tree_page=tree_page, + prefix_len=prefix_len, + extend_len=extend_len, + ), + expected, + ) + class TestMambaTrackGrid(unittest.TestCase): def _grid(self, *, interval: int, tree_page: int, chunk: int = CHUNK) -> int: From 57f762d205f1f014a1079a94ca7f70d420f25f70 Mon Sep 17 00:00:00 2001 From: ispobock Date: Wed, 16 Sep 2026 18:12:00 +0800 Subject: [PATCH 2/2] move mamba_track_seqlen_aligned comment to its assignment --- python/sglang/srt/managers/schedule_batch.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 374e67ebe5ed..1004cde857bb 100755 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -2861,6 +2861,8 @@ def _force_track_h(i: int) -> int: # names a position no page can hold. Donate only where an h snapshot exists. prefix_len = len(req.prefix_indices) seq_end = prefix_len + req.extend_range.length + # mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to + # mamba radix cache to track which seqlen this mamba state should store at. mamba_track_seqlen_aligned = (seq_end // checkpoint_grid) * checkpoint_grid mask = ( mamba_track_seqlen_aligned > prefix_len @@ -2880,9 +2882,6 @@ def _force_track_h(i: int) -> int: # we pass in mamba_track_seqlen, the actual tracked seqlen is mamba_last_track_seqlen. mamba_track_seqlen = seq_end - # mamba_track_seqlen_aligned/mamba_last_track_seqlen is actual tracked seqlen. Used to pass to - # mamba radix cache to track which seqlen this mamba state should store at. - # A coarser checkpoint grid may not be a model-state boundary, so # force retrieval from the intermediate h state in that case. mamba_track_fla_chunk_aligned = (