Skip to content

[Perf] Use SDPA for BLIP-2 Q-Former attention #55285

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Isotr0py merged 3 commits into from
Sep 6, 2026
+8 −8
Merged

[Perf] Use SDPA for BLIP-2 Q-Former attention #55285

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
a355307
Use SDPA for BLIP-2 Q-Former attention
Levius-Fubuki Sep 4, 2026
f80f219
save ci resource
Isotr0py Sep 6, 2026
37e2771
Merge branch 'main' into perf/blip2-qformer-sdpa
Isotr0py Sep 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions vllm/model_executor/models/blip2.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import (
BatchFeature,
Blip2Config,
Expand Down Expand Up @@ -141,14 +142,13 @@ def forward(

query_layer = self.transpose_for_scores(mixed_query_layer)

attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
attention_probs = torch.softmax(attention_scores * self.scaling, dim=-1)

# This is actually dropping out entire tokens to attend to, which might
# seem a bit unusual, but is taken from the original Transformer paper.
attention_probs_dropped = self.dropout(attention_probs)

context_layer = torch.matmul(attention_probs_dropped, value_layer)
context_layer = F.scaled_dot_product_attention(

@DarkLight1337 DarkLight1337 Sep 4, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this in the Transformers code? cc @hmellor

@hmellor hmellor Sep 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not really, but it could.

Blip2QFormerMultiHeadAttention in Transformers uses the ALL_ATTENTION_FUNCTIONS registry but the Blip2QFormerConfig does not specify _attn_implementation. So, it falls back to eager_attention_forward:

def eager_attention_forward(
    module: nn.Module,
    query: torch.Tensor,
    key: torch.Tensor,
    value: torch.Tensor,
    attention_mask: torch.Tensor | None,
    scaling: float,
    dropout: float = 0.0,
    **kwargs,
):
    attn_weights = torch.matmul(query, key.transpose(-1, -2)) * scaling
    if attention_mask is not None:
        attn_weights = attn_weights + attention_mask

    attn_weights = nn.functional.softmax(attn_weights, dim=-1)
    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)

    attn_output = torch.matmul(attn_weights, value)
    attn_output = attn_output.transpose(1, 2).contiguous()

    return attn_output, attn_weights

However if the config had been configured with "sdpa" then it would have called sdpa_attention_forward in Transformers.

@DarkLight1337 DarkLight1337 Sep 4, 2026 •
edited
Loading

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder whether we could simply use vllm-native MMEncoderAttention for this? @Isotr0py

@Isotr0py Isotr0py Sep 5, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we can use MMEncoderAttention since it's just a normal bidirectional attention here.

@Levius-Fubuki Levius-Fubuki Sep 5, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would MMEncoderAttention also support the Q-Former cross-attention case where q_len != kv_len (e.g. 32 vs 257)? I may be missing something, but the current Flash/Triton wrappers seem to reuse sequence metadata derived from q_len for K/V.

@Isotr0py Isotr0py Sep 6, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, let's use F.sdpa here then

query_layer,
key_layer,
value_layer,
dropout_p=self.dropout.p if self.training else 0.0,
scale=self.scaling,
)

context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
context_layer = context_layer.view(
Expand Down
Loading