[Spec] Mamba2 support in target models (#13434)

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(

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

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,

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,
         )

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,