A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 2, 13:30 UTC
[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>
Majid committed
May 28, 2026 at 08:40 UTC
c08ebebf30cdc50ddf43ed8db344b4c69028d296
4 files changed
-10
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
-2
@@ -33,7 +33,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching
33
o_ptr, # (dim, seqlen) - actually pointing to x_ptr
34
# Matrix dimensions
35
dim: tl.constexpr,
36
- seqlen: tl.int32, # cu_seqlen
36
num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
37
# Strides
38
stride_x_dim: tl.constexpr, # stride to get to next feature-value,
@@ -715,7 +714,6 @@ def causal_conv1d_fn(
714
out,
715
# Matrix dimensions
716
dim,
718
- cu_seqlen,
717
num_cache_lines,
718
# stride
719
stride_x_dim,
vllm/model_executor/layers/mamba/ops/ssd_bmm.py
-2
@@ -69,7 +69,6 @@ def _bmm_chunk_fwd_kernel(
69
out_ptr,
70
cu_chunk_seqlens_ptr,
71
# Matrix dimensions
72
- seqlen,
72
chunk_size: tl.constexpr,
73
K: tl.constexpr,
74
ngroups: tl.constexpr,
@@ -191,7 +190,6 @@ def _bmm_chunk_fwd(a, b, chunk_size, cu_chunk_seqlens, causal=False, output_dtyp
190
b_ptr=b,
191
out_ptr=out,
192
cu_chunk_seqlens_ptr=cu_chunk_seqlens,
194
- seqlen=seqlen,
193
chunk_size=chunk_size,
194
K=k,
195
ngroups=ngroups,
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
-2
@@ -163,7 +163,6 @@ def _chunk_scan_fwd_kernel(
163
chunk_size: tl.constexpr,
164
hdim: tl.constexpr,
165
dstate: tl.constexpr,
166
- seqlen,
166
nheads_ngroups_ratio: tl.constexpr,
167
# Strides
168
stride_cb_chunk: tl.int64,
@@ -482,7 +481,6 @@ def _chunk_scan_fwd(
481
chunk_size=chunk_size,
482
hdim=headdim,
483
dstate=dstate,
485
- seqlen=seqlen,
484
nheads_ngroups_ratio=nheads // ngroups,
485
stride_cb_chunk=cb.stride(0),
486
stride_cb_head=cb.stride(1),
vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
-4
@@ -35,7 +35,6 @@ def _chunk_cumsum_fwd_kernel(
35
dA_cumsum_ptr,
36
cu_chunk_seqlens_ptr,
37
# Matrix dimension
38
- seqlen,
38
nheads: tl.constexpr,
39
chunk_size: tl.constexpr,
40
dt_min: tl.constexpr,
@@ -206,7 +205,6 @@ def _chunk_state_fwd_kernel(
205
hdim: tl.constexpr,
206
dstate: tl.constexpr,
207
chunk_size: tl.constexpr,
209
- seqlen,
208
nheads_ngroups_ratio: tl.constexpr,
209
# Strides
210
stride_x_seqlen: tl.int64,
@@ -331,7 +329,6 @@ def _chunk_cumsum_fwd(
329
dt_out_ptr=dt_out,
330
dA_cumsum_ptr=dA_cumsum,
331
cu_chunk_seqlens_ptr=cu_chunk_seqlens,
334
- seqlen=seqlen,
332
nheads=nheads,
333
chunk_size=chunk_size,
334
dt_min=dt_limit[0],
@@ -389,7 +386,6 @@ def _chunk_state_fwd(
386
hdim=headdim,
387
dstate=dstate,
388
chunk_size=chunk_size,
392
- seqlen=seqlen,
389
nheads_ngroups_ratio=nheads // ngroups,
390
stride_x_seqlen=x.stride(0),
391
stride_x_head=x.stride(1),