[ExecuTorch][WebGPU] Add 4-bit weight-only quantized linear (et_vk.linear_q4gsw)

[JulianCloudNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] GPU timestamp query profiling for SDPA #20167
• [ExecuTorch][WebGPU] linear_q4gsw test suite: Llama-1B shapes + 4k/8k sweep #20227
• -> [ExecuTorch][WebGPU] Add 4-bit weight-only quantized linear (et_vk.linear_q4gsw) #20226
• [ExecuTorch][WebGPU] SDPA test suite: replay + dynamic input_pos + in-graph KV cache #20087
• [ExecuTorch][WebGPU] Add fused SDPA (sdpa_with_kv_cache) with dynamic input_pos #20086
• [ExecuTorch][WebGPU] GPU timestamp query profiling (general implementation) #20201

Adds the et_vk.linear_q4gsw operator (4-bit groupwise-symmetric weight-only linear) to the WebGPU backend: dequantize the packed int4 weight in WGSL ((q-8)*scale) and accumulate an fp32 matmul, consuming the serialized [N, K/2] uint8 weight directly (no prepack), one workgroup per output row. Mirrors the Vulkan reference (backends/vulkan/.../impl/QuantizedLinear.cpp). The dispatch carries a linear_q4gsw label for GPU-timestamp-query profiling (mirroring the SDPA kernels). The numerical test suite is in the stacked test diff.

Differential Revision: D108312283

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Unrelated Failure

As of commit bab43f6 with merge base 5526971 ():

NEW FAILURE - The following job has failed:

• pull / test-phi-3-mini-runner-linux / linux-job (gh)
  RuntimeError: Command docker exec -t 5fb5d3aad4a63d8353a77acb6d70ca46c1d107e07932fc5e25ce47ddca452d78 /exec failed with exit code 1

BROKEN TRUNK - The following job failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / android / build-android (gh) (trunk failure)
  Process completed 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.

[JulianCloudNTH]

@claude review

[claude]

Claude finished @JulianCloudNTH's task in 4m 47s —— View job

---

Claude Code Review

• Read changed files
• Review WGSL kernel
• Review C++ dispatch & resource management
• Cross-check against backend conventions (SDPA, codegen)
• Post feedback

Overall this is a clean, well-guarded first cut. Validation-before-allocation, the fp32-only byte-size checks, the K_packed == ceil(K/2) guard, and the resource-release pattern all match the existing SDPA op conventions, and the generated header's wgsl-sha256 matches the .wgsl source (verified locally), so codegen is in sync. Comments are appropriately minimal. A few correctness edge cases are worth tightening before this lands.

---

Potential bugs / correctness

1. Weight bound as array<u32> can over-read the last word when N * K_packed isn't a multiple of 4 — QuantizedLinear.cpp:198, q4gsw_linear.wgsl:45

The weight is bound as array<u32> with size = weight.nbytes = N * K_packed bytes. The shader computes t_weight[byte_idx >> 2u] with a maximum byte_idx = N*K_packed - 1, so the highest word index is (N*K_packed - 1) >> 2. When N * K_packed is not a multiple of 4, the effective array<u32> length is floor(nbytes / 4), and that final word read falls outside the bound — WebGPU robustness clamps it to 0, silently producing wrong values for the affected output column(s). This is shape-dependent (e.g. any N*K_packed ≡ 1,2,3 (mod 4)), so it can pass on aligned test shapes and fail in the wild.

Consider validating (N * K_packed) % 4 == 0, or rounding the weight binding size up to a 4-byte multiple, or indexing the weight as array<u32> only after asserting alignment. Worth a dedicated test shape with an odd N*K_packed.

Fix this →

2. scales bounds (num_groups, padded_N) are not validated against the indices the shader uses — QuantizedLinear.cpp:69-70, q4gsw_linear.wgsl:54

The shader indexes t_scales[(k / group_size) * padded_N + n], with max row (K-1)/group_size and max col N-1. The current byte-size guard only checks scales.nbytes == num_groups * padded_N * 4; it never verifies that num_groups >= ceil(K / group_size) or that padded_N >= N. If either is smaller than expected (mismatched serialization), the kernel reads OOB scales (robustness-clamped, wrong results) rather than failing loudly. Adding if (num_groups < (K + group_size - 1) / group_size || padded_N < N) throw … would catch malformed inputs at dispatch time. Note group_size is read at line 96-102, so this check belongs after that.

3. M = in_numel / K truncates silently if in_numel % K != 0 — QuantizedLinear.cpp:66

A non-multiple input numel would silently drop the remainder into a smaller M rather than erroring. A quick if (in_numel % K != 0) throw … keeps the failure mode explicit, consistent with the other malformed-dims guards.

---

Performance (non-blocking)

4. Scale reload + integer division every K iteration — q4gsw_linear.wgsl:54

(k / params.group_size) * padded_N + n recomputes an integer division and reloads the scale on every inner iteration, even though the scale only changes when k crosses a group boundary. Since this is the matmul hot loop, hoisting the scale fetch to update only at group boundaries (or iterating group-by-group) would cut a division + global load per MAC. Fine to defer if this kernel is correctness-first for now, but worth a // TODO.

---

Minor / style

• QuantizedLinear.cpp:151 declares WGPUBindGroupLayoutEntry entries[6] and the SDPA op factored this into a shared build_dispatch helper (ops/sdpa/Sdpa.cpp:150). Not required, but if more quantized-linear variants are coming in the stack, reusing that helper would avoid duplicating the bind-group/pipeline boilerplate.
• static_assert(sizeof(Q4gswParams) == 32, …) is a nice guard against accidental struct drift — good.
• Error strings consistently prefixed WebGPU linear_q4gsw: — good for triage.

None of the above are blockers except items 1–2, which I'd recommend guarding (or covering with an odd-shape test) before relying on this in production, since both fail silently via robustness clamping rather than throwing.