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 6, 18:15 UTC
[Bugfix] Fix MLA attention crash with AWQ/GPTQ quantized models (#34695)
Signed-off-by: haosdent <haosdent@gmail.com>
haosdent committed
Mar 14, 2026 at 07:25 UTC
6d53efd2a582f32b2d6e4962d67ba692b420d970
1 file changed
+10
-5
vllm/model_executor/layers/attention/mla_attention.py
+10
-5
@@ -442,6 +442,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
442
# If kv_b_proj_weight is unquantized, quantize it to mxfp4 if supported
443
self.is_aiter_triton_fp4_bmm_enabled = (
444
rocm_aiter_ops.is_fp4bmm_enabled()
445
+ and hasattr(self.kv_b_proj, "weight")
446
and self.kv_b_proj.weight.dtype == torch.bfloat16
447
)
448
@@ -2492,11 +2493,15 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2493
kv_c_normed = workspace[:toks][..., : self.kv_lora_rank]
2494
# When FP8 weights are used without FP8 prefill, kv_b_proj expects
2495
# model dtype input and will quantize internally.
2495
- if (
2496
- use_fp8_prefill
2497
- or self.kv_b_proj.weight.dtype != current_platform.fp8_dtype()
2498
- ):
2499
- kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype)
2496
+ # For quantized layers (AWQ/GPTQ) that lack a .weight attribute,
2497
+ # use params_dtype which is the expected input dtype.
2498
+ _kv_b_proj_w_dtype = (
2499
+ self.kv_b_proj.weight.dtype
2500
+ if hasattr(self.kv_b_proj, "weight")
2501
+ else self.kv_b_proj.params_dtype
2502
+ )
2503
+ if use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype():
2504
+ kv_c_normed = kv_c_normed.to(_kv_b_proj_w_dtype)
2505
2506
k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)
2507
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(