[ExecuTorch][WebGPU] Add optimized layer_norm op

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] apply_rotary_emb_hf op tests #20876
• [ExecuTorch][WebGPU] et_vk.apply_rotary_emb_hf: HuggingFace rotate-half RoPE runtime op (unblocks Qwen3) #20875
• [ExecuTorch][WebGPU] et_vk conv2d im2col-GEMM tests #20874
• [ExecuTorch][WebGPU] et_vk conv2d: route standard (groups==1) conv through an im2col tiled GEMM (1.1-2.4x) #20873
• [ExecuTorch][WebGPU] et_vk.sdpa channel-attention tests #20872
• [ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster) #20871
• [ExecuTorch][WebGPU] Migrate quantized_linear to make_compute_pipeline #20870
• [ExecuTorch][WebGPU] Migrate multi-dispatch ops to make_compute_pipeline #20869
• [ExecuTorch][WebGPU] Migrate 9 ops to make_compute_pipeline (layout:"auto") #20868
• [ExecuTorch][WebGPU] Adopt WebGPU layout:"auto" in make_compute_pipeline #20867
• [ExecuTorch][WebGPU] slice_copy native golden test #20866
• [ExecuTorch][WebGPU] slice_copy Double-arg scalar handling + view_copy/alias glue #20865
• [ExecuTorch][WebGPU] relu op tests #20864
• [ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline) #20863
• [ExecuTorch][WebGPU] max_pool2d op tests #20862
• [ExecuTorch][WebGPU] Add max_pool2d op #20861
• [ExecuTorch][WebGPU] upsample_nearest2d op tests #20860
• [ExecuTorch][WebGPU] Add upsample_nearest2d op #20859
• [ExecuTorch][WebGPU] constant_pad_nd op tests #20858
• [ExecuTorch][WebGPU] Add constant_pad_nd op #20857
• [ExecuTorch][WebGPU] addmm op tests #20856
• [ExecuTorch][WebGPU] Add optimized addmm op #20855
• [ExecuTorch][WebGPU] embedding op tests #20854
• [ExecuTorch][WebGPU] Add embedding op #20853
• [ExecuTorch][WebGPU] SDPA op tests #20852
• [ExecuTorch][WebGPU] Add optimized SDPA op #20851
• [ExecuTorch][WebGPU] conv2d op tests #20850
• [ExecuTorch][WebGPU] Add optimized conv2d op #20849
• [ExecuTorch][WebGPU] linear_fp32 op tests #20848
• [ExecuTorch][WebGPU] Add optimized linear_fp32 op #20847
• [ExecuTorch][WebGPU] layer_norm op tests #20846
• -> [ExecuTorch][WebGPU] Add optimized layer_norm op #20845
• [ExecuTorch][WebGPU] gelu op tests #20844
• [ExecuTorch][WebGPU] Add optimized gelu op #20843
• [ExecuTorch][WebGPU] Shared pipeline + dispatch-grid + vec4-align helpers (op-suite foundation) #20842

Adds native_layer_norm to the WebGPU backend. LayerNorm is pervasive in the BART decoder and the Florence-2 DaViT vision encoder (and the Whisper/Voxtral hidden dims 768/1024/1280/3072), so it is required to delegate those transformer stacks end-to-end.

Problem — The backend had no aten.native_layer_norm.default kernel. LayerNorm also needs a numerically robust mean/variance (a naive E[x^2]-E[x]^2 cancels badly on large-mean/small-variance activations), and some graphs (the group_norm LN-reframe) pass None for weight and bias, so the affine step must be optional.

Solution

• Before: aten.native_layer_norm.default was unsupported.
• After: native_layer_norm.wgsl computes per-row mean, rstd, and the normalized (optionally affine) output in a single kernel launch — one workgroup per row — writing all three outputs (out, mean, rstd) of the op's ValueList.

Implementation

• Single-pass, numerically robust mean+variance via Chan et al.'s parallel Welford: each of the 64 threads folds its strided vec4<f32> slice of the row into a running (n, mean, M2), then a shared-memory tree reduction pairwise-merges the per-thread triples — no E[x^2]-E[x]^2 cancellation.
• t_in/t_out/t_weight/t_bias are viewed as array<vec4<f32>> over the row width for wide loads/stores; the second pass re-streams the row applying (x-mean)*rstd and the affine only when has_affine == 1.
• Rows past the 65535 per-dimension ceiling are handled by a near-square 2D workgroup grid: utils::compute_row_dispatch_grid computes count_x/count_y and a stride_x override constant, and the shader decodes the flat row index as wid.y * stride_x + wid.x.
• Weight/bias are optional: when either is None the handler binds dummy storage via utils::make_optional_binding and sets has_affine == 0; the pipeline is built with the shared utils::make_compute_pipeline, epsilon read via utils::scalar_or.
• Mirrors Vulkan backends/vulkan/runtime/graph/ops/impl/NativeLayerNorm.cpp (same arg layout and [out, mean, rstd] ValueList; the WebGPU kernel additionally supports the no-affine path, which the Vulkan handler rejects).

Constraints — fp32 only (byte-size must equal numel * sizeof(float)); normalizes over the last dim only; the last dim (row_width) must be a multiple of 4 for the vec4 view (all in-scope hidden dims are); non-empty input; 2D dispatch grid for the row count.

Co-authored-with: Claude Code.

Differential Revision: D110836681

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (1 Unrelated Failure)

As of commit 8d54200 with merge base aceeb40 ():

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest-nxp-neutron / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

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

[github-actions]

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.