[Perf] Use SDPA for BLIP-2 Q-Former attention - #55285
Conversation
Signed-off-by: levius <2114377220@qq.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Team Run ID: 📒 Files selected for processing (2)
Included review availability: Your plan provides up to 10 included reviews per hour; 9 remain after this review. 📝 SummarySummary by CodeRabbit
Walkthrough
ChangesQ-Former SDPA attention
Estimated code review effort: 2 (Simple) | ~10 minutes Merge Risk: ⚪ Minimal · up to Q-Former attention now uses PyTorch scaled dot-product attention while preserving scaling, output layout, and training-only dropout behavior. Self- and cross-attention parity and dropout coverage support merge readiness with no identified current risk. 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1🛠️ Fix failing CI checks 💡
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
| attention_probs_dropped = self.dropout(attention_probs) | ||
|
|
||
| context_layer = torch.matmul(attention_probs_dropped, value_layer) | ||
| context_layer = F.scaled_dot_product_attention( |
There was a problem hiding this comment.
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_weightsHowever if the config had been configured with "sdpa" then it would have called sdpa_attention_forward in Transformers.
There was a problem hiding this comment.
I wonder whether we could simply use vllm-native MMEncoderAttention for this? @Isotr0py
There was a problem hiding this comment.
I think we can use MMEncoderAttention since it's just a normal bidirectional attention here.
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
I see, let's use F.sdpa here then
There was a problem hiding this comment.
Please remove this redundant test
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
|
/ci run |
github-actions
Bot
added
the
ready
|
✅ Triggered Buildkite CI #87415 for commit |
|
Hi @Isotr0py, the full Buildkite CI has passed. It looks like the earlier |
|
Retrying |
Signed-off-by: levius <2114377220@qq.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: Jyotirmoy Roy <jyotirmoyroy649@gmail.com>
Purpose
BLIP-2 Q-Former attention currently materializes the attention scores, applies
scaling and softmax separately, and launches a second matrix multiplication:
This PR uses standard
torch.nn.functional.scaled_dot_product_attentionforboth Q-Former self-attention and cross-attention. It preserves the existing
scale and training dropout behavior while allowing PyTorch to dispatch fused
attention kernels.
This removes the score-sized attention intermediate and the separate scale,
softmax, dropout, and value-matmul launches. It keeps one portable PyTorch path
with no custom kernel, architecture gate, device branch, or device-specific
tuning.
I searched the open PRs for BLIP-2, Q-Former, and SDPA changes and found no
active PR addressing this attention path.
Test Plan
Lint
RTX 4090 default eager encoder benchmark
The performance harness constructs the real
Blip2ForConditionalGenerationclass with BLIP-2-compatible dimensions andinvokes its real
embed_multimodalboundary. Both variants use identical,finite, deterministic weights.
Hardware/software: NVIDIA RTX 4090, CUDA 13.0, PyTorch 2.13.0+cu130.
Timing uses 25 warmups, 200 iterations, and five alternating repeats.
embed_multimodalembed_multimodalembed_multimodalembed_multimodalThe Q-Former geometric-mean speedup is 1.48x for FP16 and 1.43x for BF16.
All eight rows passed numerical, finite-value, repeatability, shape, and
input-mutation validation with no performance regression.
Q-Former peak extra allocation decreased from 1,579,008 to 1,237,504 bytes at
batch 1, and from approximately 7.3 MiB to 5.1 MiB at batch 4.
Profiler results replace 80
aten::bmmcalls and 40 separate softmax calls perten Q-Former invocations with 40
aten::scaled_dot_product_attentioncalls.All candidate calls dispatch to
aten::_flash_attention_forward.The harness, raw CSV, gate result, and profiler summary are available in the
benchmark artifact Gist.
Test Result
ruff-check,ruff-format, andgit diff --checkpassed.Limitations
These results come from one RTX 4090 run with deterministic weights, not a
full pretrained BLIP-2 checkpoint or end-to-end request throughput benchmark.
Current BLIP-2 does not implement
SupportsEncoderCudaGraph, and runtime hooksconfirmed that Q-Former executes on the default eager multimodal encoder path.
No CUDA Graph performance claim is made. If BLIP-2 gains encoder CUDA Graph
support later, batch-1 short-key shapes should be benchmarked again before
enabling capture.
No performance claim is made for other architectures, devices, or PyTorch
versions. Full upstream cross-backend CI still requires a maintainer
readyor
verifiedlabel.AI assistance
AI assistance was used during implementation, testing, benchmarking, and PR
drafting. The human contributor reviewed and takes responsibility for the
change and evidence.