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. |
|
Jie Chen (@jchen10), PTAL, thanks. <style> </style>
|
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>
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
| if (node == nullptr || (node->OpType() != "MatMul" && node->OpType() != "Gemm") || node->InputDefs().size() < 2) { | ||
| return nullptr; | ||
| } |
| 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>
|
lgtm Jiajia Qin (@qjia7) PTAL |
| InlinedVector<int64_t> a_dims, b_dims; | ||
| const NodeArg* a_source = GetProjectionSource(graph, a); | ||
| 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]; |
There was a problem hiding this comment.
C1 [P1]: Exclude inferred batch dimensions from this shortcut. This remains reproducible at this head: shared dynamic input [B,64,64], Q weights [64,64], K/V weights [64,128], and identical [-1,64,4,16] head splits pass the projection-source and literal-shape checks, but produce batches 1/2/2 at B=1. The original MatMuls legally broadcast and return [2,64,64]; MHA rejects them with "shall have same dim 0 (batch_size)." I confirmed both contracts with executable ONNX/CPU models, independently of the earlier discussion. Please reject an inferred batch here unless equality is independently proved, and add this dynamic unequal-width non-fusion regression. The current test builder always uses [0,0,N,H], so its dynamic tests cannot exercise this case.
| } | ||
| scale = is_div ? 1.0f / *value : *value; | ||
| // Also catches a tiny divisor whose reciprocal overflows. | ||
| if (!std::isfinite(scale)) { |
There was a problem hiding this comment.
C2 [P2]: Reject a computed zero scale before emitting MHA. A Div initializer of +inf or -inf passes the nonzero-input check, and its reciprocal is finite zero, so this check also accepts it. However, WebGPU MHA treats scale == 0 as the default 1/sqrt(head_size) (for example, flash_attention.cc:509), not an actual zero multiplier. For finite logits the original division yields uniform attention, whereas the replacement computes ordinary scaled attention. An executable CPU semantic comparison with BSD shape [1,32,64] and four heads produced a maximum absolute output difference of 1.428273; this was a contract comparison, not a run of the new WebGPU transformer. Please reject non-finite scalar inputs and zero computed scales, retain the reciprocal-overflow check, and add positive/negative infinite-divisor non-fusion regressions.
| transformers.emplace_back(std::make_unique<LayerNormFusion>(cpu_acl_cuda_dml_webgpu_eps, level)); | ||
| transformers.emplace_back(std::make_unique<SimplifiedLayerNormFusion>(cpu_cuda_eps)); | ||
| transformers.emplace_back(std::make_unique<AttentionFusion>(cpu_acl_cuda_dml_eps)); | ||
| if (WebGpuEpHasContribKernel(execution_providers, "MultiHeadAttention", logger)) { |
There was a problem hiding this comment.
S1 [P2]: Account for the configured Q4 eligibility requirement. With session configuration ep.webgpuexecutionprovider.kvCacheQuantizationBits=4, an otherwise matching graph with 64 queries, four heads, and head size 24 passes the transformer's multiple-of-eight and storage guards. The emitted MHA calls ApplyFlashAttention, which reads that setting and rejects non-power-of-two Q4 head sizes (flash_attention.cc:1155-1163). Thus this optimization turns a previously ordinary unquantized decomposition into a runtime INVALID_ARGUMENT, even without cache outputs. This is an exact-head source trace, not a GPU reproduction. Please propagate enough configuration to reject incompatible candidates, or conservatively disable this fusion when Q4 is configured, and add a configured head-size-24 non-fusion regression plus a compatible power-of-two positive case.
Sharing a projection source and the head-split reshape shape did not prove equal batches: a -1 in the shape absorbs the projection width, so [B, S, D] projected to D for Q and 2 * D for K/V and split with [-1, S, N, H] gives batches B and 2 * B. MatMul broadcasts that, while MultiHeadAttention rejects it. When the inferred batch dims are unknown, now also require equal projection widths (the static output dim of the 2D weight) and identical reshape shapes. Both Reshapes then provably see the same input shape and produce the same output shape. This needs no shape info for the shared source, which dynamic exports often lack. Add the reviewer's unequal-width regression, plus fused [-1, S, N, H] with equal widths and [0, -1, N, H] cases. The test builder now takes the leading head-split dims and the K/V projection width. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Dividing by +/-inf passed the nonzero check and gave a finite zero scale. WebGPU MultiHeadAttention reads scale 0 as the default 1/sqrt(head_size), whereas the original graph zeroes the logits and attends uniformly, so the fusion silently changed results. Reject non-finite Div/Mul scalars and a zero computed scale, keeping the reciprocal-overflow check. Add +inf and -inf divisor non-fusion regressions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
With ep.webgpuexecutionprovider.kvCacheQuantizationBits set to 4 or 8, WebGPU MultiHeadAttention quantizes K/V even without cache inputs or outputs, and Q4 rejects head sizes that are not powers of two. Fusing an exact attention decomposition then either loses precision or fails at run time (e.g. head size 24 under Q4 returns INVALID_ARGUMENT), while the unfused graph runs. The setting is a session config entry, so skip registering the fusion whenever it enables quantization. Add WebGPU session tests: head size 24 under Q4 and the default shape under Q8 stay unfused and run, while the same graph with quantization unset or "0" fuses and runs. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts: # onnxruntime/core/optimizer/graph_transformer_utils.cc
A rank-1 bias broadcasts a width-1 weight output to the bias length, so equal weight widths did not prove equal head-split input shapes: [D, 1] weights with Q/K/V biases of D/2D/2D and [-1, S, N, H] splits give batches B and 2 * B. Take the projection width as max(weight width, bias length), which is the width after the Add under broadcasting rules. It is computed from the constants rather than the inferred Add output shape, which is missing for dynamic-axis exports such as Depth-Anything and would disable the fusion there. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
# Conflicts: # onnxruntime/core/optimizer/graph_transformer_utils.cc
There was a problem hiding this comment.
🟡 Changes recommended
FP16 fusion can produce incorrect results because folded scales pass the float32 finiteness check but overflow FP16; an FP16 range guard and tiny-divisor regression test are needed.
3 open issues
1 resolved since last review
Review effort: Balanced
| if (q_type == nullptr || !q_type->has_tensor_type() || | ||
| (q_type->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && | ||
| q_type->tensor_type().elem_type() != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16)) { | ||
| return false; | ||
| } |
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.