From 8f89381fc6b2d54591a7a560e20ee5211ce1ac33 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Tue, 12 May 2026 14:46:22 +0200 Subject: [PATCH] [Hybrid] Warmup Mamba2 SSD kernel (#39822) Signed-off-by: Thomas Parnell Co-authored-by: Claude Opus 4.6 --- vllm/config/model.py | 4 +- .../layers/mamba/gdn_linear_attn.py | 3 +- .../layers/mamba/mamba_mixer2.py | 107 +++++++++++++++++- 3 files changed, 110 insertions(+), 4 deletions(-) diff --git a/vllm/config/model.py b/vllm/config/model.py index 22a1f6a429..193f48a545 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -1345,7 +1345,7 @@ class ModelConfig: ) raise AssertionError(f"Unsupported block type: {block_type}") - def get_mamba_chunk_size(self) -> int | None: + def get_mamba_chunk_size(self) -> int: """ Returns the mamba chunk size if it exists """ @@ -1356,7 +1356,7 @@ class ModelConfig: chunk_size = getattr(self.hf_text_config, "chunk_size", None) # Since Mamba1 does not have a chunk notion - # we use a default chunk size of 1024. + # we use a default chunk size of 2048. if chunk_size is None: chunk_size = 2048 diff --git a/vllm/model_executor/layers/mamba/gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn_linear_attn.py index 4493ee7836..011a18ec7b 100644 --- a/vllm/model_executor/layers/mamba/gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn_linear_attn.py @@ -397,6 +397,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): ) self.chunk_gated_delta_rule = ChunkGatedDeltaRule() + self._prefill_kernels_warmed_up = False self.enable_packed_recurrent_decode = ( envs.VLLM_ENABLE_FLA_PACKED_RECURRENT_DECODE ) @@ -905,7 +906,7 @@ class GatedDeltaNetAttention(PluggableLayer, MambaBase): which has fixed kernel parameters (no autotuning), so only the prefill (chunked) path needs warming up. """ - if hasattr(self, "_prefill_kernels_warmed_up"): + if self._prefill_kernels_warmed_up: return self._prefill_kernels_warmed_up = True diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index 05cb9aadfe..674134f373 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -14,6 +14,7 @@ from vllm.distributed import ( tensor_model_parallel_all_reduce, ) from vllm.forward_context import ForwardContext, get_forward_context +from vllm.logger import init_logger from vllm.model_executor.custom_op import CustomOp, PluggableLayer from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -54,6 +55,8 @@ from vllm.v1.attention.backend import AttentionMetadata from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadata from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum +logger = init_logger(__name__) + # Added by the IBM Team, 2024 @@ -473,6 +476,8 @@ class MambaMixer2(MambaBase, PluggableLayer): intermediate_size, n_groups, self.use_rms_norm, eps=rms_norm_eps ) + self._ssd_kernels_warmed_up = False + # - get hidden_states, B and C after depthwise convolution. self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split( hidden_states_B_C, @@ -557,6 +562,104 @@ class MambaMixer2(MambaBase, PluggableLayer): return output + def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None: + """Run a minimal SSD forward pass to trigger Triton autotuning + while GPU memory is still plentiful (before SSM cache allocation). + """ + if self._ssd_kernels_warmed_up: + return + self._ssd_kernels_warmed_up = True + logger.info_once("Warming up Mamba2 SSD Triton kernels...") + + device = projected_states.device + dtype = projected_states.dtype + + nheads = self.num_heads // self.tp_size + ngroups = self.n_groups // self.tp_size + headdim = self.head_dim + dstate = self.ssm_state_size + + if self.model_config is None: + return + chunk_size = self.model_config.get_mamba_chunk_size() + + # Triton's autotuner includes tensor dtypes in its cache key, + # so state_dtype must match what real inference uses. + _, ssm_state_dtype = self.get_state_dtype() + + # SSD kernel autotune keys depend on dtype and head dimensions, + # not on sequence length or batch size, so a single shape suffices. + seqlen = chunk_size + batch = 1 + nchunks = seqlen // chunk_size # = 1 + + x = torch.randn(seqlen, nheads, headdim, device=device, dtype=dtype) + dt = torch.randn(seqlen, nheads, device=device, dtype=dtype) + B = torch.randn(seqlen, ngroups, dstate, device=device, dtype=dtype) + C = torch.randn(seqlen, ngroups, dstate, device=device, dtype=dtype) + cu_seqlens = torch.tensor([0, seqlen], device=device, dtype=torch.int32) + cu_chunk_seqlens = torch.tensor( + [i * chunk_size for i in range(nchunks + 1)], + device=device, + dtype=torch.int32, + ) + last_chunk_indices = torch.tensor( + [nchunks - 1], device=device, dtype=torch.int32 + ) + seq_idx = torch.zeros(nchunks, device=device, dtype=torch.int32) + out = torch.empty(seqlen, nheads, headdim, device=device, dtype=dtype) + + # 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): + initial_states = ( + torch.randn( + batch, + nheads, + headdim, + dstate, + device=device, + dtype=ssm_state_dtype, + ) + if use_initial_states + else None + ) + try: + mamba_chunk_scan_combined_varlen( + x=x, + dt=dt, + A=self.A, + B=B, + C=C, + chunk_size=chunk_size, + cu_seqlens=cu_seqlens, + cu_chunk_seqlens=cu_chunk_seqlens, + last_chunk_indices=last_chunk_indices, + seq_idx=seq_idx, + out=out, + D=self.D, + z=None, + dt_bias=self.dt_bias, + initial_states=initial_states, + dt_softplus=True, + dt_limit=(0.0, float("inf")), + state_dtype=ssm_state_dtype, + ) + 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.", + self.prefix, + use_initial_states, + exc_info=True, + ) + + logger.debug("Mamba2 SSD kernel warmup completed for layer %s", self.prefix) + torch.accelerator.empty_cache() + def conv_ssm_forward( self, projected_states: torch.Tensor, @@ -608,7 +711,9 @@ class MambaMixer2(MambaBase, PluggableLayer): num_decode_tokens = attn_metadata.num_decode_tokens if attn_metadata is None: - # profile run + # V1 profile run -- warm up SSD kernels so that autotuning + # completes before SSM cache allocation. + self._warmup_ssd_kernels(projected_states) hidden_states_B_C = ( hidden_states_B_C.transpose(0, 1).clone().transpose(0, 1) ).contiguous()