Repository navigation
[WebGPU] Enable SimplifiedLayerNormFusion for the WebGPU EP - #34011
Open
Danny On The Air (dannyota) wants to merge 1 commit into
Open
Danny On The Air (dannyota) wants to merge 1 commit into
Danny On The Air (dannyota) wants to merge 1 commit into
Conversation
WebGPU registers SimplifiedLayerNormalization, but the fusion only ran for CPU and CUDA, so decomposed RMSNorm in WebGPU-assigned graphs stayed as Pow/ReduceMean/Add/Sqrt/Div/Mul. Run the fusion for WebGPU, including the Cast patterns already allowed for CUDA. For WebGPU nodes, fuse only when X and scale are both fp32 or both fp16, the scale covers the full last dimension, and the build is not a training build, since the kernel and the downstream MatMulNBits fusions do not handle the other cases. Also fix scale selection for rank-one X on every EP: with Mul(scale, normalized) both operands are 1D and the fusion could pick the normalized output as scale, leaving the fused node with an unresolved input and failing session initialization. Signed-off-by: Danny <dannyontheair@outlook.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The fusion can change normalization results for keepdims=0 reductions and Cast patterns, including loss of input rounding and FP16 overflow.
3 open issues
What changed in this PR
Enables decomposed RMSNorm fusion for the WebGPU execution provider, addressing the SimplifiedLayerNorm portion of #29841, and fixes rank-one scale selection across providers.
Changes:
- Registers WebGPU for
SimplifiedLayerNormFusion, including Cast patterns. - Adds WebGPU type, shape, and training-build restrictions.
- Adds optimizer regressions and end-to-end WebGPU correctness tests.
| File | Description |
|---|---|
onnxruntime/test/providers/webgpu/layer_norm_test.cc |
Adds FP32, FP16, and rank-one execution tests. |
onnxruntime/test/optimizer/graph_transform_test_layernorm.cc |
Tests fusion eligibility, rejected cases, registration, and scale selection. |
onnxruntime/core/optimizer/layer_norm_fusion.cc |
Enables WebGPU patterns, adds eligibility guards, and fixes rank-one scale selection. |
onnxruntime/core/optimizer/graph_transformer_utils.cc |
Adds WebGPU to the fusion’s compatible providers. |
Review effort: Balanced
| NodeArg* x_input = has_leading_cast ? graph.GetNode(p_pow_input_node->Index())->MutableInputDefs()[0] | ||
| : pow_node.MutableInputDefs()[0]; | ||
|
|
||
| if (pow_node.GetExecutionProviderType() == kWebGpuExecutionProvider) { |
| !scale_type->has_tensor_type()) { | ||
| continue; | ||
| } | ||
| const auto element_type = x_type->tensor_type().elem_type(); |
Comment on lines
+865
to
+867
| if ((element_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT && | ||
| element_type != ONNX_NAMESPACE::TensorProto_DataType_FLOAT16) || | ||
| element_type != scale_type->tensor_type().elem_type()) { |
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
WebGPU registers
SimplifiedLayerNormalization, butSimplifiedLayerNormFusiononly ran for CPU and CUDA, so decomposed RMSNorm (Pow/ReduceMean/Add/Sqrt/Div/Mul) in WebGPU-assigned graphs was never fused.Mul(scale, normalized)both operands are 1D, and the fusion could pick the normalized output as scale, leaving the fused node with an unresolved input. On main this makes a CPU session fail to initialize withInvalid model. Node input ... is not a graph input.Testing (macOS arm64, Metal,
--use_webgpu)onnxruntime_test_all: SimplifiedLayerNorm, LayerNorm, SkipLayerNorm, GQA pre-norm and MatMulNBits MLP/QKV fusion tests pass (91; 4 CUDA-only skipped). New tests cover WebGPU fp32/fp16 with and without Casts, the rejected cases, rank-one X with both Mul operand orders on CPU and WebGPU, and the production transformer list.onnxruntime_provider_test: WebGPU layer norm tests pass (147), including new end-to-end fp32, fp16 and rank-one tests that check the fused node runs on WebGPU and matches CPU.SimplifiedLayerNormFusionCpuRankOneScaleFirstfails on main.Testing (Linux x64, Intel Lunar Lake integrated GPU, Vulkan/Mesa ANV)
onnxruntime_provider_testWebGPU layer norm tests pass (147).GraphTransformationTests.MatMulNBitsMlpFusionMatchesUnfusedSimplifiedQuickGeluWebGpuResultsfails on this GPU (2 elements just over tolerance), identically on main, so it is unrelated to this change.Motivation and Context
Fixes #29841. The other fusions listed there (Gelu, FastGelu, BiasGelu, LayerNorm, SkipLayerNorm) are already enabled for WebGPU on main.