diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 50288f3e2..3513bd02e 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 0513e4180..70d864bbe 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/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 = ( - sum(len(v) for v in value) // FLA_CHUNK_SIZE - ) * FLA_CHUNK_SIZE + MAMBA_CACHE_V2_CHUNK_SIZE = max(FLA_CHUNK_SIZE, self.page_size) + mamba_cache_v2_chunk_aligned_seqlen = ( + 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 diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index bcdad4e3e..35ed102ed 100644 --- a/python/sglang/srt/server_args.py +++ b/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 - ), f"Page size for hybrid GDN model must be divisible by {FLA_CHUNK_SIZE}, got {self.page_size}" + 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=}" elif not self.disable_radix_cache: logger.warning(