Add rms_norm op (#19893) by JulianCloudNTH · Pull Request #19893 · pytorch/executorch

JulianCloudNTH · 2026-05-30T01:07:54Z

Summary:

Adds the et_vk.rms_norm.default operator to the WebGPU backend: a WGSL compute shader using a cooperative tree reduction, one workgroup per row. The shader mirrors the Vulkan implementation (backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp, backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl); indexing assumes contiguous fp32 inputs. The handler fails loud (throws, mirroring Vulkan's VK_CHECK_COND) on invalid shape/dtype/dispatch-limit conditions, and defaults eps to the float32 machine epsilon.

The weight constant is uploaded via the named-data path added in the parent diff.

Differential Revision: D106887028

pytorch-bot · 2026-05-30T01:07:58Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19893

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 158a749 with merge base 3b3f621 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

linux-foundation-easycla · 2026-05-30T01:08:00Z

✅ login: JulianCloudNTH / name: Julian Ng-Thow-Hing (158a749)
❌ The email address for the commit (7286c79) is not linked to the GitHub account, preventing the EasyCLA check. Consult this Help Article and GitHub Help to resolve. (To view the commit's email address, add .patch at the end of this PR page's URL.) For further assistance with EasyCLA, please visit our EasyCLA portal and chat with our support bot.

meta-codesync · 2026-05-30T01:08:03Z

@JulianCloudNTH has exported this pull request. If you are a Meta employee, you can view the originating Diff in D106887028.

github-actions · 2026-05-30T01:08:48Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

Summary: Adds the `et_vk.rms_norm.default` operator to the WebGPU backend (a WGSL compute shader using a cooperative tree reduction, one workgroup per row), fixes constant upload so the op's weight loads correctly, and wires the backend into CI. The Vulkan serializer that the WebGPU backend reuses stores every non-empty constant (e.g. the rms_norm weight) in the PTE's named-data map with `offset == UINT64_MAX` and a `named_key`, rather than inline in the VK00 blob. `WebGPUGraph::build` previously handled only inline constants, so the weight was never uploaded and the op returned all zeros. `build` now also fetches named-data constants via `NamedDataMap::get_data`, mirroring the path `VulkanBackend` already uses. `aten.add` was unaffected since it has no constant tensors. The shader mirrors the Vulkan implementation (`backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp`, `backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl`); indexing assumes contiguous fp32 inputs. The handler fails loud (throws, mirroring Vulkan's `VK_CHECK_COND`) on invalid shape/dtype/dispatch-limit conditions, and defaults `eps` to the float32 machine epsilon. Also adds a simple x86 Linux CI job, mirroring the Vulkan delegate: `backends/test/suite/flows/webgpu.py` plus a `WebGPUTester`, run by `oss/.github/workflows/test-backend-webgpu.yml` on SwiftShader (a software Vulkan adapter, via `wgpu-native`, minimal dependencies, no GPU). Two fixes were needed for SwiftShader's downlevel limits: request the adapter's full `requiredLimits` at device creation (software adapters default storage-buffer limits to 0), and lower the `add` op `workgroup_size` from 256 to 64 (256 exceeded SwiftShader's 128-invocation cap; the Vulkan delegate uses 64). Differential Revision: D106887028

SS-JIA

Review automatically exported from Phabricator review in Meta.

Summary: The Vulkan serializer that the WebGPU backend reuses stores every non-empty constant in the PTE's named-data map with `offset == UINT64_MAX` and a `named_key`, rather than inline in the VK00 blob. `WebGPUGraph::build` previously handled only inline constants, so a delegated op's constant weights were never uploaded and the op produced all zeros. `build` now also fetches named-data constants via `NamedDataMap::get_data`, mirroring the path `VulkanBackend` already uses. `aten.add` was unaffected since it has no constant tensors; the first consumer is the `rms_norm` op in the child diff. Differential Revision: D107288998

SS-JIA

Review automatically exported from Phabricator review in Meta.

Summary: Pull Request resolved: pytorch#19893 Adds the `et_vk.rms_norm.default` operator to the WebGPU backend: a WGSL compute shader using a cooperative tree reduction, one workgroup per row. The shader mirrors the Vulkan implementation (`backends/vulkan/runtime/graph/ops/impl/RmsNorm.cpp`, `backends/vulkan/runtime/graph/ops/glsl/rms_norm_buffer.glsl`); indexing assumes contiguous fp32 inputs. The handler fails loud (throws, mirroring Vulkan's `VK_CHECK_COND`) on invalid shape/dtype/dispatch-limit conditions, and defaults `eps` to the float32 machine epsilon. The weight constant is uploaded via the named-data path added in the parent diff. Differential Revision: D106887028

psiddh · 2026-06-02T20:04:07Z

+  uniform_desc.size = sizeof(RmsNormParams);
+  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  uniform_desc.mappedAtCreation = true;
+  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);


would this buffer be freed eventually ? When should this be released ?

psiddh · 2026-06-02T20:06:36Z

+      kRmsNormWorkgroupSize == 64,
+      "must match @workgroup_size and WG_SIZE in rms_norm.wgsl");
+  if (num_rows > 65535u) {
+    throw std::runtime_error(


if we are throwing exception here, should we release / free all objects that have been allocated thus far such as (uniform buffer, shader module, etc...)

psiddh · 2026-06-02T20:17:56Z

+                if (buf.ok() && buf->size() >= tensor.nbytes) {
+                  wgpuQueueWriteBuffer(
+                      queue_, tensor.buffer, 0, buf->data(), tensor.nbytes);
+                  buf->Free();


Should we free this here ? does this step "named_data_map->get_data" allocate new memory at Ln174 ?

psiddh · 2026-06-02T20:19:20Z

+                      queue_, tensor.buffer, 0, buf->data(), tensor.nbytes);
+                  buf->Free();
+                } else {
+                  throw std::runtime_error(


nits: The error message conflates "key not found" and "buffer undersized"
into one string. Splitting into two branches with sizes/error codes would make debugging ?

JulianCloudNTH requested review from kirklandsign and larryliu0820 as code owners May 30, 2026 01:07

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label May 30, 2026

meta-codesync Bot added fb-exported meta-exported labels May 30, 2026

mergennachin requested a review from psiddh June 1, 2026 15:34

mergennachin changed the title ~~Add rms_norm op + named-data upload~~ [WebGPU] Add rms_norm op + named-data upload Jun 1, 2026

mergennachin requested review from SS-JIA and digantdesai June 1, 2026 15:35

meta-codesync Bot changed the title ~~[WebGPU] Add rms_norm op + named-data upload~~ Add rms_norm op + named-data upload + x86 CI Jun 2, 2026

JulianCloudNTH force-pushed the export-D106887028 branch 2 times, most recently from 1abcd20 to 38e33e5 Compare June 2, 2026 07:23

meta-codesync Bot changed the title ~~Add rms_norm op + named-data upload + x86 CI~~ Add rms_norm op + named-data upload + x86 CI (#19893) Jun 2, 2026

JulianCloudNTH force-pushed the export-D106887028 branch from 38e33e5 to d6f278e Compare June 2, 2026 07:24

SS-JIA requested changes Jun 2, 2026

View reviewed changes

SS-JIA approved these changes Jun 2, 2026

View reviewed changes

meta-codesync Bot changed the title ~~Add rms_norm op + named-data upload + x86 CI (#19893)~~ Add rms_norm op Jun 2, 2026

JulianCloudNTH force-pushed the export-D106887028 branch 2 times, most recently from ef77d1d to cc65e02 Compare June 2, 2026 19:56

SS-JIA requested changes Jun 2, 2026

View reviewed changes

JulianCloudNTH force-pushed the export-D106887028 branch from cc65e02 to ef77d1d Compare June 2, 2026 19:56

meta-codesync Bot changed the title ~~Add rms_norm op~~ Add rms_norm op (#19893) Jun 2, 2026

JulianCloudNTH force-pushed the export-D106887028 branch from ef77d1d to 158a749 Compare June 2, 2026 20:00

psiddh reviewed Jun 2, 2026

View reviewed changes

JulianCloudNTH closed this Jun 2, 2026

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add rms_norm op (#19893)#19893

Add rms_norm op (#19893)#19893
JulianCloudNTH wants to merge 2 commits into
pytorch:mainfrom
JulianCloudNTH:export-D106887028

JulianCloudNTH commented May 30, 2026 •

edited by meta-codesync Bot

Loading

Uh oh!

pytorch-bot Bot commented May 30, 2026 •

edited

Loading

Uh oh!

linux-foundation-easycla Bot commented May 30, 2026 •

edited

Loading

Uh oh!

meta-codesync Bot commented May 30, 2026

Uh oh!

github-actions Bot commented May 30, 2026

Uh oh!

SS-JIA left a comment

Uh oh!

SS-JIA left a comment

Uh oh!

psiddh Jun 2, 2026 •

edited

Loading

Uh oh!

psiddh Jun 2, 2026

Uh oh!

psiddh Jun 2, 2026

Uh oh!

psiddh Jun 2, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

JulianCloudNTH commented May 30, 2026 • edited by meta-codesync Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19893

✅ No Failures

Uh oh!

linux-foundation-easycla Bot commented May 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

meta-codesync Bot commented May 30, 2026

Uh oh!

github-actions Bot commented May 30, 2026

This PR needs a release notes: label

Uh oh!

SS-JIA left a comment

Choose a reason for hiding this comment

Uh oh!

SS-JIA left a comment

Choose a reason for hiding this comment

Uh oh!

psiddh Jun 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

psiddh Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

psiddh Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

psiddh Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

JulianCloudNTH commented May 30, 2026 •

edited by meta-codesync Bot

Loading

pytorch-bot Bot commented May 30, 2026 •

edited

Loading

linux-foundation-easycla Bot commented May 30, 2026 •

edited

Loading

This PR needs a `release notes:` label

psiddh Jun 2, 2026 •

edited

Loading