[Perf] Use SDPA for BLIP-2 Q-Former attention (#55285)
Signed-off-by: levius <2114377220@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
Sun committed
Sep 6, 2026 at 15:04 UTC
a1541f5742a29864a80087af313ad460066a1524
1 file changed
+8
-8
vllm/model_executor/models/blip2.py
+8
-8
@@ -6,6 +6,7 @@ from typing import Annotated, Literal, TypeAlias
6
7
import torch
8
import torch.nn as nn
9
+import torch.nn.functional as F
10
from transformers import (
11
BatchFeature,
12
Blip2Config,
@@ -141,14 +142,13 @@ class Blip2QFormerMultiHeadAttention(nn.Module):
142
143
query_layer = self.transpose_for_scores(mixed_query_layer)
144
144
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
145
- attention_probs = torch.softmax(attention_scores * self.scaling, dim=-1)
146
-
147
- # This is actually dropping out entire tokens to attend to, which might
148
- # seem a bit unusual, but is taken from the original Transformer paper.
149
- attention_probs_dropped = self.dropout(attention_probs)
150
-
151
- context_layer = torch.matmul(attention_probs_dropped, value_layer)
145
+ context_layer = F.scaled_dot_product_attention(
146
+ query_layer,
147
+ key_layer,
148
+ value_layer,
149
+ dropout_p=self.dropout.p if self.training else 0.0,
150
+ scale=self.scaling,
151
+ )
152
153
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
154
context_layer = context_layer.view(