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 6e23a4639..aa40482da 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -745,8 +745,7 @@ class Mamba2AttnBackend(MambaAttnBackendBase): def init_forward_metadata(self, forward_batch: ForwardBatch): metadata = self._forward_metadata(forward_batch) self.forward_metadata = Mamba2Metadata.prepare_mixed( - metadata.query_start_loc, - metadata.mamba_cache_indices, + metadata, self.mamba_chunk_size, forward_batch, ) @@ -762,8 +761,12 @@ class Mamba2AttnBackend(MambaAttnBackendBase): spec_info: Optional[Union[EagleDraftInput, EagleVerifyInput]], ): metadata = self._capture_metadata(bs, req_pool_indices, forward_mode, spec_info) + draft_token_num = spec_info.draft_token_num if spec_info is not None else 1 self.forward_metadata = Mamba2Metadata.prepare_decode( - metadata.query_start_loc, metadata.mamba_cache_indices, seq_lens + metadata, + seq_lens, + is_target_verify=forward_mode.is_target_verify(), + draft_token_num=draft_token_num, ) def init_forward_metadata_replay_cuda_graph( @@ -780,8 +783,12 @@ class Mamba2AttnBackend(MambaAttnBackendBase): metadata = self._replay_metadata( bs, req_pool_indices, forward_mode, spec_info, seq_lens_cpu ) + draft_token_num = spec_info.draft_token_num if spec_info is not None else 1 self.forward_metadata = Mamba2Metadata.prepare_decode( - metadata.query_start_loc, metadata.mamba_cache_indices, seq_lens + metadata, + seq_lens, + is_target_verify=forward_mode.is_target_verify(), + draft_token_num=draft_token_num, ) def forward( diff --git a/python/sglang/srt/layers/attention/mamba/mamba.py b/python/sglang/srt/layers/attention/mamba/mamba.py index fcaaf2900..6eac7f298 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba.py +++ b/python/sglang/srt/layers/attention/mamba/mamba.py @@ -12,7 +12,6 @@ from sglang.srt.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, ) -from sglang.srt.distributed.utils import divide from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata from sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated import Mixer2RMSNormGated from sglang.srt.layers.attention.mamba.ops import ( @@ -401,10 +400,15 @@ class MambaMixer2(torch.nn.Module): num_prefills = metadata.num_prefills # request count num_decodes = metadata.num_decodes # token count (=request) + num_decode_tokens = ( + num_decodes * metadata.draft_token_num + if metadata.is_target_verify + else num_decodes + ) num_prefill_tokens = metadata.num_prefill_tokens # token count has_prefill = num_prefills > 0 has_decode = num_decodes > 0 - num_actual_tokens = num_prefill_tokens + num_decodes + num_actual_tokens = num_prefill_tokens + num_decode_tokens assert num_actual_tokens == projected_states.shape[0] # NOTE: V0 put prefill before decode @@ -412,12 +416,12 @@ class MambaMixer2(torch.nn.Module): # Split along token dimension hidden_states_B_C_p, hidden_states_B_C_d = torch.split( hidden_states_B_C, - [num_prefill_tokens, num_decodes], + [num_prefill_tokens, num_decode_tokens], dim=0, ) dt_p, dt_d = torch.split( dt, - [num_prefill_tokens, num_decodes], + [num_prefill_tokens, num_decode_tokens], dim=0, ) # Split along batch dimension @@ -441,7 +445,7 @@ class MambaMixer2(torch.nn.Module): ) preallocated_ssm_out_p, preallocated_ssm_out_d = torch.split( preallocated_ssm_out, - [num_prefill_tokens, num_decodes], + [num_prefill_tokens, num_decode_tokens], dim=0, ) @@ -520,20 +524,52 @@ class MambaMixer2(torch.nn.Module): # Process decode requests if has_decode: + is_target_verify = metadata.is_target_verify + # 2. Convolution sequence transformation - ccu = ( - causal_conv1d_update - if not use_triton_causal_conv - else causal_conv1d_update_triton - ) - hidden_states_B_C_d = ccu( - hidden_states_B_C_d, - conv_state, - conv_weights, - self.conv1d.bias, - self.activation, - conv_state_indices=state_indices_tensor_d, - ) + if is_target_verify: + assert ( + use_triton_causal_conv + ), "Speculative decoding requires use_triton_causal_conv=True for intermediate state support" + assert isinstance( + layer_cache, MambaPool.SpeculativeState + ), "layer_cache must be SpeculativeState for speculative decoding" + draft_token_num = metadata.draft_token_num + + # Reshape for batch processing + hidden_states_B_C_d_reshaped = hidden_states_B_C_d.view( + num_decodes, draft_token_num, -1 + ).transpose(1, 2) + + hidden_states_B_C_d_processed = causal_conv1d_update_triton( + hidden_states_B_C_d_reshaped, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=state_indices_tensor_d[:num_decodes], + intermediate_conv_window=layer_cache.intermediate_conv_window[0], + retrieve_next_token=metadata.retrieve_next_token, + retrieve_next_sibling=metadata.retrieve_next_sibling, + retrieve_parent_token=metadata.retrieve_parent_token, + ) + hidden_states_B_C_d = hidden_states_B_C_d_processed.transpose( + 1, 2 + ).view(num_decode_tokens, -1) + else: + ccu = ( + causal_conv1d_update + if not use_triton_causal_conv + else causal_conv1d_update_triton + ) + hidden_states_B_C_d = ccu( + hidden_states_B_C_d, + conv_state, + conv_weights, + self.conv1d.bias, + self.activation, + conv_state_indices=state_indices_tensor_d, + ) hidden_states_d, B_d, C_d = split_hidden_states_B_C_fn(hidden_states_B_C_d) @@ -553,24 +589,55 @@ class MambaMixer2(torch.nn.Module): -1, self.num_heads // self.tp_size, self.head_dim ) - # - the hidden is reshaped into (bs, num_heads, head_dim) - # - layer_state.ssm_state's slots will be selected - # using state_indices_tensor_d - # NOTE: final output is an in-place update of out tensor - selective_state_update( - ssm_state, - hidden_states_d, - dt_d, - A_d, - B_d, - C_d, - D_d, - z=None, - dt_bias=dt_bias, - dt_softplus=True, - state_batch_indices=state_indices_tensor_d, - out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim), - ) + if is_target_verify: + selective_state_update( + ssm_state, + hidden_states_d.view( + num_decodes, + draft_token_num, + self.num_heads // self.tp_size, + self.head_dim, + ), + dt_d.view( + num_decodes, + draft_token_num, + self.num_heads // self.tp_size, + self.head_dim, + ), + A_d, + B_d.view(num_decodes, draft_token_num, n_groups, -1), + C_d.view(num_decodes, draft_token_num, n_groups, -1), + D_d, + z=None, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_indices_tensor_d[:num_decodes], + out=preallocated_ssm_out_d.view( + num_decodes, + draft_token_num, + self.num_heads // self.tp_size, + self.head_dim, + ), + disable_state_update=True, + intermediate_states_buffer=layer_cache.intermediate_ssm, + cache_steps=draft_token_num, + retrieve_parent_token=metadata.retrieve_parent_token, + ) + else: + selective_state_update( + ssm_state, + hidden_states_d, + dt_d, + A_d, + B_d, + C_d, + D_d, + z=None, + dt_bias=dt_bias, + dt_softplus=True, + state_batch_indices=state_indices_tensor_d, + out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim), + ) # 4. gated MLP # GatedRMSNorm internally applying SiLU to the gate diff --git a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py index 2994e091e..47705fb7a 100644 --- a/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py +++ b/python/sglang/srt/layers/attention/mamba/mamba2_metadata.py @@ -30,6 +30,8 @@ class ForwardMetadata: retrieve_next_token: Optional[torch.Tensor] = None retrieve_next_sibling: Optional[torch.Tensor] = None retrieve_parent_token: Optional[torch.Tensor] = None + is_target_verify: bool = False + draft_token_num: int = 1 @dataclass(kw_only=True) @@ -141,31 +143,45 @@ class Mamba2Metadata(ForwardMetadata): @staticmethod def prepare_decode( - query_start_loc: torch.Tensor, - mamba_cache_indices: torch.Tensor, + forward_metadata: ForwardMetadata, seq_lens: torch.Tensor, + *, + is_target_verify: bool, + draft_token_num: int, ) -> "Mamba2Metadata": """This path is run during CUDA graph capture, i.e. decode only, so `num_prefills` is 0""" return Mamba2Metadata( - query_start_loc=query_start_loc, - mamba_cache_indices=mamba_cache_indices, + query_start_loc=forward_metadata.query_start_loc, + mamba_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, num_decodes=len(seq_lens), num_prefills=0, num_prefill_tokens=0, + is_target_verify=is_target_verify, + draft_token_num=draft_token_num, ) @classmethod def prepare_mixed( cls, - query_start_loc: torch.Tensor, - mamba_cache_indices: torch.Tensor, + forward_metadata: ForwardMetadata, chunk_size: int, forward_batch: ForwardBatch, ) -> "Mamba2Metadata": """This path cannot run with CUDA graph, as it contains extend requests.""" if forward_batch.extend_num_tokens is None: + draft_token_num = ( + forward_batch.spec_info.draft_token_num + if forward_batch.spec_info is not None + else 1 + ) return cls.prepare_decode( - query_start_loc, mamba_cache_indices, forward_batch.seq_lens + forward_metadata, + forward_batch.seq_lens, + is_target_verify=forward_batch.forward_mode.is_target_verify(), + draft_token_num=draft_token_num, ) num_prefills = len(forward_batch.extend_seq_lens) num_prefill_tokens = forward_batch.extend_num_tokens @@ -176,7 +192,7 @@ class Mamba2Metadata(ForwardMetadata): has_initial_states = context_lens_tensor > 0 prep_initial_states = torch.any(has_initial_states[:num_prefills]).item() - query_start_loc = query_start_loc[: num_prefills + 1] + query_start_loc = forward_metadata.query_start_loc[: num_prefills + 1] seq_idx = torch.repeat_interleave( torch.arange( num_prefills, dtype=torch.int32, device=query_start_loc.device @@ -197,12 +213,22 @@ class Mamba2Metadata(ForwardMetadata): ) ) + draft_token_num = ( + getattr(forward_batch.spec_info, "draft_token_num", 1) + if forward_batch.spec_info is not None + else 1 + ) return Mamba2Metadata( query_start_loc=query_start_loc, - mamba_cache_indices=mamba_cache_indices, + mamba_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, num_prefills=num_prefills, num_prefill_tokens=num_prefill_tokens, num_decodes=num_decodes, + is_target_verify=forward_batch.forward_mode.is_target_verify(), + draft_token_num=draft_token_num, mixed_metadata=cls.MixedMetadata( has_initial_states=has_initial_states, prep_initial_states=prep_initial_states, diff --git a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py index 5aca69397..97039d675 100644 --- a/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py +++ b/python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py @@ -42,7 +42,21 @@ else: @triton.heuristics( {"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])} ) -@triton.jit +@triton.heuristics( + { + "CACHE_INTERMEDIATE_STATES": lambda args: args["intermediate_states_buffer"] + is not None + } +) +@triton.heuristics( + { + "HAS_EAGLE_TREE_CUSTOM_ATTN_MASK": lambda args: args[ + "retrieve_parent_token_ptr" + ] + is not None + } +) +@triton.jit(do_not_specialize=["T"]) def _selective_scan_update_kernel( # Pointers to matrices state_ptr, @@ -57,8 +71,12 @@ def _selective_scan_update_kernel( out_ptr, state_batch_indices_ptr, pad_slot_id, + intermediate_states_buffer, + cache_steps, + retrieve_parent_token_ptr, # Matrix dimensions batch, + T, nheads, dim, dstate, @@ -69,9 +87,11 @@ def _selective_scan_update_kernel( stride_state_dim, stride_state_dstate, stride_x_batch, + stride_x_T, stride_x_head, stride_x_dim, stride_dt_batch, + stride_dt_T, stride_dt_head, stride_dt_dim, stride_dt_bias_head, @@ -80,19 +100,25 @@ def _selective_scan_update_kernel( stride_A_dim, stride_A_dstate, stride_B_batch, + stride_B_T, stride_B_group, stride_B_dstate, stride_C_batch, + stride_C_T, stride_C_group, stride_C_dstate, stride_D_head, stride_D_dim, stride_z_batch, + stride_z_T, stride_z_head, stride_z_dim, stride_out_batch, + stride_out_T, stride_out_head, stride_out_dim, + stride_retrieve_parent_token_batch, + stride_retrieve_parent_token_T, # Meta-parameters DT_SOFTPLUS: tl.constexpr, TIE_HDIM: tl.constexpr, @@ -101,6 +127,9 @@ def _selective_scan_update_kernel( HAS_D: tl.constexpr, HAS_Z: tl.constexpr, HAS_STATE_BATCH_INDICES: tl.constexpr, + DISABLE_STATE_UPDATE: tl.constexpr, + CACHE_INTERMEDIATE_STATES: tl.constexpr, + HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: tl.constexpr, BLOCK_SIZE_DSTATE: tl.constexpr, ): pid_m = tl.program_id(axis=0) @@ -133,67 +162,124 @@ def _selective_scan_update_kernel( state_ptrs = state_ptr + ( offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate ) - x_ptrs = x_ptr + offs_m * stride_x_dim - dt_ptrs = dt_ptr + offs_m * stride_dt_dim + + mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) + if HAS_STATE_BATCH_INDICES: + mask &= state_batch_idx != pad_slot_id + state = tl.load(state_ptrs, mask=mask, other=0.0).to(tl.float32) + if HAS_DT_BIAS: dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim if HAS_D: D_ptr += pid_h * stride_D_head - A_ptrs = A_ptr + ( - offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate - ) - B_ptrs = B_ptr + offs_n * stride_B_dstate - C_ptrs = C_ptr + offs_n * stride_C_dstate - if HAS_D: D_ptrs = D_ptr + offs_m * stride_D_dim - if HAS_Z: - z_ptrs = z_ptr + offs_m * stride_z_dim - out_ptrs = out_ptr + offs_m * stride_out_dim - mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) - if HAS_STATE_BATCH_INDICES: - mask &= state_batch_idx != pad_slot_id - state = tl.load(state_ptrs, mask=mask, other=0.0) + A_ptrs = A_ptr + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate - x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) - if not TIE_HDIM: - dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) - if HAS_DT_BIAS: - dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) - if DT_SOFTPLUS: - dt = softplus(dt) - A = tl.load( - A_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0 - ).to(tl.float32) - dA = tl.exp(A * dt[:, None]) - else: - dt = tl.load(dt_ptr).to(tl.float32) - if HAS_DT_BIAS: - dt += tl.load(dt_bias_ptr).to(tl.float32) - if DT_SOFTPLUS: - dt = softplus(dt) - A = tl.load(A_ptr).to(tl.float32) - dA = tl.exp(A * dt) # scalar, not a matrix + cache_idx = -1 + if CACHE_INTERMEDIATE_STATES: + if HAS_STATE_BATCH_INDICES: + cache_idx = state_batch_idx + else: + cache_idx = pid_b - B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) - C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) - if HAS_D: - D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) - if HAS_Z: - z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + current_step_idx = 0 + for _ in range(T): + if HAS_EAGLE_TREE_CUSTOM_ATTN_MASK: + if current_step_idx != 0 and cache_idx >= 0: + parent_ptr = ( + retrieve_parent_token_ptr + + pid_b * stride_retrieve_parent_token_batch + + current_step_idx * stride_retrieve_parent_token_T + ) + parent_step_idx = tl.load(parent_ptr).to(tl.int32) - dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt - state = state * dA + dB * x[:, None] + if parent_step_idx >= 0 and parent_step_idx < T: + step_offset = parent_step_idx * nheads * dim * dstate + cache_ptr = ( + intermediate_states_buffer + + cache_idx * cache_steps * nheads * dim * dstate + + step_offset + + pid_h * dim * dstate + + offs_m[:, None] * dstate + + offs_n[None, :] + ) + state = tl.load(cache_ptr, mask=mask, other=0.0).to(tl.float32) - mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) - if HAS_STATE_BATCH_INDICES: - mask &= state_batch_idx != pad_slot_id - tl.store(state_ptrs, state, mask=mask) - out = tl.sum(state * C[None, :], axis=1) - if HAS_D: - out += x * D - if HAS_Z: - out *= z * tl.sigmoid(z) - tl.store(out_ptrs, out, mask=offs_m < dim) + x_ptrs = x_ptr + offs_m * stride_x_dim + dt_ptrs = dt_ptr + offs_m * stride_dt_dim + B_ptrs = B_ptr + offs_n * stride_B_dstate + C_ptrs = C_ptr + offs_n * stride_C_dstate + if HAS_Z: + z_ptrs = z_ptr + offs_m * stride_z_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + + x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if not TIE_HDIM: + dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_DT_BIAS: + dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if DT_SOFTPLUS: + dt = softplus(dt) + A = tl.load( + A_ptrs, + mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), + other=0.0, + ).to(tl.float32) + dA = tl.exp(A * dt[:, None]) + else: + dt = tl.load(dt_ptr).to(tl.float32) + if HAS_DT_BIAS: + dt += tl.load(dt_bias_ptr).to(tl.float32) + if DT_SOFTPLUS: + dt = softplus(dt) + A = tl.load(A_ptr).to(tl.float32) + dA = tl.exp(A * dt) # scalar, not a matrix + + B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) + C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32) + if HAS_D: + D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + if HAS_Z: + z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32) + + dB = B[None, :] * dt[:, None] if not TIE_HDIM else B * dt + state = state * dA + dB * x[:, None] + + if CACHE_INTERMEDIATE_STATES: + if HAS_STATE_BATCH_INDICES: + if state_batch_idx != pad_slot_id: + cache_ptr_base = ( + intermediate_states_buffer + + state_batch_idx * cache_steps * nheads * dim * dstate + + current_step_idx * nheads * dim * dstate + + pid_h * dim * dstate + ) + cache_ptrs = cache_ptr_base + ( + offs_m[:, None] * dstate + offs_n[None, :] + ) + tl.store( + cache_ptrs, state.to(cache_ptrs.dtype.element_ty), mask=mask + ) + + out = tl.sum(state * C[None, :], axis=1) + if HAS_D: + out += x * D + if HAS_Z: + out *= z * tl.sigmoid(z) + tl.store(out_ptrs, out, mask=offs_m < dim) + + current_step_idx += 1 + + x_ptr += stride_x_T + dt_ptr += stride_dt_T + B_ptr += stride_B_T + C_ptr += stride_C_T + out_ptr += stride_out_T + if HAS_Z: + z_ptr += stride_z_T + + if not DISABLE_STATE_UPDATE: + tl.store(state_ptrs, state.to(state_ptrs.dtype.element_ty), mask=mask) def selective_state_update( @@ -210,14 +296,18 @@ def selective_state_update( state_batch_indices=None, pad_slot_id=PAD_SLOT_ID, out=None, + disable_state_update=False, + intermediate_states_buffer=None, + cache_steps=None, + retrieve_parent_token=None, ): """ Argument: state: (batch, dim, dstate) or (batch, nheads, dim, dstate) - x: (batch, dim) or (batch, nheads, dim) + x: (batch, dim) or (batch, nheads, dim) for single-token or (batch, T, nheads, dim) for multi-token dt: (batch, dim) or (batch, nheads, dim) A: (dim, dstate) or (nheads, dim, dstate) - B: (batch, dstate) or (batch, ngroups, dstate) + B: (batch, dstate) or (batch, ngroups, dstate) for single-token or (batch, T, ngroups, dstate) for multi-token C: (batch, dstate) or (batch, ngroups, dstate) D: (dim,) or (nheads, dim) z: (batch, dim) or (batch, nheads, dim) @@ -230,37 +320,54 @@ def selective_state_update( indices 0 and 3 out: Preallocated ssm output tensor. Assume same shape as x. In-place updated. + disable_state_update: If True, don't write back to state (for speculative verify) + intermediate_states_buffer: Buffer to cache intermediate states + cache_steps: Total number of steps in the buffer + retrieve_parent_token: (batch, T) tensor of parent token indices for EAGLE tree attention """ if state.dim() == 3: state = state.unsqueeze(1) if x.dim() == 2: x = x.unsqueeze(1) + if x.dim() == 3: + x = x.unsqueeze(1) if dt.dim() == 2: dt = dt.unsqueeze(1) + if dt.dim() == 3: + dt = dt.unsqueeze(1) if A.dim() == 2: A = A.unsqueeze(0) if B.dim() == 2: B = B.unsqueeze(1) + if B.dim() == 3: + B = B.unsqueeze(1) if C.dim() == 2: C = C.unsqueeze(1) + if C.dim() == 3: + C = C.unsqueeze(1) if D is not None and D.dim() == 1: D = D.unsqueeze(0) - if z is not None and z.dim() == 2: - z = z.unsqueeze(1) + if z is not None: + if z.dim() == 2: + z = z.unsqueeze(1) + if z.dim() == 3: + z = z.unsqueeze(1) if dt_bias is not None and dt_bias.dim() == 1: dt_bias = dt_bias.unsqueeze(0) if out.dim() == 2: out = out.unsqueeze(1) + if out.dim() == 3: + out = out.unsqueeze(1) _, nheads, dim, dstate = state.shape - batch = x.shape[0] + batch, T, _, _ = x.shape - assert x.shape == (batch, nheads, dim) + assert x.shape == (batch, T, nheads, dim) assert dt.shape == x.shape assert A.shape == (nheads, dim, dstate) - ngroups = B.shape[1] + ngroups = B.shape[2] assert nheads % ngroups == 0, "nheads must be divisible by ngroups" - assert B.shape == (batch, ngroups, dstate) + assert B.shape == (batch, T, ngroups, dstate) assert C.shape == B.shape if D is not None: assert D.shape == (nheads, dim) @@ -273,7 +380,11 @@ def selective_state_update( assert out.shape == x.shape grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) - z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) + z_strides = ( + (z.stride(0), z.stride(1), z.stride(2), z.stride(3)) + if z is not None + else (0, 0, 0, 0) + ) # We don't want autotune since it will overwrite the state # We instead tune by hand. BLOCK_SIZE_M, num_warps = ( @@ -291,6 +402,13 @@ def selective_state_update( and dt.stride(-1) == 0 and dt_bias.stride(-1) == 0 ) + + retrieve_parent_token_strides = ( + (retrieve_parent_token.stride(0), retrieve_parent_token.stride(1)) + if retrieve_parent_token is not None + else (0, 0) + ) + with torch.cuda.device(x.device.index): _selective_scan_update_kernel[grid]( state, @@ -305,7 +423,11 @@ def selective_state_update( out, state_batch_indices, pad_slot_id, + intermediate_states_buffer, + cache_steps if cache_steps is not None else 0, + retrieve_parent_token, batch, + T, nheads, dim, dstate, @@ -317,9 +439,11 @@ def selective_state_update( x.stride(0), x.stride(1), x.stride(2), + x.stride(3), dt.stride(0), dt.stride(1), dt.stride(2), + dt.stride(3), *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else 0, A.stride(0), A.stride(1), @@ -327,18 +451,25 @@ def selective_state_update( B.stride(0), B.stride(1), B.stride(2), + B.stride(3), C.stride(0), C.stride(1), C.stride(2), + C.stride(3), *(D.stride(0), D.stride(1)) if D is not None else 0, z_strides[0], z_strides[1], z_strides[2], + z_strides[3], out.stride(0), out.stride(1), out.stride(2), + out.stride(3), + retrieve_parent_token_strides[0], + retrieve_parent_token_strides[1], dt_softplus, tie_hdim, BLOCK_SIZE_M, + DISABLE_STATE_UPDATE=disable_state_update, num_warps=num_warps, ) diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index 41879abab..0add45539 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/python/sglang/srt/speculative/eagle_worker.py @@ -733,7 +733,10 @@ class EAGLEWorker(TpModelWorker): ] logits_output.hidden_states = logits_output.hidden_states[res.accepted_indices] - if self.target_worker.model_runner.hybrid_gdn_config is not None: + if ( + 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, diff --git a/test/srt/models/test_nvidia_nemotron_nano_v2.py b/test/srt/models/test_nvidia_nemotron_nano_v2.py index d0b3ab117..e29be59d3 100644 --- a/test/srt/models/test_nvidia_nemotron_nano_v2.py +++ b/test/srt/models/test_nvidia_nemotron_nano_v2.py @@ -24,5 +24,57 @@ class TestNvidiaNemotronNanoV2NVFP4(GSM8KMixin, CustomTestCase): other_args = ["--max-mamba-cache-size", "256"] +class TestNvidiaNemotronNanoV2SpeculativeDecoding(GSM8KMixin, CustomTestCase): + accuracy = 0.87 + model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + other_args = [ + "--speculative-algorithm", + "STANDALONE", + "--speculative-num-steps", + "2", + "--speculative-eagle-topk", + "3", + "--speculative-num-draft-tokens", + "5", + "--speculative-draft-model-path", + "meta-llama/Llama-3.2-1B", + "--speculative-draft-load-format", + "dummy", + "--max-running-requests", + "8", + "--max-total-tokens", + "2048", + "--json-model-override-args", + '{"vocab_size": 131072}', + ] + + +class TestNvidiaNemotronNanoV2SpeculativeDecodingBF16Cache(GSM8KMixin, CustomTestCase): + accuracy = 0.87 + model = "nvidia/NVIDIA-Nemotron-Nano-9B-v2" + other_args = [ + "--speculative-algorithm", + "STANDALONE", + "--speculative-num-steps", + "2", + "--speculative-eagle-topk", + "3", + "--speculative-num-draft-tokens", + "5", + "--speculative-draft-model-path", + "meta-llama/Llama-3.2-1B", + "--speculative-draft-load-format", + "dummy", + "--max-running-requests", + "8", + "--max-total-tokens", + "2048", + "--json-model-override-args", + '{"vocab_size": 131072}', + "--mamba-ssm-dtype", + "bfloat16", + ] + + if __name__ == "__main__": unittest.main()