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 3, 22:45 UTC
[Bugfix][Mamba] IMA in causal_conv1d kernel for long sequences (#41617)
Signed-off-by: vensen <vensenmu@gmail.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Vensen committed
May 10, 2026 at 19:38 UTC
215e2f7990d9bb8788555a49036002e69ce14eaa
1 file changed
+4
-4
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+4
-4
@@ -37,7 +37,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
37
num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
38
# Strides
39
stride_x_dim: tl.constexpr, # stride to get to next feature-value,
40
- stride_x_token: tl.constexpr, # stride to get to next token (same feature-index, same sequence-index)
40
+ stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index)
41
stride_w_dim: tl.constexpr, # stride to get to next dim-axis value
42
stride_w_width: tl.constexpr, # stride to get to next width-axis value
43
stride_istate_seq: tl.constexpr,
@@ -45,7 +45,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching
45
stride_istate_token: tl.constexpr,
46
stride_cache_indices: tl.constexpr,
47
stride_o_dim: tl.constexpr,
48
- stride_o_token: tl.constexpr,
48
+ stride_o_token: tl.int64,
49
stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M
50
# others
51
pad_slot_id: tl.constexpr,
@@ -769,7 +769,7 @@ def _causal_conv1d_update_kernel(
769
# Strides
770
stride_x_seq: tl.constexpr,
771
stride_x_dim: tl.constexpr,
772
- stride_x_token: tl.constexpr,
772
+ stride_x_token: tl.int64,
773
stride_w_dim: tl.constexpr,
774
stride_w_width: tl.constexpr,
775
stride_conv_state_seq: tl.constexpr,
@@ -778,7 +778,7 @@ def _causal_conv1d_update_kernel(
778
stride_state_indices: tl.constexpr,
779
stride_o_seq: tl.constexpr,
780
stride_o_dim: tl.constexpr,
781
- stride_o_token: tl.constexpr,
781
+ stride_o_token: tl.int64,
782
# others
783
null_block_id: tl.constexpr,
784
# Meta-parameters