support page size large than 64 for mamba radix cache (#16657)

Co-authored-by: Hanming Lu <hanming@x.ai>

python/sglang/srt/managers/schedule_batch.py

@@ -1621,7 +1621,10 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
         mamba_track_indices_cpu: List[int],
         mamba_track_seqlens_cpu: List[int],
     ):
-        mask = (req.extend_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE > 0
+        MAMBA_CACHE_V2_CHUNK_SIZE = max(
+            FLA_CHUNK_SIZE, self.token_to_kv_pool_allocator.page_size
+        )
+        mask = req.extend_input_len >= MAMBA_CACHE_V2_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()
@@ -1640,7 +1643,8 @@ 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_input_len // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE
+                + (req.extend_input_len // MAMBA_CACHE_V2_CHUNK_SIZE)
+                * MAMBA_CACHE_V2_CHUNK_SIZE
             )
             req.mamba_next_track_idx = (
                 self.req_to_token_pool.get_mamba_ping_pong_other_idx(
@@ -1652,7 +1656,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
                 # is within the current extend batch.
                 branching_seqlen_aligned_mask = (
                     req.mamba_branching_seqlen - len(req.prefix_indices)
-                ) % FLA_CHUNK_SIZE == 0
+                ) % MAMBA_CACHE_V2_CHUNK_SIZE == 0
                 if (
                     req.mamba_branching_seqlen > len(req.prefix_indices)
                     and req.mamba_branching_seqlen < mamba_track_seqlen

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -955,11 +955,14 @@ 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:
-            fla_chunk_aligned_seqlen = (
+            MAMBA_CACHE_V2_CHUNK_SIZE = max(FLA_CHUNK_SIZE, self.page_size)
-                sum(len(v) for v in value) // FLA_CHUNK_SIZE
+            mamba_cache_v2_chunk_aligned_seqlen = (
-            ) * FLA_CHUNK_SIZE
+                sum(len(v) for v in value) // MAMBA_CACHE_V2_CHUNK_SIZE
+            ) * MAMBA_CACHE_V2_CHUNK_SIZE
             mamba_branching_seqlen = (
-                fla_chunk_aligned_seqlen if fla_chunk_aligned_seqlen > 0 else None
+                mamba_cache_v2_chunk_aligned_seqlen
+                if mamba_cache_v2_chunk_aligned_seqlen > 0
+                else None
             )
         else:
             mamba_branching_seqlen = None

python/sglang/srt/server_args.py

@@ -1471,8 +1471,10 @@ class ServerArgs:
                         self.mamba_track_interval % self.page_size == 0
                     ), f"mamba_track_interval {self.mamba_track_interval} must be divisible by page_size {self.page_size}"
                     assert (
-                        FLA_CHUNK_SIZE % self.page_size == 0
+                        max(FLA_CHUNK_SIZE, self.page_size)
-                    ), f"Page size for hybrid GDN model must be divisible by {FLA_CHUNK_SIZE}, got {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=}"
 
             elif not self.disable_radix_cache:
                 logger.warning(