[ExecuTorch][WebGPU] Register-tile the q4gsw quantized-linear kernel

[JulianCloudNTH]

Stack from ghstack (oldest at bottom):

• -> [ExecuTorch][WebGPU] Register-tile the q4gsw quantized-linear kernel #20490

Register-tile the et_vk.linear_q4gsw GEMM — up to 3.4x faster prefill (M4 Pro, M=128).

Problem: et_vk.linear_q4gsw (4-bit weight-only, W4A16) computes out[m,n] = bias[n] + sum_k input[m,k] * (nibble(weight,n,k)-8) * scale[k/group_size, n] in a single dispatch over a raw [N, K/2] 4-bit weight (2 nibbles/byte, +8-shifted symmetric, groupwise scales). The shipped kernel was naive: one workgroup per output row m, threads striding N, a scalar K-loop per (m,n). For an M-row (prefill) input it re-extracts every dequantized weight M times (once per row) and re-reads each input value once per output column — redundant memory traffic that dominates the prefill GEMM.

Solution: a register-tiled GEMM where each thread owns a TM x TN = 4x4 output tile, so both weights and inputs are loaded once per tile instead of once per element.

• Before: weight (n,k) dequantized once per (m,n) (extracted Mx for prefill); input[m,k] re-read once per output column n.
• After: weight (n,k) dequantized ONCE and reused across the TM rows of the tile (weight reads drop ~TMx); each input[m,k] loaded once per k into a register and reused across the TN columns (input reads drop ~TNx).

Implementation:

• New loop nest in q4gsw_linear.wgsl: per k, hoist the TM input values into registers, then for each of the TN columns dequantize the weight once and accumulate into the 4x4 register tile.
• Host dispatch changes from M workgroups to ceil(M/TM)*ceil(N/TN) tiles over wg_size threads; wg_size is computed before the count so the dispatch is still validated against device limits before any allocation.
• Tile-edge lanes (n0+nl >= N or m0+ml >= M) clamp their weight/scale/input index to the last valid element (the never-stored overhang is harmless), since WGSL out-of-bounds reads are implementation-defined. Mirrors the Vulkan tiled GEMM q4gsw_linear_gemm__w_4x8.glsl's min(..., N-1) clamp.
• Deliberate deviations from the Vulkan kernel (recorded in DESIGN_DECISIONS): a 4x4 tile (vs Vulkan 4M x 8N) for a conservative register budget; the RAW [N,K/2] layout with scalar nibble unpack and NO W_4X8 prepack / NO wide vec4<u32> loads (prior on-device measurement found wide loads regress on this GPU); a 1D-flattened tile index (the backend is 1D-dispatch only).

Constraints: bindings, Params, the weight layout, and the single-dispatch structure are unchanged; the dequant index math is copied verbatim from the naive kernel, so the result is a floating-point accumulation reorder equal to the naive output to fp-rounding. The M=1 decode GEMV path and host M-based routing are a separate follow-up.

Authored with assistance from Claude Code.

Differential Revision: D109250327

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

• [ROCm] MI350 CI runner label rename: rebase PRs using old linux.rocm.gpu.gfx950.* labels

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.