[Perf] Add do_not_specialize to Mamba SSD chunk kernels (#43803)

Signed-off-by: Majid Taheri Andani <tahemaji@amazon.com>
Co-authored-by: Majid Taheri Andani <tahemaji@amazon.com>

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,

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,

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

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