From c08ebebf30cdc50ddf43ed8db344b4c69028d296 Mon Sep 17 00:00:00 2001 From: Majid Date: Thu, 28 May 2026 08:40:02 -0700 Subject: [PATCH] [Perf] Add do_not_specialize to Mamba SSD chunk kernels (#43803) Signed-off-by: Majid Taheri Andani Co-authored-by: Majid Taheri Andani --- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 2 -- vllm/model_executor/layers/mamba/ops/ssd_bmm.py | 2 -- vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py | 2 -- vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py | 4 ---- 4 files changed, 10 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index d87a7638533..f7c237ca2db 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -33,7 +33,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, - seqlen: tl.int32, # cu_seqlen num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines # Strides stride_x_dim: tl.constexpr, # stride to get to next feature-value, @@ -715,7 +714,6 @@ def causal_conv1d_fn( out, # Matrix dimensions dim, - cu_seqlen, num_cache_lines, # stride stride_x_dim, diff --git a/vllm/model_executor/layers/mamba/ops/ssd_bmm.py b/vllm/model_executor/layers/mamba/ops/ssd_bmm.py index 9b5901c383e..7cf5acf2fab 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_bmm.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_bmm.py @@ -69,7 +69,6 @@ def _bmm_chunk_fwd_kernel( out_ptr, cu_chunk_seqlens_ptr, # Matrix dimensions - seqlen, chunk_size: tl.constexpr, K: tl.constexpr, ngroups: tl.constexpr, @@ -191,7 +190,6 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp b_ptr=b, out_ptr=out, cu_chunk_seqlens_ptr=cu_chunk_seqlens, - seqlen=seqlen, chunk_size=chunk_size, K=k, ngroups=ngroups, diff --git a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py index e5e73625f86..bd554ede5ac 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py @@ -163,7 +163,6 @@ def _chunk_scan_fwd_kernel( chunk_size: tl.constexpr, hdim: tl.constexpr, dstate: tl.constexpr, - seqlen, nheads_ngroups_ratio: tl.constexpr, # Strides stride_cb_chunk: tl.int64, @@ -482,7 +481,6 @@ def _chunk_scan_fwd( chunk_size=chunk_size, hdim=headdim, dstate=dstate, - seqlen=seqlen, nheads_ngroups_ratio=nheads // ngroups, stride_cb_chunk=cb.stride(0), stride_cb_head=cb.stride(1), diff --git a/vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py b/vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py index 8402d5291e6..79b47985d09 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py @@ -35,7 +35,6 @@ def _chunk_cumsum_fwd_kernel( dA_cumsum_ptr, cu_chunk_seqlens_ptr, # Matrix dimension - seqlen, nheads: tl.constexpr, chunk_size: tl.constexpr, dt_min: tl.constexpr, @@ -206,7 +205,6 @@ def _chunk_state_fwd_kernel( hdim: tl.constexpr, dstate: tl.constexpr, chunk_size: tl.constexpr, - seqlen, nheads_ngroups_ratio: tl.constexpr, # Strides stride_x_seqlen: tl.int64, @@ -331,7 +329,6 @@ def _chunk_cumsum_fwd( dt_out_ptr=dt_out, dA_cumsum_ptr=dA_cumsum, cu_chunk_seqlens_ptr=cu_chunk_seqlens, - seqlen=seqlen, nheads=nheads, chunk_size=chunk_size, dt_min=dt_limit[0], @@ -389,7 +386,6 @@ def _chunk_state_fwd( hdim=headdim, dstate=dstate, chunk_size=chunk_size, - seqlen=seqlen, nheads_ngroups_ratio=nheads // ngroups, stride_x_seqlen=x.stride(0), stride_x_head=x.stride(1),