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

[JulianCloudNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] rms_norm: add a vec4 kernel for 4-aligned row widths #20458
• [ExecuTorch][WebGPU] q4gsw: route M==1 decode to a cooperative GEMV #20457
• -> [ExecuTorch][WebGPU] Register-tile the q4gsw quantized-linear kernel #20456

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.
@exported-using-ghexport

Differential Revision: D109250327

Differential Revision: D109250327

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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

⏳ 117 Pending, 1 Unrelated Failure

As of commit b3965f2 with merge base 68bb668 ():

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

• pull / test-llama-runner-linux (fp32, xnnpack+custom+quantize_kv, linux.2xlarge, executorch-ubuntu-22.04... / linux-job (gh) (matched linux rule in flaky-rules.json)
  The runner has received a shutdown signal. This can happen when the runner service is stopped, or a manually started runner is canceled.

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.

[JulianCloudNTH]

@claude review if there are any bugs

[claude]

Claude finished @JulianCloudNTH's task in 2m 10s —— View job

---

Code Review: Register-tiled q4gsw kernel

I reviewed the shader (q4gsw_linear.wgsl / q4gsw_linear_wgsl.h) and the host dispatch (QuantizedLinear.cpp) against the base branch.

• Read the diff and supporting helpers (WebGPUUtils.h)
• Verify correctness of tile indexing, edge clamping, and dispatch
• Check for bugs

Verdict: No correctness bugs found ✅

The transformation is sound. Specifically I verified:

• Tile mapping is consistent host↔shader. Host dispatches ceil(M/4)*ceil(N/4) threads (total_tiles); shader recomputes the same nrt*nct from params.M/N with TM=TN=4 matching kQ4gswTileM/N. wg_size is computed once and used for both compute_1d_workgroup_count and the wg_size override constant — consistent. The total_tiles > UINT32_MAX pre-check before the cast is correct.
• Edge clamping is safe. n_eff = min(n0+nl, N-1) and m_eff = min(m0+ml, M-1) keep all t_weight/t_scales/t_input reads in bounds for overhang tiles, and the final store is guarded by m < M && n < N, so the clamped/overhang lanes accumulate into acc slots that are never written out. M/N == 0 are already rejected upstream, so N-1u/M-1u can't underflow.
• Math unchanged. out[m,n] = Σ_k input[m,k]·(q·scale) is preserved; only loop nesting changed. The scale index (k/group_size)*padded_N + n_eff stays in bounds (n_eff ≤ N-1 < padded_N, k/group_size < num_groups).
• No barriers/shared memory, so the early return on gid.x >= tiles is safe.

Performance notes (not bugs)

1. Input row is re-read TN (=4) times per k. t_input[m_eff*K + k] depends only on ml/k, but it's read inside the nl loop nested over ml → 16 loads per (tile, k) for only 4 distinct values. Hoisting the TM input values into registers before the nl loop would cut input reads ~4×, mirroring what the PR already does for weights:

   var in_reg: array<f32, TM>;
   for (var ml = 0u; ml < TM; ml++) {
     in_reg[ml] = t_input[min(m0+ml, params.M-1u) * params.K + k];
   }
   for (var nl = 0u; nl < TN; nl++) { /* ... use in_reg[ml] ... */ }

   Fix this →

2. Decode (M=1) now does ~4× redundant work. With M=1, nrt=1 and m_eff clamps all 4 ml lanes to row 0, so each tile computes the row 4× (3 discarded). The PR body says decode is "unchanged," but it's actually somewhat more MAC/input work now (likely hidden if memory-bound). Fine to leave for the follow-up GEMV PR, just noting it isn't a no-op.

Minor

• q4gsw_ceil_div is a single-use 1-LOC helper; per the repo's "no trivial single-use helpers" guideline it could be inlined, but the named version reads clearly enough — judgment call.

Both perf items are optional and the second is explicitly deferred per the PR description, so nothing blocks this change from a correctness standpoint.
· gh/JulianCloudNTH/51/head