From b5a42233a85201a228873cc1ce08948a2e4b2d3f Mon Sep 17 00:00:00 2001
From: Julian Ng-Thow-Hing <juliannth@meta.com>
Date: Wed, 15 Jul 2026 10:01:49 -0700
Subject: [PATCH 1/3] [ExecuTorch][WebGPU] Make shader workgroup sizes
 runtime-configurable
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Pull Request resolved: https://github.com/pytorch/executorch/pull/20954

Instead of hard-coding each compute shader’s workgroup size, make it a WGSL `override wg_size` set at pipeline creation via a host `WGPUConstantEntry` (mirrors the landed `add` op). One WGSL per kernel; each distinct size becomes its own Dawn-cached pipeline — so the size can be tuned per device/shape WITHOUT generating another shader variant.

Key changes:
- `rms_norm.wgsl` / `sdpa_softmax.wgsl` / `q4gsw_linear_coop4_bicol.wgsl` (the free-knob reduction + decode-GEMV kernels) — `const WG*` -> `override wg_size`; shared arrays, tree-reduction strides, and `@workgroup_size` all use it; handlers wire it from `clamp_workgroup_size(...)`.
- `QuantizedLinear.cpp` — split GEMV out of `fixed_wg` so it takes the override; steel (256) + shmem (64) stay fixed.
- `q4gsw_linear_gemm_shmem.wgsl` — its `override wg_size` was a false knob (geometry hard-locked to 64, never host-wired) -> plain `const`.
- Geometry-lock notes on the kernels that stay fixed (`q4gsw_linear_gemm_steel.wgsl`, `sdpa_fd_decode/sdpa_fd_split.wgsl`, `sdpa_fd_reduce.wgsl`) — their size is bound to tile/head-dim geometry, not a runtime knob.
- Regenerated the affected `_wgsl.h` headers.

Behavior is identical at the default 64: `clamp_workgroup_size` only reduces toward `maxComputeInvocationsPerWorkgroup` (spec floor 256), so 64 is emitted; dispatch counts are unchanged.

Co-authored-with: Claude Code.
ghstack-source-id: 403260281
@exported-using-ghexport

Differential Revision: [D112060863](https://our.internmc.facebook.com/intern/diff/D112060863/)
---
 backends/webgpu/runtime/WebGPUUtils.h         | 10 +++++++++
 .../ops/quantized_linear/QuantizedLinear.cpp  | 12 ++++++++---
 .../q4gsw_linear_coop4_bicol.wgsl             | 11 +++++-----
 .../q4gsw_linear_coop4_bicol_wgsl.h           | 13 ++++++------
 .../q4gsw_linear_gemm_shmem.wgsl              |  3 ++-
 .../q4gsw_linear_gemm_shmem_wgsl.h            |  5 +++--
 .../q4gsw_linear_gemm_steel.wgsl              |  1 +
 ..._linear_gemm_steel_half_pwdq_f16acc_wgsl.h |  3 ++-
 .../q4gsw_linear_gemm_steel_half_pwdq_wgsl.h  |  3 ++-
 .../q4gsw_linear_gemm_steel_half_wgsl.h       |  3 ++-
 .../q4gsw_linear_gemm_steel_wgsl.h            |  3 ++-
 .../webgpu/runtime/ops/rms_norm/RmsNorm.cpp   | 14 +++++++++++--
 .../webgpu/runtime/ops/rms_norm/rms_norm.wgsl | 19 +++++++++--------
 .../runtime/ops/rms_norm/rms_norm_vec4_wgsl.h | 17 ++++++++-------
 .../runtime/ops/rms_norm/rms_norm_wgsl.h      | 15 ++++++-------
 backends/webgpu/runtime/ops/sdpa/Sdpa.cpp     |  6 ++++--
 .../webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl | 19 +++++++++--------
 .../runtime/ops/sdpa/sdpa_softmax_wgsl.h      | 21 ++++++++++---------
 .../ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl    |  1 +
 .../ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h  |  3 ++-
 .../ops/sdpa_fd_decode/sdpa_fd_split.wgsl     |  1 +
 .../sdpa_fd_decode/sdpa_fd_split_half_wgsl.h  |  3 ++-
 .../ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h   |  3 ++-
 23 files changed, 118 insertions(+), 71 deletions(-)

diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h
index 4326fde8a19..7be275db178 100644
--- a/backends/webgpu/runtime/WebGPUUtils.h
+++ b/backends/webgpu/runtime/WebGPUUtils.h
@@ -48,6 +48,16 @@ inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
   return desired;
 }
 
+// Clamp to device limit, then floor to pow2 (reduction kernels halve stride).
+inline uint32_t clamp_workgroup_size_pow2(WGPUDevice device, uint32_t desired) {
+  uint32_t v = clamp_workgroup_size(device, desired);
+  uint32_t p = 1u;
+  while (p <= (v >> 1u)) {
+    p <<= 1u;
+  }
+  return p;
+}
+
 struct WgCount {
   uint32_t x;
   uint32_t y;
diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
index b0728764310..548442f0334 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
+++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
@@ -254,6 +254,11 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const uint32_t wg_size =
       utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX);
   const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u);
+  // GEMV (bicol) is a pow2 tree reduction; compute its size only when used.
+  const uint32_t gemv_wg_size = use_gemv
+      ? utils::clamp_workgroup_size_pow2(
+            device, kQ4gswLinearCoop4BicolWorkgroupSizeX)
+      : 0u;
   // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible.
   const bool use_steel = !use_gemv && steel_supported(device) &&
       steel_workgroup_count(device, M, N, K) > 0u;
@@ -372,16 +377,17 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   WGPUPipelineLayout pipeline_layout =
       wgpuDeviceCreatePipelineLayout(device, &pl_desc);
 
+  // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed.
+  const bool fixed_wg = use_steel || use_shmem_gemm;
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
-  wg_size_constant.value = static_cast<double>(wg_size);
+  wg_size_constant.value =
+      static_cast<double>(use_gemv ? gemv_wg_size : wg_size);
 
   WGPUComputePipelineDescriptor pipeline_desc = {};
   pipeline_desc.layout = pipeline_layout;
   pipeline_desc.compute.module = shader;
   pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
-  // Only tiled GEMM overrides wg_size; GEMV/shmem (64) + steel (256) are fixed.
-  const bool fixed_wg = use_gemv || use_steel || use_shmem_gemm;
   pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u;
   pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant;
   WGPUComputePipeline pipeline =
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl
index 69174269f38..0d33c5a3a65 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl
@@ -17,10 +17,11 @@ struct Params {
 @group(0) @binding(5) var<uniform> params: Params;
 
 // Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
-const WG: u32 = 64u;
-var<workgroup> partial: array<vec2<f32>, WG>;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
+var<workgroup> partial: array<vec2<f32>, wg_size>;
 
-@compute @workgroup_size(WG, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(num_workgroups) ngrp: vec3<u32>,
@@ -69,12 +70,12 @@ fn main(
         acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
         acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
       }
-      w = w + WG;
+      w = w + wg_size;
     }
 
     partial[lid.x] = vec2<f32>(acc0, acc1);
     workgroupBarrier();
-    var s: u32 = WG >> 1u;
+    var s: u32 = wg_size >> 1u;
     loop {
       if (s == 0u) {
         break;
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h
index d6c51bcb893..864165da136 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_coop4_bicol.wgsl - DO NOT EDIT.
-// wgsl-sha256: 44c53ca7879f7a3382a45703a67cbc6fc221ecbcada58f8c413df50abbc6e544
+// wgsl-sha256: ef677c2771c3d4a73704f4725bb490174149b971bc62e696595b8d13941c979e
 inline constexpr const char* kQ4gswLinearCoop4BicolWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
@@ -34,10 +34,11 @@ struct Params {
 @group(0) @binding(5) var<uniform> params: Params;
 
 // Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
-const WG: u32 = 64u;
-var<workgroup> partial: array<vec2<f32>, WG>;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
+var<workgroup> partial: array<vec2<f32>, wg_size>;
 
-@compute @workgroup_size(WG, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(num_workgroups) ngrp: vec3<u32>,
@@ -86,12 +87,12 @@ fn main(
         acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
         acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
       }
-      w = w + WG;
+      w = w + wg_size;
     }
 
     partial[lid.x] = vec2<f32>(acc0, acc1);
     workgroupBarrier();
-    var s: u32 = WG >> 1u;
+    var s: u32 = wg_size >> 1u;
     loop {
       if (s == 0u) {
         break;
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl
index a7384f36115..e68af88ad4c 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl
@@ -16,7 +16,8 @@ struct Params {
 }
 @group(0) @binding(5) var<uniform> params: Params;
 
-override wg_size: u32 = 64u;
+// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob).
+const wg_size: u32 = 64u;
 
 // Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
 const WG_M: u32 = 32u;   // output rows per workgroup
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h
index 64bbd2cd15f..3255a05220d 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_shmem.wgsl - DO NOT EDIT.
-// wgsl-sha256: 0a1219d3b6781315a21066089fca3f92235587e8af8eb734185f35ea4bfc8a52
+// wgsl-sha256: f0180cad41701807b34912b57d9715c02588c9447783bbf691f93836168fe981
 inline constexpr const char* kQ4gswLinearGemmShmemWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
@@ -33,7 +33,8 @@ struct Params {
 }
 @group(0) @binding(5) var<uniform> params: Params;
 
-override wg_size: u32 = 64u;
+// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob).
+const wg_size: u32 = 64u;
 
 // Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
 const WG_M: u32 = 32u;   // output rows per workgroup
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl
index 4d7ab0b1d1e..f08cb39bd5c 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl
@@ -34,6 +34,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<${buffer_scalar_type(DTYPE)}, 1024>;   // BM*BK
 var<workgroup> Bs: array<${buffer_scalar_type(DTYPE)}, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h
index efefd7edce1..6d23ea7ac76 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce
+// wgsl-sha256: fe8b71afc5634e7db5607deb58f30c6517f93e1d66370ea017bcd9071201c6ca
 inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"(
 enable f16;
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@@ -50,6 +50,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<f16, 1024>;   // BM*BK
 var<workgroup> Bs: array<f16, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h
index 46057de2340..e684b00a228 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: 1f916bcd30dbbbcc7eca37e795ecc26e3c72e645ccd2c361fa0ac4e66f1a174a
+// wgsl-sha256: 1b3f5e384e09bde50c66d43db4797e01e5610d70af235355dbc4966643bf9103
 inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"(
 enable f16;
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@@ -50,6 +50,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<f16, 1024>;   // BM*BK
 var<workgroup> Bs: array<f16, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h
index ae03f63fb5c..d4c78406582 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: 00cdd2f2fb98a5c7343d16fdf7e59f1b840e180cec3f82bf9b569513c0a45396
+// wgsl-sha256: d77542e5975fab44c58644c771390e89598ce82ab25e38d57150b581b4904acd
 inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"(
 enable f16;
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@@ -50,6 +50,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<f16, 1024>;   // BM*BK
 var<workgroup> Bs: array<f16, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h
index 71b0f45bdf7..db7a7689a6f 100644
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: dd771b9ab096410f3ad0d9259bef7816e41330a434325ea28baa5abbfb2841d2
+// wgsl-sha256: dcb63b04d002899acfed88d273fe658618650d5dcbc108540cef086bb97f5bce
 inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
@@ -49,6 +49,7 @@ struct Params {
 const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
 var<workgroup> As: array<f32, 1024>;   // BM*BK
 var<workgroup> Bs: array<f32, 1024>;   // BK*BN
+// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
 @compute @workgroup_size(16, 16)
 fn main(@builtin(workgroup_id) wid: vec3<u32>,
         @builtin(local_invocation_id) lid: vec3<u32>) {
diff --git a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
index c07e58e8ec8..44e417f6dcc 100644
--- a/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
+++ b/backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
@@ -178,11 +178,21 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   WGPUPipelineLayout pipeline_layout =
       wgpuDeviceCreatePipelineLayout(device, &pl_desc);
 
+  // Runtime-overridable workgroup size (mirrors add op); clamp only reduces.
+  // Pow2 required: the kernel halves the reduction stride (wg_size / 2u).
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size_pow2(device, kRmsNormWorkgroupSizeX);
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
   // Create compute pipeline
   WGPUComputePipelineDescriptor pipeline_desc = {};
   pipeline_desc.layout = pipeline_layout;
   pipeline_desc.compute.module = shader;
   pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
   WGPUComputePipeline pipeline =
       wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
 
@@ -217,10 +227,10 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   // One workgroup per row (kRmsNormWorkgroupSizeX threads cooperate per row)
   static_assert(
       kRmsNormWorkgroupSizeX == 64,
-      "must match @workgroup_size and WG_SIZE in rms_norm.wgsl");
+      "kRmsNormWorkgroupSizeX must match override wg_size default (64)");
   static_assert(
       kRmsNormVec4WorkgroupSizeX == 64,
-      "must match @workgroup_size and WG_SIZE in rms_norm_vec4.wgsl");
+      "kRmsNormVec4WorkgroupSizeX must match override wg_size default (64)");
   const size_t dispatch_idx =
       graph.add_dispatch({pipeline, bind_group, num_rows});
 
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
index 1234174981e..11302434ec2 100644
--- a/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
@@ -12,13 +12,14 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-const WG_SIZE: u32 = 64u;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
 
-var<workgroup> shared_sum: array<${accum_scalar_type(DTYPE)}, WG_SIZE>;
+var<workgroup> shared_sum: array<${accum_scalar_type(DTYPE)}, wg_size>;
 
 fn reduce_shared(worker_id: u32) {
   workgroupBarrier();
-  var stride: u32 = WG_SIZE / 2u;
+  var stride: u32 = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -32,10 +33,10 @@ fn reduce_shared(worker_id: u32) {
 }
 
 $if VEC == 4:
-  // vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4
+  // vec4 variant of rms_norm: each lane strides by wg_size over rw4 = row_width/4
   // texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq
   // divides by it (not rw4). The host selects this only when row_width % 4 == 0.
-@compute @workgroup_size(64, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>) {
@@ -64,7 +65,7 @@ fn main(
         local_sq_sum = local_sq_sum + dot(vec4<f32>(v), vec4<f32>(v));
       $else:
         local_sq_sum = local_sq_sum + dot(v, v);
-      x4 = x4 + WG_SIZE;
+      x4 = x4 + wg_size;
     }
   $else:
     var x: u32 = worker_id;
@@ -77,7 +78,7 @@ fn main(
         local_sq_sum = local_sq_sum + f32(v) * f32(v);
       $else:
         local_sq_sum = local_sq_sum + v * v;
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
 
   shared_sum[worker_id] = local_sq_sum;
@@ -96,7 +97,7 @@ fn main(
         t_out[base4 + x4] = vec4<f16>(vec4<f32>(t_in[base4 + x4]) * rstd * vec4<f32>(t_weight[x4]));
       $else:
         t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4];
-      x4 = x4 + WG_SIZE;
+      x4 = x4 + wg_size;
     }
   $else:
     x = worker_id;
@@ -110,6 +111,6 @@ fn main(
         t_out[base + x] = f16(f32(v) * rstd * f32(w));
       $else:
         t_out[base + x] = v * rstd * w;
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
 }
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
index 02213e3d7b0..2e17ac02aa0 100644
--- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from rms_norm.wgsl - DO NOT EDIT.
-// wgsl-sha256: 4c0ba56708bf125a7ec6ea3c51d1288e05ac00a8e2cfa10e38e9a208e230b8df
+// wgsl-sha256: 62fdfe03fc67eb44fa17ca5e91e433b3a45a96c76151777856d6f557fd829919
 inline constexpr const char* kRmsNormVec4WGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<vec4<f32>>;
 @group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
@@ -27,13 +27,14 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-const WG_SIZE: u32 = 64u;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
 
-var<workgroup> shared_sum: array<f32, WG_SIZE>;
+var<workgroup> shared_sum: array<f32, wg_size>;
 
 fn reduce_shared(worker_id: u32) {
   workgroupBarrier();
-  var stride: u32 = WG_SIZE / 2u;
+  var stride: u32 = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -46,10 +47,10 @@ fn reduce_shared(worker_id: u32) {
   }
 }
 
-// vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4
+// vec4 variant of rms_norm: each lane strides by wg_size over rw4 = row_width/4
 // texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq
 // divides by it (not rw4). The host selects this only when row_width % 4 == 0.
-@compute @workgroup_size(64, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>) {
@@ -71,7 +72,7 @@ fn main(
     }
     let v = t_in[base4 + x4];
     local_sq_sum = local_sq_sum + dot(v, v);
-    x4 = x4 + WG_SIZE;
+    x4 = x4 + wg_size;
   }
 
   shared_sum[worker_id] = local_sq_sum;
@@ -86,7 +87,7 @@ fn main(
       break;
     }
     t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4];
-    x4 = x4 + WG_SIZE;
+    x4 = x4 + wg_size;
   }
 }
 )";
diff --git a/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h b/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h
index 5d9fc236e91..841b146b2ef 100644
--- a/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h
+++ b/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from rms_norm.wgsl - DO NOT EDIT.
-// wgsl-sha256: 340dcbf3c06dc311e70bef953c1e9cbbdf4121fe177eedd3253549e614b55069
+// wgsl-sha256: 55bc862d64451f5fe8f530114f0318dfbc724017314091418a1fe41b180be7c6
 inline constexpr const char* kRmsNormWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_in: array<f32>;
@@ -27,13 +27,14 @@ struct Params {
 }
 @group(0) @binding(3) var<uniform> params: Params;
 
-const WG_SIZE: u32 = 64u;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
 
-var<workgroup> shared_sum: array<f32, WG_SIZE>;
+var<workgroup> shared_sum: array<f32, wg_size>;
 
 fn reduce_shared(worker_id: u32) {
   workgroupBarrier();
-  var stride: u32 = WG_SIZE / 2u;
+  var stride: u32 = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -46,7 +47,7 @@ fn reduce_shared(worker_id: u32) {
   }
 }
 
-@compute @workgroup_size(64, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>) {
@@ -67,7 +68,7 @@ fn main(
     }
     let v = t_in[base + x];
     local_sq_sum = local_sq_sum + v * v;
-    x = x + WG_SIZE;
+    x = x + wg_size;
   }
 
   shared_sum[worker_id] = local_sq_sum;
@@ -84,7 +85,7 @@ fn main(
     let v = t_in[base + x];
     let w = t_weight[x];
     t_out[base + x] = v * rstd * w;
-    x = x + WG_SIZE;
+    x = x + wg_size;
   }
 }
 )";
diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp
index 50321ba4bdf..f3345c89c77 100644
--- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp
+++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp
@@ -188,7 +188,7 @@ void build_dispatch(
   WGPUPipelineLayout pipeline_layout =
       wgpuDeviceCreatePipelineLayout(device, &pl_desc);
 
-  // QK/AV/update_cache have an `override wg_size`; softmax (0) keeps a const.
+  // All callers pass an override wg_size; a 0 would keep the shader default.
   WGPUConstantEntry wg_size_constant = {};
   wg_size_constant.key = {"wg_size", WGPU_STRLEN};
   wg_size_constant.value = static_cast<double>(wg_size);
@@ -529,6 +529,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     // One workgroup per (h,s) row; wg_size 1 keeps the device dispatch check.
     const utils::WgCount wgc = utils::compute_2d_workgroup_count(
         device, static_cast<uint32_t>(Hq * S), 1, "softmax");
+    const uint32_t sm_wg =
+        utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX);
     SoftmaxParams p = make_softmax_params(Hq, S, context_len);
     WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p));
     BufferBinding bindings[2] = {
@@ -542,7 +544,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         sizeof(p),
         wgc.x,
         wgc.y,
-        0,
+        sm_wg,
         true,
         "sdpa_softmax");
     softmax_buf = ubuf;
diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl
index f0e349ceaff..724e427073a 100644
--- a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl
+++ b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax.wgsl
@@ -9,15 +9,16 @@ struct Params {
 }
 @group(0) @binding(2) var<uniform> params: Params;
 
-const WG_SIZE: u32 = 64u;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
 
 // WGSL forbids literal -inf; a large finite negative inits the running max.
 const NEG_INF: f32 = -1.0e30;
 
-var<workgroup> shared_max: array<f32, WG_SIZE>;
-var<workgroup> shared_sum: array<f32, WG_SIZE>;
+var<workgroup> shared_max: array<f32, wg_size>;
+var<workgroup> shared_sum: array<f32, wg_size>;
 
-@compute @workgroup_size(WG_SIZE, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>,
@@ -41,14 +42,14 @@ fn main(
         break;
       }
       local_max = max(local_max, t_in[base + x]);
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
   }
   shared_max[worker_id] = local_max;
 
   // Reduce max. workgroupBarrier() calls are in uniform control flow.
   workgroupBarrier();
-  var stride: u32 = WG_SIZE / 2u;
+  var stride: u32 = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -70,13 +71,13 @@ fn main(
         break;
       }
       local_sum = local_sum + exp(t_in[base + x] - row_max);
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
   }
   shared_sum[worker_id] = local_sum;
 
   workgroupBarrier();
-  stride = WG_SIZE / 2u;
+  stride = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -98,7 +99,7 @@ fn main(
         break;
       }
       t_out[base + x] = exp(t_in[base + x] - row_max) * inv;
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
   }
 }
diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h
index f1e7396fafe..077f173a546 100644
--- a/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h
+++ b/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from sdpa_softmax.wgsl - DO NOT EDIT.
-// wgsl-sha256: bdd45cf344663533b243153200c507f41a90295751924e70452abcd5da4cdd5a
+// wgsl-sha256: 774a372c0eae85b6656618408e40fff4b8af5647ea5353b51d3a34b9b6a7ac8a
 inline constexpr const char* kSdpaSoftmaxWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_in: array<f32>;
@@ -26,15 +26,16 @@ struct Params {
 }
 @group(0) @binding(2) var<uniform> params: Params;
 
-const WG_SIZE: u32 = 64u;
+// wg_size must be a power of two (the tree reduction halves the stride).
+override wg_size: u32 = 64u;
 
 // WGSL forbids literal -inf; a large finite negative inits the running max.
 const NEG_INF: f32 = -1.0e30;
 
-var<workgroup> shared_max: array<f32, WG_SIZE>;
-var<workgroup> shared_sum: array<f32, WG_SIZE>;
+var<workgroup> shared_max: array<f32, wg_size>;
+var<workgroup> shared_sum: array<f32, wg_size>;
 
-@compute @workgroup_size(WG_SIZE, 1, 1)
+@compute @workgroup_size(wg_size, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
     @builtin(local_invocation_id) lid: vec3<u32>,
@@ -58,14 +59,14 @@ fn main(
         break;
       }
       local_max = max(local_max, t_in[base + x]);
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
   }
   shared_max[worker_id] = local_max;
 
   // Reduce max. workgroupBarrier() calls are in uniform control flow.
   workgroupBarrier();
-  var stride: u32 = WG_SIZE / 2u;
+  var stride: u32 = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -87,13 +88,13 @@ fn main(
         break;
       }
       local_sum = local_sum + exp(t_in[base + x] - row_max);
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
   }
   shared_sum[worker_id] = local_sum;
 
   workgroupBarrier();
-  stride = WG_SIZE / 2u;
+  stride = wg_size / 2u;
   loop {
     if (stride == 0u) {
       break;
@@ -115,7 +116,7 @@ fn main(
         break;
       }
       t_out[base + x] = exp(t_in[base + x] - row_max) * inv;
-      x = x + WG_SIZE;
+      x = x + wg_size;
     }
   }
 }
diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl
index 60f86d217ff..6ba784a5d5a 100644
--- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl
+++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce.wgsl
@@ -19,6 +19,7 @@ const NEG_INF: f32 = -1.0e30;
 var<workgroup> sh_w: array<f32, MAX_SPLITS>;
 
 // FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize.
+// 64 bound to MAX_D_PER_LANE output coverage (D<=WG_SIZE*2=128); not a knob.
 @compute @workgroup_size(64, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h
index 6e5a515c065..e6eafe901ef 100644
--- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h
+++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_reduce_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from sdpa_fd_reduce.wgsl - DO NOT EDIT.
-// wgsl-sha256: 3debe9b52adece82b31a067a876105121afb4e09cb31d95596ea699d1a179d18
+// wgsl-sha256: 776e9605b309f7b5d796f28b70512ed030c7178075e6b3f843176b00494a03a6
 inline constexpr const char* kSdpaFdReduceWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_part_o: array<f32>;
@@ -36,6 +36,7 @@ const NEG_INF: f32 = -1.0e30;
 var<workgroup> sh_w: array<f32, MAX_SPLITS>;
 
 // FlashDecoding pass 2: online-softmax merge of the per-split partials, then normalize.
+// 64 bound to MAX_D_PER_LANE output coverage (D<=WG_SIZE*2=128); not a knob.
 @compute @workgroup_size(64, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl
index da67489cc4d..f5a003b4562 100644
--- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl
+++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl
@@ -28,6 +28,7 @@ var<workgroup> sh_s: array<f32, WG_SIZE>;
 var<workgroup> sh_red: array<f32, WG_SIZE>;
 
 // FlashDecoding pass 1: per-(head,split) unnormalized softmax partial.
+// 64 sizes sh_s/sh_red + MAX_D_PER_LANE output coverage (D<=128); not a knob.
 @compute @workgroup_size(64, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h
index bc69a444edb..a0898512905 100644
--- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h
+++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from sdpa_fd_split.wgsl - DO NOT EDIT.
-// wgsl-sha256: 147fc5775b76f3626dc934b2df5e34148bef12c345789332f098d13143bb646e
+// wgsl-sha256: 596f7376a2f15f788827b921b293de91541f53758d5ee7fe858bf4a72ab102df
 inline constexpr const char* kSdpaFdSplitHalfWGSL = R"(
 enable f16;
 @group(0) @binding(0) var<storage, read_write> t_part_o: array<f32>;
@@ -44,6 +44,7 @@ var<workgroup> sh_s: array<f32, WG_SIZE>;
 var<workgroup> sh_red: array<f32, WG_SIZE>;
 
 // FlashDecoding pass 1: per-(head,split) unnormalized softmax partial.
+// 64 sizes sh_s/sh_red + MAX_D_PER_LANE output coverage (D<=128); not a knob.
 @compute @workgroup_size(64, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,
diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h
index 6b23e6fd36e..cb362c9cc82 100644
--- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h
+++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_wgsl.h
@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from sdpa_fd_split.wgsl - DO NOT EDIT.
-// wgsl-sha256: 44bee2691f85839b638a51971a8966bcf8efbb72b8acc44a4c3e991f31cc8fcb
+// wgsl-sha256: b5a926169f085c870ebb551510722dad5ee10abaf97f355eab2635268923c34c
 inline constexpr const char* kSdpaFdSplitWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_part_o: array<f32>;
 @group(0) @binding(1) var<storage, read_write> t_part_ml: array<f32>;
@@ -43,6 +43,7 @@ var<workgroup> sh_s: array<f32, WG_SIZE>;
 var<workgroup> sh_red: array<f32, WG_SIZE>;
 
 // FlashDecoding pass 1: per-(head,split) unnormalized softmax partial.
+// 64 sizes sh_s/sh_red + MAX_D_PER_LANE output coverage (D<=128); not a knob.
 @compute @workgroup_size(64, 1, 1)
 fn main(
     @builtin(workgroup_id) wid: vec3<u32>,

From 52057858e390a1b938bfa37caf739d506315781b Mon Sep 17 00:00:00 2001
From: Julian Ng-Thow-Hing <juliannth@meta.com>
Date: Wed, 15 Jul 2026 10:01:50 -0700
Subject: [PATCH 2/3] [ExecuTorch][WebGPU] Remove the orphan q4gsw_linear_coop4
 shader

Pull Request resolved: https://github.com/pytorch/executorch/pull/20955

`q4gsw_linear_coop4.wgsl` is an unwired dead shader: no handler `#include`s its generated `q4gsw_linear_coop4_wgsl.h`, and it is superseded by `q4gsw_linear_coop4_bicol.wgsl` (the live GEMV decode kernel that `QuantizedLinear.cpp` uses). Remove the shader and its generated header.

Key changes:
- Deleted `q4gsw_linear_coop4.wgsl` and `q4gsw_linear_coop4_wgsl.h`.
- Dropped a stale reference to the removed kernel in `test_quantized_linear.py` and `test_webgpu_native.cpp` comments.

No build/CMake change: `WEBGPU_SRCS` lists only `.cpp`, and shaders wire via the `gen_wgsl_headers.py` glob (no static list), so removing both files leaves codegen drift-clean.

Co-authored-with: Claude Code.
ghstack-source-id: 403097256
@exported-using-ghexport

Differential Revision: [D112060303](https://our.internmc.facebook.com/intern/diff/D112060303/)
---
 .../quantized_linear/q4gsw_linear_coop4.wgsl  |  87 --------------
 .../q4gsw_linear_coop4_wgsl.h                 | 111 ------------------
 .../webgpu/test/ops/test_quantized_linear.py  |   4 +-
 backends/webgpu/test/test_webgpu_native.cpp   |   3 +-
 4 files changed, 3 insertions(+), 202 deletions(-)
 delete mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl
 delete mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h

diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl
deleted file mode 100644
index af6f661279d..00000000000
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4.wgsl
+++ /dev/null
@@ -1,87 +0,0 @@
-@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
-@group(0) @binding(1) var<storage, read> t_input: array<f32>;
-@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
-@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
-@group(0) @binding(4) var<storage, read> t_bias: array<f32>;
-
-struct Params {
-  M: u32,
-  N: u32,
-  K: u32,
-  K_packed: u32,
-  group_size: u32,
-  padded_N: u32,
-  has_bias: u32,
-  _pad: u32,
-}
-@group(0) @binding(5) var<uniform> params: Params;
-
-// Cooperative-over-K GEMV with u32-batched coalesced weight loads (64 lanes).
-const WG: u32 = 64u;
-var<workgroup> partial: array<f32, WG>;
-
-@compute @workgroup_size(WG, 1, 1)
-fn main(
-    @builtin(workgroup_id) wid: vec3<u32>,
-    @builtin(num_workgroups) ngrp: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>) {
-  let total = params.M * params.N;
-  let stride = ngrp.x;
-  let num_words = params.K >> 3u;          // K / 8 words per row
-  let row_words = params.K_packed >> 2u;   // u32s per weight row (= K/8)
-  var idx = wid.x;
-  loop {
-    if (idx >= total) {
-      break;
-    }
-    let m = idx / params.N;
-    let n = idx % params.N;
-    let in_base = m * params.K;
-    let wbase = n * row_words;
-
-    var acc: f32 = 0.0;
-    var w: u32 = lid.x;
-    loop {
-      if (w >= num_words) {
-        break;
-      }
-      let word = t_weight[wbase + w];
-      let k0 = w << 3u; // first K of this word
-      let scale = t_scales[(k0 / params.group_size) * params.padded_N + n];
-      let ib = in_base + k0;
-      // 4 bytes, low+high nibble each -> 8 consecutive K.
-      for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) {
-        let byte = (word >> (bi * 8u)) & 0xFFu;
-        let lo = f32(i32(byte & 0x0Fu) - 8);
-        let hi = f32(i32((byte >> 4u) & 0x0Fu) - 8);
-        let kk = bi << 1u;
-        acc = acc + t_input[ib + kk] * lo * scale;
-        acc = acc + t_input[ib + kk + 1u] * hi * scale;
-      }
-      w = w + WG;
-    }
-
-    partial[lid.x] = acc;
-    workgroupBarrier();
-    var s: u32 = WG >> 1u;
-    loop {
-      if (s == 0u) {
-        break;
-      }
-      if (lid.x < s) {
-        partial[lid.x] = partial[lid.x] + partial[lid.x + s];
-      }
-      workgroupBarrier();
-      s = s >> 1u;
-    }
-    if (lid.x == 0u) {
-      var o = partial[0];
-      if (params.has_bias != 0u) {
-        o = o + t_bias[n];
-      }
-      t_out[idx] = o;
-    }
-    workgroupBarrier();
-    idx = idx + stride;
-  }
-}
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h
deleted file mode 100644
index 7bacedfcb17..00000000000
--- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_wgsl.h
+++ /dev/null
@@ -1,111 +0,0 @@
-/*
- * Copyright (c) Meta Platforms, Inc. and affiliates.
- * All rights reserved.
- *
- * This source code is licensed under the BSD-style license found in the
- * LICENSE file in the root directory of this source tree.
- */
-
-#pragma once
-
-#include <cstdint>
-
-namespace executorch::backends::webgpu {
-
-// @generated from q4gsw_linear_coop4.wgsl - DO NOT EDIT.
-// wgsl-sha256: 3031886e68c375e617dfb263da39c492c6de4d8c1fb4073d70b18823a3e6a4fe
-inline constexpr const char* kQ4gswLinearCoop4WGSL = R"(
-@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
-@group(0) @binding(1) var<storage, read> t_input: array<f32>;
-@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
-@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
-@group(0) @binding(4) var<storage, read> t_bias: array<f32>;
-
-struct Params {
-  M: u32,
-  N: u32,
-  K: u32,
-  K_packed: u32,
-  group_size: u32,
-  padded_N: u32,
-  has_bias: u32,
-  _pad: u32,
-}
-@group(0) @binding(5) var<uniform> params: Params;
-
-// Cooperative-over-K GEMV with u32-batched coalesced weight loads (64 lanes).
-const WG: u32 = 64u;
-var<workgroup> partial: array<f32, WG>;
-
-@compute @workgroup_size(WG, 1, 1)
-fn main(
-    @builtin(workgroup_id) wid: vec3<u32>,
-    @builtin(num_workgroups) ngrp: vec3<u32>,
-    @builtin(local_invocation_id) lid: vec3<u32>) {
-  let total = params.M * params.N;
-  let stride = ngrp.x;
-  let num_words = params.K >> 3u;          // K / 8 words per row
-  let row_words = params.K_packed >> 2u;   // u32s per weight row (= K/8)
-  var idx = wid.x;
-  loop {
-    if (idx >= total) {
-      break;
-    }
-    let m = idx / params.N;
-    let n = idx % params.N;
-    let in_base = m * params.K;
-    let wbase = n * row_words;
-
-    var acc: f32 = 0.0;
-    var w: u32 = lid.x;
-    loop {
-      if (w >= num_words) {
-        break;
-      }
-      let word = t_weight[wbase + w];
-      let k0 = w << 3u; // first K of this word
-      let scale = t_scales[(k0 / params.group_size) * params.padded_N + n];
-      let ib = in_base + k0;
-      // 4 bytes, low+high nibble each -> 8 consecutive K.
-      for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) {
-        let byte = (word >> (bi * 8u)) & 0xFFu;
-        let lo = f32(i32(byte & 0x0Fu) - 8);
-        let hi = f32(i32((byte >> 4u) & 0x0Fu) - 8);
-        let kk = bi << 1u;
-        acc = acc + t_input[ib + kk] * lo * scale;
-        acc = acc + t_input[ib + kk + 1u] * hi * scale;
-      }
-      w = w + WG;
-    }
-
-    partial[lid.x] = acc;
-    workgroupBarrier();
-    var s: u32 = WG >> 1u;
-    loop {
-      if (s == 0u) {
-        break;
-      }
-      if (lid.x < s) {
-        partial[lid.x] = partial[lid.x] + partial[lid.x + s];
-      }
-      workgroupBarrier();
-      s = s >> 1u;
-    }
-    if (lid.x == 0u) {
-      var o = partial[0];
-      if (params.has_bias != 0u) {
-        o = o + t_bias[n];
-      }
-      t_out[idx] = o;
-    }
-    workgroupBarrier();
-    idx = idx + stride;
-  }
-}
-)";
-
-inline constexpr uint32_t kQ4gswLinearCoop4WorkgroupSizeX = 64;
-inline constexpr uint32_t kQ4gswLinearCoop4WorkgroupSizeY = 1;
-inline constexpr uint32_t kQ4gswLinearCoop4WorkgroupSizeZ = 1;
-
-} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py
index 72945c37d6d..54040a54bd7 100644
--- a/backends/webgpu/test/ops/test_quantized_linear.py
+++ b/backends/webgpu/test/ops/test_quantized_linear.py
@@ -57,8 +57,8 @@ class Q4gswConfig:
     # decode GEMV: the handler routes M==1 -> bicol, so each reads its own per-
     # column scale (col0/col1) across many K-groups (down_proj: 256 groups). q4gsw
     # requires N % 8 == 0 (torchao pads N for the scale layout), so odd-N / N=1 are
-    # not exportable -- bicol's has1 odd-N guard is defensive (mirrors coop4's
-    # general-N robustness) and unreachable through this op.
+    # not exportable -- bicol's has1 odd-N guard is defensive and unreachable
+    # through this op.
     # M>1 prefill: prefer the steel GEMM (K%16==0) on a >=256-invocation device
     # (e.g. lvp); else shmem (K>=4096 or N>=2048) or register-tiled (SwiftShader
     # caps at 128). Same fp64 golden regardless of which kernel runs.
diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp
index fbdfbd09076..2a3e1a9ce77 100644
--- a/backends/webgpu/test/test_webgpu_native.cpp
+++ b/backends/webgpu/test/test_webgpu_native.cpp
@@ -284,8 +284,7 @@ const Q4gswConfig kQ4gswConfigs[] = {
     // The M==1 configs above (q/kv/gate/down_proj) exercise the bicol 2-col
     // decode GEMV (handler routes M==1 -> bicol; each reads its own per-column
     // scale over 64-256 K-groups). q4gsw requires N % 8 == 0, so odd-N is not
-    // exportable; bicol's has1 odd-N guard is defensive (mirrors coop4
-    // general-N robustness).
+    // exportable; bicol's has1 odd-N guard is defensive.
     // M>1: steel GEMM on a >=256-invocation device (K%16==0), else shmem/tiled.
     {"steel", 96, 2048, 256, 1e-4f, 1e-3f, true, false}, // steel-isolating
     // Same shape as "steel" run under the f16-multiply steel kernel; the f16

From 42acaff7b21503eae974d4bbf9e4bb2e0e79117b Mon Sep 17 00:00:00 2001
From: Julian Ng-Thow-Hing <juliannth@meta.com>
Date: Wed, 15 Jul 2026 10:01:50 -0700
Subject: [PATCH 3/3] [ExecuTorch][WebGPU] Test: rms_norm workgroup size is
 runtime-configurable
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Pull Request resolved: https://github.com/pytorch/executorch/pull/20956

Add a native gtest that locks the wg-size configurability guarantee: the same rms_norm WGSL, run at override `wg_size` 64 and 128, must produce the same output. `Module::forward` always uses the handler-clamped size (64), so the test builds the compute pipeline directly and sets `wg_size` via a `WGPUConstantEntry`, mirroring the runtime pipeline setup + `WebGPUGraph::copy_outputs` readback.

Key changes:
- `test/test_webgpu_native.cpp` — `run_rms_norm_at_wg()` standalone-pipeline helper + `RmsNormWorkgroupSizeConfigurable` test (element-wise agree within abs 1e-4 / rel 1e-3, plus a non-zero sanity guard).

Self-contained (no `.pte`/export), so it runs inside `webgpu_native_test` under `test_webgpu_native_ci.sh` with no new wiring. Absolute rms_norm correctness stays covered by the model-driven golden tests; this adds the "128 matches 64" guarantee.

Co-authored-with: Claude Code.
ghstack-source-id: 403097259
@exported-using-ghexport

Differential Revision: [D112060302](https://our.internmc.facebook.com/intern/diff/D112060302/)
---
 backends/webgpu/test/test_webgpu_native.cpp | 232 ++++++++++++++++++++
 1 file changed, 232 insertions(+)

diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp
index 2a3e1a9ce77..5ba4a95f976 100644
--- a/backends/webgpu/test/test_webgpu_native.cpp
+++ b/backends/webgpu/test/test_webgpu_native.cpp
@@ -6,9 +6,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/backends/webgpu/runtime/WebGPUCompat.h>
 #include <executorch/backends/webgpu/runtime/WebGPUDelegateHeader.h>
 #include <executorch/backends/webgpu/runtime/WebGPUDevice.h>
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/ops/rms_norm/rms_norm_wgsl.h>
 #include <executorch/extension/module/module.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/backend/backend_options_map.h>
@@ -21,6 +23,7 @@
 #include <cstdint>
 #include <cstdio>
 #include <cstdlib>
+#include <cstring>
 #include <exception>
 #include <memory>
 #include <string>
@@ -256,6 +259,185 @@ bool sdpa_within_tol(
   return ok;
 }
 
+// Matches the WGSL Params struct in rms_norm.wgsl (16-byte aligned).
+struct RmsNormProbeParams {
+  uint32_t num_rows;
+  uint32_t row_width;
+  float epsilon;
+  uint32_t _pad;
+};
+
+struct WgMapData {
+  WGPUMapAsyncStatus status = WGPUMapAsyncStatus_Error;
+};
+void wg_map_cb(
+    WGPUMapAsyncStatus status,
+    WGPUStringView /*message*/,
+    void* userdata1,
+    void* /*userdata2*/) {
+  static_cast<WgMapData*>(userdata1)->status = status;
+}
+
+// Run the rms_norm scalar kernel at an explicit override wg_size and map the
+// output back. Module::forward can't set a second workgroup size (the handler
+// always clamps to 64), so the pipeline is built directly here; the
+// map/readback mirrors WebGPUGraph::copy_outputs.
+std::vector<float> run_rms_norm_at_wg(
+    const WebGPUContext& ctx,
+    uint32_t wg_size,
+    const std::vector<float>& input,
+    const std::vector<float>& weight,
+    uint32_t num_rows,
+    uint32_t row_width,
+    float epsilon) {
+  WGPUDevice device = ctx.device;
+  const uint64_t out_bytes =
+      static_cast<uint64_t>(num_rows) * row_width * sizeof(float);
+  const uint64_t in_bytes = static_cast<uint64_t>(input.size()) * sizeof(float);
+  const uint64_t w_bytes = static_cast<uint64_t>(weight.size()) * sizeof(float);
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kRmsNormWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry bgl_entries[4] = {};
+  bgl_entries[0].binding = 0;
+  bgl_entries[0].visibility = WGPUShaderStage_Compute;
+  bgl_entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  bgl_entries[1].binding = 1;
+  bgl_entries[1].visibility = WGPUShaderStage_Compute;
+  bgl_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  bgl_entries[2].binding = 2;
+  bgl_entries[2].visibility = WGPUShaderStage_Compute;
+  bgl_entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  bgl_entries[3].binding = 3;
+  bgl_entries[3].visibility = WGPUShaderStage_Compute;
+  bgl_entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = bgl_entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_const = {};
+  wg_const.key = {"wg_size", WGPU_STRLEN};
+  wg_const.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipe_desc = {};
+  pipe_desc.layout = pl;
+  pipe_desc.compute.module = shader;
+  pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipe_desc.compute.constantCount = 1;
+  pipe_desc.compute.constants = &wg_const;
+  WGPUComputePipeline pipe =
+      wgpuDeviceCreateComputePipeline(device, &pipe_desc);
+
+  WGPUBufferDescriptor out_bd = {};
+  out_bd.size = out_bytes;
+  out_bd.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopySrc;
+  WGPUBuffer out_buf = wgpuDeviceCreateBuffer(device, &out_bd);
+
+  WGPUBufferDescriptor in_bd = {};
+  in_bd.size = in_bytes;
+  in_bd.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst;
+  WGPUBuffer in_buf = wgpuDeviceCreateBuffer(device, &in_bd);
+  wgpuQueueWriteBuffer(ctx.queue, in_buf, 0, input.data(), in_bytes);
+
+  WGPUBufferDescriptor w_bd = {};
+  w_bd.size = w_bytes;
+  w_bd.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst;
+  WGPUBuffer w_buf = wgpuDeviceCreateBuffer(device, &w_bd);
+  wgpuQueueWriteBuffer(ctx.queue, w_buf, 0, weight.data(), w_bytes);
+
+  RmsNormProbeParams params = {num_rows, row_width, epsilon, 0u};
+  WGPUBufferDescriptor p_bd = {};
+  p_bd.size = sizeof(RmsNormProbeParams);
+  p_bd.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  WGPUBuffer p_buf = wgpuDeviceCreateBuffer(device, &p_bd);
+  wgpuQueueWriteBuffer(ctx.queue, p_buf, 0, &params, sizeof(params));
+
+  WGPUBufferDescriptor stg_bd = {};
+  stg_bd.size = out_bytes;
+  stg_bd.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst;
+  WGPUBuffer staging = wgpuDeviceCreateBuffer(device, &stg_bd);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = out_buf;
+  bg_entries[0].size = out_bytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = in_buf;
+  bg_entries[1].size = in_bytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = w_buf;
+  bg_entries[2].size = w_bytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = p_buf;
+  bg_entries[3].size = sizeof(RmsNormProbeParams);
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bg = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device, nullptr);
+  WGPUComputePassDescriptor pass_desc = {};
+  WGPUComputePassEncoder pass =
+      wgpuCommandEncoderBeginComputePass(enc, &pass_desc);
+  wgpuComputePassEncoderSetPipeline(pass, pipe);
+  wgpuComputePassEncoderSetBindGroup(pass, 0, bg, 0, nullptr);
+  wgpuComputePassEncoderDispatchWorkgroups(pass, num_rows, 1, 1);
+  wgpuComputePassEncoderEnd(pass);
+  wgpuComputePassEncoderRelease(pass);
+  wgpuCommandEncoderCopyBufferToBuffer(enc, out_buf, 0, staging, 0, out_bytes);
+  WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, nullptr);
+  wgpuQueueSubmit(ctx.queue, 1, &cmd);
+  wgpuCommandBufferRelease(cmd);
+  wgpuCommandEncoderRelease(enc);
+
+  WgMapData cb = {};
+  WGPUBufferMapCallbackInfo cb_info = {};
+  cb_info.mode = WGPUCallbackMode_WaitAnyOnly;
+  cb_info.callback = wg_map_cb;
+  cb_info.userdata1 = &cb;
+  WGPUFuture fut =
+      wgpuBufferMapAsync(staging, WGPUMapMode_Read, 0, out_bytes, cb_info);
+  const WGPUWaitStatus wait = webgpu_wait(ctx.instance, fut);
+
+  std::vector<float> result(static_cast<size_t>(num_rows) * row_width);
+  bool ok = false;
+  if (wait == WGPUWaitStatus_Success &&
+      cb.status == WGPUMapAsyncStatus_Success) {
+    const void* mapped = wgpuBufferGetConstMappedRange(staging, 0, out_bytes);
+    std::memcpy(result.data(), mapped, out_bytes);
+    wgpuBufferUnmap(staging);
+    ok = true;
+  }
+
+  wgpuBufferRelease(staging);
+  wgpuBufferRelease(p_buf);
+  wgpuBufferRelease(w_buf);
+  wgpuBufferRelease(in_buf);
+  wgpuBufferRelease(out_buf);
+  wgpuBindGroupRelease(bg);
+  wgpuComputePipelineRelease(pipe);
+  wgpuPipelineLayoutRelease(pl);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuShaderModuleRelease(shader);
+
+  if (!ok) {
+    throw std::runtime_error("rms_norm wg-size probe: output map failed");
+  }
+  return result;
+}
+
 // linear_q4gsw sweep config; mirrors CONFIGS in test_quantized_linear.py.
 struct Q4gswConfig {
   const char* name;
@@ -1414,6 +1596,56 @@ TEST(WebGPUNative, QueryPoolRoundtrip) {
 }
 #endif // WGPU_BACKEND_ENABLE_PROFILING
 
+// The override wg_size must not change results: run the rms_norm scalar kernel
+// at wg_size 64 and 128 (the handler always clamps to 64, so 128 is only
+// reachable via a direct pipeline) and require element-wise agreement. Absolute
+// rms_norm correctness is covered by the model-driven golden tests; this locks
+// the runtime-configurability guarantee (same WGSL, different size -> same
+// output).
+TEST(WebGPUNative, RmsNormWorkgroupSizeConfigurable) {
+  const WebGPUContext* ctx = get_default_webgpu_context();
+  if (ctx == nullptr || ctx->device == nullptr) {
+    GTEST_SKIP() << "no WebGPU device";
+  }
+  constexpr uint32_t num_rows = 3, row_width = 256;
+  constexpr float epsilon = 1e-5f;
+  std::vector<float> input(static_cast<size_t>(num_rows) * row_width);
+  std::vector<float> weight(row_width);
+  for (size_t i = 0; i < input.size(); i++) {
+    input[i] = std::sin(0.1f * static_cast<float>(i)) * 2.0f + 0.5f;
+  }
+  for (uint32_t j = 0; j < row_width; j++) {
+    weight[j] = 0.5f + 0.001f * static_cast<float>(j);
+  }
+
+  const std::vector<float> out64 =
+      run_rms_norm_at_wg(*ctx, 64, input, weight, num_rows, row_width, epsilon);
+  const std::vector<float> out128 = run_rms_norm_at_wg(
+      *ctx, 128, input, weight, num_rows, row_width, epsilon);
+  ASSERT_EQ(out64.size(), static_cast<size_t>(num_rows) * row_width);
+  ASSERT_EQ(out128.size(), out64.size());
+
+  double sumsq = 0.0;
+  for (float v : out64) {
+    sumsq += static_cast<double>(v) * static_cast<double>(v);
+  }
+  float max_abs = 0.0f, max_rel = 0.0f;
+  const bool consistent = sdpa_within_tol(
+      out64.data(),
+      out128.data(),
+      static_cast<int>(out64.size()),
+      &max_abs,
+      &max_rel);
+  printf(
+      "  rms_norm wg64-vs-wg128: max_abs=%e max_rel=%e sumsq=%f\n",
+      max_abs,
+      max_rel,
+      sumsq);
+  EXPECT_GT(sumsq, 0.0) << "wg64 output is all zero (kernel did not run)";
+  EXPECT_TRUE(consistent) << "wg64 vs wg128 differ beyond tol (abs " << max_abs
+                          << " rel " << max_rel << ")";
+}
+
 TEST(WebGPUNative, UpdateCache) {
   if (g_update_cache_model_path.empty()) {
     GTEST_SKIP() << "WEBGPU_TEST_UPDATE_CACHE_MODEL not set";