Add rms_norm op + named-data upload#19893
Conversation
Summary: Adds the `et_vk.rms_norm.default` operator to the WebGPU backend (WGSL compute shader using a cooperative tree reduction, one workgroup per row) and fixes constant upload so the op's weight loads correctly. 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 and CPU reference mirror 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. Differential Revision: D106887028
JulianCloudNTH
requested review from
kirklandsign and
larryliu0820
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/19893
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit b484870 with merge base 5395f20 ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
|
meta-cla
Bot
added
the
CLA Signed
|
@JulianCloudNTH has exported this pull request. If you are a Meta employee, you can view the originating Diff in D106887028. |
This PR needs a
|
1 participant
Summary:
Adds the
et_vk.rms_norm.defaultoperator to the WebGPU backend (WGSL compute shader using a cooperative tree reduction, one workgroup per row) and fixes constant upload so the op's weight loads correctly.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_MAXand anamed_key, rather than inline in the VK00 blob.WebGPUGraph::buildpreviously handled only inline constants, so the weight was never uploaded and the op returned all zeros.buildnow also fetches named-data constants viaNamedDataMap::get_data, mirroring the pathVulkanBackendalready uses.aten.addwas unaffected since it has no constant tensors.The shader and CPU reference mirror 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.Differential Revision: D106887028