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 14, 05:30 UTC
[bugfix] fix MHA for models like OpenGVLab/InternVL3_5-38B (#25146)
Signed-off-by: Yan Ma <yan.ma@intel.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>
Yan Ma committed
Sep 19, 2025 at 16:45 UTC
a684c0124cb8ac04984b6fd621d99e1463016eac
1 file changed
+5
-3
vllm/attention/layer.py
+5
-3
@@ -430,9 +430,11 @@ class MultiHeadAttention(nn.Module):
430
key: torch.Tensor,
431
value: torch.Tensor,
432
) -> torch.Tensor:
433
- """Input shape: batch_size x seq_len x hidden_size"""
434
- # TODO(Isotr0py): Use existing backend implementations and support FA3
435
- bsz, q_len, _ = query.size()
433
+ """Input shape:
434
+ (batch_size x seq_len x hidden_size) or
435
+ (batch_size x seq_len x num_heads x head_size)
436
+ """
437
+ bsz, q_len = query.size()[:2]
438
kv_len = key.size(1)
439
440
query = query.view(bsz, q_len, self.num_heads, self.head_size)