[ExecuTorch][WebGPU] Make shader workgroup sizes runtime-configurable

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] Test: rms_norm workgroup size is runtime-configurable #20956
• [ExecuTorch][WebGPU] Remove the orphan q4gsw_linear_coop4 shader #20955
• -> [ExecuTorch][WebGPU] Make shader workgroup sizes runtime-configurable #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.
@exported-using-ghexport

Differential Revision: D112060863

Differential Revision: D112060863

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 fb95a32 with merge base e500fee ():
💚 Looks good so far! There are no failures yet. 💚

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.

[JCNTH]

@claude Review this PR

[claude]

Claude finished @JCNTH's task in 2m 54s —— View job

---

PR Review: WebGPU runtime-configurable workgroup sizes

• Read the changed WGSL shaders and host handlers
• Verify override wiring & clamp_workgroup_size semantics
• Check regenerated _wgsl.h headers match source
• Review geometry-lock notes for fixed kernels
• Post review feedback

This is a clean, well-scoped refactor. The const WG* → override wg_size conversion is applied consistently across the three genuinely-tunable kernels, the host wiring mirrors the landed add op, and the geometry-lock comments on the fixed kernels are accurate and genuinely helpful for future readers. The false-knob fix on q4gsw_linear_gemm_shmem.wgsl (override → const, since geometry is hard-locked to 64 and it was never host-wired) is a nice correctness cleanup. Behavior is preserved at the default 64. Headers are regenerated 1:1 with their sources (sha256 updated).

I reviewed statically (no WebGPU device build here). A few notes below — nothing blocking.

---

Observations

1. The power-of-two invariant is documented but not enforced (latent hazard)

Each tree-reduction kernel now carries:

// wg_size must be a power of two (the tree reduction halves the stride).
override wg_size: u32 = 64u;

This is correct — stride = wg_size / 2u (or >> 1u) drops elements if wg_size isn't a power of two. The value flows from clamp_workgroup_size:

return std::min(desired, limits.maxComputeInvocationsPerWorkgroup);

Today this is always safe: the default is 64, the spec floor for maxComputeInvocationsPerWorkgroup is 256, so clamp is a no-op and 64 is emitted. But the invariant is only preserved by that coincidence. If the default is ever raised, or a device reports a non-power-of-two limit below the default, std::min could hand a tree-reduction kernel a broken size with no guardrail. Since the PR's stated goal is "the size can be tuned per device/shape", this is worth hardening now rather than after a future tuning change silently corrupts a reduction.

Suggestion: a debug assert / comment at the call sites for the reduction kernels (rms_norm, sdpa_softmax, coop4_bicol) noting the pow2 requirement, or a small clamp_workgroup_size_pow2 variant that floors to the nearest power of two. The plain-dispatch kernels (add/mul/cat/…) don't have this constraint, so a shared clamp helper shouldn't unconditionally force pow2.

2. gemv_wg_size is computed unconditionally — QuantizedLinear.cpp:258

const uint32_t gemv_wg_size =
    utils::clamp_workgroup_size(device, kQ4gswLinearCoop4BicolWorkgroupSizeX);

This runs a second wgpuDeviceGetLimits on every call even for the M>1 GEMM paths that never use it. Negligible (pipeline-build time, cheap query), but it could be sunk into the use_gemv branch of the ternary if you'd rather not query twice.

3. Sdpa.cpp softmax: 0 → sm_wg — confirmed behavior-preserving

The softmax dispatch now passes a clamped 64 instead of the sentinel 0. build_dispatch still guards with if (wg_size != 0) (Sdpa.cpp:200), so the constant is now supplied and the shader's override wg_size is set explicitly rather than falling back to its default. Since both resolve to 64 and compute_2d_workgroup_count is independent of wg_size (one workgroup per row), dispatch counts are unchanged. Good.

4. Override-sized workgroup arrays — e.g. var<workgroup> shared_sum: array<f32, wg_size>

WGSL permits an override-expression as a var<workgroup> array element count (this is the one place override array sizes are allowed), so this is legal and, per the PR body, mirrors the already-landed add op. No issue.

---

Verdict

Looks good to merge. The only item I'd genuinely consider addressing is (1) — it's a latent footgun that the "make it tunable" framing invites, even though every path is safe at the current default. (2) and (3) are informational.
· branch gh/JCNTH/86/head