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 14, 05:30 UTC
[Bugfix] fixes the causal_conv1d_update kernel update non-speculative decoding cases (#24680)
Signed-off-by: Tao He <linzhu.ht@alibaba-inc.com> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Tao He committed
Sep 12, 2025 at 09:16 UTC
880c741bb6d755c1da25d51e10985b4fa06671bd
1 file changed
+8
-5
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+8
-5
@@ -720,15 +720,15 @@ def _causal_conv1d_update_kernel(
720
# STEP 2: assume state_len > seqlen
721
idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
722
723
- # The conv_state updates works in a sliding window manner,
724
- # at each forward pass, the tokens are shift by 1, so we
723
+ # With speculative decoding, the conv_state updates works in a sliding
724
+ # window manner, at each forward pass, the tokens are shift by 1, so we
725
# load since idx_tokens + 1.
726
conv_state_ptrs_source = (
727
conv_state_ptr + (conv_state_batch_coord * stride_conv_state_seq) +
728
conv_state_token_offset * stride_conv_state_tok +
729
(idx_feats * stride_conv_state_dim)[None, :] +
730
- ((idx_tokens + 1) * stride_conv_state_tok)[:, None]
731
- ) # [BLOCK_M, BLOCK_N]
730
+ ((idx_tokens + (1 if IS_SPEC_DECODING else seqlen)) *
731
+ stride_conv_state_tok)[:, None]) # [BLOCK_M, BLOCK_N]
732
mask = ((conv_state_batch_coord < num_cache_lines)
733
& ((idx_tokens + seqlen) < state_len)[:, None]
734
& (idx_feats < dim)[None, :])
@@ -924,7 +924,10 @@ def causal_conv1d_update(
924
)
925
stride_state_indices = conv_state_indices.stride(
926
0) if conv_state_indices is not None else 0
927
- state_len = width - 1 + (seqlen - 1) # effective state_len needed
927
+ if num_accepted_tokens is not None:
928
+ state_len = width - 1 + (seqlen - 1) # effective state_len needed
929
+ else:
930
+ state_len = width - 1
931
np2_statelen = triton.next_power_of_2(state_len)
932
933
def grid(META):