diff --git a/docs/advanced_features/server_arguments.md b/docs/advanced_features/server_arguments.md index 95355f730..c2c8b1682 100644 --- a/docs/advanced_features/server_arguments.md +++ b/docs/advanced_features/server_arguments.md @@ -303,6 +303,8 @@ Please consult the documentation below and [server_args.py](https://github.com/s | `--max-mamba-cache-size` | The maximum size of the mamba cache. | `None` | Type: int | | `--mamba-ssm-dtype` | The data type of the SSM states in mamba cache. | `float32` | `float32`, `bfloat16` | | `--mamba-full-memory-ratio` | The ratio of mamba state memory to full kv cache memory. | `0.2` | Type: float | +| `--mamba-scheduler-strategy` | The strategy to use for mamba scheduler. `auto` currently defaults to `no_buffer`. 1. `no_buffer` does not support overlap scheduler due to not allocating extra mamba state buffers. Branching point caching support is feasible but not implemented. 2. `extra_buffer` supports overlap schedule by allocating extra mamba state buffers to track mamba state for caching (mamba state usage per running req becomes `2x` for non-spec; `1+(1/(2+speculative_num_draft_tokens))x` for spec dec (e.g. 1.16x if speculative_num_draft_tokens==4)). 2a. `extra_buffer` is strictly better for non-KV-cache-bound cases; for KV-cache-bound cases, the tradeoff depends on whether enabling overlap outweighs reduced max running requests. 2b. mamba caching at radix cache branching point is strictly better than non-branch but requires kernel support (currently only FLA backend), currently only extra_buffer supports branching. | `auto` | `auto`, `no_buffer`, `extra_buffer` | +| `--mamba-track-interval` | The interval (in tokens) to track the mamba state during decode. Only used when `--mamba-scheduler-strategy` is `extra_buffer`. Must be divisible by page_size if set, and must be >= speculative_num_draft_tokens when using speculative decoding. | `256` | Type: int | ## Args for multi-item scoring | Argument | Description | Defaults | Options | diff --git a/python/sglang/srt/disaggregation/decode.py b/python/sglang/srt/disaggregation/decode.py index b92ffce1d..f462f3301 100644 --- a/python/sglang/srt/disaggregation/decode.py +++ b/python/sglang/srt/disaggregation/decode.py @@ -155,8 +155,16 @@ class HybridMambaDecodeReqToTokenPool(HybridReqToTokenPool): pre_alloc_size=pre_alloc_size, ) self.enable_memory_saver = enable_memory_saver + self.enable_mamba_extra_buffer = ( + False # TODO: add PD support for mamba cache extra_buffer + ) self._init_mamba_pool( - size + pre_alloc_size, cache_params, device, speculative_num_draft_tokens + size=size + pre_alloc_size, + mamba_spec_state_size=size + pre_alloc_size, + cache_params=cache_params, + device=device, + enable_mamba_extra_buffer=self.enable_mamba_extra_buffer, + speculative_num_draft_tokens=speculative_num_draft_tokens, ) def clear(self): diff --git a/python/sglang/srt/layers/attention/fla/chunk.py b/python/sglang/srt/layers/attention/fla/chunk.py index 21d93ac00..fc1206135 100644 --- a/python/sglang/srt/layers/attention/fla/chunk.py +++ b/python/sglang/srt/layers/attention/fla/chunk.py @@ -67,7 +67,7 @@ def chunk_gated_delta_rule_fwd( cu_seqlens=cu_seqlens, ) if SUPPRESS_LEVEL < 3: - return g, o, A, final_state, None, None, None + return g, o, A, final_state, None, h, None elif SUPPRESS_LEVEL >= 3: return g, o, A, final_state, w, h, v_new @@ -108,7 +108,7 @@ class ChunkGatedDeltaRuleFunction(torch.autograd.Function): output_final_state=output_final_state, cu_seqlens=cu_seqlens, ) - return o.to(q.dtype), final_state + return o.to(q.dtype), final_state, h @torch.compiler.disable @@ -224,7 +224,7 @@ def chunk_gated_delta_rule( ) if scale is None: scale = k.shape[-1] ** -0.5 - o, final_state = ChunkGatedDeltaRuleFunction.apply( + o, final_state, h = ChunkGatedDeltaRuleFunction.apply( q, k, v, @@ -238,4 +238,4 @@ def chunk_gated_delta_rule( ) if head_first: o = rearrange(o, "b t h ... -> b h t ...") - return o, final_state + return o, final_state, h diff --git a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py index dbae27c16..e926fc4d3 100644 --- a/python/sglang/srt/layers/attention/fla/chunk_delta_h.py +++ b/python/sglang/srt/layers/attention/fla/chunk_delta_h.py @@ -16,6 +16,7 @@ from sglang.srt.layers.attention.fla.op import exp, safe_exp from sglang.srt.layers.attention.fla.utils import is_nvidia_hopper NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8, 16] +CHUNK_SIZE = 64 # @triton.autotune( @@ -274,16 +275,15 @@ def chunk_gated_delta_rule_fwd_h( gk: Optional[torch.Tensor] = None, initial_state: Optional[torch.Tensor] = None, output_final_state: bool = False, - chunk_size: int = 64, # SY: remove this argument and force chunk size 64? save_new_value: bool = True, cu_seqlens: Optional[torch.LongTensor] = None, -) -> Tuple[torch.Tensor, torch.Tensor]: +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: B, T, Hg, K, V = *k.shape, u.shape[-1] H = u.shape[-2] - BT = chunk_size + BT = CHUNK_SIZE chunk_indices = ( - prepare_chunk_indices(cu_seqlens, chunk_size) + prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) if cu_seqlens is not None else None ) diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_recurrent.py index 92c236edd..058662280 100644 --- a/python/sglang/srt/layers/attention/fla/fused_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_recurrent.py @@ -353,6 +353,7 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( cu_seqlens, scale, intermediate_states_buffer, + intermediate_state_indices, cache_steps, retrieve_parent_token_ptr, stride_retrieve_parent_token_seq: tl.constexpr, @@ -431,7 +432,7 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel( # Prepare intermediate state cache variables if enabled cache_idx = -1 if CACHE_INTERMEDIATE_STATES: - cache_idx = tl.load(h0_indices + i_n) + cache_idx = tl.load(intermediate_state_indices + i_n) step_idx = 0 for _ in range(0, T): @@ -532,6 +533,7 @@ def fused_recurrent_gated_delta_rule_update_fwd( disable_state_update: bool = False, disable_output_calculation: bool = False, intermediate_states_buffer: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, cache_steps: Optional[int] = None, retrieve_parent_token: Optional[torch.Tensor] = None, ) -> torch.Tensor: @@ -574,6 +576,7 @@ def fused_recurrent_gated_delta_rule_update_fwd( cu_seqlens=cu_seqlens, scale=scale, intermediate_states_buffer=intermediate_states_buffer, + intermediate_state_indices=intermediate_state_indices, cache_steps=0 if cache_steps is None else cache_steps, retrieve_parent_token_ptr=retrieve_parent_token, stride_retrieve_parent_token_seq=stride_retrieve_parent_token_seq, @@ -588,13 +591,13 @@ def fused_recurrent_gated_delta_rule_update_fwd( BK=BK, BV=BV, USE_INITIAL_STATE=initial_state_source is not None, - IS_BETA_HEADWISE=beta.ndim == v.ndim, - USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, IS_VARLEN=cu_seqlens is not None, - DISABLE_STATE_UPDATE=disable_state_update, - DISABLE_OUTPUT_CALCULATION=disable_output_calculation, CACHE_INTERMEDIATE_STATES=intermediate_states_buffer is not None, HAS_EAGLE_TREE_CUSTOM_ATTN_MASK=retrieve_parent_token is not None, + IS_BETA_HEADWISE=beta.ndim == v.ndim, + USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel, + DISABLE_STATE_UPDATE=disable_state_update, + DISABLE_OUTPUT_CALCULATION=disable_output_calculation, num_warps=num_warps, num_stages=num_stages, ) @@ -621,6 +624,7 @@ class FusedRecurrentUpdateFunction(torch.autograd.Function): disable_state_update: bool = False, disable_output_calculation: bool = False, intermediate_states_buffer: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, cache_steps: Optional[int] = None, retrieve_parent_token: Optional[torch.Tensor] = None, ): @@ -638,6 +642,7 @@ class FusedRecurrentUpdateFunction(torch.autograd.Function): disable_state_update=disable_state_update, disable_output_calculation=disable_output_calculation, intermediate_states_buffer=intermediate_states_buffer, + intermediate_state_indices=intermediate_state_indices, cache_steps=cache_steps, retrieve_parent_token=retrieve_parent_token, ) @@ -668,6 +673,7 @@ def fused_recurrent_gated_delta_rule_update( disable_state_update: bool = False, disable_output_calculation: bool = False, intermediate_states_buffer: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, cache_steps: Optional[int] = None, retrieve_parent_token: Optional[torch.Tensor] = None, ) -> torch.Tensor: @@ -677,14 +683,17 @@ def fused_recurrent_gated_delta_rule_update( f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`." f"Please flatten variable-length inputs before processing." ) - if ( - initial_state_source is not None - and initial_state_indices.shape[0] != len(cu_seqlens) - 1 - ): - raise ValueError( - f"The number of initial states is expected to be equal to the number of input sequences, " - f"i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}." - ) + if initial_state_source is not None: + if initial_state_indices.shape[0] != len(cu_seqlens) - 1: + raise ValueError( + f"The number of initial states is expected to be equal to the number of input sequences, " + f"i.e., {len(cu_seqlens) - 1} rather than {initial_state_indices.shape[0]}." + ) + if initial_state_indices.shape[0] != intermediate_state_indices.shape[0]: + raise ValueError( + f"The number of intermediate state indices is expected to be equal to the number of input sequences, " + f"i.e., {initial_state_indices.shape[0]} != {intermediate_state_indices.shape[0]}." + ) if scale is None: scale = k.shape[-1] ** -0.5 else: @@ -705,6 +714,7 @@ def fused_recurrent_gated_delta_rule_update( disable_state_update, disable_output_calculation, intermediate_states_buffer, + intermediate_state_indices, cache_steps, retrieve_parent_token, ) 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 aa40482da..20fbb29f8 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -1,10 +1,13 @@ from typing import Optional, Union import torch +import triton +import triton.language as tl from einops import rearrange from sglang.srt.layers.attention.base_attn_backend import AttentionBackend from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule +from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating from sglang.srt.layers.attention.fla.fused_recurrent import ( fused_recurrent_gated_delta_rule_update, @@ -57,6 +60,108 @@ elif is_npu(): causal_conv1d_update = causal_conv1d_update_npu +# Kernel to track mamba states if needed based on track mask +@triton.jit +def track_mamba_state_if_needed_kernel( + conv_states_ptr, + ssm_states_ptr, + cache_indices_ptr, + mamba_track_mask_ptr, + mamba_track_indices_ptr, + conv_state_stride_0, # stride for first dimension (batch/pool index) + ssm_state_stride_0, # stride for first dimension (batch/pool index) + conv_state_numel_per_row: tl.constexpr, # total elements per row + ssm_state_numel_per_row: tl.constexpr, # total elements per row + BLOCK_SIZE: tl.constexpr, +): + """ + Track conv_states and ssm_states rows based on track mask. + + This kernel replaces a Python loop that copies state tensors for mamba attention. + For each batch element, if the track mask is True, it copies the entire row from + the source index (cache_indices[i]) to the destination index (mamba_track_indices[i]). + + Grid: (batch_size,) + Each block handles one batch element, using multiple threads to copy data in parallel. + """ + batch_idx = tl.program_id(0) + + # Load the copy mask for this batch element + track_mask = tl.load(mamba_track_mask_ptr + batch_idx) + + # Early exit if we don't need to track + if not track_mask: + return + + # Load source and destination indices + src_idx = tl.load(cache_indices_ptr + batch_idx) + dst_idx = tl.load(mamba_track_indices_ptr + batch_idx) + + # Copy conv_states + # Each thread handles BLOCK_SIZE elements + for offset in range(0, conv_state_numel_per_row, BLOCK_SIZE): + element_indices = offset + tl.arange(0, BLOCK_SIZE) + mask = element_indices < conv_state_numel_per_row + + src_ptr = conv_states_ptr + src_idx * conv_state_stride_0 + element_indices + dst_ptr = conv_states_ptr + dst_idx * conv_state_stride_0 + element_indices + + data = tl.load(src_ptr, mask=mask, other=0.0) + tl.store(dst_ptr, data, mask=mask) + + # Copy ssm_states + for offset in range(0, ssm_state_numel_per_row, BLOCK_SIZE): + element_indices = offset + tl.arange(0, BLOCK_SIZE) + mask = element_indices < ssm_state_numel_per_row + + src_ptr = ssm_states_ptr + src_idx * ssm_state_stride_0 + element_indices + dst_ptr = ssm_states_ptr + dst_idx * ssm_state_stride_0 + element_indices + + data = tl.load(src_ptr, mask=mask, other=0.0) + tl.store(dst_ptr, data, mask=mask) + + +def track_mamba_states_if_needed( + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + mamba_track_mask: torch.Tensor, + mamba_track_indices: torch.Tensor, + batch_size: int, +): + """ + Track mamba states using Triton kernel for better performance. + + Args: + conv_states: Convolution states tensor [pool_size, ...] + ssm_states: SSM states tensor [pool_size, ...] + cache_indices: Source indices for each batch element [batch_size] + mamba_track_mask: Boolean mask indicating which elements to track [batch_size] + mamba_track_indices: Indices to track for each batch element [batch_size] + batch_size: Number of batch elements + """ + conv_state_numel_per_row = conv_states[0].numel() + ssm_state_numel_per_row = ssm_states[0].numel() + + # Choose BLOCK_SIZE based on the size of the data + BLOCK_SIZE = 1024 + + # Launch kernel with batch_size blocks + grid = (batch_size,) + track_mamba_state_if_needed_kernel[grid]( + conv_states, + ssm_states, + cache_indices, + mamba_track_mask, + mamba_track_indices, + conv_states.stride(0), + ssm_states.stride(0), + conv_state_numel_per_row, + ssm_state_numel_per_row, + BLOCK_SIZE, + ) + + class MambaAttnBackendBase(AttentionBackend): def __init__(self, model_runner: ModelRunner): super().__init__() @@ -71,6 +176,7 @@ class MambaAttnBackendBase(AttentionBackend): self.retrieve_parent_token_list = [] self.cached_cuda_graph_decode_query_start_loc: torch.Tensor = None self.cached_cuda_graph_verify_query_start_loc: torch.Tensor = None + self.conv_states_shape: tuple[int, int] = None def _forward_metadata(self, forward_batch: ForwardBatch): bs = forward_batch.batch_size @@ -78,6 +184,15 @@ class MambaAttnBackendBase(AttentionBackend): retrieve_next_token = None retrieve_next_sibling = None retrieve_parent_token = None + track_conv_indices = None + track_ssm_h_src = None + track_ssm_h_dst = None + track_ssm_final_src = None + track_ssm_final_dst = None + + mamba_cache_indices = self.req_to_token_pool.get_mamba_indices( + forward_batch.req_pool_indices + ) if forward_batch.forward_mode.is_decode_or_idle(): query_start_loc = torch.arange( @@ -108,22 +223,160 @@ class MambaAttnBackendBase(AttentionBackend): forward_batch.extend_start_loc[-1] + forward_batch.extend_seq_lens[-1] ) + if ( + forward_batch.mamba_track_mask is not None + and forward_batch.mamba_track_mask.any() + ): + track_conv_indices = self._init_track_conv_indices( + query_start_loc, forward_batch + ) + + ( + track_ssm_h_src, + track_ssm_h_dst, + track_ssm_final_src, + track_ssm_final_dst, + ) = self._init_track_ssm_indices(mamba_cache_indices, forward_batch) else: raise ValueError(f"Invalid forward mode: {forward_batch.forward_mode=}") - mamba_cache_indices = self.req_to_token_pool.get_mamba_indices( - forward_batch.req_pool_indices - ) + return ForwardMetadata( query_start_loc=query_start_loc, mamba_cache_indices=mamba_cache_indices, retrieve_next_token=retrieve_next_token, retrieve_next_sibling=retrieve_next_sibling, retrieve_parent_token=retrieve_parent_token, + track_conv_indices=track_conv_indices, + track_ssm_h_src=track_ssm_h_src, + track_ssm_h_dst=track_ssm_h_dst, + track_ssm_final_src=track_ssm_final_src, + track_ssm_final_dst=track_ssm_final_dst, ) def init_forward_metadata(self, forward_batch: ForwardBatch): self.forward_metadata = self._forward_metadata(forward_batch) + def _init_track_conv_indices( + self, query_start_loc: torch.Tensor, forward_batch: ForwardBatch + ): + """ + Compute indices for extracting conv states from the input sequence during extend. + + In Mamba models, the conv layer maintains a sliding window of recent inputs. + After processing a prefill chunk, we need to save the last `conv_state_len` tokens + of the processed region for prefix caching. + + The key insight is that FLA (Flash Linear Attention) processes sequences in chunks + of FLA_CHUNK_SIZE. We only track the conv state up to the last complete chunk boundary + (aligned_len). + + start_indices is the starting token index of the conv state to track in this extend batch. + indices include all pos to track in this extend batch, conv_state_len for each req that + needs to be tracked (i.e. mamba_track_mask is True) + + Returns: + indices: Tensor of shape [num_tracked_requests, conv_state_len] containing + flattened positions into the packed input tensor. + """ + conv_state_len = self.conv_states_shape[-1] + + # Calculate the end position of the last aligned chunk + lens_to_track = ( + forward_batch.mamba_track_seqlens - forward_batch.extend_prefix_lens + ) + aligned_len = (lens_to_track // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE + start_indices = query_start_loc[:-1] + aligned_len - conv_state_len + start_indices = start_indices[forward_batch.mamba_track_mask] + + # Create indices: [batch_size, conv_state_len] + indices = start_indices.unsqueeze(-1) + torch.arange( + conv_state_len, + device=self.device, + dtype=start_indices.dtype, + ) + + return indices.clamp(0, query_start_loc[-1] - 1) + + def _init_track_ssm_indices( + self, mamba_cache_indices: torch.Tensor, forward_batch: ForwardBatch + ): + """ + Compute source and destination indices for tracking SSM states for prefix caching. + + After processing a prefill, we need to save the SSM recurrent state for prefix caching. + The FLA kernel outputs intermediate hidden states `h` at each chunk boundary, + plus a `last_recurrent_state` at the end of the chunked prefill size. + + The challenge is that sequences may or may not end on a chunk boundary: + - Aligned case (len % FLA_CHUNK_SIZE == 0): In this case, FLA will store the to-cache + state in the last_recurrent_state. + - Unaligned case (len % FLA_CHUNK_SIZE != 0): The last_recurrent_state includes the + unaligned position, but we only want state up to the last chunk boundary. + We must extract from the intermediate `h` tensor at the appropriate chunk index. + + We compute the src and dst indices for all requests that need to be cached + (i.e. mamba_track_mask is True) based on the rule above. + + For example: + 1. If chunked prefill length is < 64, then only final state has value. In this case we + cache `final` state. + 2. if chunked prefill length == 64, then only final state has value. In this case we + cache pos 64, from `final` state + 3. if chunked prefill length >64 and < 128, then both h and final state have value. + We cache pos 64 from `h` state + 4. if chunked prefill length ==128, then both h and final state have value. We cache + pos 128 from `final` state. Note `h` doesn't include the pos 128. + + Returns: + track_ssm_h_src: Source indices into the packed `h` tensor (for unaligned seqs) + track_ssm_h_dst: Destination cache slot indices (for unaligned seqs) + track_ssm_final_src: Source indices into last_recurrent_state buffer (for aligned seqs) + track_ssm_final_dst: Destination cache slot indices (for aligned seqs) + """ + # Move to CPU to avoid kernel launches for masking operations + mamba_track_mask = forward_batch.mamba_track_mask.cpu() + extend_seq_lens = forward_batch.extend_seq_lens.cpu() + mamba_track_indices = forward_batch.mamba_track_indices.cpu() + mamba_cache_indices = mamba_cache_indices.cpu() + mamba_track_seqlens = forward_batch.mamba_track_seqlens.cpu() + prefix_lens = forward_batch.extend_prefix_lens.cpu() + + # Calculate the number of hidden states per request + num_h_states = (extend_seq_lens - 1) // FLA_CHUNK_SIZE + 1 + + # Calculate the starting offset for each sequence in the packed batch + track_ssm_src_offset = torch.zeros_like(num_h_states) + track_ssm_src_offset[1:] = torch.cumsum(num_h_states[:-1], dim=0) + + # Filter variables by track mask + lens_to_track = mamba_track_seqlens - prefix_lens + lens_masked = lens_to_track[mamba_track_mask] + offset_masked = track_ssm_src_offset[mamba_track_mask] + dst_masked = mamba_track_indices[mamba_track_mask] + + # Determine if the sequence ends at a chunk boundary + is_aligned = (lens_masked % FLA_CHUNK_SIZE) == 0 + + # Case 1: Aligned. Use last_recurrent_state from ssm_states. + track_ssm_final_src = mamba_cache_indices[mamba_track_mask][is_aligned] + track_ssm_final_dst = dst_masked[is_aligned] + + # Case 2: Unaligned. Use intermediate state from h. + # TODO: if support FLA_CHUNK_SIZE % page size != 0, then need to modify this + not_aligned = ~is_aligned + track_ssm_h_src = offset_masked[not_aligned] + ( + lens_masked[not_aligned] // FLA_CHUNK_SIZE + ) + track_ssm_h_dst = dst_masked[not_aligned] + + # Move back to GPU + return ( + track_ssm_h_src.to(self.device, non_blocking=True), + track_ssm_h_dst.to(self.device, non_blocking=True), + track_ssm_final_src.to(self.device, non_blocking=True), + track_ssm_final_dst.to(self.device, non_blocking=True), + ) + def init_forward_metadata_capture_cuda_graph( self, bs: int, @@ -165,7 +418,7 @@ class MambaAttnBackendBase(AttentionBackend): ) ) self.query_start_loc_list.append( - torch.empty((i + 2,), dtype=torch.int32, device=self.device) + torch.zeros((i + 2,), dtype=torch.int32, device=self.device) ) self.retrieve_next_token_list.append( torch.zeros( @@ -277,7 +530,6 @@ class MambaAttnBackendBase(AttentionBackend): # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask if forward_mode.is_target_verify() and spec_info.topk > 1: bs_without_pad = spec_info.retrive_next_token.shape[0] - # print(spec_info.retrive_next_token, spec_info.retrive_next_sibling) self.retrieve_next_token_list[bs - 1][:bs_without_pad].copy_( spec_info.retrive_next_token ) @@ -300,6 +552,70 @@ class MambaAttnBackendBase(AttentionBackend): def get_cuda_graph_seq_len_fill_value(self): return 1 # Mamba attn does not use seq lens to index kv cache + def _track_mamba_state_decode( + self, + forward_batch: ForwardBatch, + conv_states: torch.Tensor, + ssm_states: torch.Tensor, + cache_indices: torch.Tensor, + ): + """ + Track and copy Mamba conv/SSM states during decode for prefix caching. + + During decode, each token update modifies conv_states and ssm_states in-place + at positions indexed by cache_indices (the working slots). For prefix caching, + we need to copy these updated states to persistent cache slots (mamba_track_indices) + so they can be prefix cached. + + This delegates to `track_mamba_states_if_needed`, which performs: + conv_states[mamba_track_indices[i]] = conv_states[cache_indices[i]] + ssm_states[mamba_track_indices[i]] = ssm_states[cache_indices[i]] + for all requests where mamba_track_mask[i] is True. + """ + if forward_batch.mamba_track_mask is not None: + track_mamba_states_if_needed( + conv_states, + ssm_states, + cache_indices, + forward_batch.mamba_track_mask, + forward_batch.mamba_track_indices, + forward_batch.batch_size, + ) + + def _track_mamba_state_extend( + self, + forward_batch: ForwardBatch, + h: torch.Tensor, + ssm_states: torch.Tensor, + forward_metadata: ForwardMetadata, + ): + """ + Track and copy SSM states during extend for prefix caching. + + After the FLA chunked prefill kernel runs, we need to save the SSM recurrent + state at the last chunk boundary so it can be reused for prefix caching. + The source of the state depends on whether the sequence length is aligned + to FLA_CHUNK_SIZE. See `_init_track_ssm_indices` for more details on how + the source and destination indices are computed. + + Note: Conv state tracking for extend is handled separately via gather operations + using indices computed by `_init_track_conv_indices`. + """ + if ( + forward_batch.mamba_track_mask is not None + and forward_batch.mamba_track_mask.any() + ): + h = h.squeeze(0) + + if forward_metadata.track_ssm_h_src.numel() > 0: + ssm_states[forward_metadata.track_ssm_h_dst] = h[ + forward_metadata.track_ssm_h_src + ].to(ssm_states.dtype, copy=False) + if forward_metadata.track_ssm_final_src.numel() > 0: + ssm_states[forward_metadata.track_ssm_final_dst] = ssm_states[ + forward_metadata.track_ssm_final_src + ] + class KimiLinearAttnBackend(MambaAttnBackendBase): """Attention backend using Mamba kernel.""" @@ -521,6 +837,15 @@ class KimiLinearAttnBackend(MambaAttnBackendBase): class GDNAttnBackend(MambaAttnBackendBase): """Attention backend using Mamba kernel.""" + def __init__(self, model_runner: ModelRunner): + super().__init__(model_runner) + self.conv_states_shape = ( + model_runner.req_to_token_pool.mamba_pool.mamba_cache.conv[0].shape + ) + assert ( + self.conv_states_shape[-1] < FLA_CHUNK_SIZE + ), f"{self.conv_states_shape[-1]=} should be less than {FLA_CHUNK_SIZE}" + def forward_decode( self, q: torch.Tensor, @@ -593,6 +918,10 @@ class GDNAttnBackend(MambaAttnBackendBase): softplus_threshold=20.0, ) + self._track_mamba_state_decode( + forward_batch, conv_states, ssm_states, cache_indices + ) + return core_attn_out def forward_extend( @@ -622,12 +951,13 @@ class GDNAttnBackend(MambaAttnBackendBase): seq_len = kwargs["seq_len"] is_target_verify = forward_batch.forward_mode.is_target_verify() + forward_metadata = self.forward_metadata - query_start_loc = self.forward_metadata.query_start_loc - cache_indices = self.forward_metadata.mamba_cache_indices - retrieve_next_token = self.forward_metadata.retrieve_next_token - retrieve_next_sibling = self.forward_metadata.retrieve_next_sibling - retrieve_parent_token = self.forward_metadata.retrieve_parent_token + query_start_loc = forward_metadata.query_start_loc + cache_indices = forward_metadata.mamba_cache_indices + retrieve_next_token = forward_metadata.retrieve_next_token + retrieve_next_sibling = forward_metadata.retrieve_next_sibling + retrieve_parent_token = forward_metadata.retrieve_parent_token mamba_cache_params = self.req_to_token_pool.mamba2_layer_cache(layer_id) conv_states = mamba_cache_params.conv[0] @@ -643,6 +973,9 @@ class GDNAttnBackend(MambaAttnBackendBase): dtype=torch.bool, device=forward_batch.input_ids.device, ) + intermediate_state_indices = torch.arange( + cache_indices.shape[0], dtype=torch.int32, device=cache_indices.device + ) else: has_initial_states = forward_batch.extend_prefix_lens > 0 @@ -660,14 +993,30 @@ class GDNAttnBackend(MambaAttnBackendBase): activation, conv_state_indices=cache_indices[:batch_size], intermediate_conv_window=intermediate_conv_window_cache, + intermediate_state_indices=intermediate_state_indices[:batch_size], retrieve_next_token=retrieve_next_token, retrieve_next_sibling=retrieve_next_sibling, retrieve_parent_token=retrieve_parent_token, ) mixed_qkv = mixed_qkv_processed.transpose(1, 2).view(seq_len, -1) else: + mixed_qkv = mixed_qkv.transpose(0, 1) + if ( + forward_batch.mamba_track_mask is not None + and forward_batch.mamba_track_mask.any() + ): + conv_dst = forward_batch.mamba_track_indices + # Gather all slices at once: [:, track_conv_indices] -> [d, num_masked, slice_len] + # track_conv_indices is already filtered and clamped in _init_track_conv_indices + mixed_qkv_to_track = mixed_qkv[ + :, forward_metadata.track_conv_indices + ].transpose(0, 1) + # Apply mask and assign to destinations + mask_indices = forward_batch.mamba_track_mask.nonzero(as_tuple=True)[0] + conv_states[conv_dst[mask_indices]] = mixed_qkv_to_track + mixed_qkv = causal_conv1d_fn( - mixed_qkv.transpose(0, 1), + mixed_qkv, conv_weights, bias, activation=activation, @@ -710,12 +1059,13 @@ class GDNAttnBackend(MambaAttnBackendBase): use_qk_l2norm_in_kernel=True, disable_state_update=True, intermediate_states_buffer=intermediate_state_cache, + intermediate_state_indices=intermediate_state_indices, cache_steps=forward_batch.spec_info.draft_token_num, retrieve_parent_token=retrieve_parent_token, ) else: recurrent_state = ssm_states[cache_indices] - core_attn_out, last_recurrent_state = chunk_gated_delta_rule( + core_attn_out, last_recurrent_state, h = chunk_gated_delta_rule( q=query, k=key, v=value, @@ -730,6 +1080,10 @@ class GDNAttnBackend(MambaAttnBackendBase): last_recurrent_state = last_recurrent_state.to(ssm_states.dtype, copy=False) ssm_states[cache_indices] = last_recurrent_state + self._track_mamba_state_extend( + forward_batch, h, ssm_states, forward_metadata + ) + return core_attn_out @@ -965,14 +1319,23 @@ class HybridLinearAttnBackend(AttentionBackend): **kwargs, ) - def update_mamba_state_after_mtp_verify(self, accepted_indices, model): - request_number = accepted_indices.shape[0] + def update_mamba_state_after_mtp_verify( + self, + accepted_steps: torch.Tensor, + mamba_track_indices: Optional[torch.Tensor], + mamba_steps_to_track: Optional[torch.Tensor], + model, + ): + request_number = accepted_steps.shape[0] state_indices_tensor = ( self.linear_attn_backend.forward_metadata.mamba_cache_indices[ :request_number ] ) + intermediate_state_indices = torch.arange( + request_number, dtype=torch.int32, device=state_indices_tensor.device + ) mamba_caches = ( self.linear_attn_backend.req_to_token_pool.get_speculative_mamba2_params_all_layers() @@ -983,19 +1346,41 @@ class HybridLinearAttnBackend(AttentionBackend): intermediate_state_cache = mamba_caches.intermediate_ssm intermediate_conv_window_cache = mamba_caches.intermediate_conv_window[0] - # SSM state updates (chunked to reduce peak memory) - valid_mask = accepted_indices >= 0 - # Compute common indices once to avoid duplication - valid_state_indices = state_indices_tensor[valid_mask].to(torch.int64) # [N] - last_steps = accepted_indices[valid_mask].to(torch.int64) # [N] + valid_mask = accepted_steps >= 0 + dst_state_indices = state_indices_tensor[valid_mask].to(torch.int64) # [N] + src_state_indices = intermediate_state_indices[valid_mask].to( + torch.int64 + ) # [N] + last_steps = accepted_steps[valid_mask].to(torch.int64) # [N] # scatter into ssm_states at the chosen cache lines - ssm_states[:, valid_state_indices, :] = intermediate_state_cache[ - :, valid_state_indices, last_steps + ssm_states[:, dst_state_indices, :] = intermediate_state_cache[ + :, src_state_indices, last_steps ].to(ssm_states.dtype, copy=False) # Scatter into conv_states at the chosen cache lines - conv_states[:, valid_state_indices, :, :] = intermediate_conv_window_cache[ - :, valid_state_indices, last_steps + conv_states[:, dst_state_indices, :] = intermediate_conv_window_cache[ + :, src_state_indices, last_steps ].to(conv_states.dtype, copy=False) + + # Track indices used for tracking mamba states for prefix cache + if mamba_track_indices is not None: + assert mamba_steps_to_track is not None + track_mask = mamba_steps_to_track >= 0 + track_steps = mamba_steps_to_track[track_mask].to(torch.int64) # [N] + if track_steps.numel() == 0: + # No track indices to update + return + dst_track_indices = mamba_track_indices[track_mask].to(torch.int64) + src_track_indices = intermediate_state_indices[track_mask].to(torch.int64) + + # scatter into ssm_states at the chosen track states + ssm_states[:, dst_track_indices, :] = intermediate_state_cache[ + :, src_track_indices, track_steps + ].to(ssm_states.dtype, copy=False) + + # scatter into conv_states at the chosen track states + conv_states[:, dst_track_indices, :] = intermediate_conv_window_cache[ + :, src_track_indices, track_steps + ].to(conv_states.dtype, copy=False) diff --git a/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py b/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py index 08e695b75..247ae4326 100644 --- a/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py +++ b/python/sglang/srt/layers/attention/mamba/causal_conv1d_triton.py @@ -578,6 +578,7 @@ def _causal_conv1d_update_kernel( conv_state_indices_ptr, num_accepted_tokens_ptr, intermediate_conv_window_ptr, + intermediate_state_indices_ptr, retrieve_next_token_ptr, retrieve_next_sibling_ptr, retrieve_parent_token_ptr, @@ -602,6 +603,7 @@ def _causal_conv1d_update_kernel( stride_inter_step: tl.constexpr, stride_inter_dim: tl.constexpr, stride_inter_win: tl.constexpr, + stride_intermediate_state_indices: tl.constexpr, stride_retrieve_next_token_seq: tl.constexpr, stride_retrieve_next_token_token: tl.constexpr, stride_retrieve_next_sibling_seq: tl.constexpr, @@ -639,6 +641,11 @@ def _causal_conv1d_update_kernel( conv_state_batch_coord = tl.load( conv_state_indices_ptr + idx_seq * stride_state_indices ).to(tl.int64) + if SAVE_INTERMEDIATE: + intermediate_state_batch_coord = tl.load( + intermediate_state_indices_ptr + + idx_seq * stride_intermediate_state_indices + ).to(tl.int64) else: conv_state_batch_coord = idx_seq if USE_PAD_SLOT: # noqa @@ -847,7 +854,7 @@ def _causal_conv1d_update_kernel( # Layout: [seq(cache line), step, dim, win(K-1)] base_ptr = ( intermediate_conv_window_ptr - + conv_state_batch_coord * stride_inter_seq + + intermediate_state_batch_coord * stride_inter_seq + idx_token * stride_inter_step + idx_feats * stride_inter_dim ) @@ -934,7 +941,7 @@ def _causal_conv1d_update_kernel( # Layout: [seq(cache line), step, dim, win(K-1)] base_ptr = ( intermediate_conv_window_ptr - + conv_state_batch_coord * stride_inter_seq + + intermediate_state_batch_coord * stride_inter_seq + idx_token * stride_inter_step + idx_feats * stride_inter_dim ) @@ -980,6 +987,7 @@ def causal_conv1d_update( conv_state_indices: Optional[torch.Tensor] = None, num_accepted_tokens: Optional[torch.Tensor] = None, intermediate_conv_window: Optional[torch.Tensor] = None, + intermediate_state_indices: Optional[torch.Tensor] = None, retrieve_next_token: Optional[torch.Tensor] = None, retrieve_next_sibling: Optional[torch.Tensor] = None, retrieve_parent_token: Optional[torch.Tensor] = None, @@ -1040,6 +1048,8 @@ def causal_conv1d_update( assert conv_state.size(0) >= batch else: assert (batch,) == conv_state_indices.shape + assert intermediate_state_indices is not None + assert (batch,) == intermediate_state_indices.shape assert num_cache_lines >= batch assert weight.stride(1) == 1 # Need this @@ -1056,6 +1066,11 @@ def causal_conv1d_update( stride_state_indices = ( conv_state_indices.stride(0) if conv_state_indices is not None else 0 ) + stride_intermediate_state_indices = ( + intermediate_state_indices.stride(0) + if intermediate_state_indices is not None + else 0 + ) if num_accepted_tokens is not None: state_len = width - 1 + (seqlen - 1) # effective state_len needed else: @@ -1117,6 +1132,7 @@ def causal_conv1d_update( conv_state_indices, num_accepted_tokens, intermediate_conv_window if intermediate_conv_window is not None else x, + intermediate_state_indices, retrieve_next_token, retrieve_next_sibling, retrieve_parent_token, @@ -1141,6 +1157,7 @@ def causal_conv1d_update( stride_inter_step, stride_inter_dim, stride_inter_win, + stride_intermediate_state_indices, stride_retrieve_next_token_seq, stride_retrieve_next_token_token, stride_retrieve_next_sibling_seq, diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 47705fb7a..5eeb2b65e 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -27,9 +27,17 @@ from sglang.srt.model_executor.forward_batch_info import ForwardBatch class ForwardMetadata: query_start_loc: torch.Tensor mamba_cache_indices: torch.Tensor + # For topk > 1 eagle retrieve_next_token: Optional[torch.Tensor] = None retrieve_next_sibling: Optional[torch.Tensor] = None retrieve_parent_token: Optional[torch.Tensor] = None + # For prefill radix cache + track_conv_indices: Optional[torch.Tensor] = None + track_ssm_h_src: Optional[torch.Tensor] = None + track_ssm_h_dst: Optional[torch.Tensor] = None + track_ssm_final_src: Optional[torch.Tensor] = None + track_ssm_final_dst: Optional[torch.Tensor] = None + is_target_verify: bool = False draft_token_num: int = 1 diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5d0c44437..87eb98758 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -57,6 +57,7 @@ from sglang.srt.disaggregation.decode_schedule_batch_mixin import ( from sglang.srt.disaggregation.utils import DisaggregationMode from sglang.srt.distributed.parallel_state import get_tensor_model_parallel_rank from sglang.srt.environ import envs +from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.mem_cache.allocator import ( BaseTokenToKVPoolAllocator, SWATokenToKVPoolAllocator, @@ -543,6 +544,14 @@ class Req: # Memory pool info self.req_pool_idx: Optional[int] = None self.mamba_pool_idx: Optional[torch.Tensor] = None # shape (1) + self.mamba_ping_pong_track_buffer: Optional[torch.Tensor] = None # shape (2) + self.mamba_next_track_idx: Optional[int] = None # 0 or 1 + self.mamba_last_track_seqlen: Optional[int] = ( + None # seq len of the last cached mamba state + ) + # the branching point seqlen to track mamba state. If set, given by prefix match, + # it will be the tracked seqlen in the ping pong buffer for the right prefill pass. + self.mamba_branching_seqlen: Optional[int] = None # Check finish self.tokenizer = None @@ -824,11 +833,13 @@ class Req: self.last_node, self.last_host_node, self.host_hit_length, + self.mamba_branching_seqlen, ) = ( match_result.device_indices, match_result.last_device_node, match_result.last_host_node, match_result.host_hit_length, + match_result.mamba_branching_seqlen, ) self.cache_protected_len = len(self.prefix_indices) @@ -1027,6 +1038,10 @@ class Req: self.extend_logprob_start_len = 0 self.is_chunked = 0 self.mamba_pool_idx = None + self.mamba_ping_pong_track_buffer = None + self.mamba_next_track_idx = None + self.mamba_last_track_seqlen = None + self.mamba_branching_seqlen = None self.already_computed = 0 self.kv_allocated_len = 0 self.kv_committed_len = 0 @@ -1115,6 +1130,11 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): out_cache_loc: torch.Tensor = None # shape: [b], int64 output_ids: torch.Tensor = None # shape: [b], int64 + # For hybrid GDN prefix cache + mamba_track_indices: torch.Tensor = None # shape: [b], int64 + mamba_track_mask: torch.Tensor = None # shape: [b], bool + mamba_track_seqlens: torch.Tensor = None # shape: [b], int64 + # For multimodal inputs multimodal_inputs: Optional[List] = None @@ -1380,6 +1400,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): input_embeds = [] extend_input_logprob_token_ids = [] multimodal_inputs = [] + mamba_track_mask_cpu = [] + mamba_track_indices_cpu = [] + mamba_track_seqlens_cpu = [] for i, (req, seq_len, pre_len) in enumerate(zip(reqs, seq_lens, prefix_lens)): req.req_pool_idx = req_pool_indices[i] @@ -1403,6 +1426,14 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): req.already_computed = seq_len req.is_retracted = False + if get_global_server_args().enable_mamba_extra_buffer(): + self._mamba_radix_cache_v2_req_prepare_for_extend( + req, + mamba_track_mask_cpu, + mamba_track_indices_cpu, + mamba_track_seqlens_cpu, + ) + # Compute the relative logprob_start_len in an extend batch # # Key variables: @@ -1512,6 +1543,23 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.extend_logprob_start_lens = [r.extend_logprob_start_len for r in reqs] self.extend_input_logprob_token_ids = extend_input_logprob_token_ids + if get_global_server_args().enable_mamba_extra_buffer(): + self.mamba_track_indices = torch.tensor( + mamba_track_indices_cpu, + dtype=torch.int64, + device=self.device, + ) + self.mamba_track_mask = torch.tensor( + mamba_track_mask_cpu, + dtype=torch.bool, + device=self.device, + ) + self.mamba_track_seqlens = torch.tensor( + mamba_track_seqlens_cpu, + dtype=torch.int64, + device=self.device, + ) + if self.model_config.is_encoder_decoder: self.prepare_encoder_info_extend(input_ids, seq_lens) @@ -1521,6 +1569,60 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.model_config.vocab_size, ) + def _mamba_radix_cache_v2_req_prepare_for_extend( + self, + req: Req, + mamba_track_mask_cpu: List[bool], + 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_track_mask_cpu.append(mask) + mamba_track_indices_cpu.append( + req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx].item() + ) + mamba_track_seqlen = -1 + if mask: + # mamba_track_seqlen is used to calculate the indices to track in + # hybrid_linear_attn_backend's _init_track_ssm_indices. Due to the + # fact that the ssm state between aligned and non-aligned are retrieved differently, + # if 1) last pos and 2) is aligned, then retrieved from the last_recurrent_state, + # 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_input_len + # 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 // FLA_CHUNK_SIZE) * FLA_CHUNK_SIZE + ) + req.mamba_next_track_idx = ( + self.req_to_token_pool.get_mamba_ping_pong_other_idx( + req.mamba_next_track_idx + ) + ) + if req.mamba_branching_seqlen is not None: + # track branching point in this forward if the branching point + # is within the current extend batch. + branching_seqlen_aligned_mask = ( + req.mamba_branching_seqlen - len(req.prefix_indices) + ) % FLA_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 + 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) + def prepare_for_split_prefill(self): self.prepare_for_extend() # For split prefill, we need to set the forward mode to SPLIT_PREFILL @@ -1786,6 +1888,24 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.orig_seq_lens.add_(1) self.seq_lens_sum += bs + if get_global_server_args().enable_mamba_extra_buffer(): + self.mamba_track_indices = torch.tensor( + [ + req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx] + for req in self.reqs + ], + dtype=torch.int64, + device=self.device, + ) + self.mamba_track_mask = torch.tensor( + [ + sl % get_global_server_args().mamba_track_interval == 0 + for sl in self.seq_lens_cpu + ], + dtype=torch.bool, + device=self.device, + ) + def maybe_wait_verify_done(self): if self.is_v2_eagle: draft_input: EagleDraftInput = self.spec_info @@ -1842,6 +1962,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.out_cache_loc = None self.seq_lens_sum = self.seq_lens.sum().item() self.output_ids = self.output_ids[keep_indices_device] + self.mamba_track_indices = None + self.mamba_track_mask = None + self.mamba_track_seqlens = None self.return_logprob = any(req.return_logprob for req in self.reqs) if self.return_logprob: self.top_logprobs_nums = [self.top_logprobs_nums[i] for i in keep_indices] @@ -1889,6 +2012,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): self.seq_lens_sum += other.seq_lens_sum if self.output_ids is not None: self.output_ids = torch.cat([self.output_ids, other.output_ids]) + self.mamba_track_indices = None + self.mamba_track_mask = None + self.mamba_track_seqlens = None if self.return_logprob and other.return_logprob: self.top_logprobs_nums.extend(other.top_logprobs_nums) self.token_ids_logprobs.extend(other.token_ids_logprobs) @@ -1982,6 +2108,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): dllm_config=self.dllm_config, reqs=self.reqs, has_grammar=self.has_grammar, + mamba_track_indices=self.mamba_track_indices, + mamba_track_mask=self.mamba_track_mask, + mamba_track_seqlens=self.mamba_track_seqlens, ) def copy(self): @@ -2003,6 +2132,9 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin): is_prefill_only=self.is_prefill_only, seq_lens_cpu=self.seq_lens_cpu, enable_overlap=self.enable_overlap, + mamba_track_indices=self.mamba_track_indices, + mamba_track_mask=self.mamba_track_mask, + mamba_track_seqlens=self.mamba_track_seqlens, ) def _is_available_size_sufficient(self, num_tokens: int) -> bool: @@ -2104,3 +2236,8 @@ class ModelWorkerBatch: # FIXME(lsyin): remove this after fully overlap grammar reqs: Optional[List[Req]] = None has_grammar: bool = False + + # For mamba state tracking + mamba_track_indices: Optional[torch.Tensor] = None # shape: [b], int64 + mamba_track_mask: Optional[torch.Tensor] = None # shape: [b], bool + mamba_track_seqlens: Optional[torch.Tensor] = None # shape: [b], int64 diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 0f1176990..3bafdaa3c 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/python/sglang/srt/managers/schedule_policy.py @@ -363,7 +363,7 @@ class PrefillAdder: self.is_hybrid_swa = isinstance( self.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator ) - self.is_ssm_radix_cache = isinstance(self.tree_cache, MambaRadixCache) + self.is_hybrid_ssm_cache = isinstance(self.tree_cache, MambaRadixCache) self.priority_scheduling_preemption_threshold = ( priority_scheduling_preemption_threshold @@ -389,7 +389,7 @@ class PrefillAdder: self.token_to_kv_pool_allocator.swa_available_size() + self.tree_cache.swa_evictable_size(), ) - elif self.is_ssm_radix_cache: + elif self.is_hybrid_ssm_cache: available_and_evictable = ( self.token_to_kv_pool_allocator.available_size() + self.tree_cache.full_evictable_size() @@ -411,7 +411,7 @@ class PrefillAdder: self.token_to_kv_pool_allocator.swa_available_size() + self.tree_cache.swa_evictable_size(), ) - elif self.is_ssm_radix_cache: + elif self.is_hybrid_ssm_cache: available_and_evictable = ( self.token_to_kv_pool_allocator.available_size() + self.tree_cache.full_evictable_size() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 084b4a747..0f9272dd5 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -405,7 +405,7 @@ class Scheduler( # Hybrid memory pool self.is_hybrid_swa = self.tp_worker.is_hybrid_swa - self.is_ssm_model = ( + self.is_hybrid_ssm = ( self.tp_worker.model_runner.hybrid_gdn_config is not None or self.tp_worker.model_runner.mamba2_config is not None ) @@ -772,6 +772,7 @@ class Scheduler( eviction_policy=server_args.radix_eviction_policy, enable_metrics=self.enable_metrics, enable_kv_cache_events=self.enable_kv_cache_events, + enable_mamba_extra_buffer=server_args.enable_mamba_extra_buffer(), ) if ( @@ -808,7 +809,7 @@ class Scheduler( self.tree_cache = SWARadixCache( params=params, sliding_window_size=self.sliding_window_size ) - elif self.is_ssm_model: + elif self.is_hybrid_ssm: from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache self.tree_cache = MambaRadixCache(params) diff --git a/python/sglang/srt/managers/scheduler_metrics_mixin.py b/python/sglang/srt/managers/scheduler_metrics_mixin.py index a2549719b..6fe5fdc81 100644 --- a/python/sglang/srt/managers/scheduler_metrics_mixin.py +++ b/python/sglang/srt/managers/scheduler_metrics_mixin.py @@ -112,7 +112,7 @@ class SchedulerMetricsMixin: f"full token usage: {full_token_usage:.2f}, " f"swa token usage: {swa_token_usage:.2f}, " ) - elif self.is_ssm_model: + elif self.is_hybrid_ssm: ( full_num_used, _, @@ -166,7 +166,7 @@ class SchedulerMetricsMixin: self.stats.token_usage = token_usage if self.is_hybrid_swa: self.stats.swa_token_usage = swa_token_usage - if self.is_ssm_model: + if self.is_hybrid_ssm: self.stats.mamba_usage = mamba_usage self.stats.num_queue_reqs = len(self.waiting_queue) self.stats.num_grammar_queue_reqs = len(self.grammar_queue) @@ -238,7 +238,7 @@ class SchedulerMetricsMixin: f"#swa token: {swa_num_used}, " f"swa token usage: {swa_token_usage:.2f}, " ) - elif self.is_ssm_model: + elif self.is_hybrid_ssm: ( full_num_used, mamba_used, @@ -315,7 +315,7 @@ class SchedulerMetricsMixin: self.stats.token_usage = token_usage if self.is_hybrid_swa: self.stats.swa_token_usage = swa_token_usage - if self.is_ssm_model: + if self.is_hybrid_ssm: self.stats.mamba_usage = mamba_usage self.stats.decode_sum_seq_lens = batch.seq_lens_cpu.sum().item() self.stats.gen_throughput = self.last_gen_throughput @@ -402,7 +402,7 @@ class SchedulerMetricsMixin: if self.is_hybrid_swa: full_num_used, swa_num_used, *_ = self._get_swa_token_info() num_tokens = max(full_num_used, swa_num_used) - elif self.is_ssm_model: + elif self.is_hybrid_ssm: num_tokens = self._get_mamba_token_info()[0] else: num_tokens = self._get_token_info()[0] diff --git a/python/sglang/srt/managers/scheduler_output_processor_mixin.py b/python/sglang/srt/managers/scheduler_output_processor_mixin.py index c48f5f893..b3144c5cd 100644 --- a/python/sglang/srt/managers/scheduler_output_processor_mixin.py +++ b/python/sglang/srt/managers/scheduler_output_processor_mixin.py @@ -21,6 +21,7 @@ from sglang.srt.managers.schedule_batch import ( ScheduleBatch, ) from sglang.srt.mem_cache.common import release_kv_cache +from sglang.srt.server_args import get_global_server_args from sglang.srt.tracing.trace import trace_slice, trace_slice_batch, trace_slice_end if TYPE_CHECKING: @@ -267,6 +268,7 @@ class SchedulerOutputProcessorMixin: next_token_ids = result.next_token_ids.tolist() accept_lens = result.accept_lens.tolist() result.num_accepted_tokens = sum(accept_lens) - len(batch.reqs) + result.accept_length_per_req_cpu = [x - 1 for x in accept_lens] predict_tokens = [] stride = self.draft_worker.speculative_num_draft_tokens @@ -359,6 +361,9 @@ class SchedulerOutputProcessorMixin: req.output_ids.extend(next_token_id) new_accepted_len = len(next_token_id) + # Update Mamba last track seqlen + self._mamba_prefix_cache_update(req, batch, result, i) + req.check_finished(new_accepted_len) if req.finished(): @@ -424,6 +429,31 @@ class SchedulerOutputProcessorMixin: ): self.log_decode_stats(can_run_cuda_graph, running_batch=batch) + def _mamba_prefix_cache_update( + self, req: Req, batch: ScheduleBatch, result: GenerationBatchResult, i: int + ) -> None: + seq_len = len(req.origin_input_ids) + len(req.output_ids) - 1 + if req.mamba_ping_pong_track_buffer is not None: + mamba_track_interval = get_global_server_args().mamba_track_interval + if batch.spec_algorithm.is_none() and seq_len % mamba_track_interval == 0: + # for non-spec decode, we update mamba_last_track_seqlen at the end of each track interval + req.mamba_next_track_idx = 1 - req.mamba_next_track_idx + req.mamba_last_track_seqlen = seq_len + elif ( + not batch.spec_algorithm.is_none() + and result.accept_length_per_req_cpu is not None + ): + # for spec decode, update mamba_last_track_seqlen if this iteration crosses a track interval + actual_seq_len = req.seqlen - 1 + if ( + actual_seq_len // mamba_track_interval + != (actual_seq_len - result.accept_length_per_req_cpu[i]) + // mamba_track_interval + ): + req.mamba_last_track_seqlen = ( + actual_seq_len // mamba_track_interval * mamba_track_interval + ) + def _process_input_token_logprobs( self, req: Req, input_token_logprobs: List ) -> None: diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index acfaa5178..b6fa4bde3 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py @@ -121,9 +121,22 @@ class SchedulerRuntimeCheckerMixin: full_num_used != self.tree_cache.full_protected_size() or mamba_num_used != self.tree_cache.mamba_protected_size() ) + free_full_pages = set( + self.token_to_kv_pool_allocator.free_pages.tolist() + + self.token_to_kv_pool_allocator.release_pages.tolist() + ) + cached_full_pages = set(self.tree_cache.all_values_flatten().tolist()) + expected_full_pages = set(range(1, self.token_to_kv_pool_allocator.size + 1)) + leaked_full_pages = expected_full_pages - free_full_pages - cached_full_pages + free_mamba_pages = set(self.req_to_token_pool.mamba_pool.free_slots.tolist()) + cached_mamba_pages = set(self.tree_cache.all_mamba_values_flatten().tolist()) + expected_mamba_pages = set(range(self.req_to_token_pool.mamba_pool.size)) + leaked_mamba_pages = ( + expected_mamba_pages - free_mamba_pages - cached_mamba_pages + ) token_msg = ( f"{full_available_size=}, {full_evictable_size=}, {self.token_to_kv_pool_allocator.size=}, {self.tree_cache.full_protected_size()=}\n" - f"{mamba_available_size=}, {mamba_evictable_size=}, {self.req_to_token_pool.mamba_pool.size=}, {self.tree_cache.mamba_protected_size()=}\n" + f"{mamba_available_size=}, {mamba_evictable_size=}, {self.req_to_token_pool.mamba_pool.size=}, {self.tree_cache.mamba_protected_size()=}, leaked_full_pages={leaked_full_pages if len(leaked_full_pages) > 0 else None}, leaked_mamba_pages={leaked_mamba_pages if len(leaked_mamba_pages) > 0 else None}\n" ) return memory_leak, token_msg @@ -207,7 +220,7 @@ class SchedulerRuntimeCheckerMixin: def check_memory(self: Scheduler): if self.is_hybrid_swa: memory_leak, token_msg = self._check_hybrid_memory() - elif self.is_ssm_model and isinstance(self.tree_cache, MambaRadixCache): + elif self.is_hybrid_ssm and isinstance(self.tree_cache, MambaRadixCache): memory_leak, token_msg = self._check_mamba_memory() else: memory_leak, token_msg = self._check_radix_cache_memory() @@ -242,7 +255,7 @@ class SchedulerRuntimeCheckerMixin: ) = self._get_swa_token_info() num_used = max(full_num_used, swa_num_used) token_usage = max(full_token_usage, swa_token_usage) - elif self.is_ssm_model: + elif self.is_hybrid_ssm: ( num_used, _, @@ -281,7 +294,7 @@ class SchedulerRuntimeCheckerMixin: def check_tree_cache(self: Scheduler): if (self.is_hybrid_swa and isinstance(self.tree_cache, SWARadixCache)) or ( - self.is_ssm_model and isinstance(self.tree_cache, MambaRadixCache) + self.is_hybrid_ssm and isinstance(self.tree_cache, MambaRadixCache) ): self.tree_cache.sanity_check() @@ -344,7 +357,7 @@ class SchedulerWatchdog: # Print batch size and memory pool info to check whether there are de-sync issues. if self.scheduler.is_hybrid_swa: _, info_msg = self.scheduler._check_hybrid_memory() - elif self.scheduler.is_ssm_model and isinstance( + elif self.scheduler.is_hybrid_ssm and isinstance( self.scheduler.tree_cache, MambaRadixCache ): _, info_msg = self.scheduler._check_mamba_memory() diff --git a/python/sglang/srt/managers/utils.py b/python/sglang/srt/managers/utils.py index 5df932295..06469a061 100644 --- a/python/sglang/srt/managers/utils.py +++ b/python/sglang/srt/managers/utils.py @@ -24,7 +24,8 @@ class GenerationBatchResult: logits_output: Optional[LogitsProcessorOutput] = None pp_hidden_states_proxy_tensors: Optional[PPProxyTensors] = None next_token_ids: Optional[torch.Tensor] = None - num_accepted_tokens: Optional[int] = None + num_accepted_tokens: int = 0 + accept_length_per_req_cpu: Optional[List[int]] = None can_run_cuda_graph: bool = False # For output processing diff --git a/python/sglang/srt/mem_cache/cache_init_params.py b/python/sglang/srt/mem_cache/cache_init_params.py index 06ca57521..c1258deb3 100644 --- a/python/sglang/srt/mem_cache/cache_init_params.py +++ b/python/sglang/srt/mem_cache/cache_init_params.py @@ -24,3 +24,5 @@ class CacheInitParams: enable_metrics: bool = False enable_kv_cache_events: bool = False + + enable_mamba_extra_buffer: bool = False diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index 48bbb3605..f8bbac4fe 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/python/sglang/srt/mem_cache/common.py @@ -19,6 +19,10 @@ from sglang.srt.utils.common import ceil_align if TYPE_CHECKING: from sglang.srt.managers.schedule_batch import Req, ScheduleBatch +# Needs 2 + 1 slots for mamba request with prefix cache. 2 for ping pong cache, 1 for running mamba state. +MAMBA_STATE_PER_REQ_PREFIX_CACHE = 3 +MAMBA_STATE_PER_REQ_NO_CACHE = 1 + logger = logging.getLogger(__name__) @@ -300,9 +304,15 @@ def alloc_req_slots( """Allocate request slots from the pool.""" if isinstance(req_to_token_pool, HybridReqToTokenPool): mamba_available_size = req_to_token_pool.mamba_pool.available_size() - if mamba_available_size < num_reqs: + factor = ( + MAMBA_STATE_PER_REQ_PREFIX_CACHE + if isinstance(tree_cache, MambaRadixCache) + else MAMBA_STATE_PER_REQ_NO_CACHE + ) + mamba_state_needed = num_reqs * factor + if mamba_available_size < mamba_state_needed: if tree_cache is not None and isinstance(tree_cache, MambaRadixCache): - mamba_num = max(0, num_reqs - mamba_available_size) + mamba_num = max(0, mamba_state_needed - mamba_available_size) tree_cache.evict_mamba(mamba_num) req_pool_indices = req_to_token_pool.alloc(num_reqs, reqs) else: diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 8ab4ac9a8..0513e4180 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/python/sglang/srt/mem_cache/mamba_radix_cache.py @@ -21,16 +21,24 @@ The radix tree data structure for managing the hybrid (full and Mamba) KV cache. import heapq from collections import defaultdict +from functools import partial from typing import TYPE_CHECKING, List, Optional, Tuple import torch from numpy import float64 -from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator +from sglang.srt.distributed import get_tensor_model_parallel_rank +from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE +from sglang.srt.mem_cache.allocator import ( + PagedTokenToKVPoolAllocator, + TokenToKVPoolAllocator, +) from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult +from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool from sglang.srt.mem_cache.radix_cache import ( RadixKey, _key_match_page_size1, + _key_match_paged, get_child_key, ) @@ -252,6 +260,30 @@ class LRUList: return False return node.id in self.cache + def pretty_print(self, tree_cache: Optional["MambaRadixCache"] = None): + """ + Pretty print the lru list + """ + msg = f"{self.mamba=} LRU list: " + x_lru = self._get_lru() + while x_lru is not None and x_lru.id in self.cache: + msg += f"[{x_lru.id}] {x_lru.last_access_time:f} -> " + x_lru = getattr(x_lru, self.prv) + print(msg) + + if not tree_cache: + return + msg = f"{self.mamba=} Nodes (sorted by last_access_time): " + if self.mamba: + nodes = tree_cache._collect_nontombstone_nodes() + else: + nodes = tree_cache._collect_all_nodes() + heapq.heapify(nodes) + while len(nodes): + x = heapq.heappop(nodes) + msg += f"[{x.id}] {x.last_access_time:f} -> " + print(msg) + # Note: this is expensive, only use for debug def sanity_check_evictable_size(self): """ @@ -292,9 +324,13 @@ class LRUList: if x == tree_cache.root_node: # root node is not in the lru list continue + assert ( + x_lru is not None and x_lru.id in self.cache + ), f"Incorrect LRU list, x_lru is None or not in cache: {x_lru=}, {x.id=}" + assert ( x == x_lru - ), f"Incorrect LRU list, {self.mamba=}, x: {x.id=} != x_lru: {x_lru.id=}" + ), f"Incorrect LRU list, {self.mamba=}, x: {x.id=} != x_lru: {x_lru.id=}, {x.last_access_time=}, {x_lru.last_access_time=}" assert ( x_lru.full_lock_ref == 0 ), f"x_lru should not be locked when idle, {x_lru.full_lock_ref=}, {x_lru.id=}" @@ -314,22 +350,33 @@ class LRUList: evictable_size == lru_list_evictable_size ), f"{self.mamba=}, total nodes: {total_nodes}, total lru: {total_lru}, evictable size: {evictable_size} != lru list evictable size: {lru_list_evictable_size}" except Exception as e: - msg = f"Mamba Radix tree sanity check failed, ping @yizhang2077: {e}" - logger.error(msg) - raise Exception(msg) + if get_tensor_model_parallel_rank() == 0: + msg = f"Mamba Radix tree sanity check failed, ping @yizhang2077: {e}" + logger.error(msg) + tree_cache.pretty_print() + tree_cache.full_lru_list.pretty_print(tree_cache) + tree_cache.mamba_lru_list.pretty_print(tree_cache) + raise Exception(msg) class MambaRadixCache(BasePrefixCache): def __init__(self, params: CacheInitParams): - assert isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator) - self.req_to_token_pool = params.req_to_token_pool + assert isinstance( + params.token_to_kv_pool_allocator, TokenToKVPoolAllocator + ) or isinstance(params.token_to_kv_pool_allocator, PagedTokenToKVPoolAllocator) + self.req_to_token_pool: HybridReqToTokenPool = params.req_to_token_pool self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator - assert ( - params.page_size == 1 - ), "Only support page_size=1 in mamba radix cache now." self.page_size = params.page_size self.disable = params.disable + self.enable_mamba_extra_buffer = params.enable_mamba_extra_buffer + + if not self.enable_mamba_extra_buffer: + assert ( + self.page_size == 1 + ), f"Page size must be 1 for MambaRadixCache v1, got {self.page_size}" + else: + logger.info(f"Mamba extra_buffer is enabled.") if self.token_to_kv_pool_allocator: self.device = self.token_to_kv_pool_allocator.device @@ -339,15 +386,19 @@ class MambaRadixCache(BasePrefixCache): if params.enable_metrics: self.init_metrics_collector() - self.key_match_fn = _key_match_page_size1 - self.get_child_key_fn = get_child_key + if self.page_size == 1: + self.key_match_fn = _key_match_page_size1 + self.get_child_key_fn = get_child_key + else: + self.key_match_fn = partial(_key_match_paged, page_size=self.page_size) + self.get_child_key_fn = partial(get_child_key, page_size=self.page_size) self.reset() ##### Public API ##### def reset(self) -> None: self.root_node = TreeNode() - self.root_node.key = [] + self.root_node.key = RadixKey([], None) self.root_node.value = [] self.root_node.full_lock_ref = 1 self.root_node.mamba_lock_ref = 1 @@ -384,7 +435,7 @@ class MambaRadixCache(BasePrefixCache): last_host_node=self.root_node, ) - value, last_node = self._match_prefix_helper(key) + value, last_node, mamba_branching_seqlen = self._match_prefix_helper(key) # copy mamba state to req local space if cow is true if cow_mamba and last_node.mamba_value is not None: @@ -415,18 +466,26 @@ class MambaRadixCache(BasePrefixCache): device_indices=value, last_device_node=last_node, last_host_node=last_node, + mamba_branching_seqlen=mamba_branching_seqlen, ) def insert(self, key: RadixKey, value=None, mamba_value=None) -> Tuple[int, bool]: if self.disable: - return 0 + return 0, False if value is None: value = torch.tensor([x for x in key.token_ids], dtype=torch.int64) return self._insert_helper(self.root_node, key, value, mamba_value) - def cache_finished_req(self, req: Req, is_insert: bool = True): + def cache_finished_req(self, req: Req, is_insert: bool = True) -> None: """Cache request when it finishes.""" + # for abort with prefix cache hit and before alloc is called + if req.req_pool_idx is None: + if req.mamba_pool_idx is not None: + self.req_to_token_pool.mamba_pool.free(req.mamba_pool_idx.unsqueeze(-1)) + req.mamba_pool_idx = None + return + kv_committed_len = req.pop_committed_kv_cache() if self.disable: @@ -442,56 +501,135 @@ class MambaRadixCache(BasePrefixCache): req.req_pool_idx, :kv_committed_len ] - page_aligned_len = len(kv_indices) - page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) - - # Radix Cache takes one ref in memory pool - # insert the token_ids and kv_indices into the radix tree - # Note: the insert function already frees the overlapped kv_indices - mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone() - if is_insert: + cache_len = ( + req.mamba_last_track_seqlen + if self.enable_mamba_extra_buffer + else len(token_ids) + ) + if cache_len is None: + cache_len = 0 + if cache_len != len(token_ids): + cache_end_idx = max(cache_len, req.cache_protected_len) + self.token_to_kv_pool_allocator.free(kv_indices[cache_end_idx:]) + token_ids = token_ids[:cache_len] + kv_indices = kv_indices[:cache_len] + + if self.page_size != 1: + page_aligned_len = len(kv_indices) // self.page_size * self.page_size + page_aligned_kv_indices = kv_indices[:page_aligned_len].to( + dtype=torch.int64, copy=True + ) + else: + page_aligned_len = len(kv_indices) + page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) + + assert ( + cache_len == page_aligned_len + ), f"It is required {cache_len=}, {page_aligned_len=}, {kv_committed_len=}, {len(req.origin_input_ids)=}, {len(req.output_ids)=} ping @yizhang2077 if you see this" + + # Radix Cache takes one ref in memory pool + # insert the token_ids and kv_indices into the radix tree + if self.enable_mamba_extra_buffer: + mamba_ping_pong_track_buffer_to_keep = ( + self.req_to_token_pool.get_mamba_ping_pong_other_idx( + req.mamba_next_track_idx + ) + ) + mamba_value = ( + req.mamba_ping_pong_track_buffer[ + mamba_ping_pong_track_buffer_to_keep + ] + .unsqueeze(-1) + .clone() + ) + else: + mamba_value = req.mamba_pool_idx.unsqueeze(-1).clone() + mamba_ping_pong_track_buffer_to_keep = None + new_prefix_len, mamba_exist = self.insert( RadixKey(token_ids[:page_aligned_len], req.extra_key), page_aligned_kv_indices, mamba_value, ) + self.token_to_kv_pool_allocator.free( - kv_indices[len(req.prefix_indices) : new_prefix_len] + kv_indices[req.cache_protected_len : new_prefix_len] ) else: - self.token_to_kv_pool_allocator.free( - kv_indices[len(req.prefix_indices) : page_aligned_len] - ) + self.token_to_kv_pool_allocator.free(kv_indices[req.cache_protected_len :]) mamba_exist = True - if req.req_pool_idx is not None: - self.req_to_token_pool.free(req.req_pool_idx, free_mamba_cache=mamba_exist) - self.dec_lock_ref(req.last_node) - else: # for abort case - self.req_to_token_pool.mamba_pool.free(mamba_value) + if mamba_exist: + mamba_ping_pong_track_buffer_to_keep = None + + free_mamba_cache = True if self.enable_mamba_extra_buffer else mamba_exist + + self.req_to_token_pool.free( + req.req_pool_idx, + free_mamba_cache=free_mamba_cache, + mamba_ping_pong_track_buffer_to_keep=mamba_ping_pong_track_buffer_to_keep, + ) + + self.dec_lock_ref(req.last_node) def cache_unfinished_req(self, req: Req, chunked=False) -> None: """Cache request when it is unfinished.""" - if self.disable: + + def _skip_cache_unfinished_req(req: Req) -> None: kv_indices = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(req.fill_ids) ] + # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - req.prefix_indices = kv_indices + req.prefix_indices = kv_indices.to(dtype=torch.int64, copy=True) return token_ids = req.fill_ids - kv_indices = self.req_to_token_pool.req_to_token[ + cache_len = ( + req.mamba_last_track_seqlen + if self.enable_mamba_extra_buffer + else len(token_ids) + ) + if self.disable or cache_len is None: + return _skip_cache_unfinished_req(req) + + kv_indices_orig = self.req_to_token_pool.req_to_token[ req.req_pool_idx, : len(token_ids) ] - page_aligned_len = len(kv_indices) - page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) + # kv_indices is the kv indices to be cached + kv_indices = kv_indices_orig[:cache_len] + if self.page_size != 1: + page_aligned_len = len(kv_indices) // self.page_size * self.page_size + page_aligned_kv_indices = kv_indices[:page_aligned_len].to( + dtype=torch.int64, copy=True + ) + else: + page_aligned_len = len(kv_indices) + page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True) + + assert page_aligned_len == len( + kv_indices + ), f"page_aligned_len != len(kv_indices), {page_aligned_len=}, {len(kv_indices)=}, {cache_len=}, {self.page_size=}, {FLA_CHUNK_SIZE=}" + page_aligned_token_ids = token_ids[:page_aligned_len] - mamba_value = self.req_to_token_pool.get_mamba_indices( - req.req_pool_idx - ).unsqueeze(-1) + if self.enable_mamba_extra_buffer: + # copy from the ping pong track buffer + mamba_ping_pong_track_buffer_to_keep = ( + self.req_to_token_pool.get_mamba_ping_pong_other_idx( + req.mamba_next_track_idx + ) + ) + mamba_value = ( + req.mamba_ping_pong_track_buffer[mamba_ping_pong_track_buffer_to_keep] + .unsqueeze(-1) + .clone() + ) + else: + mamba_value = self.req_to_token_pool.get_mamba_indices( + req.req_pool_idx + ).unsqueeze(-1) # radix tree mamba value is forked from req space mamba_value_forked = self.req_to_token_pool.mamba_pool.fork_from(mamba_value) @@ -508,7 +646,7 @@ class MambaRadixCache(BasePrefixCache): mamba_value_forked, ) self.token_to_kv_pool_allocator.free( - kv_indices[len(req.prefix_indices) : new_prefix_len] + kv_indices[req.cache_protected_len : new_prefix_len] ) # there is a mamba cache in radix cache, release it if mamba_exist: @@ -526,21 +664,28 @@ class MambaRadixCache(BasePrefixCache): if not mamba_exist: assert torch.equal(new_last_node.mamba_value, mamba_value_forked) - assert len(req.prefix_indices) <= len( + assert ( + req.cache_protected_len <= len(new_indices) + self.page_size - 1 + ), f"{req.cache_protected_len=}, {len(new_indices)=}, {len(page_aligned_token_ids)=}, {mamba_exist=}" + assert new_prefix_len <= len( new_indices - ), f"{req.prefix_indices=}, {new_indices=}" - assert new_prefix_len <= len(new_indices), f"{new_prefix_len=}, {new_indices=}" + ), f"{new_prefix_len=}, {len(new_indices)=}" self.req_to_token_pool.write( - (req.req_pool_idx, slice(len(req.prefix_indices), len(new_indices))), - new_indices[len(req.prefix_indices) :], + (req.req_pool_idx, slice(req.cache_protected_len, len(new_indices))), + new_indices[req.cache_protected_len :], ) self.dec_lock_ref(req.last_node) self.inc_lock_ref(new_last_node) # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later - req.prefix_indices = new_indices + # NOTE: this is needed for both page_size == 1 and page_size > 1 + req.prefix_indices = torch.cat( + [new_indices, kv_indices_orig[len(new_indices) :]] + ) + req.cache_protected_len = len(new_indices) + req.mamba_last_track_seqlen = None req.last_node = new_last_node def pretty_print(self) -> None: @@ -670,7 +815,7 @@ class MambaRadixCache(BasePrefixCache): It unlocks the mamba_lock_ref for current node if its mamba_value exists. """ if self.disable: - return + return None if node.mamba_value is not None: assert ( @@ -692,6 +837,8 @@ class MambaRadixCache(BasePrefixCache): node = node.parent def sanity_check(self): + if self.disable: + return self.full_lru_list.sanity_check(self) self.mamba_lru_list.sanity_check(self) @@ -734,13 +881,25 @@ class MambaRadixCache(BasePrefixCache): _dfs_helper(child) _dfs_helper(self.root_node) - return torch.cat(values) + return torch.cat(values) if len(values) > 0 else torch.tensor([]) + + def all_mamba_values_flatten(self) -> torch.Tensor: + values = [] + + def _dfs_helper(node: TreeNode): + if node.mamba_value is not None: + values.append(node.mamba_value) + for _, child in node.children.items(): + _dfs_helper(child) + + _dfs_helper(self.root_node) + return torch.cat(values) if len(values) > 0 else torch.tensor([]) ##### Internal Helper Functions ##### def _match_prefix_helper( self, key: RadixKey - ) -> Tuple[List[torch.Tensor], TreeNode]: + ) -> Tuple[List[torch.Tensor], TreeNode, Optional[int]]: """ Mamba prefix matching helper. It factors in the sliding window size such that the matched node is guaranteed to either 1. connected to root without mamba tombstone, @@ -750,7 +909,7 @@ class MambaRadixCache(BasePrefixCache): node = self.root_node child_key = self.get_child_key_fn(key) - value = [] + value: List[torch.Tensor] = [] best_value_len = 0 best_last_node = node while len(key) > 0 and child_key in node.children.keys(): @@ -793,7 +952,19 @@ class MambaRadixCache(BasePrefixCache): ) node_update = node_update.parent - return value[:best_value_len], best_last_node + # 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_branching_seqlen = ( + fla_chunk_aligned_seqlen if fla_chunk_aligned_seqlen > 0 else None + ) + else: + mamba_branching_seqlen = None + + return value[:best_value_len], best_last_node, mamba_branching_seqlen def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode: # new_node -> child diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index 271a2eb34..ce9c8e35f 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -152,6 +152,7 @@ class MambaPool: self, *, size: int, + spec_state_size: int, cache_params: BaseLinearStateParams, device: str, enable_memory_saver: bool = False, @@ -198,7 +199,7 @@ class MambaPool: intermediate_ssm_state_cache = torch.zeros( size=( num_mamba_layers, - size + 1, + spec_state_size + 1, speculative_num_draft_tokens, temporal_state_shape[0], temporal_state_shape[1], @@ -213,7 +214,7 @@ class MambaPool: torch.zeros( size=( num_mamba_layers, - size + 1, + spec_state_size + 1, speculative_num_draft_tokens, conv_shape[0], conv_shape[1], @@ -267,6 +268,10 @@ class MambaPool: select_index = self.free_slots[:need_size] self.free_slots = self.free_slots[need_size:] + # clear at alloc time + for i in range(len(self.mamba_cache.conv)): + self.mamba_cache.conv[i][:, select_index] = 0 + self.mamba_cache.temporal[:, select_index] = 0 return select_index @@ -274,9 +279,6 @@ class MambaPool: if free_index.numel() == 0: return self.free_slots = torch.cat((self.free_slots, free_index)) - for i in range(len(self.mamba_cache.conv)): - self.mamba_cache.conv[i][:, free_index] = 0 - self.mamba_cache.temporal[:, free_index] = 0 def clear(self): # Zero the entire mamba cache before resetting free_slots @@ -333,10 +335,12 @@ class HybridReqToTokenPool(ReqToTokenPool): *, size: int, mamba_size: int, + mamba_spec_state_size: int, max_context_len: int, device: str, enable_memory_saver: bool, cache_params: BaseLinearStateParams, + enable_mamba_extra_buffer: bool, speculative_num_draft_tokens: int = None, ): super().__init__( @@ -345,23 +349,32 @@ class HybridReqToTokenPool(ReqToTokenPool): device=device, enable_memory_saver=enable_memory_saver, ) + self.mamba_ping_pong_track_buffer_size = ( + 2 if speculative_num_draft_tokens is None else 1 + ) + self.enable_mamba_extra_buffer = enable_mamba_extra_buffer self.enable_memory_saver = enable_memory_saver self._init_mamba_pool( size=mamba_size, + mamba_spec_state_size=mamba_spec_state_size, cache_params=cache_params, device=device, + enable_mamba_extra_buffer=enable_mamba_extra_buffer, speculative_num_draft_tokens=speculative_num_draft_tokens, ) def _init_mamba_pool( self, size: int, + mamba_spec_state_size: int, cache_params: BaseLinearStateParams, device: str, + enable_mamba_extra_buffer: bool, speculative_num_draft_tokens: int = None, ): self.mamba_pool = MambaPool( size=size, + spec_state_size=mamba_spec_state_size, cache_params=cache_params, device=device, enable_memory_saver=self.enable_memory_saver, @@ -373,32 +386,67 @@ class HybridReqToTokenPool(ReqToTokenPool): self.req_index_to_mamba_index_mapping: torch.Tensor = torch.zeros( size, dtype=torch.int32, device=self.device ) + if enable_mamba_extra_buffer: + self.req_index_to_mamba_ping_pong_track_buffer_mapping: torch.Tensor = ( + torch.zeros( + (size, self.mamba_ping_pong_track_buffer_size), + dtype=torch.int32, + device=self.device, + ) + ) # For chunk prefill req, we do not need to allocate mamba cache, # We could use allocated mamba cache instead. - def alloc( - self, need_size: int, reqs: Optional[List[Req]] = None - ) -> Optional[List[int]]: + def alloc(self, need_size: int, reqs: Optional[List["Req"]]) -> Optional[List[int]]: + assert reqs is not None select_index = super().alloc(need_size) if select_index == None: return None mamba_index = [] + mamba_ping_pong_track_buffer_list = [] for req in reqs: mid = None if req.mamba_pool_idx is not None: # for radix cache mid = req.mamba_pool_idx else: - mid = self.mamba_pool.alloc(1)[0] + mid = self.mamba_pool.alloc(1) + assert ( + mid is not None + ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size. {mid=}, {self.mamba_pool.size=}, {self.mamba_pool.available_size()=}, {len(reqs)=}" + mid = mid[0] req.mamba_pool_idx = mid - if mid is not None: - mamba_index.append(mid) + mamba_index.append(mid) + if self.enable_mamba_extra_buffer: + if req.mamba_ping_pong_track_buffer is None: + req.mamba_ping_pong_track_buffer = self.mamba_pool.alloc( + self.mamba_ping_pong_track_buffer_size + ) + assert ( + req.mamba_ping_pong_track_buffer is not None + ), "Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." + req.mamba_next_track_idx = 0 + mamba_ping_pong_track_buffer_list.append( + req.mamba_ping_pong_track_buffer.tolist() + ) assert len(select_index) == len( mamba_index - ), f"Not enough space for mamba cache, try to increase --max-mamba-cache-size." + ), f"Not enough space for mamba cache, try to increase --mamba-full-memory-ratio or --max-mamba-cache-size." + if self.enable_mamba_extra_buffer: + assert len(select_index) == len( + mamba_ping_pong_track_buffer_list + ), f"Not enough space for mamba ping pong idx, try to increase --mamba-full-memory-ratio." self.req_index_to_mamba_index_mapping[select_index] = torch.tensor( mamba_index, dtype=torch.int32, device=self.device ) + if self.enable_mamba_extra_buffer: + self.req_index_to_mamba_ping_pong_track_buffer_mapping[select_index] = ( + torch.tensor( + mamba_ping_pong_track_buffer_list, + dtype=torch.int32, + device=self.device, + ) + ) return select_index def get_mamba_indices(self, req_indices: torch.Tensor) -> torch.Tensor: @@ -411,8 +459,19 @@ class HybridReqToTokenPool(ReqToTokenPool): def get_speculative_mamba2_params_all_layers(self) -> MambaPool.SpeculativeState: return self.mamba_pool.get_speculative_mamba2_params_all_layers() + def get_mamba_ping_pong_other_idx(self, mamba_next_track_idx: int) -> int: + if self.mamba_ping_pong_track_buffer_size == 2: + return 1 - mamba_next_track_idx + else: + return mamba_next_track_idx + # For chunk prefill, we can not free mamba cache, we need use it in the future - def free(self, free_index: Union[int, List[int]], free_mamba_cache: bool = True): + def free( + self, + free_index: Union[int, List[int]], + free_mamba_cache: bool = True, + mamba_ping_pong_track_buffer_to_keep: Optional[int] = None, + ): if isinstance(free_index, (int,)): free_index = [free_index] super().free(free_index) @@ -420,9 +479,31 @@ class HybridReqToTokenPool(ReqToTokenPool): mamba_index = self.req_index_to_mamba_index_mapping[free_index] self.mamba_pool.free(mamba_index) + if self.enable_mamba_extra_buffer: + mamba_ping_pong_track_buffer_to_free = ( + self.req_index_to_mamba_ping_pong_track_buffer_mapping[ + free_index + ].squeeze(0) + ) + if mamba_ping_pong_track_buffer_to_keep is not None: + assert mamba_ping_pong_track_buffer_to_keep in [ + 0, + 1, + ], f"mamba_ping_pong_track_buffer_to_keep must be 0 or 1, {mamba_ping_pong_track_buffer_to_keep=}" + idx_to_free = list(range(self.mamba_ping_pong_track_buffer_size)) + idx_to_free.remove(mamba_ping_pong_track_buffer_to_keep) + mamba_ping_pong_track_buffer_to_free = ( + mamba_ping_pong_track_buffer_to_free[idx_to_free] + ) + self.mamba_pool.free(mamba_ping_pong_track_buffer_to_free) + def clear(self): + logger.info("Reset HybridReqToTokenPool") super().clear() self.mamba_pool.clear() + self.req_index_to_mamba_index_mapping.zero_() + if self.enable_mamba_extra_buffer: + self.req_index_to_mamba_ping_pong_track_buffer_mapping.zero_() class KVCache(abc.ABC): @@ -1133,6 +1214,9 @@ class HybridLinearKVPool(KVCache): cache_v, ) + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor): + self.full_kv_pool.move_kv_cache(tgt_loc, src_loc) + def get_v_head_dim(self): return self.full_kv_pool.get_value_buffer(0).shape[-1] diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index bbc59c16e..2d96431a3 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/cuda_graph_runner.py @@ -321,6 +321,11 @@ class CudaGraphRunner: num_tokens_per_bs=self.num_tokens_per_bs, ) + enable_mamba_track = ( + self.model_runner.server_args.enable_mamba_extra_buffer() + and self.model_runner.spec_algorithm.is_none() + ) + if self.require_gathered_buffer: assert self.require_mlp_tp_gather or self.require_attn_tp_gather self.buffers: GraphInputBuffers = GraphInputBuffers.create( @@ -338,6 +343,7 @@ class CudaGraphRunner: encoder_len_fill_value=self.encoder_len_fill_value, num_tokens_per_bs=self.num_tokens_per_bs, cache_loc_dtype=self._cache_loc_dtype(), + enable_mamba_track=enable_mamba_track, ) self.tbo_plugin = TboCudaGraphRunnerPlugin() @@ -537,7 +543,7 @@ class CudaGraphRunner: def capture_one_batch_size( self, bs: int, forward: Callable, stream_idx: Optional[int] = None ): - buffers = self.buffers + buffers: GraphInputBuffers = self.buffers graph = self._create_device_graph() stream = self.stream num_tokens = bs * self.num_tokens_per_bs @@ -611,6 +617,18 @@ class CudaGraphRunner: else: lora_ids = None + # mamba state tracking + mamba_track_indices = ( + buffers.mamba_track_indices[:bs] + if buffers.mamba_track_indices is not None + else None + ) + mamba_track_mask = ( + buffers.mamba_track_mask[:bs] + if buffers.mamba_track_mask is not None + else None + ) + if stream_idx is None: attn_backend = self.model_runner.attn_backend else: @@ -631,6 +649,9 @@ class CudaGraphRunner: attn_backend=attn_backend, out_cache_loc=out_cache_loc, seq_lens_sum=seq_lens.sum().item(), + mamba_track_indices=mamba_track_indices, + mamba_track_mask=mamba_track_mask, + mamba_track_seqlens=None, # Prefill only encoder_lens=encoder_lens, return_logprob=False, positions=positions, diff --git a/python/sglang/srt/model_executor/forward_batch_info.py b/python/sglang/srt/model_executor/forward_batch_info.py index b8f9d3a65..dbdeea8c7 100644 --- a/python/sglang/srt/model_executor/forward_batch_info.py +++ b/python/sglang/srt/model_executor/forward_batch_info.py @@ -249,6 +249,12 @@ class ForwardBatch: # The indices of output tokens in the token_to_kv_pool_swa # TODO(shiyang, biao): integrate out_cache_loc_swa into multiple attention backends out_cache_loc_swa: Optional[torch.Tensor] = None + # The indices to track mamba state with + mamba_track_indices: Optional[torch.Tensor] = None # shape: [b], int64 + # The mask to track mamba state if needed + mamba_track_mask: Optional[torch.Tensor] = None # shape: [b], bool + # The seqlens to track mamba state if masked, prefill only. + mamba_track_seqlens: Optional[torch.Tensor] = None # shape: [b], int64 # Optional seq_lens on cpu seq_lens_cpu: Optional[torch.Tensor] = None @@ -398,6 +404,9 @@ class ForwardBatch: req_pool_indices=batch.req_pool_indices, seq_lens=batch.seq_lens, out_cache_loc=batch.out_cache_loc, + mamba_track_indices=batch.mamba_track_indices, + mamba_track_mask=batch.mamba_track_mask, + mamba_track_seqlens=batch.mamba_track_seqlens, mm_inputs=batch.multimodal_inputs, encoder_cached=batch.encoder_cached, encoder_lens=batch.encoder_lens, @@ -881,6 +890,16 @@ class ForwardBatch: if self.encoder_lens is not None: self.encoder_lens = self._pad_tensor_to_size(self.encoder_lens, bs) self.positions = self._pad_tensor_to_size(self.positions, num_tokens) + if self.mamba_track_indices is not None: + self.mamba_track_indices = self._pad_tensor_to_size( + self.mamba_track_indices, bs + ) + if self.mamba_track_mask is not None: + self.mamba_track_mask = self._pad_tensor_to_size(self.mamba_track_mask, bs) + if self.mamba_track_seqlens is not None: + self.mamba_track_seqlens = self._pad_tensor_to_size( + self.mamba_track_seqlens, bs + ) if self.mrope_positions is not None: self.mrope_positions = self._pad_tensor_to_size(self.mrope_positions, bs) diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index b070028e6..f4468a70c 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/python/sglang/srt/model_executor/input_buffers.py @@ -25,6 +25,8 @@ class GraphInputBuffers: num_token_non_padded: torch.Tensor custom_mask: torch.Tensor next_token_logits_buffer: torch.Tensor + mamba_track_indices: Optional[torch.Tensor] + mamba_track_mask: Optional[torch.Tensor] global_num_tokens_gpu: torch.Tensor global_num_tokens_for_logprob_gpu: torch.Tensor encoder_lens: Optional[torch.Tensor] @@ -48,6 +50,7 @@ class GraphInputBuffers: encoder_len_fill_value: int, num_tokens_per_bs: int, cache_loc_dtype: torch.dtype, + enable_mamba_track: bool, ) -> "GraphInputBuffers": with torch.device(device): input_ids = torch.zeros((max_num_token,), dtype=torch.int64) @@ -66,6 +69,14 @@ class GraphInputBuffers: (max_num_token, vocab_size), dtype=torch.float, ) + mamba_track_indices = ( + torch.zeros((max_bs,), dtype=torch.int64) + if enable_mamba_track + else None + ) + mamba_track_mask = ( + torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None + ) if pp_size > 1: pp_proxy_tensors = { @@ -111,6 +122,8 @@ class GraphInputBuffers: num_token_non_padded=num_token_non_padded, custom_mask=custom_mask, next_token_logits_buffer=next_token_logits_buffer, + mamba_track_indices=mamba_track_indices, + mamba_track_mask=mamba_track_mask, encoder_lens=encoder_lens, global_num_tokens_gpu=global_num_tokens_gpu, global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu, @@ -134,6 +147,10 @@ class GraphInputBuffers: if bs != raw_bs: self.seq_lens.fill_(seq_len_fill_value) self.out_cache_loc.zero_() + if self.mamba_track_indices is not None: + self.mamba_track_indices.zero_() + if self.mamba_track_mask is not None: + self.mamba_track_mask.fill_(False) # Common inputs self.input_ids[:raw_num_token].copy_(forward_batch.input_ids) @@ -142,6 +159,17 @@ class GraphInputBuffers: self.out_cache_loc[:raw_num_token].copy_(forward_batch.out_cache_loc) self.positions[:raw_num_token].copy_(forward_batch.positions) + if ( + self.mamba_track_indices is not None + and forward_batch.mamba_track_indices is not None + ): + self.mamba_track_indices[:raw_bs].copy_(forward_batch.mamba_track_indices) + if ( + self.mamba_track_mask is not None + and forward_batch.mamba_track_mask is not None + ): + self.mamba_track_mask[:raw_bs].copy_(forward_batch.mamba_track_mask) + seq_lens_cpu: Optional[torch.Tensor] = None if forward_batch.seq_lens_cpu is not None: if bs != raw_bs: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 81d2dd86b..a7e267b74 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -238,8 +238,10 @@ def add_chunked_prefix_cache_attention_backend(backend_name): # Detect stragger ranks in model loading UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480 # leave more time for post data processing -# the ratio of mamba cache pool size to max_running_requests, it will be safe when it is larger than 2 (yizhang2077) +# the ratio of mamba cache pool size to max_running_requests MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3 +MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP = 2 +MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP = 1 logger = logging.getLogger(__name__) @@ -1446,14 +1448,9 @@ class ModelRunner: server_args = self.server_args assert config is not None - speculativa_ratio = ( - 0 - if server_args.speculative_num_draft_tokens is None - else server_args.speculative_num_draft_tokens - ) if ( server_args.disable_radix_cache - or config.mamba2_cache_params.mamba_cache_per_req == 0 + or server_args.max_mamba_cache_size is not None ): # with disable radix cache, sets the max_mamba_cache_size based on the max_running_requests if server_args.max_mamba_cache_size is None: @@ -1461,7 +1458,25 @@ class ModelRunner: server_args.max_mamba_cache_size = server_args.max_running_requests else: server_args.max_mamba_cache_size = 512 + server_args.max_mamba_cache_size = server_args.max_mamba_cache_size // ( + server_args.dp_size if server_args.enable_dp_attention else 1 + ) else: + assert config.mamba2_cache_params.mamba_cache_per_req > 0 + # reserve the memory for the intermediate mamba states used for spec dec + if not self.spec_algorithm.is_none(): + assert server_args.speculative_num_draft_tokens is not None + assert server_args.max_running_requests is not None + + mamba_state_intermediate_size = ( + config.mamba2_cache_params.mamba_cache_per_req + * server_args.max_running_requests + * server_args.speculative_num_draft_tokens + ) + total_rest_memory = total_rest_memory - ( + mamba_state_intermediate_size / (1 << 30) + ) + # allocate the memory based on the ratio between mamba state memory vs. full kv cache memory # solve the equations: # 1. mamba_state_memory + full_kv_cache_memory == total_rest_memory @@ -1475,21 +1490,22 @@ class ModelRunner: server_args.max_mamba_cache_size = int( (mamba_state_memory_raw * (1 << 30)) // config.mamba2_cache_params.mamba_cache_per_req - // (1 + speculativa_ratio) ) - if self.hybrid_gdn_config is not None: - server_args.max_mamba_cache_size = server_args.max_mamba_cache_size // ( - server_args.dp_size if server_args.enable_dp_attention else 1 - ) mamba_state_memory = ( server_args.max_mamba_cache_size * config.mamba2_cache_params.mamba_cache_per_req - * (1 + speculativa_ratio) / (1 << 30) ) return total_rest_memory - mamba_state_memory + @property + def qwen3_next_config(self): + config = self.model_config.hf_config + if isinstance(config, Qwen3NextConfig): + return config + return None + @property def hybrid_gdn_config(self): config = self.model_config.hf_config @@ -1683,11 +1699,18 @@ class ModelRunner: ) if self.mambaish_config is not None: - ratio = ( - MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO - if not self.server_args.disable_radix_cache - else 1 - ) + additional_ratio = 0 + if ( + self.server_args.enable_mamba_extra_buffer() + and not self.spec_algorithm.is_none() + ): + additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_NO_OVERLAP + else: + additional_ratio = MAMBA_CACHE_V2_ADDITIONAL_RATIO_OVERLAP + if self.server_args.disable_radix_cache: + ratio = 1 + else: + ratio = MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO + additional_ratio max_num_reqs = min( max_num_reqs, self.server_args.max_mamba_cache_size // ratio ) @@ -1789,11 +1812,13 @@ class ModelRunner: self.req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=self.server_args.max_mamba_cache_size, + mamba_spec_state_size=max_num_reqs, max_context_len=self.model_config.context_len + extra_max_context_len, device=self.device, enable_memory_saver=self.server_args.enable_memory_saver, cache_params=config.mamba2_cache_params, + enable_mamba_extra_buffer=self.server_args.enable_mamba_extra_buffer(), speculative_num_draft_tokens=self.server_args.speculative_num_draft_tokens, ) else: diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index a669287b1..719f72374 100644 --- a/python/sglang/srt/server_args.py +++ b/python/sglang/srt/server_args.py @@ -30,6 +30,7 @@ import orjson from sglang.srt.connector import ConnectorType from sglang.srt.environ import ToolStrictLevel, envs from sglang.srt.function_call.function_call_parser import FunctionCallParser +from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE from sglang.srt.lora.lora_registry import LoRARef from sglang.srt.parser.reasoning_parser import ReasoningParser from sglang.srt.utils.common import ( @@ -185,6 +186,8 @@ FP8_GEMM_RUNNER_BACKEND_CHOICES = [ MAMBA_SSM_DTYPE_CHOICES = ["float32", "bfloat16"] +mamba_scheduler_strategy_CHOICES = ["auto", "no_buffer", "extra_buffer"] + # Allow external code to add more choices def add_load_format_choices(choices): @@ -466,6 +469,8 @@ class ServerArgs: max_mamba_cache_size: Optional[int] = None mamba_ssm_dtype: str = "float32" mamba_full_memory_ratio: float = 0.9 + mamba_scheduler_strategy: str = "auto" + mamba_track_interval: int = 256 # Hierarchical cache enable_hierarchical_cache: bool = False @@ -737,6 +742,10 @@ class ServerArgs: self.random_seed = random.randint(0, 1 << 30) if self.mm_process_config is None: self.mm_process_config = {} + if self.mamba_scheduler_strategy == "auto": + # TODO: when extra_buffer is more verified, we can set the default path based on + # [overlap, non-overlap] + self.mamba_scheduler_strategy = "no_buffer" # Handle ModelScope model downloads if get_bool_env_var("SGLANG_USE_MODELSCOPE"): @@ -1333,12 +1342,6 @@ class ServerArgs: f"{model_arch}" ) elif model_arch in ["Qwen3NextForCausalLM"]: - if not self.disable_radix_cache: - logger.warning( - "Disabling overlap schedule since MambaRadixCache is not compatible with " - "overlap schedule currently, try to use --disable-radix-cache if overlap schedule is necessary" - ) - self.disable_overlap_schedule = True if is_sm100_supported(): quantization_config = getattr(hf_config, "quantization_config", None) quant_method = ( @@ -1372,15 +1375,52 @@ class ServerArgs: ) self.disable_radix_cache = True self.disable_overlap_schedule = False + + # Mamba radix cache v2 + if self.enable_mamba_extra_buffer(): + assert ( + is_cuda() + ), "Mamba extra_buffer is only supported on CUDA devices with FLA backend" + assert ( + self.disaggregation_mode == "null" + ), "Mamba extra_buffer is not compatible with disaggregation mode yet." + if self.speculative_num_draft_tokens is not None: + assert ( + self.mamba_track_interval >= self.speculative_num_draft_tokens + ), f"mamba_track_interval {self.mamba_track_interval} must be greater than or equal to speculative_num_draft_tokens {self.speculative_num_draft_tokens}" + + if self.page_size is not None: + assert ( + 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}" + + if self.speculative_algorithm is not None: + logger.info( + f"Disable overlap schedule for {model_arch} model speculative decoding." + ) + self.disable_overlap_schedule = True + elif not self.disable_radix_cache: + logger.warning( + "Disabling overlap schedule since MambaRadixCache no_buffer is not compatible with " + "overlap schedule currently, try to use --mamba-scheduler-strategy extra_buffer to enable overlap schedule" + ) + self.disable_overlap_schedule = True + elif model_arch in [ "FalconH1ForCausalLM", "JetNemotronForCausalLM", "JetVLMForConditionalGeneration", ]: + assert ( + not self.enable_mamba_extra_buffer() + ), f"mamba extra_buffer is not supported for {model_arch} model" if not self.disable_radix_cache: logger.warning( - "Disabling overlap schedule since MambaRadixCache is not compatible with " - "overlap schedule currently, try to use --disable-radix-cache if overlap schedule is necessary" + "Disabling overlap schedule since mamba no_buffer is not compatible with " + "overlap schedule, try to use --disable-radix-cache if overlap schedule is necessary" ) self.disable_overlap_schedule = True if is_sm100_supported(): @@ -3535,6 +3575,19 @@ class ServerArgs: default=ServerArgs.mamba_full_memory_ratio, help="The ratio of mamba state memory to full kv cache memory.", ) + parser.add_argument( + "--mamba-scheduler-strategy", + type=str, + choices=mamba_scheduler_strategy_CHOICES, + default=ServerArgs.mamba_scheduler_strategy, + help="The strategy to use for mamba radix cache.", + ) + parser.add_argument( + "--mamba-track-interval", + type=int, + default=ServerArgs.mamba_track_interval, + help="The interval to track the mamba state during decode.", + ) # Hierarchical cache parser.add_argument( @@ -4326,6 +4379,9 @@ class ServerArgs: model_config = self.get_model_config() return model_config.attention_arch == AttentionArch.MLA + def enable_mamba_extra_buffer(self) -> bool: + return self.mamba_scheduler_strategy == "extra_buffer" + def check_server_args(self): # Check parallel size constraints assert ( diff --git a/python/sglang/srt/speculative/eagle_info.py b/python/sglang/srt/speculative/eagle_info.py index 6566e0ce9..75fc1070f 100644 --- a/python/sglang/srt/speculative/eagle_info.py +++ b/python/sglang/srt/speculative/eagle_info.py @@ -147,6 +147,16 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin): bs, ) + if get_global_server_args().enable_mamba_extra_buffer(): + batch.mamba_track_indices = torch.tensor( + [ + req.mamba_ping_pong_track_buffer[req.mamba_next_track_idx] + for req in batch.reqs + ], + dtype=torch.int64, + device=batch.device, + ) + def generate_attn_arg_prefill( self, req_pool_indices: torch.Tensor, diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 8109640e8..239e8136b 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -316,6 +316,7 @@ class EAGLEWorker(TpModelWorker): logits_output=logits_output, next_token_ids=verify_output.verified_id, num_accepted_tokens=sum(verify_output.accept_length_per_req_cpu), + accept_length_per_req_cpu=verify_output.accept_length_per_req_cpu, can_run_cuda_graph=can_run_cuda_graph, ) @@ -669,6 +670,7 @@ class EAGLEWorker(TpModelWorker): pass def verify(self, batch: ScheduleBatch, spec_info: EagleVerifyInput): + seq_lens_pre_verify = batch.seq_lens.clone() spec_info.prepare_for_verify(batch, self.page_size) spec_info.num_tokens_per_batch = self.speculative_num_steps + 1 batch.return_hidden_states = False @@ -743,44 +745,8 @@ class EAGLEWorker(TpModelWorker): self.target_worker.model_runner.hybrid_gdn_config is not None or self.target_worker.model_runner.mamba2_config is not None ): - accepted_length = ( - torch.tensor( - res.accept_length_per_req_cpu, - device=logits_output.hidden_states.device, - dtype=torch.int64, - ) - + 1 - ) - - # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask - # res.accepted_indices.shape[0] > 0 skips DP attn idle batch - if spec_info.topk > 1 and res.accepted_indices.shape[0] > 0: - # accepted_indices=[0,2,3,4,5,7,9,10,11], accepted_length=[4, 3, 2], cumulative_accepted_lengths=[4, 7, 9] - # first_token_indices_per_req=prepend(0, accepted_indices[cumulative_accepted_lengths[:-1]]) = [0, 5, 10] - # last_token_indices_per_req=accepted_indices[cumulative_accepted_lengths - 1] = [4, 9, 11] (last token ID of each req) - # max_relative_indices_per_req = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches - cumulative_accepted_lengths = torch.cumsum(accepted_length, dim=0) - req_start_positions = torch.cat( - [ - torch.zeros( - 1, - dtype=cumulative_accepted_lengths.dtype, - device=cumulative_accepted_lengths.device, - ), - cumulative_accepted_lengths[:-1], - ] - ) - first_token_indices_per_req = res.accepted_indices[req_start_positions] - last_token_indices_per_req = res.accepted_indices[ - cumulative_accepted_lengths - 1 - ] - max_relative_indices_per_req = ( - last_token_indices_per_req - first_token_indices_per_req - ) - else: - max_relative_indices_per_req = accepted_length - 1 - self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify( - max_relative_indices_per_req, self.target_worker.model_runner.model + self._mamba_verify_update( + batch, res, logits_output, spec_info, seq_lens_pre_verify ) if batch.return_logprob: @@ -794,6 +760,85 @@ class EAGLEWorker(TpModelWorker): return logits_output, res, model_worker_batch, can_run_cuda_graph + def _mamba_verify_update( + self, + batch: ScheduleBatch, + res: EagleVerifyOutput, + logits_output: LogitsProcessorOutput, + spec_info: EagleVerifyInput, + seq_lens_pre_verify: torch.Tensor, + ): + accepted_length = ( + torch.tensor( + res.accept_length_per_req_cpu, + device=logits_output.hidden_states.device, + dtype=torch.int64, + ) + + 1 + ) + cumulative_accepted_lengths = torch.cumsum(accepted_length, dim=0) + # prepend 0 to the cumulative_accepted_lengths + accepted_indices_start = torch.cat( + [ + torch.zeros( + 1, + dtype=cumulative_accepted_lengths.dtype, + device=cumulative_accepted_lengths.device, + ), + cumulative_accepted_lengths[:-1], + ] + ) + accepted_indices_offset = torch.arange( + 0, + len(batch.seq_lens) * batch.spec_info.draft_token_num, + step=batch.spec_info.draft_token_num, + dtype=accepted_indices_start.dtype, + device=accepted_indices_start.device, + ) + + # If topk > 1, we need to use retrieve_next_token and retrieve_next_sibling to handle the eagle tree custom attention mask + # res.accepted_indices.shape[0] > 0 skips DP attn idle batch + if spec_info.topk > 1 and res.accepted_indices.shape[0] > 0: + # accepted_indices=[0,2,3,4,5,7,9,10,11], accepted_length=[4, 3, 2], cumulative_accepted_lengths=[4, 7, 9] + # first_token_indices_per_req=prepend(0, accepted_indices[cumulative_accepted_lengths[:-1]]) = [0, 5, 10] + # last_token_indices_per_req=accepted_indices[cumulative_accepted_lengths - 1] = [4, 9, 11] (last token ID of each req) + # max_relative_indices_per_req = [4,4,1]; those are the per-req spec-decoding step offsets that contain the correct mamba caches + # first_token_indices_per_req = res.accepted_indices[accepted_indices_start] + accepted_steps = ( + res.accepted_indices[cumulative_accepted_lengths - 1] + - accepted_indices_offset + ) + else: + accepted_steps = accepted_length - 1 + + if batch.mamba_track_indices is not None: + # If after verify, the request's seq_lens has crossed a mamba track interval, + # we need to update the mamba state for the request at the crossing point. + mamba_track_interval = self.server_args.mamba_track_interval + to_track_mask = ( + seq_lens_pre_verify // mamba_track_interval + != batch.seq_lens // mamba_track_interval + ) + tracking_point = ( + batch.seq_lens // mamba_track_interval * mamba_track_interval + ) + to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0) + mamba_steps_to_track = torch.where( + to_track_mask, + res.accepted_indices[to_track_ith + accepted_indices_start] + - accepted_indices_offset, + -1, + ) + else: + mamba_steps_to_track = None + + self.target_worker.model_runner.attn_backend.update_mamba_state_after_mtp_verify( + accepted_steps=accepted_steps, + mamba_track_indices=batch.mamba_track_indices, + mamba_steps_to_track=mamba_steps_to_track, + model=self.target_worker.model_runner.model, + ) + def add_logprob_values( self, batch: ScheduleBatch, diff --git a/python/sglang/test/kl_test_utils.py b/python/sglang/test/kl_test_utils.py index 0161d0273..116f0ad7e 100644 --- a/python/sglang/test/kl_test_utils.py +++ b/python/sglang/test/kl_test_utils.py @@ -1,6 +1,7 @@ import inspect import json import os +import random import numpy as np import requests @@ -79,8 +80,9 @@ def get_input_ids( break text = format_longbench_v2_example(example) tokens = tokenizer.encode(text) - # Truncate to max_tokens - input_ids.append(tokens[:max_prompt_tokens]) + # Truncate to a random length between 0.5x and 1.5x of max_prompt_tokens + truncate_len = int(max_prompt_tokens * random.uniform(0.5, 1.5)) + input_ids.append(tokens[:truncate_len]) # Save to local cache with open(cache_file, "w") as f: diff --git a/test/srt/models/test_qwen3_next_models.py b/test/srt/models/test_qwen3_next_models.py index 436585caf..fd0c6c41a 100644 --- a/test/srt/models/test_qwen3_next_models.py +++ b/test/srt/models/test_qwen3_next_models.py @@ -1,6 +1,8 @@ import unittest from types import SimpleNamespace +import requests + from sglang.srt.utils import kill_process_tree from sglang.test.few_shot_gsm8k import run_eval from sglang.test.kl_test_utils import ( @@ -16,7 +18,22 @@ from sglang.test.test_utils import ( QWEN3_NEXT_MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct" -ACC_THRESHOLDS = {QWEN3_NEXT_MODEL: {"kl_div": 0.01, "gsm8k": 0.93}} +ACC_THRESHOLDS = { + QWEN3_NEXT_MODEL: {"kl_div": 0.008, "gsm8k": 0.93}, +} + + +def send_request_helper(base_url: str, text: str): + response = requests.post( + base_url + "/generate", + json={ + "text": text, + "sampling_params": { + "max_new_tokens": 1, + }, + }, + ) + return response.json() class TestQwen3Next(CustomTestCase): @@ -33,6 +50,10 @@ class TestQwen3Next(CustomTestCase): "4", "--chunked-prefill-size", "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", ], ) @@ -61,8 +82,8 @@ class TestQwen3Next(CustomTestCase): self.base_url, ACC_THRESHOLDS, self.model, - max_samples=16, - max_new_tokens=256, + max_samples=32, + max_new_tokens=512, ) def test_input_output_logprobs_match_decode_cache_hit(self): @@ -70,10 +91,37 @@ class TestQwen3Next(CustomTestCase): self.base_url, ACC_THRESHOLDS, self.model, - max_samples=16, - max_new_tokens=256, + max_samples=32, + max_new_tokens=512, ) + def test_prefix_cache_branching(self): + print("running test_prefix_cache_branching") + requests.get(self.base_url + "/flush_cache") + branching_pos = 257 + text_prefix = "hi" * branching_pos + suffix_list = ["this" * 256, "here" * 256, "that" * 256] + cache_hit_list = [False, False, True] + + # First request only prefill the entire sequence + # Second request won't have cache hit, but will cache the branching point + # Third request will have cache hit on the branching point + for i, (suffix, cache_hit) in enumerate( + zip(suffix_list, cache_hit_list, strict=True) + ): + result = send_request_helper(self.base_url, text_prefix + suffix) + cached_tokens = result["meta_info"]["cached_tokens"] + if cache_hit: + expected_cached_tokens = branching_pos // 64 * 64 + assert ( + cached_tokens == expected_cached_tokens + ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" + else: + assert ( + cached_tokens == 0 + ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" + print("test_prefix_cache_branching passed") + class TestQwen3NextMTP(CustomTestCase): @classmethod @@ -98,6 +146,10 @@ class TestQwen3NextMTP(CustomTestCase): "0.8", "--tp", "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "no_buffer", ], ) @@ -121,6 +173,24 @@ class TestQwen3NextMTP(CustomTestCase): metrics["accuracy"], ACC_THRESHOLDS[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, + self.model, + max_samples=32, + max_new_tokens=512, + ) + + def test_input_output_logprobs_match_decode_cache_hit(self): + test_input_output_logprobs_match_decode_cache_hit_helper( + self.base_url, + ACC_THRESHOLDS, + self.model, + max_samples=32, + max_new_tokens=512, + ) + class TestQwen3NextMTPTopk(CustomTestCase): @classmethod @@ -145,6 +215,12 @@ class TestQwen3NextMTPTopk(CustomTestCase): "0.8", "--tp", "4", + "--chunked-prefill-size", + "2048", + "--mamba-scheduler-strategy", + "extra_buffer", + "--mamba-track-interval", + "128", ], ) @@ -168,6 +244,51 @@ class TestQwen3NextMTPTopk(CustomTestCase): metrics["accuracy"], ACC_THRESHOLDS[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, + self.model, + max_samples=32, + max_new_tokens=512, + ) + + def test_input_output_logprobs_match_decode_cache_hit(self): + test_input_output_logprobs_match_decode_cache_hit_helper( + self.base_url, + ACC_THRESHOLDS, + self.model, + max_samples=32, + max_new_tokens=512, + ) + + def test_prefix_cache_branching(self): + print("running test_prefix_cache_branching") + requests.get(self.base_url + "/flush_cache") + branching_pos = 257 + text_prefix = "hi" * branching_pos + suffix_list = ["this" * 256, "here" * 256, "that" * 256] + cache_hit_list = [False, False, True] + + # First request only prefill the entire sequence + # Second request won't have cache hit, but will cache the branching point + # Third request will have cache hit on the branching point + for i, (suffix, cache_hit) in enumerate( + zip(suffix_list, cache_hit_list, strict=True) + ): + result = send_request_helper(self.base_url, text_prefix + suffix) + cached_tokens = result["meta_info"]["cached_tokens"] + if cache_hit: + expected_cached_tokens = branching_pos // 64 * 64 + assert ( + cached_tokens == expected_cached_tokens + ), f"{i=}, {cache_hit=}, {cached_tokens=} is not equal to {expected_cached_tokens=}, {branching_pos=}" + else: + assert ( + cached_tokens == 0 + ), f"{i=}, {cache_hit=}, {cached_tokens=} is not 0" + print("test_prefix_cache_branching passed") + class TestQwen3NextPiecewiseCudaGraph(CustomTestCase): diff --git a/test/srt/run_suite.py b/test/srt/run_suite.py index 4f79828e9..1f50ef0ff 100644 --- a/test/srt/run_suite.py +++ b/test/srt/run_suite.py @@ -146,7 +146,7 @@ suites = { TestFile("test_eagle_dp_attention.py", 200), ], "per-commit-4-gpu": [ - TestFile("models/test_qwen3_next_models.py", 472), + TestFile("models/test_qwen3_next_models.py", 590), TestFile("test_gpt_oss_4gpu.py", 300), TestFile("test_local_attn.py", 411), TestFile("test_multi_instance_release_memory_occupation.py", 64), diff --git a/test/srt/test_mamba_unittest.py b/test/srt/test_mamba_unittest.py index d72cad94e..d3cea4209 100644 --- a/test/srt/test_mamba_unittest.py +++ b/test/srt/test_mamba_unittest.py @@ -81,10 +81,12 @@ class TestMamba(unittest.TestCase): req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=mamba_cache_size, + mamba_spec_state_size=max_num_reqs, max_context_len=max_context_len, device=device, enable_memory_saver=False, cache_params=mamba2_cache_params, + enable_mamba_extra_buffer=False, speculative_num_draft_tokens=3, ) @@ -159,10 +161,12 @@ class TestMamba(unittest.TestCase): req_to_token_pool = HybridReqToTokenPool( size=max_num_reqs, mamba_size=mamba_cache_size, + mamba_spec_state_size=max_num_reqs, max_context_len=max_context_len, device=device, enable_memory_saver=False, cache_params=mamba2_cache_params, + enable_mamba_extra_buffer=False, speculative_num_draft_tokens=3, ) # setup kv pool