Skip to content

[WebGPU] Fix SkipSimplifiedLayerNormalization bias#28427

Merged
guschmue merged 3 commits into
microsoft:mainfrom
xenova:fix-28424
May 12, 2026
Merged

[WebGPU] Fix SkipSimplifiedLayerNormalization bias#28427
guschmue merged 3 commits into
microsoft:mainfrom
xenova:fix-28424

Conversation

@xenova
Copy link
Copy Markdown
Contributor

@xenova xenova commented May 8, 2026

Description

The op spec for com.microsoft.SkipSimplifiedLayerNormalization puts the optional 1-D bias at input slot 3 (the simplified variant has no beta), but onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.cc hardcodes beta = Input(3); bias = Input(4) regardless of the simplified template parameter — so the user-supplied bias is silently consumed as beta and the kernel reads slot 4 (out of range for this op) for the actual bias. The CPU EP handles the slot correctly; only WebGPU diverges (and produces NaN in larger graphs where slot 4's memory is non-zero).

Motivation and Context

Closes #28424

@xenova xenova changed the title Fix 28424 [WebGPU] Fix SkipSimplifiedLayerNormalization bias May 8, 2026
@tianleiwu tianleiwu requested a review from Copilot May 11, 2026 04:47
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

Fixes a WebGPU-specific input indexing bug in com.microsoft.SkipSimplifiedLayerNormalization where the optional 1-D bias was read from the wrong input slot (and could lead to incorrect results/NaNs), bringing WebGPU behavior in line with the operator schema and other EPs.

Changes:

  • Correct WebGPU SkipSimplifiedLayerNormalization optional input handling so bias is read from slot 3 (and beta is absent for the simplified variant).
  • Strengthen WebGPU shader program caching by including hasBeta/hasBias in the program cache hint to avoid cache collisions across optional-input combinations.
  • Add a regression unit test that exercises SkipSimplifiedLayerNormalization with bias in FP16.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File Description
onnxruntime/contrib_ops/webgpu/bert/skip_layer_norm.cc Fixes simplified-vs-non-simplified beta/bias input slot selection and updates shader cache hint to account for optional inputs.
onnxruntime/test/contrib_ops/skiplayernorm_op_test.cc Adds an FP16 regression test validating correct bias handling for SkipSimplifiedLayerNormalization.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

tianleiwu
tianleiwu previously approved these changes May 11, 2026
@guschmue guschmue enabled auto-merge (squash) May 11, 2026 15:34
guschmue
guschmue previously approved these changes May 11, 2026
@guschmue guschmue added the ep:WebGPU ort-web webgpu provider label May 11, 2026
@guschmue
Copy link
Copy Markdown
Contributor

DirectML EP is not happy with the new unit test.
We don't maintain that EP anymore and don't want to touch it.
Maybe easiest way out is to disable the test for DirectML.

@xenova
Copy link
Copy Markdown
Contributor Author

xenova commented May 11, 2026

Any pointers on how to disable it only for DirectML? 👀

@guschmue
Copy link
Copy Markdown
Contributor

yeap, copy this line:

bool enable_dml = (nullptr != DefaultDmlExecutionProvider().get()) && !disable_dml;

auto-merge was automatically disabled May 12, 2026 17:18

Head branch was pushed to by a user without write access

@xenova xenova dismissed stale reviews from guschmue and tianleiwu via 936cafb May 12, 2026 17:18
@xenova
Copy link
Copy Markdown
Contributor Author

xenova commented May 12, 2026

thanks 👍 936cafb should hopefully fix that.

@guschmue guschmue enabled auto-merge (squash) May 12, 2026 17:37
@guschmue guschmue merged commit fafe564 into microsoft:main May 12, 2026
92 of 171 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

ep:WebGPU ort-web webgpu provider

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[WebGPU] SkipSimplifiedLayerNormalization reads bias from the wrong input slot

4 participants