[ExecuTorch][WebGPU] f16-accumulate (pwdqf16acc) steel q4gsw prefill GEMM

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] Test coverage for the f16-accumulate steel GEMM #20801
• -> [ExecuTorch][WebGPU] f16-accumulate (pwdqf16acc) steel q4gsw prefill GEMM #20800
• [ExecuTorch][WebGPU] Test coverage for the packed-word-dequant f16 steel GEMM #20799
• [ExecuTorch][WebGPU] Packed-word-dequant f16 steel q4gsw prefill GEMM #20798
• [ExecuTorch][WebGPU] Reuse pool for single-op scratch (-86.7 MiB) #20797
• [ExecuTorch][WebGPU] Test coverage for the f16 KV cache #20773
• [ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF) #20772
• [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

+46-56% end-to-end prefill tok/s over the shipped f16 steel q4gsw GEMM (Apple M4 Pro / Chrome Canary), behind a new opt-in build flag; perplexity held (13.32 -> 13.37, +0.05).

Problem: the f16 steel prefill GEMM (and its packed-word-dequant variant pwdq) accumulates its 4x4 register tile in f32. WebLLM/MLC accumulate in f16, which halves the accumulator footprint and raises occupancy — the largest remaining prefill gap vs MLC on Apple.

Solution: an f16-accumulate variant of the pwdq kernel — identical staging (dequant-once + hoisted scale) and 64x64/256-thread/BK=16 geometry, but the 4x4 accumulator is kept in f16 with fma() (mirrors MLC's array<f16,16> reduction) and cast to f32 in the epilogue for the f32 output/bias.

Before (pwdq): f16 multiply, f32 accumulate.
After (pwdqf16acc): f16 multiply, f16 accumulate, f32 epilogue.

Implementation:

• New q4gsw_linear_gemm_steel_half_pwdqf16acc.wgsl (+ generated header) + a new CMake option EXECUTORCH_WEBGPU_STEEL_F16ACC (default OFF, mirrors EXECUTORCH_WEBGPU_STEEL_F16) -> WGPU_BACKEND_STEEL_F16ACC.
• When enabled, overrides the f32-accumulate steel kernels for M>1 prefill whenever the device negotiated shader-f16 and group_size % BK == 0; else the f32-accumulate pwdq / half / f32 kernels run (fail-closed).

Constraints: LOSSY — f16 accumulation over the full K (up to 8192) is not bit-exact, so it ships on a perplexity bar, not a bit-exact gate (as MLC does): measured Llama-3.2-1B int4 perplexity 13.32 -> 13.37 (+0.05) on the real prefill path. Default-OFF keeps upstream builds on the strict f32-accumulate golden; the option is opt-in for latency-sensitive deployments.

Co-authored-with: Claude Code.

Differential Revision: D111163606

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 5 New Failures, 1 Unrelated Failure

As of commit d29f101 with merge base 4af91c3 ():

NEW FAILURES - The following jobs have failed:

• Lint / lintrunner (gh)
  >>> Lint for backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h:
• pull / android / build-android (gh)
  ImportError: /opt/hostedtoolcache/Python/3.11.15/x64/lib/python3.11/site-packages/executorch/extension/training/pybindings/_training_lib.cpython-311-x86_64-linux-gnu.so: undefined symbol: _ZN3c104impl3cow23materialize_cow_storageERNS_11StorageImplE
• pull / test-binary-size-linux-gcc / linux-job (gh)
  RuntimeError: Command docker exec -t ae714300fccb1c15d4ecdc749d39084c07333f4db8c21a1a3d8262dc246fd9d3 /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 9e7f290ac6fed934ec69ff6cb629f8ad37e80dc090befcefc65359647ea7f476 /exec failed with exit code 1
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 962b74b782c409a74eba174c36141114b4eefbb9f87a64475835cefe1baee87e /exec failed with exit code 1

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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.