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 15, 01:00 UTC
[BugFix] fix 3 issues: (1) using metadata for causal-conv1d, (2) indexing overflow in v1 vLLM, and (3) init_states in v0 (#20838)
Signed-off-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com> Co-authored-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com>
Tuan, Hoang-Trong committed
Jul 15, 2025 at 16:08 UTC
f29fd8a7f8b2f20ff9f89062a6474ee668ee8e5f
2 files changed
+14
-9
vllm/model_executor/layers/mamba/mamba_mixer2.py
+11
-5
@@ -573,8 +573,8 @@ class MambaMixer2(MambaBase, CustomOp):
573
x = hidden_states_B_C_p.transpose(
574
0, 1) # this is the form that causal-conv see
575
if mamba2_metadata.cu_seqlen is None:
576
- mamba2_metadata = update_metadata(
577
- x, attn_metadata.query_start_loc, mamba2_metadata)
576
+ mamba2_metadata = update_metadata(x, query_start_loc_p,
577
+ mamba2_metadata)
578
hidden_states_B_C_p = causal_conv1d_fn(
579
x,
580
conv_weights,
@@ -583,6 +583,7 @@ class MambaMixer2(MambaBase, CustomOp):
583
conv_states=conv_state,
584
has_initial_state=has_initial_states_p,
585
cache_indices=state_indices_tensor_p,
586
+ metadata=mamba2_metadata,
587
query_start_loc=query_start_loc_p).transpose(
588
0, 1)[:num_prefill_tokens]
589
@@ -593,9 +594,14 @@ class MambaMixer2(MambaBase, CustomOp):
594
initial_states = None
595
if (has_initial_states_p is not None and prep_initial_states):
596
# making a copy of the states
596
- initial_states = torch.where(
597
- has_initial_states_p[:, None, None, None],
598
- ssm_state[state_indices_tensor_p], 0)
597
+ if envs.VLLM_USE_V1:
598
+ initial_states = torch.where(
599
+ has_initial_states_p[:, None, None, None],
600
+ ssm_state[state_indices_tensor_p], 0)
601
+ else:
602
+ initial_states = torch.where(
603
+ has_initial_states_p[:num_prefills, None, None, None],
604
+ ssm_state[state_indices_tensor_p], 0)
605
606
scan_output, varlen_state = mamba_chunk_scan_combined(
607
hidden_states_p.view(1, num_prefill_tokens,
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+3
-4
@@ -55,7 +55,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching
55
IS_CONTINUOUS_BATCHING: tl.constexpr,
56
USE_PAD_SLOT: tl.constexpr,
57
NP2_STATELEN: tl.constexpr,
58
- DECODE_SEQLEN: tl.constexpr,
58
BLOCK_M: tl.constexpr,
59
BLOCK_N: tl.constexpr,
60
):
@@ -416,7 +415,7 @@ def causal_conv1d_fn(
415
activation = "silu"
416
417
args = None
419
- out = torch.zeros_like(x)
418
+ out = torch.empty_like(x)
419
if metadata is not None:
420
cu_seqlen = metadata.cu_seqlen
421
nums_dict = metadata.nums_dict
@@ -607,7 +606,6 @@ def causal_conv1d_fn(
606
IS_CONTINUOUS_BATCHING=cache_indices is not None,
607
USE_PAD_SLOT=pad_slot_id is not None,
608
NP2_STATELEN=np2_statelen,
610
- DECODE_SEQLEN=1,
609
#launch_cooperative_grid=True
610
BLOCK_M=8,
611
BLOCK_N=256,
@@ -665,7 +663,8 @@ def _causal_conv1d_update_kernel(
663
664
if IS_CONTINUOUS_BATCHING:
665
# mask = idx_seq < batch
668
- conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq)
666
+ conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq).to(
667
+ tl.int64)
668
else:
669
conv_state_batch_coord = idx_seq
670
if USE_PAD_SLOT: # noqa