From 215e2f7990d9bb8788555a49036002e69ce14eaa Mon Sep 17 00:00:00 2001 From: Vensen Date: Sun, 10 May 2026 19:38:28 +0700 Subject: [PATCH] [Bugfix][Mamba] IMA in causal_conv1d kernel for long sequences (#41617) Signed-off-by: vensen Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- vllm/model_executor/layers/mamba/ops/causal_conv1d.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 1160105ad1..d87a763853 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -37,7 +37,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching 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, - stride_x_token: tl.constexpr, # stride to get to next token (same feature-index, same sequence-index) + stride_x_token: tl.int64, # stride to get to next token (same feature-index, same sequence-index) stride_w_dim: tl.constexpr, # stride to get to next dim-axis value stride_w_width: tl.constexpr, # stride to get to next width-axis value stride_istate_seq: tl.constexpr, @@ -45,7 +45,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching stride_istate_token: tl.constexpr, stride_cache_indices: tl.constexpr, stride_o_dim: tl.constexpr, - stride_o_token: tl.constexpr, + stride_o_token: tl.int64, stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M # others pad_slot_id: tl.constexpr, @@ -769,7 +769,7 @@ def _causal_conv1d_update_kernel( # Strides stride_x_seq: tl.constexpr, stride_x_dim: tl.constexpr, - stride_x_token: tl.constexpr, + stride_x_token: tl.int64, stride_w_dim: tl.constexpr, stride_w_width: tl.constexpr, stride_conv_state_seq: tl.constexpr, @@ -778,7 +778,7 @@ def _causal_conv1d_update_kernel( stride_state_indices: tl.constexpr, stride_o_seq: tl.constexpr, stride_o_dim: tl.constexpr, - stride_o_token: tl.constexpr, + stride_o_token: tl.int64, # others null_block_id: tl.constexpr, # Meta-parameters