From 1b20fa908902170d6afe5023be0d8e11071a0158 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Mon, 2 Nov 2026 12:29:43 +0800 Subject: [PATCH 1/7] [WebGPU] Fuse unmasked ViT attention into MultiHeadAttention 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 --- .../core/optimizer/graph_transformer_utils.cc | 3 + .../optimizer/multi_head_attention_fusion.cc | 290 ++++++++++++++++++ .../optimizer/multi_head_attention_fusion.h | 42 +++ .../test/optimizer/graph_transform_test.cc | 184 +++++++++++ 4 files changed, 519 insertions(+) create mode 100644 onnxruntime/core/optimizer/multi_head_attention_fusion.cc create mode 100644 onnxruntime/core/optimizer/multi_head_attention_fusion.h diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index 64372c52e1a20..a1062969f90c2 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -57,6 +57,7 @@ #include "core/optimizer/layer_norm_fusion.h" #include "core/optimizer/matmul_activation_fusion.h" #include "core/optimizer/matmul_add_fusion.h" +#include "core/optimizer/multi_head_attention_fusion.h" #include "core/optimizer/group_query_attention_pre_norm_fusion.h" #include "core/optimizer/matmul_nbits_qkv_fusion.h" #include "core/optimizer/matmul_nbits_mlp_fusion.h" @@ -415,6 +416,8 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_webgpu_eps, level)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps)); + transformers.emplace_back(std::make_unique( + InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc new file mode 100644 index 0000000000000..2b9505551c22e --- /dev/null +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -0,0 +1,290 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "core/optimizer/multi_head_attention_fusion.h" + +#include +#include +#include + +#include "core/framework/tensorprotoutils.h" +#include "core/graph/graph_utils.h" +#include "core/optimizer/initializer.h" +#include "core/optimizer/utils.h" + +namespace onnxruntime { + +namespace { + +constexpr std::initializer_list kTransposeVersions = {1, 13, 21, 23, 24, 25}; +constexpr std::initializer_list kReshapeVersions = {5, 13, 14, 19, 21, 23, 24, 25}; +constexpr std::initializer_list kMatMulVersions = {1, 9, 13}; +constexpr std::initializer_list kBinaryVersions = {7, 13, 14}; + +bool IsTransposeWithPerm(const Node& node, const std::vector& perm) { + return graph_utils::IsSupportedOptypeVersionAndDomain(node, "Transpose", kTransposeVersions) && + optimizer_utils::IsAttributeWithExpectedValues(node, "perm", perm); +} + +bool IsReshapeWithoutAllowZero(const Node& node) { + return graph_utils::IsSupportedOptypeVersionAndDomain(node, "Reshape", kReshapeVersions) && + node.InputDefs().size() == 2 && + !optimizer_utils::IsAttributeWithExpectedValue(node, "allowzero", static_cast(1)); +} + +bool GetConstantShape(const Graph& graph, const Node& reshape, InlinedVector& shape) { + return optimizer_utils::AppendTensorFromInitializer(graph, *reshape.InputDefs()[1], shape, true); +} + +// Reads a scalar float/float16 constant. +std::optional GetScalarConstant(const Graph& graph, const NodeArg& arg) { + if (!optimizer_utils::IsScalar(arg)) { + return std::nullopt; + } + const auto* tensor_proto = graph_utils::GetConstantInitializer(graph, arg.Name()); + if (tensor_proto == nullptr) { + return std::nullopt; + } + Initializer init{graph, *tensor_proto, graph.ModelPath()}; + if (init.size() != 1) { + return std::nullopt; + } + switch (tensor_proto->data_type()) { + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT: + return *init.data(); + case ONNX_NAMESPACE::TensorProto_DataType_FLOAT16: + return static_cast(*init.data()); + default: + return std::nullopt; + } +} + +// Matches `Reshape[*,*,N,H] -> Transpose(perm)` feeding `consumer` at `input_index`. +const Node* MatchHeadSplit(const Graph& graph, const Node& consumer, int input_index, const std::vector& perm, + int64_t& num_heads, int64_t& head_size, const Node*& reshape) { + const Node* transpose = graph_utils::GetInputNode(consumer, input_index); + if (transpose == nullptr || !IsTransposeWithPerm(*transpose, perm) || + !optimizer_utils::CheckOutputEdges(graph, *transpose, 1)) { + return nullptr; + } + + reshape = graph_utils::GetInputNode(*transpose, 0); + if (reshape == nullptr || !IsReshapeWithoutAllowZero(*reshape)) { + return nullptr; + } + + InlinedVector shape; + if (!GetConstantShape(graph, *reshape, shape) || shape.size() != 4 || shape[2] <= 0 || shape[3] <= 0) { + return nullptr; + } + + num_heads = shape[2]; + head_size = shape[3]; + return transpose; +} + +bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet& compatible_eps, + NodeArg*& bsd_shape_arg, const logging::Logger& logger) { + if (!graph_utils::IsSupportedOptypeVersionAndDomain(softmax, "Softmax", {1, 11, 13}) || + !graph_utils::IsSupportedProvider(softmax, compatible_eps) || + !optimizer_utils::CheckOutputEdges(graph, softmax, 1)) { + return false; + } + + // Before opset 13 Softmax coerces the input to 2D at `axis` (default 1), so only the last axis of the 4D score + // tensor matches MultiHeadAttention. + const auto* axis_attr = graph_utils::GetNodeAttribute(softmax, "axis"); + const int64_t axis = axis_attr != nullptr ? axis_attr->i() : (softmax.SinceVersion() >= 13 ? -1 : 1); + if (axis != -1 && axis != 3) { + return false; + } + + float scale = 1.0f; + const Node* scale_node = nullptr; + const Node* qk_matmul = graph_utils::GetInputNode(softmax, 0); + if (qk_matmul != nullptr && + (graph_utils::IsSupportedOptypeVersionAndDomain(*qk_matmul, "Div", kBinaryVersions) || + graph_utils::IsSupportedOptypeVersionAndDomain(*qk_matmul, "Mul", kBinaryVersions))) { + scale_node = qk_matmul; + const bool is_div = scale_node->OpType() == "Div"; + if (!optimizer_utils::CheckOutputEdges(graph, *scale_node, 1)) { + return false; + } + // For Mul the scalar may be either operand; for Div it must be the divisor. + int scalar_index = 1; + std::optional value = GetScalarConstant(graph, *scale_node->InputDefs()[1]); + if (!value.has_value() && !is_div) { + scalar_index = 0; + value = GetScalarConstant(graph, *scale_node->InputDefs()[0]); + } + if (!value.has_value() || *value == 0.0f || !std::isfinite(*value)) { + return false; + } + scale = is_div ? 1.0f / *value : *value; + qk_matmul = graph_utils::GetInputNode(*scale_node, 1 - scalar_index); + } + + if (qk_matmul == nullptr || !graph_utils::IsSupportedOptypeVersionAndDomain(*qk_matmul, "MatMul", kMatMulVersions) || + !optimizer_utils::CheckOutputEdges(graph, *qk_matmul, 1)) { + return false; + } + + Node* pv_matmul = graph.GetNode(softmax.OutputNodesBegin()->Index()); + if (pv_matmul == nullptr || !graph_utils::IsSupportedOptypeVersionAndDomain(*pv_matmul, "MatMul", kMatMulVersions) || + pv_matmul->InputDefs()[0] != softmax.OutputDefs()[0] || + !optimizer_utils::CheckOutputEdges(graph, *pv_matmul, 1)) { + return false; + } + + Node* out_transpose = graph.GetNode(pv_matmul->OutputNodesBegin()->Index()); + if (out_transpose == nullptr || !IsTransposeWithPerm(*out_transpose, {0, 2, 1, 3}) || + !optimizer_utils::CheckOutputEdges(graph, *out_transpose, 1)) { + return false; + } + + // The output Reshape is kept and now consumes the [B, S, N*H] MultiHeadAttention output instead of the + // [B, S, N, H] tensor. A 0 in its shape copies the input dim at that index, which is only equivalent for the + // leading B and S dims. + const Node* out_reshape = graph.GetNode(out_transpose->OutputNodesBegin()->Index()); + InlinedVector out_shape; + if (out_reshape == nullptr || !IsReshapeWithoutAllowZero(*out_reshape) || + out_reshape->InputDefs()[0] != out_transpose->OutputDefs()[0] || + !GetConstantShape(graph, *out_reshape, out_shape)) { + return false; + } + for (size_t i = 2; i < out_shape.size(); ++i) { + if (out_shape[i] == 0) { + return false; + } + } + + int64_t q_heads = 0, q_head_size = 0, k_heads = 0, k_head_size = 0, v_heads = 0, v_head_size = 0; + const Node *q_reshape = nullptr, *k_reshape = nullptr, *v_reshape = nullptr; + const Node* q_transpose = MatchHeadSplit(graph, *qk_matmul, 0, {0, 2, 1, 3}, q_heads, q_head_size, q_reshape); + const Node* k_transpose = MatchHeadSplit(graph, *qk_matmul, 1, {0, 2, 3, 1}, k_heads, k_head_size, k_reshape); + const Node* v_transpose = MatchHeadSplit(graph, *pv_matmul, 1, {0, 2, 1, 3}, v_heads, v_head_size, v_reshape); + if (q_transpose == nullptr || k_transpose == nullptr || v_transpose == nullptr) { + return false; + } + if (q_heads != k_heads || q_heads != v_heads || q_head_size != k_head_size || q_head_size != v_head_size) { + LOGS(logger, VERBOSE) << "MultiHeadAttentionFusion[" << softmax.Name() << "]: mismatched Q/K/V head shapes"; + return false; + } + // The WebGPU MultiHeadAttention kernel takes its FlashAttention path for head sizes that are a multiple of 8 on + // every adapter. The fallback path is slower than the unfused graph and does not support graph capture. + if (q_head_size % 8 != 0) { + return false; + } + // A short query leaves only a small score tensor to avoid, and MultiHeadAttention sends queries shorter than 32 to + // the split-reduce decode kernel, which is slower than the unfused graph for long keys. + const auto* q_bnsh_shape = q_transpose->OutputDefs()[0]->Shape(); + if (q_bnsh_shape != nullptr && q_bnsh_shape->dim_size() == 4 && utils::HasDimValue(q_bnsh_shape->dim(2)) && + q_bnsh_shape->dim(2).dim_value() < 32) { + return false; + } + + const auto* q_type = q_reshape->OutputDefs()[0]->TypeAsProto(); + 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; + } + + const auto& provider = softmax.GetExecutionProviderType(); + for (const Node* node : {qk_matmul, scale_node, static_cast(pv_matmul), + static_cast(out_transpose), q_transpose, k_transpose, v_transpose, + q_reshape, k_reshape, v_reshape, out_reshape}) { + if (node != nullptr && node->GetExecutionProviderType() != provider) { + return false; + } + } + + if (bsd_shape_arg == nullptr) { + ONNX_NAMESPACE::TensorProto shape_proto; + shape_proto.set_name(graph.GenerateNodeArgName("mha_fusion_bsd_shape")); + shape_proto.set_data_type(ONNX_NAMESPACE::TensorProto_DataType_INT64); + shape_proto.add_dims(3); + const std::array bsd_shape{0, 0, -1}; + utils::SetRawDataInTensorProto(shape_proto, bsd_shape.data(), bsd_shape.size() * sizeof(int64_t)); + bsd_shape_arg = &graph_utils::AddInitializerWithOrtValue(graph, shape_proto); + } + + // Only the element type is set; shape inference fills in the shapes when the graph is resolved. + ONNX_NAMESPACE::TypeProto tensor_type; + tensor_type.mutable_tensor_type()->set_elem_type(q_type->tensor_type().elem_type()); + + auto add_to_bsd = [&](const Node& reshape) -> Node& { + NodeArg* bsnh = const_cast(reshape.OutputDefs()[0]); + NodeArg& bsd = graph.GetOrCreateNodeArg(graph.GenerateNodeArgName(bsnh->Name() + "_bsd"), &tensor_type); + Node& node = graph.AddNode(graph.GenerateNodeName(reshape.Name() + "/MhaBsd"), "Reshape", + "Fold heads into the hidden dim for MultiHeadAttention", {bsnh, bsd_shape_arg}, {&bsd}); + node.SetExecutionProviderType(provider); + graph.AddEdge(reshape.Index(), node.Index(), 0, 0); + return node; + }; + Node& q_bsd = add_to_bsd(*q_reshape); + Node& k_bsd = add_to_bsd(*k_reshape); + Node& v_bsd = add_to_bsd(*v_reshape); + + NodeArg& mha_output = graph.GetOrCreateNodeArg(graph.GenerateNodeArgName("mha_fusion_output"), &tensor_type); + Node& mha = graph.AddNode(graph.GenerateNodeName(softmax.Name() + "/MultiHeadAttention"), "MultiHeadAttention", + "Fused scaled dot-product attention", + {q_bsd.MutableOutputDefs()[0], k_bsd.MutableOutputDefs()[0], v_bsd.MutableOutputDefs()[0]}, + {&mha_output}, nullptr, kMSDomain); + mha.AddAttribute("num_heads", q_heads); + mha.AddAttribute("scale", scale); + mha.SetExecutionProviderType(provider); + graph.AddEdge(q_bsd.Index(), mha.Index(), 0, 0); + graph.AddEdge(k_bsd.Index(), mha.Index(), 0, 1); + graph.AddEdge(v_bsd.Index(), mha.Index(), 0, 2); + + graph_utils::ReplaceDownstreamNodeInput(graph, *out_transpose, 0, mha, 0); + + InlinedVector nodes_to_remove{q_transpose->Index(), k_transpose->Index(), v_transpose->Index(), + qk_matmul->Index(), softmax.Index(), pv_matmul->Index(), + out_transpose->Index()}; + if (scale_node != nullptr) { + nodes_to_remove.push_back(scale_node->Index()); + } + graph.NotifyNodeReplacement(nodes_to_remove, mha.Index()); + for (NodeIndex index : nodes_to_remove) { + Node* node = graph.GetNode(index); + graph_utils::RemoveNodeOutputEdges(graph, *node); + graph.RemoveNode(index); + } + + LOGS(logger, VERBOSE) << "MultiHeadAttentionFusion: fused attention at " << mha.Name(); + return true; +} + +} // namespace + +Status MultiHeadAttentionFusion::ApplyImpl(Graph& graph, bool& modified, int graph_level, + const logging::Logger& logger) const { + GraphViewer graph_viewer(graph); + const auto& node_topology_list = graph_viewer.GetNodesInTopologicalOrder(); + + NodeArg* bsd_shape_arg = nullptr; + int fused_count = 0; + for (auto node_index : node_topology_list) { + Node* node = graph.GetNode(node_index); + if (node == nullptr) { + continue; + } + + ORT_RETURN_IF_ERROR(Recurse(*node, modified, graph_level, logger)); + + if (TryFuse(graph, *node, GetCompatibleExecutionProviders(), bsd_shape_arg, logger)) { + ++fused_count; + modified = true; + } + } + + if (fused_count > 0) { + LOGS(logger, INFO) << "Total fused MultiHeadAttention node count: " << fused_count; + } + + return Status::OK(); +} + +} // namespace onnxruntime diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.h b/onnxruntime/core/optimizer/multi_head_attention_fusion.h new file mode 100644 index 0000000000000..e313b1817cd75 --- /dev/null +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.h @@ -0,0 +1,42 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/optimizer/graph_transformer.h" + +namespace onnxruntime { + +/** +@Class MultiHeadAttentionFusion + +Fuses the unmasked scaled-dot-product attention emitted by ViT-style exporters (e.g. HuggingFace +ViT / DINOv2) with separate Q, K and V projections into a com.microsoft MultiHeadAttention node: + + Q -> Reshape[*,*,N,H] -> Transpose(0,2,1,3) --+ + MatMul -> Div|Mul(scalar)? -> Softmax(-1) + K -> Reshape[*,*,N,H] -> Transpose(0,2,3,1) --+ | + MatMul -> Transpose(0,2,1,3) + V -> Reshape[*,*,N,H] -> Transpose(0,2,1,3) ------------------------------------+ -> Reshape(out) + +is rewritten to + + Q -> Reshape[*,*,N,H] -> Reshape[0,0,-1] --+ + K -> Reshape[*,*,N,H] -> Reshape[0,0,-1] --+-> MultiHeadAttention(num_heads=N, scale) -> Reshape(out) + V -> Reshape[*,*,N,H] -> Reshape[0,0,-1] --+ + +The added Reshapes only fold the trailing (N, H) dims back to the BSD layout MultiHeadAttention expects, so they are +free for EPs whose Reshape aliases its input. This keeps the [B, N, S, S] score tensor out of memory, which dominates +the attention cost for long vision sequences. +*/ +class MultiHeadAttentionFusion : public GraphTransformer { + public: + explicit MultiHeadAttentionFusion( + const InlinedHashSet& compatible_execution_providers = {}) noexcept + : GraphTransformer("MultiHeadAttentionFusion", compatible_execution_providers) { + } + + Status ApplyImpl(Graph& graph, bool& modified, int graph_level, const logging::Logger& logger) const override; +}; + +} // namespace onnxruntime diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 875b28279671b..3698eaab7fa80 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -8,6 +8,7 @@ #include #include #include +#include #include #include @@ -30,6 +31,7 @@ #include "core/graph/model.h" #include "core/mlas/inc/mlas_q4.h" #include "core/optimizer/attention_fusion.h" +#include "core/optimizer/multi_head_attention_fusion.h" #include "core/optimizer/bias_dropout_fusion.h" #include "core/optimizer/bias_gelu_fusion.h" #include "core/optimizer/bias_softmax_fusion.h" @@ -9516,6 +9518,188 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaProjectionRewriteFa CheckMobileClipAttentionUnfusedMatMulGraph)); } +#if !defined(DISABLE_CONTRIB_OPS) +enum class VitAttentionScaleOp { + Div, + Mul, +}; + +struct VitAttentionTestConfig { + VitAttentionScaleOp scale_op = VitAttentionScaleOp::Div; + // Unset means no Softmax axis attribute. + std::optional softmax_axis = 3; + std::vector k_perm = {0, 2, 3, 1}; + int64_t seq_len = 64; + int64_t head_size = 16; +}; + +// Unmasked ViT/DINOv2 self-attention with separate Q, K and V projections: +// x -> MatMul(Wq/Wk/Wv) -> Reshape[0,0,N,H] -> Transpose -> MatMul(Q, K^T) -> Div|Mul -> Softmax +// -> MatMul(V) -> Transpose(0,2,1,3) -> Reshape[0,0,N*H] +static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttentionTestConfig& config = {}) { + constexpr int64_t num_heads = 4; + const int64_t head_size = config.head_size; + const int64_t hidden_size = num_heads * head_size; + + auto* input = builder.MakeInput({1, config.seq_len, hidden_size}, -1.0f, 1.0f); + auto* split_heads_shape = builder.Make1DInitializer({0, 0, num_heads, head_size}); + auto* merge_heads_shape = builder.Make1DInitializer({0, 0, hidden_size}); + const float sqrt_head_size = std::sqrt(static_cast(head_size)); + auto* scale = builder.MakeScalarInitializer( + config.scale_op == VitAttentionScaleOp::Div ? sqrt_head_size : 1.0f / sqrt_head_size); + + auto add_head_split = [&](const std::vector& perm) { + auto* weight = builder.MakeInitializer({hidden_size, hidden_size}, -0.2f, 0.2f); + auto* projected = builder.MakeIntermediate(); + auto* bsnh = builder.MakeIntermediate(); + auto* transposed = builder.MakeIntermediate(); + builder.AddNode("MatMul", {input, weight}, {projected}); + builder.AddNode("Reshape", {projected, split_heads_shape}, {bsnh}); + builder.AddNode("Transpose", {bsnh}, {transposed}).AddAttribute("perm", perm); + return transposed; + }; + auto* q = add_head_split({0, 2, 1, 3}); + auto* k = add_head_split(config.k_perm); + auto* v = add_head_split({0, 2, 1, 3}); + + auto* qk = builder.MakeIntermediate(); + auto* scaled = builder.MakeIntermediate(); + auto* probs = builder.MakeIntermediate(); + auto* context = builder.MakeIntermediate(); + auto* context_bsnh = builder.MakeIntermediate(); + auto* output = builder.MakeOutput(); + builder.AddNode("MatMul", {q, k}, {qk}); + builder.AddNode(config.scale_op == VitAttentionScaleOp::Div ? "Div" : "Mul", {qk, scale}, {scaled}); + auto& softmax = builder.AddNode("Softmax", {scaled}, {probs}); + if (config.softmax_axis.has_value()) { + softmax.AddAttribute("axis", *config.softmax_axis); + } + builder.AddNode("MatMul", {probs, v}, {context}); + builder.AddNode("Transpose", {context}, {context_bsnh}).AddAttribute("perm", std::vector{0, 2, 1, 3}); + builder.AddNode("Reshape", {context_bsnh, merge_heads_shape}, {output}); +} + +static Status AssignVitAttentionEp(Graph& graph, const char* ep) { + for (auto& node : graph.Nodes()) { + node.SetExecutionProviderType(ep); + } + return Status::OK(); +} + +static Status CheckVitAttentionFused(Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(op_to_count["com.microsoft.MultiHeadAttention"] == 1, "expected one MultiHeadAttention"); + ORT_RETURN_IF_NOT(op_to_count["Softmax"] == 0, "Softmax was not fused"); + ORT_RETURN_IF_NOT(op_to_count["Transpose"] == 0, "head Transposes were not removed"); + ORT_RETURN_IF_NOT(op_to_count["MatMul"] == 3, "only the Q/K/V projections should remain"); + ORT_RETURN_IF_NOT(op_to_count["Div"] == 0 && op_to_count["Mul"] == 0, "the scale was not folded"); + + for (const auto& node : graph.Nodes()) { + if (node.OpType() == "MultiHeadAttention") { + const auto& attrs = node.GetAttributes(); + ORT_RETURN_IF_NOT(attrs.at("num_heads").i() == 4, "unexpected num_heads"); + ORT_RETURN_IF_NOT(std::abs(attrs.at("scale").f() - 0.25f) < 1e-6f, "unexpected scale"); + ORT_RETURN_IF_NOT(node.GetExecutionProviderType() == kWebGpuExecutionProvider, "unexpected EP"); + } + } + return Status::OK(); +} + +static Status CheckVitAttentionUnfused(Graph& graph) { + auto op_to_count = CountOpsInGraph(graph); + ORT_RETURN_IF_NOT(op_to_count["com.microsoft.MultiHeadAttention"] == 0, "unexpected MultiHeadAttention"); + ORT_RETURN_IF_NOT(op_to_count["Softmax"] == 1, "Softmax should remain"); + return Status::OK(); +} + +static Status TestVitAttentionFusion(const VitAttentionTestConfig& config, int opset, const char* ep, + const std::function& check, const logging::Logger& logger) { + return TestGraphTransformer( + [&config](ModelTestBuilder& builder) { BuildVitAttentionTestCase(builder, config); }, opset, logger, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}), + TransformerLevel::Level2, 1, [ep](Graph& graph) { return AssignVitAttentionEp(graph, ep); }, check); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDivScale) { + // Opset 11 matches exported ViT/DINOv2 models, where Softmax needs the explicit last-axis attribute. + ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, kWebGpuExecutionProvider, CheckVitAttentionFused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitMulScaleDefaultAxis) { + // From opset 13 Softmax defaults to the last axis. + ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Mul, std::nullopt}, 13, kWebGpuExecutionProvider, + CheckVitAttentionFused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitOpset11DefaultAxisNotFused) { + // Before opset 13 the default Softmax axis is 1, which normalizes over heads and keys together. + ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, std::nullopt}, 11, kWebGpuExecutionProvider, + CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitUntransposedKeyNotFused) { + // seq_len == head_size keeps Q x K a valid MatMul although it no longer computes Q x K^T. + ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, 3, {0, 2, 1, 3}, 32, 32}, 11, kWebGpuExecutionProvider, + CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitShortQueryNotFused) { + // Queries shorter than 32 would run the split-reduce decode kernel. + ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, 3, {0, 2, 3, 1}, 31}, 11, kWebGpuExecutionProvider, + CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitHeadSizeNotMultipleOf8NotFused) { + // Such head sizes can miss the FlashAttention path, whose fallback does not support graph capture. + ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, 3, {0, 2, 3, 1}, 64, 12}, 11, + kWebGpuExecutionProvider, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitOtherEpNotFused) { + ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, kCpuExecutionProvider, CheckVitAttentionUnfused, *logger_)); +} + +// Runs the graph on WebGPU with and without the fusion through the default session optimizers and compares the +// outputs. A sequence length of at least 32 takes MultiHeadAttention down the dense FlashAttention prefill path; 97 +// also leaves a partial K/V tile. Each session gets its own EP instance. +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitWebGpuNumerics) { + if (!DefaultWebGpuExecutionProvider()) { + GTEST_SKIP() << "WebGPU execution provider is not available"; + } + + const std::unordered_map domain_to_version{{kOnnxDomain, 11}, {kMSDomain, 1}}; + Model model("VitAttention", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), + domain_to_version, {}, *logger_); + ModelTestBuilder builder(model.MainGraph()); + BuildVitAttentionTestCase(builder, {VitAttentionScaleOp::Div, 3, {0, 2, 3, 1}, 97}); + builder.SetGraphOutputs(); + ASSERT_STATUS_OK(model.MainGraph().Resolve()); + std::string model_data; + model.ToProto().SerializeToString(&model_data); + + auto run = [&](bool fuse, std::vector& fetches) { + InferenceSessionWrapper session{SessionOptions{}, GetEnvironment()}; + ASSERT_STATUS_OK(session.RegisterExecutionProvider(DefaultWebGpuExecutionProvider())); + std::istringstream model_istream(model_data); + ASSERT_STATUS_OK(session.Load(model_istream)); + if (!fuse) { + ASSERT_STATUS_OK(session.FilterEnabledOptimizers({"MultiHeadAttentionFusion"})); + } + ASSERT_STATUS_OK(session.Initialize()); + auto op_to_count = CountOpsInGraph(session.GetGraph()); + EXPECT_EQ(op_to_count["com.microsoft.MultiHeadAttention"], fuse ? 1 : 0); + ASSERT_STATUS_OK(session.Run(RunOptions{}, builder.feeds_, builder.output_names_, &fetches)); + }; + + std::vector unfused, fused; + ASSERT_NO_FATAL_FAILURE(run(false, unfused)); + ASSERT_NO_FATAL_FAILURE(run(true, fused)); + ASSERT_EQ(fused.size(), 1u); + auto [result, message] = CompareOrtValue(fused[0], unfused[0], 1e-4, 1e-4, false); + EXPECT_EQ(result, COMPARE_RESULT::SUCCESS) << message; +} +#endif // !defined(DISABLE_CONTRIB_OPS) + // Current-opset regression tests for fusion optimizers. // These construct minimal graphs at the current ONNX opset and verify the optimizer fires. // If ONNX bumps an op's since_version, the optimizer's version list will miss it and the test fails. From 30c79878f4ea901cc7482e565c1aad74d0d8dad3 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Tue, 3 Nov 2026 12:36:36 +0800 Subject: [PATCH 2/7] Clarify why the output Reshape can be kept in MultiHeadAttentionFusion Co-Authored-By: Claude Opus 5.5 --- onnxruntime/core/optimizer/multi_head_attention_fusion.cc | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc index 2b9505551c22e..7e25c5dd7c16f 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -142,9 +142,10 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet return false; } - // The output Reshape is kept and now consumes the [B, S, N*H] MultiHeadAttention output instead of the - // [B, S, N, H] tensor. A 0 in its shape copies the input dim at that index, which is only equivalent for the - // leading B and S dims. + // The output Reshape is kept and now reads the [B, S, N*H] MultiHeadAttention output instead of the + // [B, S, N, H] Transpose output. Both tensors hold the same elements in the same order, so explicit dims and -1 + // keep their meaning, but a 0 copies the input dim at that index, which only matches for the two leading dims. + // The shape therefore has to be a constant without zeros past index 1. const Node* out_reshape = graph.GetNode(out_transpose->OutputNodesBegin()->Index()); InlinedVector out_shape; if (out_reshape == nullptr || !IsReshapeWithoutAllowZero(*out_reshape) || From 6206d02d7eafae6cc059aace629d648736a621c4 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Tue, 3 Nov 2026 16:34:31 +0800 Subject: [PATCH 3/7] Address review findings in MultiHeadAttentionFusion Co-Authored-By: Claude Opus 5.5 --- .../core/optimizer/graph_transformer_utils.cc | 52 ++--- .../optimizer/multi_head_attention_fusion.cc | 78 +++++++- .../optimizer/multi_head_attention_fusion.h | 6 + .../test/optimizer/graph_transform_test.cc | 185 +++++++++++++++--- 4 files changed, 258 insertions(+), 63 deletions(-) diff --git a/onnxruntime/core/optimizer/graph_transformer_utils.cc b/onnxruntime/core/optimizer/graph_transformer_utils.cc index a1062969f90c2..5457e5f9b4b46 100644 --- a/onnxruntime/core/optimizer/graph_transformer_utils.cc +++ b/onnxruntime/core/optimizer/graph_transformer_utils.cc @@ -209,6 +209,26 @@ std::unique_ptr GenerateRuleBasedGraphTransformer( return rule_transformer; } +#if !defined(DISABLE_CONTRIB_OPS) +// Whether a registered WebGPU EP has a kernel for the com.microsoft opset-1 operator `op_type`. Fusions that emit +// WebGPU-only contrib ops check this so they never produce a node the EP cannot run. +static bool WebGpuEpHasContribKernel(const ExecutionProviders* execution_providers, std::string_view op_type, + const logging::Logger& logger) { + if (execution_providers == nullptr) { + return false; + } + const auto* webgpu_ep = execution_providers->Get(onnxruntime::kWebGpuExecutionProvider); + if (webgpu_ep == nullptr) { + return false; + } + auto registry = webgpu_ep->GetKernelRegistry(); + return registry != nullptr && + registry->TryFindKernel(onnxruntime::kWebGpuExecutionProvider, op_type, kMSDomain, 1, + KernelRegistry::TypeConstraintMap{}, logger, nullptr) + .IsOK(); +} +#endif // !defined(DISABLE_CONTRIB_OPS) + InlinedVector> GenerateTransformers( TransformerLevel level, const SessionOptions& session_options, @@ -416,8 +436,10 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_webgpu_eps, level)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps)); - transformers.emplace_back(std::make_unique( - InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); + if (WebGpuEpHasContribKernel(execution_providers, "MultiHeadAttention", logger)) { + transformers.emplace_back(std::make_unique( + InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); + } transformers.emplace_back(std::make_unique(cpu_acl_cuda_dml_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); transformers.emplace_back(std::make_unique(cpu_cuda_eps)); @@ -470,28 +492,10 @@ InlinedVector> GenerateTransformers( transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kCudaExecutionProvider, onnxruntime::kWebGpuExecutionProvider})); - bool has_matmul_nbits_mlp_kernel = false; - bool has_matmul_nbits_qkv_kernel = false; - if (execution_providers != nullptr) { - const auto* webgpu_ep = execution_providers->Get(onnxruntime::kWebGpuExecutionProvider); - if (webgpu_ep != nullptr) { - auto registry = webgpu_ep->GetKernelRegistry(); - if (registry) { - auto has_webgpu_kernel = [&](std::string_view op_type) { - return registry->TryFindKernel(onnxruntime::kWebGpuExecutionProvider, - op_type, - kMSDomain, - 1, - KernelRegistry::TypeConstraintMap{}, - logger, - nullptr) - .IsOK(); - }; - has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp"); - has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv"); - } - } - } + const bool has_matmul_nbits_mlp_kernel = + WebGpuEpHasContribKernel(execution_providers, "MatMulNBitsMlp", logger); + const bool has_matmul_nbits_qkv_kernel = + WebGpuEpHasContribKernel(execution_providers, "MatMulNBitsQkv", logger); if (has_matmul_nbits_mlp_kernel) { transformers.emplace_back(std::make_unique( InlinedHashSet{onnxruntime::kWebGpuExecutionProvider})); diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc index 7e25c5dd7c16f..ee76e8dc7e4cf 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -83,6 +83,51 @@ const Node* MatchHeadSplit(const Graph& graph, const Node& consumer, int input_i return transpose; } +// Walks up from a head-split Reshape to the tensor its projection reads: through an optional bias Add, a MatMul or +// Gemm with a constant weight, and the flattening Reshape that MatMulAddFusion puts in front of a Gemm. +const NodeArg* GetProjectionSource(const Graph& graph, const Node& head_split_reshape) { + const Node* node = graph_utils::GetInputNode(head_split_reshape, 0); + if (node != nullptr && node->OpType() == "Add" && node->InputDefs().size() == 2) { + const int data_index = graph_utils::IsConstantInitializer(graph, node->InputDefs()[1]->Name(), true) ? 0 : 1; + if (!graph_utils::IsConstantInitializer(graph, node->InputDefs()[1 - data_index]->Name(), true)) { + return nullptr; + } + node = graph_utils::GetInputNode(*node, data_index); + } + if (node == nullptr || (node->OpType() != "MatMul" && node->OpType() != "Gemm") || node->InputDefs().size() < 2 || + !graph_utils::IsConstantInitializer(graph, node->InputDefs()[1]->Name(), true)) { + return nullptr; + } + const Node* flatten = graph_utils::GetInputNode(*node, 0); + if (node->OpType() == "Gemm" && flatten != nullptr && flatten->OpType() == "Reshape") { + return flatten->InputDefs()[0]; + } + return node->InputDefs()[0]; +} + +// Whether two head-split Reshapes produce the same batch dim. That holds when both project the same tensor and keep +// the same leading reshape dims, as in self-attention; otherwise it needs statically comparable inferred dims. +bool HaveSameBatchDim(const Graph& graph, const Node& a, const Node& b) { + InlinedVector a_dims, b_dims; + const NodeArg* a_source = GetProjectionSource(graph, a); + if (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]) { + return true; + } + + const auto* a_shape = a.OutputDefs()[0]->Shape(); + const auto* b_shape = b.OutputDefs()[0]->Shape(); + if (a_shape == nullptr || b_shape == nullptr || a_shape->dim_size() != 4 || b_shape->dim_size() != 4) { + return false; + } + const auto& a_dim = a_shape->dim(0); + const auto& b_dim = b_shape->dim(0); + if (utils::HasDimValue(a_dim) && utils::HasDimValue(b_dim)) { + return a_dim.dim_value() == b_dim.dim_value(); + } + return utils::HasDimParam(a_dim) && utils::HasDimParam(b_dim) && a_dim.dim_param() == b_dim.dim_param(); +} + bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet& compatible_eps, NodeArg*& bsd_shape_arg, const logging::Logger& logger) { if (!graph_utils::IsSupportedOptypeVersionAndDomain(softmax, "Softmax", {1, 11, 13}) || @@ -167,20 +212,13 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet if (q_transpose == nullptr || k_transpose == nullptr || v_transpose == nullptr) { return false; } + // Unequal head counts also cover multi-query attention, where K/V have one head that MatMul broadcasts. if (q_heads != k_heads || q_heads != v_heads || q_head_size != k_head_size || q_head_size != v_head_size) { LOGS(logger, VERBOSE) << "MultiHeadAttentionFusion[" << softmax.Name() << "]: mismatched Q/K/V head shapes"; return false; } - // The WebGPU MultiHeadAttention kernel takes its FlashAttention path for head sizes that are a multiple of 8 on - // every adapter. The fallback path is slower than the unfused graph and does not support graph capture. - if (q_head_size % 8 != 0) { - return false; - } - // A short query leaves only a small score tensor to avoid, and MultiHeadAttention sends queries shorter than 32 to - // the split-reduce decode kernel, which is slower than the unfused graph for long keys. - const auto* q_bnsh_shape = q_transpose->OutputDefs()[0]->Shape(); - if (q_bnsh_shape != nullptr && q_bnsh_shape->dim_size() == 4 && utils::HasDimValue(q_bnsh_shape->dim(2)) && - q_bnsh_shape->dim(2).dim_value() < 32) { + // MatMul broadcasts the batch dim, but MultiHeadAttention requires Q, K and V to share it. + if (!HaveSameBatchDim(graph, *q_reshape, *k_reshape) || !HaveSameBatchDim(graph, *q_reshape, *v_reshape)) { return false; } @@ -191,6 +229,26 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet return false; } + // The remaining checks keep the fused node on the WebGPU kernel's dense FlashAttention prefill path. Its fallbacks + // are slower than the unfused graph, and the non-flash one does not support graph capture. + // FlashAttention requires a head size that is a multiple of 8 on every adapter. + if (q_head_size % 8 != 0) { + return false; + } + // Dense prefill needs K/V tiles of 16 keys to fit in the 16 KB of workgroup storage every adapter guarantees. + const int64_t element_size = + q_type->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ? 2 : 4; + if (2 * 16 * q_head_size * element_size > 16384) { + return false; + } + // Queries shorter than 32 go to the split-reduce decode kernel, and leave only a small score tensor to avoid. + // A symbolic query length cannot be checked here, so it is fused. + const auto* q_bnsh_shape = q_transpose->OutputDefs()[0]->Shape(); + if (q_bnsh_shape != nullptr && q_bnsh_shape->dim_size() == 4 && utils::HasDimValue(q_bnsh_shape->dim(2)) && + q_bnsh_shape->dim(2).dim_value() < 32) { + return false; + } + const auto& provider = softmax.GetExecutionProviderType(); for (const Node* node : {qk_matmul, scale_node, static_cast(pv_matmul), static_cast(out_transpose), q_transpose, k_transpose, v_transpose, diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.h b/onnxruntime/core/optimizer/multi_head_attention_fusion.h index e313b1817cd75..b9ca7657638cf 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.h +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.h @@ -28,6 +28,12 @@ is rewritten to The added Reshapes only fold the trailing (N, H) dims back to the BSD layout MultiHeadAttention expects, so they are free for EPs whose Reshape aliases its input. This keeps the [B, N, S, S] score tensor out of memory, which dominates the attention cost for long vision sequences. + +Besides the structural match, the fusion requires Q, K and V to share their batch dim, head count and head size, and +it only fires when the WebGPU MultiHeadAttention kernel would take its dense FlashAttention prefill path: a head size +that is a multiple of 8 and fits the guaranteed workgroup storage, and a query length that is not statically known to +be below 32. These cutoffs are specific to the WebGPU kernel, so the transformer is only meant to be registered for +the WebGPU EP. */ class MultiHeadAttentionFusion : public GraphTransformer { public: diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 3698eaab7fa80..7708cbc68df42 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -9526,41 +9526,76 @@ enum class VitAttentionScaleOp { struct VitAttentionTestConfig { VitAttentionScaleOp scale_op = VitAttentionScaleOp::Div; + // Mul only: put the scalar in input 0. + bool scale_is_first_input = false; // Unset means no Softmax axis attribute. std::optional softmax_axis = 3; std::vector k_perm = {0, 2, 3, 1}; int64_t seq_len = 64; int64_t head_size = 16; + // Fewer K/V heads than the 4 query heads makes MatMul broadcast them, as in multi-query attention. + int64_t kv_num_heads = 4; + // Projects K/V from a second input with batch 1; a larger query batch then makes MatMul broadcast them. + bool separate_kv_input = false; + int64_t q_batch = 1; + // Leaves the batch and sequence dims of the inputs unknown, as in models exported with dynamic axes. + bool dynamic_input_dims = false; + // Shape of the Reshape that merges the heads back; empty means {0, 0, hidden_size}. + std::vector merge_heads_shape; + bool fp16 = false; }; // Unmasked ViT/DINOv2 self-attention with separate Q, K and V projections: // x -> MatMul(Wq/Wk/Wv) -> Reshape[0,0,N,H] -> Transpose -> MatMul(Q, K^T) -> Div|Mul -> Softmax // -> MatMul(V) -> Transpose(0,2,1,3) -> Reshape[0,0,N*H] +// The fp16 variant uses constant data and is only meant for graph checks. static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttentionTestConfig& config = {}) { constexpr int64_t num_heads = 4; const int64_t head_size = config.head_size; const int64_t hidden_size = num_heads * head_size; - auto* input = builder.MakeInput({1, config.seq_len, hidden_size}, -1.0f, 1.0f); - auto* split_heads_shape = builder.Make1DInitializer({0, 0, num_heads, head_size}); - auto* merge_heads_shape = builder.Make1DInitializer({0, 0, hidden_size}); - const float sqrt_head_size = std::sqrt(static_cast(head_size)); - auto* scale = builder.MakeScalarInitializer( - config.scale_op == VitAttentionScaleOp::Div ? sqrt_head_size : 1.0f / sqrt_head_size); + auto make_input = [&](int64_t batch) -> NodeArg* { + if (config.dynamic_input_dims) { + return builder.MakeInput(std::optional>{{-1, -1, hidden_size}}); + } + const std::vector shape{batch, config.seq_len, hidden_size}; + if (config.fp16) { + return builder.MakeInput(std::optional>{shape}); + } + return builder.MakeInput(shape, -1.0f, 1.0f); + }; + auto make_weight = [&](int64_t output_size) -> NodeArg* { + const std::vector shape{hidden_size, output_size}; + if (config.fp16) { + return builder.MakeInitializer( + shape, std::vector(static_cast(hidden_size * output_size), MLFloat16(0.01f))); + } + return builder.MakeInitializer(shape, -0.2f, 0.2f); + }; - auto add_head_split = [&](const std::vector& perm) { - auto* weight = builder.MakeInitializer({hidden_size, hidden_size}, -0.2f, 0.2f); + NodeArg* q_input = make_input(config.q_batch); + NodeArg* kv_input = config.separate_kv_input ? make_input(1) : q_input; + const float sqrt_head_size = std::sqrt(static_cast(head_size)); + const float scale_value = config.scale_op == VitAttentionScaleOp::Div ? sqrt_head_size : 1.0f / sqrt_head_size; + NodeArg* scale = config.fp16 ? builder.MakeScalarInitializer(MLFloat16(scale_value)) + : builder.MakeScalarInitializer(scale_value); + NodeArg* merge_heads_shape = builder.Make1DInitializer( + config.merge_heads_shape.empty() ? std::vector{0, 0, hidden_size} : config.merge_heads_shape); + + auto add_head_split = [&](NodeArg* x, int64_t heads, const std::vector& perm) { + auto* weight = make_weight(heads * head_size); + auto* split_heads_shape = builder.Make1DInitializer({0, 0, heads, head_size}); auto* projected = builder.MakeIntermediate(); auto* bsnh = builder.MakeIntermediate(); auto* transposed = builder.MakeIntermediate(); - builder.AddNode("MatMul", {input, weight}, {projected}); + builder.AddNode("MatMul", {x, weight}, {projected}); builder.AddNode("Reshape", {projected, split_heads_shape}, {bsnh}); builder.AddNode("Transpose", {bsnh}, {transposed}).AddAttribute("perm", perm); return transposed; }; - auto* q = add_head_split({0, 2, 1, 3}); - auto* k = add_head_split(config.k_perm); - auto* v = add_head_split({0, 2, 1, 3}); + auto* q = add_head_split(q_input, num_heads, {0, 2, 1, 3}); + auto* k = add_head_split(kv_input, config.kv_num_heads, config.k_perm); + auto* v = add_head_split(kv_input, config.kv_num_heads, {0, 2, 1, 3}); auto* qk = builder.MakeIntermediate(); auto* scaled = builder.MakeIntermediate(); @@ -9569,7 +9604,13 @@ static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttent auto* context_bsnh = builder.MakeIntermediate(); auto* output = builder.MakeOutput(); builder.AddNode("MatMul", {q, k}, {qk}); - builder.AddNode(config.scale_op == VitAttentionScaleOp::Div ? "Div" : "Mul", {qk, scale}, {scaled}); + if (config.scale_op == VitAttentionScaleOp::Div) { + builder.AddNode("Div", {qk, scale}, {scaled}); + } else if (config.scale_is_first_input) { + builder.AddNode("Mul", {scale, qk}, {scaled}); + } else { + builder.AddNode("Mul", {qk, scale}, {scaled}); + } auto& softmax = builder.AddNode("Softmax", {scaled}, {probs}); if (config.softmax_axis.has_value()) { softmax.AddAttribute("axis", *config.softmax_axis); @@ -9612,51 +9653,135 @@ static Status CheckVitAttentionUnfused(Graph& graph) { return Status::OK(); } -static Status TestVitAttentionFusion(const VitAttentionTestConfig& config, int opset, const char* ep, - const std::function& check, const logging::Logger& logger) { +static Status AssignVitAttentionWebGpuEp(Graph& graph) { + return AssignVitAttentionEp(graph, kWebGpuExecutionProvider); +} + +static Status TestVitAttentionFusion(const VitAttentionTestConfig& config, int opset, + const std::function& check, const logging::Logger& logger, + const std::function& assign_ep = AssignVitAttentionWebGpuEp) { return TestGraphTransformer( [&config](ModelTestBuilder& builder) { BuildVitAttentionTestCase(builder, config); }, opset, logger, std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}), - TransformerLevel::Level2, 1, [ep](Graph& graph) { return AssignVitAttentionEp(graph, ep); }, check); + TransformerLevel::Level2, 1, assign_ep, check); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDivScale) { // Opset 11 matches exported ViT/DINOv2 models, where Softmax needs the explicit last-axis attribute. - ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, kWebGpuExecutionProvider, CheckVitAttentionFused, *logger_)); + ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, CheckVitAttentionFused, *logger_)); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitMulScaleDefaultAxis) { // From opset 13 Softmax defaults to the last axis. - ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Mul, std::nullopt}, 13, kWebGpuExecutionProvider, - CheckVitAttentionFused, *logger_)); + VitAttentionTestConfig config; + config.scale_op = VitAttentionScaleOp::Mul; + config.softmax_axis = std::nullopt; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 13, CheckVitAttentionFused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitMulScaleFirstInput) { + VitAttentionTestConfig config; + config.scale_op = VitAttentionScaleOp::Mul; + config.scale_is_first_input = true; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionFused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitFloat16) { + VitAttentionTestConfig config; + config.fp16 = true; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionFused, *logger_)); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitOpset11DefaultAxisNotFused) { // Before opset 13 the default Softmax axis is 1, which normalizes over heads and keys together. - ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, std::nullopt}, 11, kWebGpuExecutionProvider, - CheckVitAttentionUnfused, *logger_)); + VitAttentionTestConfig config; + config.softmax_axis = std::nullopt; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitUntransposedKeyNotFused) { // seq_len == head_size keeps Q x K a valid MatMul although it no longer computes Q x K^T. - ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, 3, {0, 2, 1, 3}, 32, 32}, 11, kWebGpuExecutionProvider, - CheckVitAttentionUnfused, *logger_)); + VitAttentionTestConfig config; + config.k_perm = {0, 2, 1, 3}; + config.seq_len = 32; + config.head_size = 32; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitMultiQueryNotFused) { + VitAttentionTestConfig config; + config.kv_num_heads = 1; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitBroadcastBatchNotFused) { + VitAttentionTestConfig config; + config.separate_kv_input = true; + config.q_batch = 2; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDynamicDimsSelfAttention) { + // The inferred batch dims are unknown, but Q, K and V project the same tensor, so they share them. + VitAttentionTestConfig config; + config.dynamic_input_dims = true; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionFused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDynamicDimsCrossAttentionNotFused) { + // With unknown dims and a separate K/V input, nothing shows that the batch dims match. + VitAttentionTestConfig config; + config.dynamic_input_dims = true; + config.separate_kv_input = true; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitZeroInMergedHeadDimNotFused) { + // The output Reshape would read the MultiHeadAttention output, whose dim 2 is N*H rather than N. + VitAttentionTestConfig config; + config.merge_heads_shape = {0, 0, 0, -1}; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitShortQueryNotFused) { // Queries shorter than 32 would run the split-reduce decode kernel. - ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, 3, {0, 2, 3, 1}, 31}, 11, kWebGpuExecutionProvider, - CheckVitAttentionUnfused, *logger_)); + VitAttentionTestConfig config; + config.seq_len = 31; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitHeadSizeNotMultipleOf8NotFused) { // Such head sizes can miss the FlashAttention path, whose fallback does not support graph capture. - ASSERT_STATUS_OK(TestVitAttentionFusion({VitAttentionScaleOp::Div, 3, {0, 2, 3, 1}, 64, 12}, 11, - kWebGpuExecutionProvider, CheckVitAttentionUnfused, *logger_)); + VitAttentionTestConfig config; + config.head_size = 12; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitLargeFloatHeadNotFused) { + // fp32 K/V tiles for a head size of 256 exceed the guaranteed 16 KB of workgroup storage. + VitAttentionTestConfig config; + config.head_size = 256; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitOtherEpNotFused) { - ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, kCpuExecutionProvider, CheckVitAttentionUnfused, *logger_)); + ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, CheckVitAttentionUnfused, *logger_, [](Graph& graph) { + return AssignVitAttentionEp(graph, kCpuExecutionProvider); + })); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitMixedEpNotFused) { + // One head-split Transpose left on the CPU EP. + ASSERT_STATUS_OK(TestVitAttentionFusion({}, 11, CheckVitAttentionUnfused, *logger_, [](Graph& graph) { + ORT_RETURN_IF_ERROR(AssignVitAttentionEp(graph, kWebGpuExecutionProvider)); + for (auto& node : graph.Nodes()) { + if (node.OpType() == "Transpose") { + node.SetExecutionProviderType(kCpuExecutionProvider); + break; + } + } + return Status::OK(); + })); } // Runs the graph on WebGPU with and without the fusion through the default session optimizers and compares the @@ -9671,7 +9796,9 @@ TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitWebGpuNumerics) { Model model("VitAttention", false, ModelMetaData(), PathString(), IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, {}, *logger_); ModelTestBuilder builder(model.MainGraph()); - BuildVitAttentionTestCase(builder, {VitAttentionScaleOp::Div, 3, {0, 2, 3, 1}, 97}); + VitAttentionTestConfig config; + config.seq_len = 97; + BuildVitAttentionTestCase(builder, config); builder.SetGraphOutputs(); ASSERT_STATUS_OK(model.MainGraph().Resolve()); std::string model_data; From 08a398efe6e967f95613c208181ab7e4c73a55df Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Wed, 4 Nov 2026 14:45:20 +0800 Subject: [PATCH 4/7] Tighten the shared-projection batch proof in MultiHeadAttentionFusion 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 --- .../optimizer/multi_head_attention_fusion.cc | 54 +++++++++------- .../test/optimizer/graph_transform_test.cc | 62 ++++++++++++++++--- 2 files changed, 87 insertions(+), 29 deletions(-) diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc index ee76e8dc7e4cf..83a5882f467ce 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -83,19 +83,31 @@ const Node* MatchHeadSplit(const Graph& graph, const Node& consumer, int input_i return transpose; } +bool HasRankAtMost(const NodeArg& arg, int rank) { + const auto* shape = arg.Shape(); + return shape != nullptr && shape->dim_size() <= rank; +} + // Walks up from a head-split Reshape to the tensor its projection reads: through an optional bias Add, a MatMul or -// Gemm with a constant weight, and the flattening Reshape that MatMulAddFusion puts in front of a Gemm. +// Gemm, and the flattening Reshape that MatMulAddFusion puts in front of a Gemm. Only projections that keep the +// leading dims of that tensor are followed: a 2D constant weight and a bias of rank 1 or less, since higher-rank +// constants could broadcast new leading dims. const NodeArg* GetProjectionSource(const Graph& graph, const Node& head_split_reshape) { const Node* node = graph_utils::GetInputNode(head_split_reshape, 0); if (node != nullptr && node->OpType() == "Add" && node->InputDefs().size() == 2) { const int data_index = graph_utils::IsConstantInitializer(graph, node->InputDefs()[1]->Name(), true) ? 0 : 1; - if (!graph_utils::IsConstantInitializer(graph, node->InputDefs()[1 - data_index]->Name(), true)) { + const NodeArg& bias = *node->InputDefs()[1 - data_index]; + if (!graph_utils::IsConstantInitializer(graph, bias.Name(), true) || !HasRankAtMost(bias, 1)) { return nullptr; } node = graph_utils::GetInputNode(*node, data_index); } - if (node == nullptr || (node->OpType() != "MatMul" && node->OpType() != "Gemm") || node->InputDefs().size() < 2 || - !graph_utils::IsConstantInitializer(graph, node->InputDefs()[1]->Name(), true)) { + if (node == nullptr || (node->OpType() != "MatMul" && node->OpType() != "Gemm") || node->InputDefs().size() < 2) { + return nullptr; + } + const NodeArg& weight = *node->InputDefs()[1]; + if (!graph_utils::IsConstantInitializer(graph, weight.Name(), true) || !weight.Shape() || + weight.Shape()->dim_size() != 2) { return nullptr; } const Node* flatten = graph_utils::GetInputNode(*node, 0); @@ -105,27 +117,27 @@ const NodeArg* GetProjectionSource(const Graph& graph, const Node& head_split_re return node->InputDefs()[0]; } -// Whether two head-split Reshapes produce the same batch dim. That holds when both project the same tensor and keep -// the same leading reshape dims, as in self-attention; otherwise it needs statically comparable inferred dims. +// Whether two head-split Reshapes produce the same batch dim. Statically comparable inferred dims decide it. When +// they are unknown (dynamic axes), it still holds if both project the same tensor through leading-dim-preserving +// projections and keep the same leading reshape dims, as in self-attention. bool HaveSameBatchDim(const Graph& graph, const Node& a, const Node& b) { - InlinedVector a_dims, b_dims; - const NodeArg* a_source = GetProjectionSource(graph, a); - if (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]) { - return true; - } - const auto* a_shape = a.OutputDefs()[0]->Shape(); const auto* b_shape = b.OutputDefs()[0]->Shape(); - if (a_shape == nullptr || b_shape == nullptr || a_shape->dim_size() != 4 || b_shape->dim_size() != 4) { - return false; - } - const auto& a_dim = a_shape->dim(0); - const auto& b_dim = b_shape->dim(0); - if (utils::HasDimValue(a_dim) && utils::HasDimValue(b_dim)) { - return a_dim.dim_value() == b_dim.dim_value(); + if (a_shape != nullptr && b_shape != nullptr && a_shape->dim_size() == 4 && b_shape->dim_size() == 4) { + const auto& a_dim = a_shape->dim(0); + const auto& b_dim = b_shape->dim(0); + if (utils::HasDimValue(a_dim) && utils::HasDimValue(b_dim)) { + return a_dim.dim_value() == b_dim.dim_value(); + } + if (utils::HasDimParam(a_dim) && utils::HasDimParam(b_dim) && a_dim.dim_param() == b_dim.dim_param()) { + return true; + } } - return utils::HasDimParam(a_dim) && utils::HasDimParam(b_dim) && a_dim.dim_param() == b_dim.dim_param(); + + InlinedVector 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]; } bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet& compatible_eps, diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index c97dbc0a5c298..54bf59832cf45 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -9740,6 +9740,10 @@ struct VitAttentionTestConfig { int64_t q_batch = 1; // Leaves the batch and sequence dims of the inputs unknown, as in models exported with dynamic axes. bool dynamic_input_dims = false; + // A non-zero value gives the Q weight this leading batch dim, which MatMul broadcasts into the Q batch. + int64_t q_weight_batch = 0; + // Adds bias Adds after the projections: Q uses this shape, K/V use their 1D hidden size. + std::vector q_bias_shape; // Shape of the Reshape that merges the heads back; empty means {0, 0, hidden_size}. std::vector merge_heads_shape; bool fp16 = false; @@ -9764,11 +9768,14 @@ static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttent } return builder.MakeInput(shape, -1.0f, 1.0f); }; - auto make_weight = [&](int64_t output_size) -> NodeArg* { - const std::vector shape{hidden_size, output_size}; + auto make_weight = [&](int64_t output_size, int64_t batch = 0) -> NodeArg* { + std::vector shape{hidden_size, output_size}; + if (batch > 0) { + shape.insert(shape.begin(), batch); + } if (config.fp16) { return builder.MakeInitializer( - shape, std::vector(static_cast(hidden_size * output_size), MLFloat16(0.01f))); + shape, std::vector(static_cast(TensorShape(shape).Size()), MLFloat16(0.01f))); } return builder.MakeInitializer(shape, -0.2f, 0.2f); }; @@ -9782,20 +9789,28 @@ static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttent NodeArg* merge_heads_shape = builder.Make1DInitializer( config.merge_heads_shape.empty() ? std::vector{0, 0, hidden_size} : config.merge_heads_shape); - auto add_head_split = [&](NodeArg* x, int64_t heads, const std::vector& perm) { - auto* weight = make_weight(heads * head_size); + auto add_head_split = [&](NodeArg* x, int64_t heads, const std::vector& perm, int64_t weight_batch, + const std::vector& bias_shape) { + auto* weight = make_weight(heads * head_size, weight_batch); auto* split_heads_shape = builder.Make1DInitializer({0, 0, heads, head_size}); auto* projected = builder.MakeIntermediate(); auto* bsnh = builder.MakeIntermediate(); auto* transposed = builder.MakeIntermediate(); builder.AddNode("MatMul", {x, weight}, {projected}); + if (!config.q_bias_shape.empty()) { + auto* bias = builder.MakeInitializer(bias_shape, -0.1f, 0.1f); + auto* biased = builder.MakeIntermediate(); + builder.AddNode("Add", {bias, projected}, {biased}); + projected = biased; + } builder.AddNode("Reshape", {projected, split_heads_shape}, {bsnh}); builder.AddNode("Transpose", {bsnh}, {transposed}).AddAttribute("perm", perm); return transposed; }; - auto* q = add_head_split(q_input, num_heads, {0, 2, 1, 3}); - auto* k = add_head_split(kv_input, config.kv_num_heads, config.k_perm); - auto* v = add_head_split(kv_input, config.kv_num_heads, {0, 2, 1, 3}); + const std::vector kv_bias_shape{config.kv_num_heads * head_size}; + auto* q = add_head_split(q_input, num_heads, {0, 2, 1, 3}, config.q_weight_batch, config.q_bias_shape); + auto* k = add_head_split(kv_input, config.kv_num_heads, config.k_perm, 0, kv_bias_shape); + auto* v = add_head_split(kv_input, config.kv_num_heads, {0, 2, 1, 3}, 0, kv_bias_shape); auto* qk = builder.MakeIntermediate(); auto* scaled = builder.MakeIntermediate(); @@ -9936,6 +9951,37 @@ TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDynamicDimsCrossAtte ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); } +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDynamicDimsWithBiasFused) { + // 1D biases keep the leading dims, so the shared projection source still proves matching batches. + VitAttentionTestConfig config; + config.dynamic_input_dims = true; + config.q_bias_shape = {64}; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionFused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitSharedSourceBatchedWeightNotFused) { + // Q, K and V project the same [1, S, D] input, but the [2, D, D] Q weight broadcasts Q to batch 2. + VitAttentionTestConfig config; + config.q_weight_batch = 2; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitSharedSourceBatchedWeightDynamicDimsNotFused) { + // As above with unknown inferred batch dims, so only the projection-source check can reject it. + VitAttentionTestConfig config; + config.dynamic_input_dims = true; + config.q_weight_batch = 2; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitSharedSourceBatchedBiasDynamicDimsNotFused) { + // A [2, 1, D] Q bias broadcasts Q to batch 2 even though all three project the same input. + VitAttentionTestConfig config; + config.dynamic_input_dims = true; + config.q_bias_shape = {2, 1, 64}; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitZeroInMergedHeadDimNotFused) { // The output Reshape would read the MultiHeadAttention output, whose dim 2 is N*H rather than N. VitAttentionTestConfig config; From 489abac6260aa70691a7400099fb05ec1331748d Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Wed, 4 Nov 2026 14:54:27 +0800 Subject: [PATCH 5/7] Cap fused head count at WebGPU's guaranteed workgroup size 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 --- .../optimizer/multi_head_attention_fusion.cc | 4 +++ .../test/optimizer/graph_transform_test.cc | 26 ++++++++++++++----- 2 files changed, 24 insertions(+), 6 deletions(-) diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc index 83a5882f467ce..f3a1c9885b228 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -243,6 +243,10 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet // The remaining checks keep the fused node on the WebGPU kernel's dense FlashAttention prefill path. Its fallbacks // are slower than the unfused graph, and the non-flash one does not support graph capture. + // The kernel rejects more heads than maxComputeInvocationsPerWorkgroup, whose guaranteed minimum is 256. + if (q_heads > 256) { + return false; + } // FlashAttention requires a head size that is a multiple of 8 on every adapter. if (q_head_size % 8 != 0) { return false; diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 54bf59832cf45..5e034849e517b 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -9732,9 +9732,11 @@ struct VitAttentionTestConfig { std::optional softmax_axis = 3; std::vector k_perm = {0, 2, 3, 1}; int64_t seq_len = 64; + int64_t num_heads = 4; int64_t head_size = 16; - // Fewer K/V heads than the 4 query heads makes MatMul broadcast them, as in multi-query attention. - int64_t kv_num_heads = 4; + // Fewer K/V heads than the query heads makes MatMul broadcast them, as in multi-query attention. + // 0 means the same as num_heads. + int64_t kv_num_heads = 0; // Projects K/V from a second input with batch 1; a larger query batch then makes MatMul broadcast them. bool separate_kv_input = false; int64_t q_batch = 1; @@ -9754,7 +9756,8 @@ struct VitAttentionTestConfig { // -> MatMul(V) -> Transpose(0,2,1,3) -> Reshape[0,0,N*H] // The fp16 variant uses constant data and is only meant for graph checks. static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttentionTestConfig& config = {}) { - constexpr int64_t num_heads = 4; + const int64_t num_heads = config.num_heads; + const int64_t kv_num_heads = config.kv_num_heads == 0 ? num_heads : config.kv_num_heads; const int64_t head_size = config.head_size; const int64_t hidden_size = num_heads * head_size; @@ -9807,10 +9810,10 @@ static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttent builder.AddNode("Transpose", {bsnh}, {transposed}).AddAttribute("perm", perm); return transposed; }; - const std::vector kv_bias_shape{config.kv_num_heads * head_size}; + const std::vector kv_bias_shape{kv_num_heads * head_size}; auto* q = add_head_split(q_input, num_heads, {0, 2, 1, 3}, config.q_weight_batch, config.q_bias_shape); - auto* k = add_head_split(kv_input, config.kv_num_heads, config.k_perm, 0, kv_bias_shape); - auto* v = add_head_split(kv_input, config.kv_num_heads, {0, 2, 1, 3}, 0, kv_bias_shape); + auto* k = add_head_split(kv_input, kv_num_heads, config.k_perm, 0, kv_bias_shape); + auto* v = add_head_split(kv_input, kv_num_heads, {0, 2, 1, 3}, 0, kv_bias_shape); auto* qk = builder.MakeIntermediate(); auto* scaled = builder.MakeIntermediate(); @@ -9989,6 +9992,17 @@ TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitZeroInMergedHeadDimN ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); } +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitTooManyHeadsNotFused) { + // The WebGPU kernel rejects more heads than maxComputeInvocationsPerWorkgroup, which is only guaranteed to be 256. + // The other guards pass: the head size is a multiple of 8 and the query is not short. + VitAttentionTestConfig config; + config.num_heads = 257; + config.head_size = 8; + config.seq_len = 32; + config.fp16 = true; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitShortQueryNotFused) { // Queries shorter than 32 would run the split-reduce decode kernel. VitAttentionTestConfig config; From be444fa3d38ba9f5b2e2c255d95c432870b8beb0 Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Wed, 4 Nov 2026 15:00:39 +0800 Subject: [PATCH 6/7] Document the Qualcomm exception to the workgroup-storage bound 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 --- onnxruntime/core/optimizer/multi_head_attention_fusion.cc | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc index f3a1c9885b228..e5b4add36984c 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -251,7 +251,9 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet if (q_head_size % 8 != 0) { return false; } - // Dense prefill needs K/V tiles of 16 keys to fit in the 16 KB of workgroup storage every adapter guarantees. + // Dense prefill needs K/V tiles of 16 keys to fit in the 16 KB of workgroup storage every adapter guarantees. That + // matches the kernel except on Qualcomm, whose extra output tile can exceed 16 KB on 16 KB adapters; the adapter is + // unknown here, and missing dense prefill only costs speed, so the bound is not tightened for that case. const int64_t element_size = q_type->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ? 2 : 4; if (2 * 16 * q_head_size * element_size > 16384) { From ff959265ce242de14b10fca57104dcab952a82ba Mon Sep 17 00:00:00 2001 From: Xinghua Cao Date: Thu, 5 Nov 2026 15:49:23 +0800 Subject: [PATCH 7/7] Address second-round review of MultiHeadAttentionFusion - 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 --- .../optimizer/multi_head_attention_fusion.cc | 39 +++++----- .../optimizer/multi_head_attention_fusion.h | 3 +- .../test/optimizer/graph_transform_test.cc | 75 ++++++++++++++++++- 3 files changed, 94 insertions(+), 23 deletions(-) diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc index e5b4add36984c..42b63a6ef0fbe 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.cc +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.cc @@ -88,10 +88,10 @@ bool HasRankAtMost(const NodeArg& arg, int rank) { return shape != nullptr && shape->dim_size() <= rank; } -// Walks up from a head-split Reshape to the tensor its projection reads: through an optional bias Add, a MatMul or -// Gemm, and the flattening Reshape that MatMulAddFusion puts in front of a Gemm. Only projections that keep the -// leading dims of that tensor are followed: a 2D constant weight and a bias of rank 1 or less, since higher-rank -// constants could broadcast new leading dims. +// Walks up from a head-split Reshape through an optional bias Add and a MatMul to the tensor the projection reads. +// Only projections that keep the leading dims of that tensor are followed: a 2D constant weight and a bias of rank 1 +// or less, since higher-rank constants could broadcast new leading dims. Gemm projections are not followed: the +// MatMulAddFusion that produces them requires static shapes, whose batch dims are compared directly. const NodeArg* GetProjectionSource(const Graph& graph, const Node& head_split_reshape) { const Node* node = graph_utils::GetInputNode(head_split_reshape, 0); if (node != nullptr && node->OpType() == "Add" && node->InputDefs().size() == 2) { @@ -102,7 +102,7 @@ const NodeArg* GetProjectionSource(const Graph& graph, const Node& head_split_re } node = graph_utils::GetInputNode(*node, data_index); } - if (node == nullptr || (node->OpType() != "MatMul" && node->OpType() != "Gemm") || node->InputDefs().size() < 2) { + if (node == nullptr || !graph_utils::IsSupportedOptypeVersionAndDomain(*node, "MatMul", kMatMulVersions)) { return nullptr; } const NodeArg& weight = *node->InputDefs()[1]; @@ -110,10 +110,6 @@ const NodeArg* GetProjectionSource(const Graph& graph, const Node& head_split_re weight.Shape()->dim_size() != 2) { return nullptr; } - const Node* flatten = graph_utils::GetInputNode(*node, 0); - if (node->OpType() == "Gemm" && flatten != nullptr && flatten->OpType() == "Reshape") { - return flatten->InputDefs()[0]; - } return node->InputDefs()[0]; } @@ -174,10 +170,14 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet scalar_index = 0; value = GetScalarConstant(graph, *scale_node->InputDefs()[0]); } - if (!value.has_value() || *value == 0.0f || !std::isfinite(*value)) { + if (!value.has_value() || *value == 0.0f) { return false; } scale = is_div ? 1.0f / *value : *value; + // Also catches a tiny divisor whose reciprocal overflows. + if (!std::isfinite(scale)) { + return false; + } qk_matmul = graph_utils::GetInputNode(*scale_node, 1 - scalar_index); } @@ -242,24 +242,29 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet } // The remaining checks keep the fused node on the WebGPU kernel's dense FlashAttention prefill path. Its fallbacks - // are slower than the unfused graph, and the non-flash one does not support graph capture. - // The kernel rejects more heads than maxComputeInvocationsPerWorkgroup, whose guaranteed minimum is 256. + // are slower than the unfused graph, and the non-flash one does not support graph capture. They mirror kernel logic + // in contrib_ops/webgpu/bert/flash_attention.cc and contrib_ops/cpu/bert/multihead_attention_helper.h, so keep them + // in sync with the functions named below. + // multihead_attention_helper::CheckInputs rejects more heads than maxComputeInvocationsPerWorkgroup, whose + // guaranteed minimum is 256. if (q_heads > 256) { return false; } - // FlashAttention requires a head size that is a multiple of 8 on every adapter. + // CanApplyFlashAttention requires a head size that is a multiple of 8 on every adapter. if (q_head_size % 8 != 0) { return false; } - // Dense prefill needs K/V tiles of 16 keys to fit in the 16 KB of workgroup storage every adapter guarantees. That - // matches the kernel except on Qualcomm, whose extra output tile can exceed 16 KB on 16 KB adapters; the adapter is + // CanApplyDensePrefill needs K/V tiles of 16 keys to fit in the 16 KB of workgroup storage every adapter + // guarantees (DensePrefillFitsWorkgroupStorage, SelectDensePrefillMaxKStep). That matches the kernel except on + // Qualcomm, whose extra output tile can exceed 16 KB on 16 KB adapters; the adapter is // unknown here, and missing dense prefill only costs speed, so the bound is not tightened for that case. const int64_t element_size = q_type->tensor_type().elem_type() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ? 2 : 4; if (2 * 16 * q_head_size * element_size > 16384) { return false; } - // Queries shorter than 32 go to the split-reduce decode kernel, and leave only a small score tensor to avoid. + // CanApplyDensePrefill sends queries shorter than 32 to the split-reduce decode kernel, and they leave only a small + // score tensor to avoid. // A symbolic query length cannot be checked here, so it is fused. const auto* q_bnsh_shape = q_transpose->OutputDefs()[0]->Shape(); if (q_bnsh_shape != nullptr && q_bnsh_shape->dim_size() == 4 && utils::HasDimValue(q_bnsh_shape->dim(2)) && @@ -291,7 +296,7 @@ bool TryFuse(Graph& graph, Node& softmax, const InlinedHashSet tensor_type.mutable_tensor_type()->set_elem_type(q_type->tensor_type().elem_type()); auto add_to_bsd = [&](const Node& reshape) -> Node& { - NodeArg* bsnh = const_cast(reshape.OutputDefs()[0]); + NodeArg* bsnh = graph.GetNode(reshape.Index())->MutableOutputDefs()[0]; NodeArg& bsd = graph.GetOrCreateNodeArg(graph.GenerateNodeArgName(bsnh->Name() + "_bsd"), &tensor_type); Node& node = graph.AddNode(graph.GenerateNodeName(reshape.Name() + "/MhaBsd"), "Reshape", "Fold heads into the hidden dim for MultiHeadAttention", {bsnh, bsd_shape_arg}, {&bsd}); diff --git a/onnxruntime/core/optimizer/multi_head_attention_fusion.h b/onnxruntime/core/optimizer/multi_head_attention_fusion.h index b9ca7657638cf..eb4f6466f5665 100644 --- a/onnxruntime/core/optimizer/multi_head_attention_fusion.h +++ b/onnxruntime/core/optimizer/multi_head_attention_fusion.h @@ -37,8 +37,7 @@ the WebGPU EP. */ class MultiHeadAttentionFusion : public GraphTransformer { public: - explicit MultiHeadAttentionFusion( - const InlinedHashSet& compatible_execution_providers = {}) noexcept + explicit MultiHeadAttentionFusion(const InlinedHashSet& compatible_execution_providers) noexcept : GraphTransformer("MultiHeadAttentionFusion", compatible_execution_providers) { } diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 5e034849e517b..813555fcb4f8f 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -9722,12 +9722,16 @@ TEST_F(GraphTransformationTests, AttentionFusionMobileClipMhaProjectionRewriteFa enum class VitAttentionScaleOp { Div, Mul, + // MatMul feeds Softmax directly. + None, }; struct VitAttentionTestConfig { VitAttentionScaleOp scale_op = VitAttentionScaleOp::Div; // Mul only: put the scalar in input 0. bool scale_is_first_input = false; + // Overrides the Div/Mul scalar, which defaults to the usual sqrt(head_size) divisor or its reciprocal. + std::optional scale_value; // Unset means no Softmax axis attribute. std::optional softmax_axis = 3; std::vector k_perm = {0, 2, 3, 1}; @@ -9786,7 +9790,8 @@ static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttent NodeArg* q_input = make_input(config.q_batch); NodeArg* kv_input = config.separate_kv_input ? make_input(1) : q_input; const float sqrt_head_size = std::sqrt(static_cast(head_size)); - const float scale_value = config.scale_op == VitAttentionScaleOp::Div ? sqrt_head_size : 1.0f / sqrt_head_size; + const float scale_value = config.scale_value.value_or( + config.scale_op == VitAttentionScaleOp::Div ? sqrt_head_size : 1.0f / sqrt_head_size); NodeArg* scale = config.fp16 ? builder.MakeScalarInitializer(MLFloat16(scale_value)) : builder.MakeScalarInitializer(scale_value); NodeArg* merge_heads_shape = builder.Make1DInitializer( @@ -9824,6 +9829,8 @@ static void BuildVitAttentionTestCase(ModelTestBuilder& builder, const VitAttent builder.AddNode("MatMul", {q, k}, {qk}); if (config.scale_op == VitAttentionScaleOp::Div) { builder.AddNode("Div", {qk, scale}, {scaled}); + } else if (config.scale_op == VitAttentionScaleOp::None) { + scaled = qk; } else if (config.scale_is_first_input) { builder.AddNode("Mul", {scale, qk}, {scaled}); } else { @@ -9845,25 +9852,50 @@ static Status AssignVitAttentionEp(Graph& graph, const char* ep) { return Status::OK(); } -static Status CheckVitAttentionFused(Graph& graph) { +static Status CheckVitAttentionFusedWithScale(Graph& graph, float expected_scale) { auto op_to_count = CountOpsInGraph(graph); ORT_RETURN_IF_NOT(op_to_count["com.microsoft.MultiHeadAttention"] == 1, "expected one MultiHeadAttention"); ORT_RETURN_IF_NOT(op_to_count["Softmax"] == 0, "Softmax was not fused"); ORT_RETURN_IF_NOT(op_to_count["Transpose"] == 0, "head Transposes were not removed"); ORT_RETURN_IF_NOT(op_to_count["MatMul"] == 3, "only the Q/K/V projections should remain"); ORT_RETURN_IF_NOT(op_to_count["Div"] == 0 && op_to_count["Mul"] == 0, "the scale was not folded"); + // Three head-split Reshapes, the three added BSD Reshapes and the output Reshape. + ORT_RETURN_IF_NOT(op_to_count["Reshape"] == 7, "unexpected Reshape count"); for (const auto& node : graph.Nodes()) { if (node.OpType() == "MultiHeadAttention") { const auto& attrs = node.GetAttributes(); ORT_RETURN_IF_NOT(attrs.at("num_heads").i() == 4, "unexpected num_heads"); - ORT_RETURN_IF_NOT(std::abs(attrs.at("scale").f() - 0.25f) < 1e-6f, "unexpected scale"); + ORT_RETURN_IF_NOT(std::abs(attrs.at("scale").f() - expected_scale) < 1e-6f, "unexpected scale"); ORT_RETURN_IF_NOT(node.GetExecutionProviderType() == kWebGpuExecutionProvider, "unexpected EP"); + // Q, K and V must be the [0, 0, -1] Reshapes of the head-split [B, S, N, H] tensors, not those tensors. + for (int i = 0; i < 3; ++i) { + const Node* bsd = graph_utils::GetInputNode(node, i); + ORT_RETURN_IF_NOT(bsd != nullptr && bsd->OpType() == "Reshape", "MultiHeadAttention input ", i, + " is not a BSD Reshape"); + InlinedVector bsd_shape; + const InlinedVector expected_bsd_shape{0, 0, -1}; + const bool is_bsd_shape = + optimizer_utils::AppendTensorFromInitializer(graph, *bsd->InputDefs()[1], bsd_shape, true) && + bsd_shape == expected_bsd_shape; + ORT_RETURN_IF_NOT(is_bsd_shape, "MultiHeadAttention input ", i, " is not reshaped to [0, 0, -1]"); + const Node* head_split = graph_utils::GetInputNode(*bsd, 0); + ORT_RETURN_IF_NOT(head_split != nullptr && head_split->OpType() == "Reshape", "MultiHeadAttention input ", i, + " does not come from a head-split Reshape"); + } } } return Status::OK(); } +static Status CheckVitAttentionFused(Graph& graph) { + return CheckVitAttentionFusedWithScale(graph, 0.25f); +} + +static Status CheckVitAttentionFusedUnscaled(Graph& graph) { + return CheckVitAttentionFusedWithScale(graph, 1.0f); +} + static Status CheckVitAttentionUnfused(Graph& graph) { auto op_to_count = CountOpsInGraph(graph); ORT_RETURN_IF_NOT(op_to_count["com.microsoft.MultiHeadAttention"] == 0, "unexpected MultiHeadAttention"); @@ -9904,6 +9936,40 @@ TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitMulScaleFirstInput) ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionFused, *logger_)); } +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitNoScale) { + VitAttentionTestConfig config; + config.scale_op = VitAttentionScaleOp::None; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionFusedUnscaled, *logger_)); +} + +// Fails when ONNX bumps the since_version of an op in the pattern beyond the version lists in +// onnxruntime/core/optimizer/multi_head_attention_fusion.cc. +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionCurrentOpsetTest) { + const int current_opset = GetCurrentOnnxOpset(); + auto check = [current_opset](Graph& graph) -> Status { + Status status = CheckVitAttentionFused(graph); + if (!status.IsOK()) { + return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "MultiHeadAttentionFusion failed at opset ", current_opset, ": ", + status.ErrorMessage(), ". Update the version lists in ", + "onnxruntime/core/optimizer/multi_head_attention_fusion.cc."); + } + return Status::OK(); + }; + // As in GeluFusionCurrentOpsetTest, allow the current opset even while it is still under development in ONNX. + ASSERT_STATUS_OK(TestGraphTransformer( + [](ModelTestBuilder& builder) { BuildVitAttentionTestCase(builder); }, current_opset, *logger_, + std::make_unique(InlinedHashSet{kWebGpuExecutionProvider}), + TransformerLevel::Level2, 1, AssignVitAttentionWebGpuEp, check, + ModelOptions{kAllowReleasedOpsetsOnly, /*strict_shape_type_inference*/ false})); +} + +TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitOverflowingScaleNotFused) { + // The reciprocal of a denormal divisor overflows to inf, so there is no finite scale attribute to emit. + VitAttentionTestConfig config; + config.scale_value = 1e-40f; + ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_)); +} + TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitFloat16) { VitAttentionTestConfig config; config.fp16 = true; @@ -9963,7 +10029,8 @@ TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitDynamicDimsWithBiasF } TEST_F(GraphTransformationTests, MultiHeadAttentionFusionVitSharedSourceBatchedWeightNotFused) { - // Q, K and V project the same [1, S, D] input, but the [2, D, D] Q weight broadcasts Q to batch 2. + // Q, K and V project the same [1, S, D] input, but the [2, D, D] Q weight broadcasts Q to batch 2. The static + // inferred batch dims (2 vs. 1) reject it before the projection-source check is consulted. VitAttentionTestConfig config; config.q_weight_batch = 2; ASSERT_STATUS_OK(TestVitAttentionFusion(config, 11, CheckVitAttentionUnfused, *logger_));