diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 5d5a64c9d..4c63913f3 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/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] diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 17ea53769..7bfa712ff 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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( - FLA_CHUNK_SIZE, self.token_to_kv_pool_allocator.page_size - ) - mask = req.extend_input_len >= MAMBA_CACHE_V2_CHUNK_SIZE + def _force_track_h(i: int) -> int: + assert i % FLA_CHUNK_SIZE == 0 + # There are 3 cases for mamba_track_seqlen passed to mamba_track_seqlens_cpu: + # 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) - * MAMBA_CACHE_V2_CHUNK_SIZE + + (req.extend_input_len // mamba_cache_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 - # because the branching point is not the last aligned position, so we need - # to retrieve its state from h. Adding 1 will give us the correct index in h, - # otherwise the calculation will retrieve the state from the last_recurrent_state, - # which is not correct. - mamba_track_seqlen = req.mamba_branching_seqlen + 1 + # 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. + # See _force_track_h() for more details. + mamba_track_seqlen = _force_track_h(req.mamba_branching_seqlen) 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) diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 70d864bbe..5c61299cf 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/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_v2_chunk_aligned_seqlen = ( - sum(len(v) for v in value) // MAMBA_CACHE_V2_CHUNK_SIZE - ) * MAMBA_CACHE_V2_CHUNK_SIZE + mamba_cache_chunk_size = get_global_server_args().mamba_cache_chunk_size + mamba_cache_chunk_aligned_seqlen = ( + sum(len(v) for v in value) // mamba_cache_chunk_size + ) * mamba_cache_chunk_size mamba_branching_seqlen = ( - mamba_cache_v2_chunk_aligned_seqlen - if mamba_cache_v2_chunk_aligned_seqlen > 0 + mamba_cache_chunk_aligned_seqlen + if mamba_cache_chunk_aligned_seqlen > 0 else None ) else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index ca32e98c4..261895b01 100644 --- a/python/sglang/srt/server_args.py +++ b/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 ( diff --git a/test/srt/models/test_qwen3_next_models.py b/test/srt/models/test_qwen3_next_models.py index fd0c6c41a..f8eff84cb 100644 --- a/test/srt/models/test_qwen3_next_models.py +++ b/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,