[ExecuTorch][WebGPU] Add f16-multiply variant of the steel q4gsw prefill GEMM

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] Test coverage for the f16-multiply steel q4gsw GEMM #20753
• -> [ExecuTorch][WebGPU] Add f16-multiply variant of the steel q4gsw prefill GEMM #20752
• [ExecuTorch][WebGPU] Steel q4gsw prefill GEMM — test coverage #20731
• [ExecuTorch][WebGPU] Add 256-thread "steel" q4gsw prefill GEMM #20730
• [ExecuTorch][WebGPU] Request shader-f16 device feature (enablement, fail-open) #20729
• [ExecuTorch][WebGPU] Unit tests for the WGSL shader-variant codegen #20728
• [ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup #20727

Enables an f16-multiply path for the 256-thread steel q4gsw prefill GEMM on devices that negotiated WebGPU shader-f16 — f16 shared memory + f16 multiply with an f32 accumulator (storage stays f32). It is opt-in and default-off, so no ExecuTorch consumer's numerics change; the default build keeps the f32 kernel and the strict f32 golden. f16 multiply / f32 accumulate (never f16-accumulate, which is numerically unstable on the target GPUs) is a measured prefill win the f32-only kernel could not express.

Key changes:

• q4gsw_linear_gemm_steel.yaml — add a half variant (SUFFIX: half) beside the float variant.
• q4gsw_linear_gemm_steel.wgsl — three $if DTYPE == "half" splits (enable f16;, an f16 dequant multiply, an f32(a * b) cast on the MAC so the accumulator stays f32); each $else is the verbatim f32 line, so the float variant regenerates byte-identical.
• q4gsw_linear_gemm_steel_half_wgsl.h — generated kQ4gswLinearGemmSteelHalfWGSL (f16 shared memory + multiply, f32 accumulate, f32 storage).
• QuantizedLinear.cpp — under WGPU_BACKEND_STEEL_F16, select the half shader only when use_steel and ctx->shader_f16_supported; else fall back to the f32 steel kernel (fail-closed). Bindings, tile, and workgroup count are identical — only shader_src differs.

The device-negotiation gate has no Vulkan analogue: WebGPU shader-f16 is an optional WGSL feature requested at device creation, unlike Vulkan's host-side dtype pick (add_dtype_suffix).

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110802531

Differential Revision: D110802531

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 4c20d00 with merge base fa49f8f ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 99104c286a0b4a4a00f2fed5359c202dfed23b8f5ee7235661ef96668614e2c2 /exec failed with exit code 1
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 192a62ce45dfc63cc50a5f5d5d33ddfb9045ccf17ade2eb95f9ee666990b4b8a /exec failed with exit code 1

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 throoughly review, check for overengineering, correctness, and areaas for modularization

[claude]

Claude finished @JCNTH's task in 4m 27s —— View job

---

Review: f16-multiply steel q4gsw prefill GEMM

• Read PR diff (QuantizedLinear.cpp, .wgsl, .yaml, generated header)
• Verify shader/binding parity and codegen mechanics
• Check correctness (f16 mul / f32 accum, enable f16 placement, fail-closed gate)
• Assess overengineering & modularization

---

Overall this is a clean, well-scoped change. The template splits are minimal and the f32 variant provably regenerates byte-identical (the committed q4gsw_linear_gemm_steel_wgsl.h is untouched, and every $else is the verbatim original line). Binding layout, Params, tile constants, and @workgroup_size are identical between variants, so reusing the same workgroup_count/bindings and only swapping shader_src is correct. Not overengineered.

Correctness — looks right

• enable f16; is emitted as the first line of the half shader (WGSL requires enable before any declaration) and is absent from the float variant. ✓
• f16 multiply / f32 accumulate is implemented as intended: acc stays array<array<f32,4>,4>, and the MAC is acc += f32(a[m] * bvec[n]) — the product is computed in f16, then widened before the f32 add. Never f16-accumulate. ✓
• The dequant f16(i32(nib) - 8) * f16(t_scales[...]) is safe on the nib side (-8..7 is exact in f16); scale/activation/product carry the expected f16 rounding, which is the deliberate tradeoff. ✓
• The gate is genuinely fail-closed: the half shader is selected only when use_steel && ctx != nullptr && ctx->shader_f16_supported, and shader_f16_supported is set only when WGPUFeatureName_ShaderF16 was required of the device (WebGPUDevice.cpp:151-153) — so enable f16; can never reach a device that rejected it. get_default_webgpu_context() is the same context the graph uses, so the negotiation state is consistent. ✓

Things worth flagging (not blockers)

1. WGPU_BACKEND_STEEL_F16 has no build wiring. grep finds the macro only at the two #ifdef sites in QuantizedLinear.cpp — nothing in any CMakeLists.txt, preset, or .bzl defines or exposes it. So today the f16 path is unreachable in a normal build; it can only be turned on by hand-defining the macro. That matches the "opt-in, default-off" intent, but consider adding a documented CMake option (e.g. EXECUTORCH_WEBGPU_STEEL_F16) so it's discoverable and testable, rather than an undocumented compile define. Fix this →

2. No numerical test for the half path. As noted in the description, the golden stays strict-f32 and the half variant isn't exercised. Acceptable while it's off-by-default, but there's currently zero coverage of the f16 kernel's numerics (a tolerance-based golden) — worth a follow-up before anyone flips it on, since f16-product overflow for large activations (f16 max ≈ 65504) is an inherent risk of this path.

3. Modularization: the 6-line #ifdef override in q4gsw_linear_impl is fine as-is; a select_steel_shader() helper would be over-abstraction for a single call site. No change recommended.

Note on verification

I could not run python backends/webgpu/scripts/gen_wgsl_headers.py --check in this sandbox (command approval), so I confirmed codegen consistency by inspection instead: the half header's structure matches render_header, the float header is unchanged in the diff, and CI's --check will catch any embedded-sha/content drift. If you want, I can add the CMake option from point (1) — say the word.
· branch gh/JCNTH/6/head