[Mamba] Fix checkpoint depth for prefixes that end off the radix page (#39115)

Co-authored-by: Ke Bao <ispobaoke@gmail.com>

python/sglang/srt/managers/schedule_batch.py

@@ -2965,7 +2965,18 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             # 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/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
+            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()
@@ -2978,14 +2989,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             # 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_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
-            )
+            mamba_track_seqlen = seq_end
 
             # A coarser checkpoint grid may not be a model-state boundary, so
             # force retrieval from the intermediate h state in that case.