Fix FlashAttention MLA prefill V unpadding (#42642)
Signed-off-by: Martin Vit <martin@voipmonitor.org> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>
Martin Vit committed
Jul 8, 2026 at 11:22 UTC
c2ecd0f888228f6583ac56392b79cfbe3626ed5e
3 files changed
+12
-4
vllm/model_executor/layers/attention/mla_attention.py
+4
@@ -2409,6 +2409,9 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2409
q, kv_c_and_k_pe_cache, attn_metadata, k_scale
2410
)
2411
2412
+ context_output = context_output[..., : self.v_head_dim]
2413
+ suffix_output = suffix_output[..., : self.v_head_dim]
2414
+
2415
output = output.view(-1, self.num_heads, self.v_head_dim)
2416
merge_attn_states(
2417
output=output,
@@ -2421,6 +2424,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
2424
elif output_scale is None:
2425
# With output_scale set, backend already wrote into `output` in place.
2426
assert isinstance(output_prefill, torch.Tensor)
2427
+ output_prefill = output_prefill[..., : self.v_head_dim]
2428
output_prefill = output_prefill.flatten(start_dim=-2)
2429
output.copy_(output_prefill)
2430
vllm/v1/attention/backends/mla/prefill/flash_attn.py
-4
@@ -192,10 +192,6 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
192
if isinstance(attn_out, tuple):
193
attn_out, lse = attn_out[0], attn_out[1]
194
195
- # Unpad output back to v_head_dim if we padded V
196
- if self.requires_v_padding:
197
- attn_out = attn_out[..., : v.shape[-1]]
198
-
195
# Remain consistent with old `flash_attn_varlen_func` where there
196
# is only one output tensor if `return_softmax_lse` is False.
197
if return_softmax_lse:
vllm/v1/attention/ops/merge_attn_states.py
+8
@@ -46,6 +46,14 @@ def merge_attn_states(
46
When provided, output must be FP8 dtype.
47
"""
48
49
+ # Both the CUDA and Triton kernels derive the suffix head stride from
50
+ # prefix_output, so suffix_output must share the same head stride.
51
+ assert prefix_output.stride(1) == suffix_output.stride(1), (
52
+ "merge_attn_states requires prefix_output and suffix_output to have "
53
+ f"matching head strides, got {prefix_output.stride(1)} and "
54
+ f"{suffix_output.stride(1)}"
55
+ )
56
+
57
# NOTE(DefTruth): Currently, custom merge_attn_states CUDA kernel
58
# does not support FP8 dtype for inputs, fallback to use Triton kernel.
59
# However, when output_scale is provided, the inputs are still BF16/FP16