Skip to content

[WebGPU] Fuse unmasked ViT attention into MultiHeadAttention #33854

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
xhcao wants to merge 8 commits into
base: main
Choose a base branch
from
Open

[WebGPU] Fuse unmasked ViT attention into MultiHeadAttention #33854

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
1b20fa9
[WebGPU] Fuse unmasked ViT attention into MultiHeadAttention
xhcao Nov 2, 2026
30c7987
Clarify why the output Reshape can be kept in MultiHeadAttentionFusion
xhcao Nov 3, 2026
6206d02
Address review findings in MultiHeadAttentionFusion
xhcao Nov 3, 2026
d4e0eb6
Merge remote-tracking branch 'upstream/main' into webgpu-vit-mha-fusion
xhcao Nov 4, 2026
08a398e
Tighten the shared-projection batch proof in MultiHeadAttentionFusion
xhcao Nov 4, 2026
489abac
Cap fused head count at WebGPU's guaranteed workgroup size
xhcao Nov 4, 2026
be444fa
Document the Qualcomm exception to the workgroup-storage bound
xhcao Nov 4, 2026
ff95926
Address second-round review of MultiHeadAttentionFusion
xhcao Nov 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
55 changes: 31 additions & 24 deletions onnxruntime/core/optimizer/graph_transformer_utils.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -209,6 +210,26 @@ std::unique_ptr<RuleBasedGraphTransformer> 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<std::unique_ptr<GraphTransformer>> GenerateTransformers(
TransformerLevel level,
const SessionOptions& session_options,
Expand Down Expand Up @@ -416,6 +437,10 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
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)) {
transformers.emplace_back(std::make_unique<MultiHeadAttentionFusion>(
InlinedHashSet<std::string_view>{onnxruntime::kWebGpuExecutionProvider}));
}
transformers.emplace_back(std::make_unique<EmbedLayerNormFusion>(cpu_acl_cuda_dml_eps));
transformers.emplace_back(std::make_unique<GatherSliceToSplitFusion>(cpu_cuda_eps));
transformers.emplace_back(std::make_unique<GatherToSliceFusion>(cpu_cuda_eps));
Expand Down Expand Up @@ -468,30 +493,12 @@ InlinedVector<std::unique_ptr<GraphTransformer>> GenerateTransformers(
transformers.emplace_back(std::make_unique<GroupQueryAttentionPreNormFusion>(
InlinedHashSet<std::string_view>{onnxruntime::kCudaExecutionProvider,
onnxruntime::kWebGpuExecutionProvider}));
bool has_matmul_nbits_qlora_kernel = false;
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_qlora_kernel = has_webgpu_kernel("MatMulNBitsQLoRA");
has_matmul_nbits_mlp_kernel = has_webgpu_kernel("MatMulNBitsMlp");
has_matmul_nbits_qkv_kernel = has_webgpu_kernel("MatMulNBitsQkv");
}
}
}
const bool has_matmul_nbits_qlora_kernel =
WebGpuEpHasContribKernel(execution_providers, "MatMulNBitsQLoRA", logger);
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_qlora_kernel) {
transformers.emplace_back(std::make_unique<MatMulNBitsQLoRAFusion>());
}
Expand Down
Loading