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][Mamba] - Fix Conv State Kernel FP32 Support (#24883)
Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
Asaf Joseph Gardin committed
Sep 18, 2025 at 15:21 UTC
66072b36dbf1707440ff43d57273d9e9974349d7
2 files changed
+14
-5
tests/models/language/generation/test_hybrid.py
+6
-3
@@ -418,7 +418,9 @@ def test_full_cuda_graph(
418
@pytest.mark.parametrize("model", FP32_STATE_MODELS)
419
@pytest.mark.parametrize("max_tokens", [64])
420
@pytest.mark.parametrize("num_logprobs", [5])
421
-def test_fp32_state(
421
+@pytest.mark.parametrize("cache_dtype_param",
422
+ ["mamba_ssm_cache_dtype", "mamba_cache_dtype"])
423
+def test_fp32_cache_state(
424
hf_runner,
425
vllm_runner,
426
example_prompts,
@@ -426,6 +428,7 @@ def test_fp32_state(
428
model: str,
429
max_tokens: int,
430
num_logprobs: int,
431
+ cache_dtype_param: str,
432
) -> None:
433
434
try:
@@ -443,13 +446,13 @@ def test_fp32_state(
446
m.setenv("VLLM_USE_V1", "0")
447
with vllm_runner(model,
448
max_num_seqs=MAX_NUM_SEQS,
446
- mamba_ssm_cache_dtype="float32") as vllm_model:
449
+ **{cache_dtype_param: "float32"}) as vllm_model:
450
vllm_v0_outputs = vllm_model.generate_greedy_logprobs(
451
example_prompts, max_tokens, num_logprobs)
452
453
with vllm_runner(model,
454
max_num_seqs=MAX_NUM_SEQS,
452
- mamba_ssm_cache_dtype="float32") as vllm_model:
455
+ **{cache_dtype_param: "float32"}) as vllm_model:
456
vllm_v1_outputs = vllm_model.generate_greedy_logprobs(
457
example_prompts, max_tokens, num_logprobs)
458
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+8
-2
@@ -415,6 +415,9 @@ def causal_conv1d_fn(
415
activation = "silu"
416
417
args = None
418
+ # Store original dtype to cast back at the end
419
+ original_x_dtype = x.dtype
420
+ x = x.to(conv_states.dtype)
421
out = torch.empty_like(x)
422
if metadata is not None:
423
cu_seqlen = metadata.cu_seqlen
@@ -613,7 +616,7 @@ def causal_conv1d_fn(
616
BLOCK_N=256,
617
num_stages=2,
618
)
616
- return out
619
+ return out.to(original_x_dtype)
620
621
622
@triton.jit()
@@ -973,6 +976,9 @@ def causal_conv1d_update(
976
activation = "silu" if activation is True else None
977
elif activation is not None:
978
assert activation in ["silu", "swish"]
979
+
980
+ original_x_dtype = x.dtype
981
+ x = x.to(conv_state.dtype)
982
unsqueeze = query_start_loc is None and x.dim() == 2
983
if unsqueeze:
984
# make it (batch, dim, seqlen) with seqlen == 1
@@ -1081,4 +1087,4 @@ def causal_conv1d_update(
1087
)
1088
if unsqueeze:
1089
out = out.squeeze(-1)
1084
- return out
1090
+ return out.to(original_x_dtype)