[Perf][Mamba] Reduce host overhead in the Mamba2 Triton SSD prefill path

vllm/model_executor/layers/mamba/mamba_mixer2.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 
+import itertools
 from collections.abc import Sequence
 from dataclasses import replace
 
@@ -658,12 +659,18 @@ def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None:
         )
         seq_idx = torch.zeros(nchunks, device=device, dtype=torch.int32)
         out = torch.empty(seqlen, nheads, headdim, device=device, dtype=dtype)
+        final_states = torch.empty(
+            batch, nheads, headdim, dstate, device=device, dtype=ssm_state_dtype
+        )
+        final_state_indices = torch.zeros(batch, device=device, dtype=torch.int32)
 
         # Two kernels (_state_passing_fwd, _chunk_scan_fwd) use
-        # HAS_INITSTATES as a constexpr, producing separate compiled
-        # binaries. Warm up both code paths so neither triggers
-        # JIT compilation during inference.
-        for use_initial_states in (False, True):
+        # HAS_INITSTATES as a constexpr, and _state_passing_fwd also
+        # HAS_FINAL_STATES, producing separate compiled binaries. Warm up
+        # all code paths so none triggers JIT compilation during inference.
+        for use_initial_states, write_final_states in itertools.product(
+            (False, True), repeat=2
+        ):
             initial_states = (
                 torch.randn(
                     batch,
@@ -696,14 +703,19 @@ def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None:
                     dt_softplus=True,
                     dt_limit=(0.0, float("inf")),
                     state_dtype=ssm_state_dtype,
+                    final_states_out=final_states if write_final_states else None,
+                    final_state_indices=(
+                        final_state_indices if write_final_states else None
+                    ),
                 )
             except Exception:
                 logger.warning(
                     "Mamba2 SSD kernel warmup failed for layer %s "
-                    "(initial_states=%s). First inference may experience "
-                    "latency spike or OOM due to autotuner.",
+                    "(initial_states=%s, final_states=%s). First inference may "
+                    "experience latency spike or OOM due to autotuner.",
                     self.prefix,
                     use_initial_states,
+                    write_final_states,
                     exc_info=True,
                 )
 
@@ -845,7 +857,10 @@ def conv_ssm_forward(
 
             # NOTE: final output is an in-place update of out tensor
             assert preallocated_ssm_out_p is not None
+            assert state_indices_tensor_p is not None
             has_checkpoints = checkpoint_chunk_idx is not None
+            # Without checkpoints, the SSD scan writes each sequence's final
+            # state straight into its ssm_state slot.
             varlen_states = mamba_chunk_scan_combined_varlen(
                 hidden_states_p.view(
                     num_prefill_tokens, self.num_heads // self.tp_size, self.head_dim
@@ -868,18 +883,15 @@ def conv_ssm_forward(
                 dt_limit=(0.0, float("inf")),
                 out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim),
                 state_dtype=ssm_state.dtype,
+                final_states_out=None if has_checkpoints else ssm_state,
+                final_state_indices=None if has_checkpoints else state_indices_tensor_p,
             )
 
-            # update ssm states
-            # - varlen state is a (num_prefills, nheads, headdim, dstate)
-            #   tensor
-            assert state_indices_tensor_p is not None
-            final_states = (
-                varlen_states[last_chunk_indices_p]
-                if has_checkpoints
-                else varlen_states
-            )
             if has_checkpoints:
+                # update ssm states
+                # - varlen state is a (num_chunks, nheads, headdim, dstate)
+                #   tensor
+                final_states = varlen_states[last_chunk_indices_p]
                 assert checkpoint_meta is not None
                 assert query_start_loc_p is not None
                 store_prefill_checkpoint(
@@ -892,7 +904,7 @@ def conv_ssm_forward(
                     query_start_loc_p,
                     state_len=self.conv_kernel_size - 1,
                 )
-            ssm_state[state_indices_tensor_p] = final_states
+                ssm_state[state_indices_tensor_p] = final_states
             if ring_start is not None and self._updates_replayssm_trackers:
                 assert prev_num_accepted is not None
                 reset_replayssm_ring_trackers(

vllm/model_executor/layers/mamba/ops/ssd_bmm.py

@@ -8,6 +8,7 @@
 
 import torch
 
+from vllm.model_executor.layers.mamba.ops.triton_helpers import launch_autotuned
 from vllm.triton_utils import tl, triton
 
 
@@ -185,26 +186,27 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp
         * triton.cdiv(chunk_size, META["BLOCK_SIZE_N"]),
         nchunks * ngroups,
     )
-    with torch.accelerator.device_index(a.device.index):
-        _bmm_chunk_fwd_kernel[grid](
-            a_ptr=a,
-            b_ptr=b,
-            out_ptr=out,
-            cu_chunk_seqlens_ptr=cu_chunk_seqlens,
-            chunk_size=chunk_size,
-            K=k,
-            ngroups=ngroups,
-            stride_a_seqlen=a.stride(0),
-            stride_a_head=a.stride(1),
-            stride_ak=a.stride(2),
-            stride_b_seqlen=b.stride(0),
-            stride_b_head=b.stride(1),
-            stride_bk=b.stride(2),
-            stride_out_chunk=out.stride(0),
-            stride_out_head=out.stride(1),
-            stride_outm=out.stride(-2),
-            stride_outn=out.stride(-1),
-            IS_CAUSAL=causal,
-            dot_dtype=dot_dtype,
-        )
+    launch_autotuned(
+        _bmm_chunk_fwd_kernel,
+        grid,
+        a_ptr=a,
+        b_ptr=b,
+        out_ptr=out,
+        cu_chunk_seqlens_ptr=cu_chunk_seqlens,
+        chunk_size=chunk_size,
+        K=k,
+        ngroups=ngroups,
+        stride_a_seqlen=a.stride(0),
+        stride_a_head=a.stride(1),
+        stride_ak=a.stride(2),
+        stride_b_seqlen=b.stride(0),
+        stride_b_head=b.stride(1),
+        stride_bk=b.stride(2),
+        stride_out_chunk=out.stride(0),
+        stride_out_head=out.stride(1),
+        stride_outm=out.stride(-2),
+        stride_outn=out.stride(-1),
+        IS_CAUSAL=causal,
+        dot_dtype=dot_dtype,
+    )
     return out

vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py

@@ -8,7 +8,10 @@
 
 from packaging import version
 
-from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
+from vllm.model_executor.layers.mamba.ops.triton_helpers import (
+    fast_exp,
+    launch_autotuned,
+)
 from vllm.platforms import current_platform
 from vllm.triton_utils import tl, triton
 
@@ -488,7 +491,9 @@ def _chunk_scan_fwd(
         else (0, 0, 0, 0)
     )
 
-    _chunk_scan_fwd_kernel[grid](
+    launch_autotuned(
+        _chunk_scan_fwd_kernel,
+        grid,
         cb_ptr=cb,
         x_ptr=x,
         z_ptr=z,

vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py

@@ -8,7 +8,10 @@
 
 import torch
 
-from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp
+from vllm.model_executor.layers.mamba.ops.triton_helpers import (
+    fast_exp,
+    launch_autotuned,
+)
 from vllm.triton_utils import tl, triton
 
 from .mamba_ssm import softplus
@@ -321,32 +324,33 @@ def _chunk_cumsum_fwd(
         nheads, nchunks, chunk_size, device=dt.device, dtype=torch.float32
     )
     grid_chunk_cs = lambda META: (nchunks, triton.cdiv(nheads, META["BLOCK_SIZE_H"]))
-    with torch.accelerator.device_index(dt.device.index):
-        _chunk_cumsum_fwd_kernel[grid_chunk_cs](
-            dt_ptr=dt,
-            A_ptr=A,
-            dt_bias_ptr=dt_bias,
-            dt_out_ptr=dt_out,
-            dA_cumsum_ptr=dA_cumsum,
-            cu_chunk_seqlens_ptr=cu_chunk_seqlens,
-            nheads=nheads,
-            chunk_size=chunk_size,
-            dt_min=dt_limit[0],
-            dt_max=dt_limit[1],
-            stride_dt_seqlen=dt.stride(0),
-            stride_dt_head=dt.stride(1),
-            stride_A_head=A.stride(0),
-            stride_dt_bias_head=dt_bias.stride(0) if dt_bias is not None else 0,
-            stride_dt_out_head=dt_out.stride(0),
-            stride_dt_out_chunk=dt_out.stride(1),
-            stride_dt_out_csize=dt_out.stride(2),
-            stride_dA_cs_head=dA_cumsum.stride(0),
-            stride_dA_cs_chunk=dA_cumsum.stride(1),
-            stride_dA_cs_csize=dA_cumsum.stride(2),
-            DT_SOFTPLUS=dt_softplus,
-            HAS_DT_BIAS=dt_bias is not None,
-            BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size),
-        )
+    launch_autotuned(
+        _chunk_cumsum_fwd_kernel,
+        grid_chunk_cs,
+        dt_ptr=dt,
+        A_ptr=A,
+        dt_bias_ptr=dt_bias,
+        dt_out_ptr=dt_out,
+        dA_cumsum_ptr=dA_cumsum,
+        cu_chunk_seqlens_ptr=cu_chunk_seqlens,
+        nheads=nheads,
+        chunk_size=chunk_size,
+        dt_min=dt_limit[0],
+        dt_max=dt_limit[1],
+        stride_dt_seqlen=dt.stride(0),
+        stride_dt_head=dt.stride(1),
+        stride_A_head=A.stride(0),
+        stride_dt_bias_head=dt_bias.stride(0) if dt_bias is not None else 0,
+        stride_dt_out_head=dt_out.stride(0),
+        stride_dt_out_chunk=dt_out.stride(1),
+        stride_dt_out_csize=dt_out.stride(2),
+        stride_dA_cs_head=dA_cumsum.stride(0),
+        stride_dA_cs_chunk=dA_cumsum.stride(1),
+        stride_dA_cs_csize=dA_cumsum.stride(2),
+        DT_SOFTPLUS=dt_softplus,
+        HAS_DT_BIAS=dt_bias is not None,
+        BLOCK_SIZE_CHUNK=triton.next_power_of_2(chunk_size),
+    )
     return dA_cumsum, dt_out
 
 
@@ -375,33 +379,34 @@ def _chunk_state_fwd(
         nchunks,
         nheads,
     )
-    with torch.accelerator.device_index(x.device.index):
-        _chunk_state_fwd_kernel[grid](
-            x_ptr=x,
-            b_ptr=B,
-            states_ptr=states,
-            dt_ptr=dt,
-            dA_cumsum_ptr=dA_cumsum,
-            cu_chunk_seqlens_ptr=cu_chunk_seqlens,
-            hdim=headdim,
-            dstate=dstate,
-            chunk_size=chunk_size,
-            nheads_ngroups_ratio=nheads // ngroups,
-            stride_x_seqlen=x.stride(0),
-            stride_x_head=x.stride(1),
-            stride_x_hdim=x.stride(2),
-            stride_b_seqlen=B.stride(0),
-            stride_b_head=B.stride(1),
-            stride_b_dstate=B.stride(2),
-            stride_states_chunk=states.stride(0),
-            stride_states_head=states.stride(1),
-            stride_states_hdim=states.stride(2),
-            stride_states_dstate=states.stride(3),
-            stride_dt_head=dt.stride(0),
-            stride_dt_chunk=dt.stride(1),
-            stride_dt_csize=dt.stride(2),
-            stride_dA_cs_head=dA_cumsum.stride(0),
-            stride_dA_cs_chunk=dA_cumsum.stride(1),
-            stride_dA_cs_csize=dA_cumsum.stride(2),
-        )
+    launch_autotuned(
+        _chunk_state_fwd_kernel,
+        grid,
+        x_ptr=x,
+        b_ptr=B,
+        states_ptr=states,
+        dt_ptr=dt,
+        dA_cumsum_ptr=dA_cumsum,
+        cu_chunk_seqlens_ptr=cu_chunk_seqlens,
+        hdim=headdim,
+        dstate=dstate,
+        chunk_size=chunk_size,
+        nheads_ngroups_ratio=nheads // ngroups,
+        stride_x_seqlen=x.stride(0),
+        stride_x_head=x.stride(1),
+        stride_x_hdim=x.stride(2),
+        stride_b_seqlen=B.stride(0),
+        stride_b_head=B.stride(1),
+        stride_b_dstate=B.stride(2),
+        stride_states_chunk=states.stride(0),
+        stride_states_head=states.stride(1),
+        stride_states_hdim=states.stride(2),
+        stride_states_dstate=states.stride(3),
+        stride_dt_head=dt.stride(0),
+        stride_dt_chunk=dt.stride(1),
+        stride_dt_csize=dt.stride(2),
+        stride_dA_cs_head=dA_cumsum.stride(0),
+        stride_dA_cs_chunk=dA_cumsum.stride(1),
+        stride_dA_cs_csize=dA_cumsum.stride(2),
+    )
     return states