[Bugfix] Only pad transformers backend `value` when it is narrower (#50060)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Jul 27, 2026 at 21:36 UTC
d223c900d85224c02f2162ee2c757a769e99f519
1 file changed
+6
-3
vllm/model_executor/models/transformers/__init__.py
+6
-3
@@ -64,12 +64,15 @@ def vllm_attention_forward(
64
head_dim_v = value.shape[-1]
65
query, key, value = (x.transpose(1, 2) for x in (query, key, value))
66
query, key, value = (x.reshape(hidden, -1) for x in (query, key, value))
67
- # Pad `value` up to the query/key head size when they differ (expanded MLA).
68
- if head_dim_v != head_dim_qk:
67
+ # Pad `value` up to the query/key head size when it is smaller (expanded
68
+ # MLA). A larger last dim just means `value` isn't split per head, e.g.
69
+ # packed grouped/multi-query projections, and needs no padding.
70
+ pad_value = head_dim_v < head_dim_qk
71
+ if pad_value:
72
value = F.pad(value.view(-1, head_dim_v), (0, head_dim_qk - head_dim_v))
73
value = value.reshape(hidden, -1)
74
attn_output = self_attn.forward(query, key, value)
72
- if head_dim_v != head_dim_qk:
75
+ if pad_value:
76
attn_output = attn_output.view(-1, head_dim_qk)[..., :head_dim_v]
77
attn_output = attn_output.reshape(hidden, -1)
78
return attn_output, None