[ExecuTorch][WebGPU] Add 256-thread "steel" q4gsw prefill GEMM

[JCNTH]

Stack from ghstack (oldest at bottom):

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

Add a shared-memory 64x64-tile prefill GEMM for et_vk.linear_q4gsw that beats the existing shmem/tiled GEMM by 1.75-3.17x and moves in-browser Llama-3.2-1B prefill from behind to ahead of llama.cpp WebGPU (+11-33%, M4 Pro / Chrome Canary).

Problem: the q4gsw prefill path (M>1) routes to either the register-tiled or the 32x32 shmem GEMM. On Apple / M4 Pro both leave the 4-bit linear ~2.65x slower than llama.cpp's simdgroup-matrix GEMM, so prefill is the one axis where the WebGPU delegate trails.

Solution:

• Before: M>1 selects use_shmem_gemm (large K/N) else the register-tiled GEMM.
• After: M>1 prefers a new use_steel branch above shmem — a 64x64 output tile computed by 256 threads (16x16) with a 4x4 register sub-tile per thread, staging a BK=16 K-slice of activations (f32) and dequantized weights into shared memory once and reusing it across the tile. shmem/tiled remain the fallback when steel is ineligible.

Implementation:

• Add q4gsw_linear_gemm_steel.wgsl (same 6 bindings + Params as the sibling q4gsw kernels) and its generated _wgsl.h.
• Add a fail-closed steel_supported (requires maxComputeInvocationsPerWorkgroup >= 256; SwiftShader caps at 128 so it falls back) and steel_workgroup_count (one workgroup per 64x64 tile; returns 0 to fall back when K % 16 != 0 — the kernel stages a full BK K-tile with no K-mask — or when the tile count exceeds the 1D dispatch limit).
• Thread use_steel through the shared compute_q4gsw_workgroup_count helper (a steel branch beside gemv/shmem) AND the add_tensor_resize_hook, so a dynamic-shape prefill recomputes the steel tile count for the live M; add steel to fixed_wg (fixed @workgroup_size(16, 16), no wg_size override).
• No direct Vulkan analogue: Vulkan's q4gsw GEMM is register-tiled with an is_gemv split and no 256-thread shmem-staged tile. This tiling is WebGPU-specific (buffer-only storage, compile-time @workgroup_size, 1D-dispatch fold); only the signed-nibble dequant mirrors the q4gsw reference.

Constraints: engine-independent — depends only on shape + device limits and participates in the dynamic-resize hook exactly like gemv/shmem. Falls back to shmem/tiled on non-256-invocation devices, odd K, or over-limit dispatch, so behavior is unchanged where steel is ineligible. Layout and numerics are unchanged (same bindings, same dequant, f32 accumulator). The f16-multiply variant is a separate follow-up.

Co-authored-with: Claude Code.
@exported-using-ghexport

Differential Revision: D110660965

Differential Revision: D110660965

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 4 New Failures, 4 Pending

As of commit 4ad3bd1 with merge base fa49f8f ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  RuntimeError: Command docker exec -t ddc3f8c99148161b844692389524e9451a13f9ce5aa4a745b1b3c78b56773f1d /exec failed with exit code 1
• pull / unittest / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_vit_model
• pull / unittest-editable / linux / linux-job (gh)
  RuntimeError: Command docker exec -t ce8c061ca2c74b84c8fdaead8e14e6334b67309a5c5286da8b7a7a162a2e9db2 /exec failed with exit code 1
• pull / unittest-editable / macos / macos-job (gh)
  export/tests/test_target_recipes.py::TestTargetRecipes::test_vit_model

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[linux-foundation-easycla]

The committers listed above are authorized under a signed CLA.

• ✅ login: JCNTH / name: Julian Ng-Thow-Hing (215344e, e7bb34e)

[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.

[JCNTH]

@claude throoughly review, check for overengineering, correctness, and areaas for modularization

[claude]

Claude finished @JCNTH's task in 2m 43s —— View job

---

Review: steel q4gsw prefill GEMM

• Read the diff (QuantizedLinear.cpp, steel.wgsl, steel_wgsl.h)
• Review correctness (dispatch, guards, kernel numerics)
• Check for overengineering
• Identify modularization opportunities

Solid, well-scoped change. Numerics and dequant mirror the shmem sibling exactly, the fail-closed device guard is the right instinct, and the dynamic-resize wiring is complete. A few things worth addressing before merge — nothing blocking correctness for the target shapes, but a couple of real inconsistencies and clear dedup wins.

---

Correctness

1. Build-time falls back gracefully, resize hard-throws — asymmetric. steel_workgroup_count returns 0 (→ fall back to shmem/tiled) when the tile count exceeds maxComputeWorkgroupsPerDimension, so at build time an over-limit shape silently picks a different kernel. But use_steel is baked into the pipeline at build. On a dynamic-shape resize, compute_q4gsw_workgroup_count re-runs with the live M and, for the same over-limit condition, throws (QuantizedLinear.cpp:125-129) instead of falling back — the kernel is already fixed, so it can't. That's defensible (you can't swap pipelines mid-resize), but the two paths reason about the same limit in opposite ways. Worth a one-line comment on the throw explaining why resize can't fall back the way the build path does, so a future reader doesn't "fix" the asymmetry.

2. maxComputeWorkgroupsPerDimension is queried three times per dispatch. On the build path for a steel dispatch: once in steel_supported (invocations), once in steel_workgroup_count (dim limit), and a third time inside compute_q4gsw_workgroup_count's steel branch. All redundant. Minor, but see modularization below — folding them removes it for free.

3. The steel branch throw is unreachable on the build path (only reachable via resize). Since use_steel is only ever true when steel_workgroup_count(...) > 0 — i.e. total <= max_count — the identical check in compute_q4gsw_workgroup_count can never fire at build time. That's fine (it exists for resize), but it means the limit is validated twice with two separate code copies that must be kept in sync.

4. Kernel numerics — verified correct. The A/B staging (ar/ac, br/bc), the K%BK==0 host guard removing the need for a K-mask, the signed-nibble dequant (f32(i32(nib) - 8) * scale), scale_row computed per shmem-K-row, the 1D→2D tile decode (bx = wid.x % nbN), and the bounds-guarded write all line up with the shmem sibling. No issues found.

---

Overengineering

5. Steel now shadows the shmem GEMM for essentially all M>1 shapes. Previously shmem was gated on K >= 4096 || N >= 2048. Now use_shmem_gemm = !use_gemv && !use_steel && (...), and steel wins for any M>1 with K%16==0 on a ≥256-invocation device. So shmem is only reachable on odd-K shapes or sub-256-invocation devices (SwiftShader) that also hit the large-K/N gate — a narrow slice. That's an intended consequence, but if the plan is for steel to be the prefill path going forward, it's worth a comment noting shmem is now a fallback-only path (and confirming it's still exercised in CI on SwiftShader, which caps at 128 and will route here). Not a request to delete it — just flag the intent so it isn't mistaken for dead code later.

Beyond that, I don't see gratuitous complexity — the kernel is a straightforward 64×64/256-thread tile and the host guards are minimal.

---

Modularization

6. Three copies of the "max workgroups per dimension" query. The block

WGPULimits limits = {};
const uint32_t max_wgs =
    wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
        limits.maxComputeWorkgroupsPerDimension > 0
    ? limits.maxComputeWorkgroupsPerDimension : 65535u;

appears in steel_workgroup_count and twice in compute_q4gsw_workgroup_count (steel + shmem branches), all identical including the 65535u fallback. Extract one helper:

uint32_t max_workgroups_per_dim(WGPUDevice device) {
  WGPULimits limits = {};
  return (wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
          limits.maxComputeWorkgroupsPerDimension > 0)
      ? limits.maxComputeWorkgroupsPerDimension : 65535u;
}

Fix this →

7. The steel and shmem branches of compute_q4gsw_workgroup_count are byte-identical except for the tile constants and the error string. Both are "one workgroup per tile, throw over dim-limit." Collapse into a small helper:

uint32_t tiled_wg_count(WGPUDevice device, uint32_t m, uint32_t n,
                        uint32_t tile_m, uint32_t tile_n,
                        const char* op_name, const char* kind) {
  const int64_t total = utils::div_up<int64_t>(m, tile_m) *
      utils::div_up<int64_t>(n, tile_n);
  if (total > static_cast<int64_t>(max_workgroups_per_dim(device))) {
    throw std::runtime_error(std::string("WebGPU ") + op_name + ": " + kind +
        " tile count exceeds the 1D dispatch limit");
  }
  return static_cast<uint32_t>(total);
}

Then both branches are one call each, and combined with #6 the routing function drops ~25 lines.

8. Consider deriving steel_workgroup_count's dim-limit check from the same helper so the "return 0 to fall back" and the "throw over limit" logic share their notion of the max, rather than each re-deriving max_count/max_wgs independently (the bug surface in #1/#3).

---

Minor / nits

• steel_workgroup_count guards K % 16 != 0 and over-limit but the total == 0u case (m==0 or n==0) also returns 0 — correct, just undocumented in the comment which only mentions K and dispatch limit.
• The generated _wgsl.h correctly carries a wgsl-sha256; confirm it was regenerated by the codegen tool rather than hand-edited (the source and embedded copy match, so this looks fine).

Overall: correct for the target shapes, good guard discipline, and the perf win is well-motivated. The main asks are the max-limit dedup (#6/#7) and a comment clarifying the build-vs-resize fallback asymmetry (#1).