iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[Bugfix] Fix MLA kv_b_proj activation dtype with Marlin FP8 (#38771)

Browse Source 
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>
This commit is contained in:
Jacob Zhang
2026-08-05 07:25:41 -07:00
committed by GitHub
co-authored by Matthew Bonanni mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent 397094da17
commit cd930c8e78
1 changed files with 24 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/model_executor/layers/attention/mla_attention.py
+24 -25
View File
@@ -344,6 +344,20 @@ def _canonicalize_sparse_mla_kv_cache_dtype(
return kv_cache_dtype
def _get_kv_b_proj_input_dtype(
kv_b_proj: ColumnParallelLinear, use_fp8_prefill: bool
) -> torch.dtype | None:
weight = getattr(kv_b_proj, "weight", None)
weight_dtype = weight.dtype if weight is not None else kv_b_proj.params_dtype
if weight_dtype == torch.int32:
return kv_b_proj.params_dtype
if weight_dtype == torch.uint8:
return None
if weight_dtype == current_platform.fp8_dtype() and not use_fp8_prefill:
return None
return weight_dtype
class MLAAttention(nn.Module, AttentionLayerBase):
"""Multi-Head Latent Attention layer.
@@ -2271,6 +2285,9 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
assert prefill_metadata.chunked_context is not None
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
self.kv_b_proj, use_fp8_prefill
)
output = None
merge_output = None
@@ -2316,21 +2333,8 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
# Extract kv_c_normed from workspace
kv_c_normed = workspace[:toks][..., : self.kv_lora_rank]
# When FP8 weights are used without FP8 prefill, kv_b_proj expects
# model dtype input and will quantize internally.
# For quantized layers (AWQ/GPTQ) that lack a .weight attribute,
# use params_dtype which is the expected input dtype.
_kv_b_proj_w_dtype = (
self.kv_b_proj.weight.dtype
if hasattr(self.kv_b_proj, "weight")
else self.kv_b_proj.params_dtype
)
# For NVFP4, weights are packed uint8 — keep input in model dtype
# since the NVFP4 linear layer quantizes internally.
if (
use_fp8_prefill or _kv_b_proj_w_dtype != current_platform.fp8_dtype()
) and _kv_b_proj_w_dtype != torch.uint8:
kv_c_normed = kv_c_normed.to(self.kv_b_proj.weight.dtype)
if kv_b_proj_input_dtype is not None:
kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
k_pe = workspace[:toks][..., self.kv_lora_rank :].unsqueeze(1)
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
@@ -2401,6 +2405,9 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
assert prefill_metadata.chunked_context.chunk_size is not None
use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()
kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(
self.kv_b_proj, use_fp8_prefill
)
output = None
merge_output = None
iters = len(prefill_metadata.chunked_context.seq_tot)
@@ -2482,16 +2489,8 @@ class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):
chunk_idx=i,
toks=toks,
)
kv_b_proj_w_dtype = (
self.kv_b_proj.weight.dtype
if hasattr(self.kv_b_proj, "weight")
else self.kv_b_proj.params_dtype
)
if (
use_fp8_prefill or kv_b_proj_w_dtype != current_platform.fp8_dtype()
) and kv_b_proj_w_dtype != torch.uint8:
kv_c_normed = kv_c_normed.to(kv_b_proj_w_dtype)
if kv_b_proj_input_dtype is not None:
kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype)
kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
-1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim