support page size large than 64 for mamba radix cache with fix (#16768)

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -34,6 +34,7 @@ from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, MambaPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
 from sglang.srt.model_executor.model_runner import ModelRunner
+from sglang.srt.server_args import get_global_server_args
 from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
 from sglang.srt.speculative.spec_info import SpecInput
 from sglang.srt.utils import is_cuda, is_npu
@@ -284,7 +285,8 @@ class MambaAttnBackendBase(AttentionBackend):
         lens_to_track = (
             forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens
         )
-        aligned_len = (lens_to_track // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
+        mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
+        aligned_len = (lens_to_track // mamba_cache_chunk_size) * mamba_cache_chunk_size
         start_indices = query_start_loc[:-1] + aligned_len - conv_state_len
         start_indices = start_indices[forward_batch.mamba_track_mask]
 

python/sglang/srt/managers/schedule_batch.py

@@ -1618,10 +1618,19 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         mamba_track_indices_cpu: List[int],
         mamba_track_seqlens_cpu: List[int],
     ):
-        MAMBA_CACHE_V2_CHUNK_SIZE = max(
+        def _force_track_h(i: int) -> int:
-            FLA_CHUNK_SIZE, self.token_to_kv_pool_allocator.page_size
+            assert i % FLA_CHUNK_SIZE == 0
-        )
+            # There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu:
-        mask = req.extend_input_len >= MAMBA_CACHE_V2_CHUNK_SIZE
+            # 1) aligned with FLA_CHUNK_SIZE-> retrieve from last_recurrent_state
+            #    a) is the last position -> retrieve from last_recurrent_state
+            #    b) is NOT the last position -> retrieve from h
+            # 2) unaligned with FLA_CHUNK_SIZE -> retrieve from h
+            # Currently, the math calculation only supports case 1a and 2. So for 1b, we need to add 1
+            # to force the math calculation to retrieve the correct mamba state from h.
+            return i + 1
+
+        mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
+        mask = req.extend_input_len >= mamba_cache_chunk_size
         mamba_track_mask_cpu.append(mask)
         mamba_track_indices_cpu.append(
             req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item()
@@ -1636,13 +1645,28 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
             # 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_input_len
-            # mamba_last_track_seqlen is actual tracked seqlen. Used to pass to
+
+            # 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_input_len // MAMBA_CACHE_V2_CHUNK_SIZE)
+                + (req.extend_input_len // mamba_cache_chunk_size)
-                * MAMBA_CACHE_V2_CHUNK_SIZE
+                * mamba_cache_chunk_size
             )
+
+            # mamba_track_fla_chunk_aligned is the aligned seqlen based on FLA_CHUNK_SIZE
+            # If mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned, which can be true when
+            # page_size > FLA_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_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
+            )
+            if mamba_track_fla_chunk_aligned != mamba_track_seqlen_aligned:
+                # We want to track mamba_track_seqlen_aligned, and it's not the last position,
+                # so we need to add 1 to the seqlen to retrieve the correct mamba state from h.
+                mamba_track_seqlen = _force_track_h(mamba_track_seqlen_aligned)
+
             req.mamba_next_track_idx = (
                 self.req_to_token_pool.get_mamba_ping_pong_other_idx(
                     req.mamba_next_track_idx
@@ -1653,18 +1677,16 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 # is within the current extend batch.
                 branching_seqlen_aligned_mask = (
                     req.mamba_branching_seqlen - len(req.prefix_indices)
-                ) % MAMBA_CACHE_V2_CHUNK_SIZE == 0
+                ) % mamba_cache_chunk_size == 0
                 if (
                     req.mamba_branching_seqlen > len(req.prefix_indices)
                     and req.mamba_branching_seqlen < mamba_track_seqlen
                     and branching_seqlen_aligned_mask
                 ):
-                    # NOTE: See the comment above for mamba_track_seqlen, the +1 is necessary
+                    # We want to track mamba_track_seqlen_aligned, and it's not the last position,
-                    # because the branching point is not the last aligned position, so we need
+                    # so we need to add 1 to the seqlen to retrieve the correct mamba state from h.
-                    # to retrieve its state from h. Adding 1 will give us the correct index in h,
+                    # See _force_track_h() for more details.
-                    # otherwise the calculation will retrieve the state from the last_recurrent_state,
+                    mamba_track_seqlen = _force_track_h(req.mamba_branching_seqlen)
-                    # which is not correct.
-                    mamba_track_seqlen = req.mamba_branching_seqlen + 1
                     mamba_track_seqlen_aligned = req.mamba_branching_seqlen
             req.mamba_last_track_seqlen = mamba_track_seqlen_aligned
         mamba_track_seqlens_cpu.append(mamba_track_seqlen)

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -41,6 +41,7 @@ from sglang.srt.mem_cache.radix_cache import (
     _key_match_paged,
     get_child_key,
 )
+from sglang.srt.server_args import get_global_server_args
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -955,13 +956,13 @@ class MambaRadixCache(BasePrefixCache):
         # Calculate the branching point. It is defined as the last aligned position that
         # does not have a mamba value.
         if len(value) > best_value_len:
-            MAMBA_CACHE_V2_CHUNK_SIZE = max(FLA_CHUNK_SIZE, self.page_size)
+            mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size
-            mamba_cache_v2_chunk_aligned_seqlen = (
+            mamba_cache_chunk_aligned_seqlen = (
-                sum(len(v) for v in value) // MAMBA_CACHE_V2_CHUNK_SIZE
+                sum(len(v) for v in value) // mamba_cache_chunk_size
-            ) * MAMBA_CACHE_V2_CHUNK_SIZE
+            ) * mamba_cache_chunk_size
             mamba_branching_seqlen = (
-                mamba_cache_v2_chunk_aligned_seqlen
+                mamba_cache_chunk_aligned_seqlen
-                if mamba_cache_v2_chunk_aligned_seqlen > 0
+                if mamba_cache_chunk_aligned_seqlen > 0
                 else None
             )
         else:

python/sglang/srt/server_args.py

@@ -1475,7 +1475,7 @@ class ServerArgs:
                         max(FLA_CHUNK_SIZE, self.page_size)
                         % min(FLA_CHUNK_SIZE, self.page_size)
                         == 0
-                    ), f"max(FLA_CHUNK_SIZE, self.page_size) % min(FLA_CHUNK_SIZE, self.page_size) should be 0, got {FLA_CHUNK_SIZE=}, {self.page_size=}"
+                    ), f"For SSM models with extra buffer, either FLA_CHUNK_SIZE or page_size must be divisible by the other, got {FLA_CHUNK_SIZE=}, {self.page_size=}"
 
             elif not self.disable_radix_cache:
                 logger.warning(
@@ -4630,6 +4630,12 @@ class ServerArgs:
     def enable_mamba_extra_buffer(self) -> bool:
         return self.mamba_scheduler_strategy == "extra_buffer"
 
+    @property
+    def mamba_cache_chunk_size(self) -> int:
+        # For mamba cache with extra buffer, the chunk size is the max of FLA_CHUNK_SIZE and page_size.
+        # It is used to determine the caching point in a sequence during prefill.
+        return max(FLA_CHUNK_SIZE, self.page_size)
+
     def check_server_args(self):
         # Check parallel size constraints
         assert (

test/srt/models/test_qwen3_next_models.py

@@ -19,6 +19,11 @@ from sglang.test.test_utils import (
 QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct"
 
 ACC_THRESHOLDS = {
+    QWEN3_NEXT_MODEL: {"kl_div": 0.0025, "gsm8k": 0.93},
+}
+
+# MTP has higher KL divergence threshold
+ACC_THRESHOLDS_MTP = {
     QWEN3_NEXT_MODEL: {"kl_div": 0.008, "gsm8k": 0.93},
 }
 
@@ -241,13 +246,13 @@ class TestQwen3NextMTPTopk(CustomTestCase):
         metrics = run_eval(args)
         print(f"{metrics=}")
         self.assertGreaterEqual(
-            metrics["accuracy"], ACC_THRESHOLDS[self.model]["gsm8k"]
+            metrics["accuracy"], ACC_THRESHOLDS_MTP[self.model]["gsm8k"]
         )
 
     def test_input_output_logprobs_match_prefill_cache_hit(self):
         test_input_output_logprobs_match_prefill_cache_hit_helper(
             self.base_url,
-            ACC_THRESHOLDS,
+            ACC_THRESHOLDS_MTP,
             self.model,
             max_samples=32,
             max_new_tokens=512,
@@ -256,7 +261,7 @@ class TestQwen3NextMTPTopk(CustomTestCase):
     def test_input_output_logprobs_match_decode_cache_hit(self):
         test_input_output_logprobs_match_decode_cache_hit_helper(
             self.base_url,
-            ACC_THRESHOLDS,
+            ACC_THRESHOLDS_MTP,
             self.model,
             max_samples=32,
             max_new_tokens=512,