Repository navigation
Conversation
Add MultiHeadAttentionFusion, which rewrites the scaled dot-product
attention emitted by ViT/DINOv2 exporters with separate Q/K/V projections
Reshape -> Transpose -> MatMul(Q, K^T) -> Div|Mul -> Softmax
-> MatMul(V) -> Transpose -> Reshape
into com.microsoft.MultiHeadAttention so the WebGPU EP runs it through
FlashAttention instead of materializing the [B, N, S, S] score tensor.
The fusion is registered at Level2 for the WebGPU EP only. It is skipped
when the head size is not a multiple of 8 (the non-flash MHA path is
slower and does not support graph capture) and when the query length is
statically known to be below 32 (the split-reduce decode kernel is slower
than the unfused graph there).
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts: # onnxruntime/core/optimizer/graph_transformer_utils.cc
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
Author
|
Jie Chen (@jchen10), PTAL, thanks. <style> </style>
|
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Missing adapter head-count limits and batch-compatibility checks allow fusion to turn valid attention graphs into fused nodes that fail at runtime.
3 open issues
What changed in this PR
Adds a Level2 WebGPU graph optimization that replaces unmasked ViT-style attention with com.microsoft.MultiHeadAttention, aiming to avoid materializing the full attention-score tensor.
Changes:
- Matches separate Q/K/V attention patterns and folds scalar scaling into MHA.
- Registers the fusion only when WebGPU exposes the required kernel.
- Adds graph-matching, rejection, and fused-versus-unfused numerical tests.
| File | Description |
|---|---|
onnxruntime/test/optimizer/graph_transform_test.cc |
Adds fusion coverage and a WebGPU numerical comparison. |
onnxruntime/core/optimizer/multi_head_attention_fusion.h |
Declares the transformer and documents its supported pattern. |
onnxruntime/core/optimizer/multi_head_attention_fusion.cc |
Implements pattern matching, eligibility checks, and graph rewriting. |
onnxruntime/core/optimizer/graph_transformer_utils.cc |
Registers the fusion and shares WebGPU kernel-availability checks. |
Review effort: Balanced
Projecting the same tensor only proves equal batches when the projection keeps the leading dims. A constant weight or bias with its own leading dims (e.g. a [2, D, D] Q weight on a [1, S, D] input) broadcasts the Q batch, so the fused MultiHeadAttention would reject the unequal batches at run time. - Decide on statically comparable inferred batch dims first, so known conflicts are rejected. - Only follow projections with a 2D constant weight and a bias of rank 1 or less when falling back to the shared-source proof. - Add regression tests for a batched Q weight (static and dynamic dims), a batched Q bias, and a 1D bias that still fuses. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The WebGPU MultiHeadAttention kernel rejects num_heads above maxComputeInvocationsPerWorkgroup, which is only guaranteed to be 256, and the adapter limit is not known at graph optimization time. A graph with more heads that passes the other guards would fuse and then fail at run time, so skip it. Make the test head count configurable and add a 257-head regression. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The 16-key K/V tile estimate matches the WebGPU dense-prefill storage calculation on every adapter except Qualcomm, whose extra output tile can exceed 16 KB on 16 KB adapters. The adapter is not known at graph optimization time, and a conservative bound covering that case would stop fusing fp32 head size 64 everywhere, while a miss only costs speed. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The fusion’s batch-equality checks incorrectly accept inferred batch dimensions and transposed-A Gemm chains, converting valid broadcastable attention graphs into fused operations that reject the unequal batches.
2 open issues
3 resolved since last review
Review effort: Balanced
Comment on lines
+105
to
+107
| if (node == nullptr || (node->OpType() != "MatMul" && node->OpType() != "Gemm") || node->InputDefs().size() < 2) { | ||
| return nullptr; | ||
| } |
Comment on lines
+139
to
+140
| return a_source != nullptr && a_source == GetProjectionSource(graph, b) && GetConstantShape(graph, a, a_dims) && | ||
| GetConstantShape(graph, b, b_dims) && a_dims[0] == b_dims[0] && a_dims[1] == b_dims[1]; |
- Drop the Gemm branch from the projection-source walk. MatMulAddFusion only emits Gemm for static shapes, whose batch dims are compared directly, so the branch was never consulted. - Check the computed scale for finiteness, which also catches a tiny divisor whose reciprocal overflows. - Require an explicit EP set in the constructor, since the guards are specific to the WebGPU kernel. - Avoid a const_cast when wiring the BSD Reshapes. - Name the kernel functions the WebGPU guards mirror. - Tests: add a current-opset regression, the unscaled MatMul -> Softmax form and an overflowing scale; check the Reshape count and that MultiHeadAttention reads the [0, 0, -1] Reshapes; clarify which guard rejects the static batched-weight case. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds
MultiHeadAttentionFusion, a Level2 graph transformer for the WebGPU EP. It is registered only when the WebGPU EP has acom.microsoft.MultiHeadAttentionkernel. It rewrites the unmasked scaled-dot-product attention that ViT/DINOv2 exporters emit with separate Q/K/V projections:into
com.microsoft.MultiHeadAttention(num_heads = N,scaletaken from the Div/Mul). Each projection gets aReshape[0,0,-1]back to BSD; WebGPU's Reshape aliases its input, so this costs nothing. The output Reshape is kept and reads the MHA output, which has the same elements in the same order.Motivation and Context
On WebGPU these ViT attention blocks run as separate kernels that write and re-read the full score tensor.