[Bugfix] Fix MLA kv_b_proj activation dtype with Marlin FP8 (#38771)
Signed-off-by: Jacob Zhang <jaco8123@gmail.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Jacob Zhang committed
Aug 5, 2026 at 07:25 UTC
cd930c8e786bcfb84397e0aa424d789ebec58433
1 file changed
+24
-25
vllm/model_executor/layers/attention/mla_attention.py
+24
-25
@@ -344,6 +344,20 @@ def _canonicalize_sparse_mla_kv_cache_dtype(
344
return kv_cache_dtype
345
346
347
+def _get_kv_b_proj_input_dtype(
348
+ kv_b_proj: ColumnParallelLinear, use_fp8_prefill: bool
349
+) -> torch.dtype | None:
350
+ weight = getattr(kv_b_proj, "weight", None)
351
+ weight_dtype = weight.dtype if weight is not None else kv_b_proj.params_dtype
352
+ if weight_dtype == torch.int32:
353
+ return kv_b_proj.params_dtype
354
+ if weight_dtype == torch.uint8:
355
+ return None
356
+ if weight_dtype == current_platform.fp8_dtype() and not use_fp8_prefill:
357
+ return None
358
+ return weight_dtype
359
+
360
+
361
class MLAAttention(nn.Module, AttentionLayerBase):
362
"""Multi-Head Latent Attention layer.
363
@@ -2271,6 +2285,9 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2285
assert prefill_metadata.chunked_context is not None
2286
2287
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2288
+ kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
2289
+ self.kv_b_proj, use_fp8_prefill
2290
+ )
2291
2292
output = None
2293
merge_output = None
@@ -2316,21 +2333,8 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2333
2334
# Extract kv_c_normed from workspace
2335
kv_c_normed = workspace[:toks][..., : self.kv_lora_rank]
2319
- # When FP8 weights are used without FP8 prefill, kv_b_proj expects
2320
- # model dtype input and will quantize internally.
2321
- # For quantized layers (AWQ/GPTQ) that lack a .weight attribute,
2322
- # use params_dtype which is the expected input dtype.
2323
- _kv_b_proj_w_dtype = (
2324
- self.kv_b_proj.weight.dtype
2325
- if hasattr(self.kv_b_proj, "weight")
2326
- else self.kv_b_proj.params_dtype
2327
- )
2328
- # For NVFP4, weights are packed uint8 — keep input in model dtype
2329
- # since the NVFP4 linear layer quantizes internally.
2330
- if (
2331
- use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype()
2332
- ) and _kv_b_proj_w_dtype != torch.uint8:
2333
- kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype)
2336
+ if kv_b_proj_input_dtype is not None:
2337
+ kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
2338
2339
k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)
2340
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
@@ -2401,6 +2405,9 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2405
assert prefill_metadata.chunked_context.chunk_size is not None
2406
2407
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
2408
+ kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
2409
+ self.kv_b_proj, use_fp8_prefill
2410
+ )
2411
output = None
2412
merge_output = None
2413
iters = len(prefill_metadata.chunked_context.seq_tot)
@@ -2482,16 +2489,8 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
2489
chunk_idx=i,
2490
toks=toks,
2491
)
2485
-
2486
- kv_b_proj_w_dtype = (
2487
- self.kv_b_proj.weight.dtype
2488
- if hasattr(self.kv_b_proj, "weight")
2489
- else self.kv_b_proj.params_dtype
2490
- )
2491
- if (
2492
- use_fp8_prefill or kv_b_proj_w_dtype != current_platform.fp8_dtype()
2493
- ) and kv_b_proj_w_dtype != torch.uint8:
2494
- kv_c_normed = kv_c_normed.to(kv_b_proj_w_dtype)
2492
+ if kv_b_proj_input_dtype is not None:
2493
+ kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
2494
2495
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
2496
-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim