[ExecuTorch][WebGPU] Packed-word-dequant f16 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

+24-28% microbench GFLOP/s (+11-16% end-to-end prefill tok/s) on the f16 steel q4gsw prefill GEMM, Apple M4 Pro / Chrome Canary — bit-exact.

Problem: the f16 steel prefill GEMM (the half variant) dequantizes weights per-nibble — each of the 256 threads extracts ONE 4-bit weight from a full u32 word, so every packed word is re-loaded ~8x and the per-column scale is re-read ~16x across the BK=16 tile. That redundant global traffic dominates the kernel on Apple.

Solution: a packed-word-dequant staging variant, bit-exact to the half kernel — identical As staging, compute loop, epilogue, and 64x64 tile / 256-thread / BK=16 geometry; only the weight-dequant staging differs.

Before (half): each thread extracts 1 nibble from a re-loaded word; scale re-read per nibble.
After (pwdq): threads [0,BN) each stage one full BK-column — load the column's two u32 words ONCE, unpack all 16 nibbles, read the per-column scale ONCE.

Implementation:

• New q4gsw_linear_gemm_steel_half_pwdq.wgsl (+ generated header); selected on the WGPU_BACKEND_STEEL_F16 path when the device negotiated shader-f16 AND group_size % BK == 0 (the hoisted scale must be constant across the BK tile), else the per-nibble half kernel.
• No new build option, dispatch, or bindings — same tile/count as steel; only the shader source differs.

Constraints: requires K % BK == 0 (the steel route already guarantees it, making K_packed = K/2 a multiple of 8 so every column is u32-aligned) and group_size % BK == 0 (all real q4 group sizes: 32/64/128). f16-multiply / f32-accumulate, so numerically identical to the half kernel.

Co-authored-with: Claude Code.

Differential Revision: D111163510

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures

As of commit 42c7fe2 with merge base 4af91c3 ():

NEW FAILURES - The following jobs have failed:

• 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 6cc7fd70b311cd8edcdc3b71cf41025f0032e262e2c8806594199e99af71871f /exec failed with exit code 1
• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t f22e2103add29c2e3ea536def6b0e77bb8de06a5fd7a4ec8b879ed91200f93af /exec failed with exit code 1
• pull / unittest / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_mv2_model
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t 228d9a07a82d04d93e23c5ccc6bb163403fad7b7d4ddf9858a1f9508b08d058d /exec failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_resnet18_model

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.