[ExecuTorch][WebGPU] Add 256-thread "steel" q4gsw prefill GEMM#20815
Merged
Conversation
Pull Request resolved: #20730 **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. ghstack-source-id: 401515156 @exported-using-ghexport Differential Revision: [D110660965](https://our.internmc.facebook.com/intern/diff/D110660965/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20815
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Pull Request resolved: #20731 **Add steel-GEMM coverage to the `et_vk.linear_q4gsw` golden sweep (stacked on the steel op diff).** The op diff routes M>1 q4gsw prefill to the new steel GEMM on a >=256-invocation device (`K % 16 == 0`), falling back to shmem/register-tiled otherwise. The existing M>1 CONFIGS (`q_proj_4k`, `gate_proj_pf`, `down_proj_pf`, `shmem_edge`) already exercise steel on such a device via the shape-discovering native sweep; this adds one small config that isolates the steel branch specifically and documents the routing. **Changes:** - `test_quantized_linear.py` / `test_webgpu_native.cpp`: add the `steel` config (M=96, K=2048, N=256) — below the shmem thresholds (K<4096, N<2048) so pre-steel it was register-tiled, which uniquely pins the steel branch; M=96 exercises the partial 64-row tile (edge masking). - Document that M>1 `K % 16 == 0` shapes prefer steel on a >=256-invocation device (lvp) and fall back on a <256 device (SwiftShader) — the same fp64 golden validates whichever kernel runs. Co-authored-with: Claude Code. ghstack-source-id: 401515169 @exported-using-ghexport Differential Revision: [D110660966](https://our.internmc.facebook.com/intern/diff/D110660966/)
…ill GEMM Pull Request resolved: #20752 **Enables an f16-multiply path for the 256-thread steel q4gsw prefill GEMM** on devices that negotiated WebGPU shader-f16 — f16 shared memory + f16 multiply with an f32 accumulator (storage stays f32). It is opt-in and default-off, so no ExecuTorch consumer's numerics change; the default build keeps the f32 kernel and the strict f32 golden. f16 multiply / f32 accumulate (never f16-accumulate, which is numerically unstable on the target GPUs) is a measured prefill win the f32-only kernel could not express. **Key changes:** - `q4gsw_linear_gemm_steel.yaml` — add a `half` variant (`SUFFIX: half`) beside the `float` variant. - `q4gsw_linear_gemm_steel.wgsl` — three `$if DTYPE == "half"` splits (`enable f16;`, an f16 dequant multiply, an `f32(a * b)` cast on the MAC so the accumulator stays f32); each `$else` is the verbatim f32 line, so the float variant regenerates byte-identical. - `q4gsw_linear_gemm_steel_half_wgsl.h` — generated `kQ4gswLinearGemmSteelHalfWGSL` (f16 shared memory + multiply, f32 accumulate, f32 storage). - `QuantizedLinear.cpp` — under `WGPU_BACKEND_STEEL_F16`, select the half shader only when `use_steel` and `ctx->shader_f16_supported`; else fall back to the f32 steel kernel (fail-closed). Bindings, tile, and workgroup count are identical — only `shader_src` differs. The device-negotiation gate has no Vulkan analogue: WebGPU shader-f16 is an optional WGSL feature requested at device creation, unlike Vulkan's host-side dtype pick (`add_dtype_suffix`). Co-authored-with: Claude Code. ghstack-source-id: 401515172 @exported-using-ghexport Differential Revision: [D110802531](https://our.internmc.facebook.com/intern/diff/D110802531/)
Pull Request resolved: #20753 **Adds numeric coverage for the f16-multiply steel q4gsw kernel** from the parent diff, which runs only in a `-DWGPU_BACKEND_STEEL_F16` build and had none. Its f16 multiply has a rounding floor (~2.3e-4) above the strict f32 golden (1e-4), so the new configs use a looser abs tolerance without touching the f32 golden. **Key changes:** - `test_quantized_linear.py` CONFIGS — add `steel_f16` (same shape as the `steel` config, exact-N tile) and `steel_f16_edge` (M=70, K=1024, N=136, partial M and N tiles). The exported `.pte` is dtype-independent, so each fixture drives whichever steel kernel the runtime build selected. - `test_webgpu_native.cpp` `kQ4gswConfigs` — matching `steel_f16`/`steel_f16_edge` entries at abs 2.3e-4 / rel 1e-3, both under `#ifdef WGPU_BACKEND_STEEL_F16`. The native sweep self-discovers each `.pte` by name and goldens against the same fp64 truth the f32 configs use. The looser abs gate reflects the f16 rounding floor (uniform in K, not an accumulate artifact); the accumulator stays f32. The default f32 build's config table and its strict 1e-4 golden are `#ifdef`-excluded and byte-unchanged. Co-authored-with: Claude Code. ghstack-source-id: 401515178 @exported-using-ghexport Differential Revision: [D110802559](https://our.internmc.facebook.com/intern/diff/D110802559/)
…t-OFF)
**Add an opt-in native f16 KV cache for WebGPU SDPA — −280 MiB device memory, decode-neutral, default-OFF (byte-identical when off).**
**Problem:** The SDPA K/V cache is stored fp32, so at long context it dominates WebGPU device memory. On a device that negotiated the shader-f16 feature the cache can be stored as native f16 (half the bytes) while attention still accumulates in f32, at no measured decode-speed cost.
**Solution:** Behind a new opt-in build flag `EXECUTORCH_WEBGPU_KV_F16` (defines `WGPU_BACKEND_KV_F16`, default OFF), and only when the device supports shader-f16 (fail-closed to f32 otherwise):
- **Before:** the K/V caches are allocated fp32 (`numel*4` bytes) and read by the f32 SDPA kernels.
- **After:** the K/V caches are allocated on a dedicated f16 buffer (`numel*2` bytes, zero-init); SDPA and FlashDecoding select the generated f16 (`_half`) kernel variants that bind the cache as f16 and widen to f32 on read — compute and accumulation are unchanged.
**Implementation:**
- Retires the 4 SDPA/decode kernels (`sdpa_compute_attn_weights`, `sdpa_compute_out`, `sdpa_fd_split`, `update_cache`) into `DTYPE`-templated variants via the landed WGSL shader-variant codegen: each `<kernel>.wgsl` + a `<kernel>.yaml` (`DTYPE: float/half`) generates BOTH the existing f32 header (regenerated byte-identical) and a new `_half` header — the f16 variant is generated, not a hand-duplicated file. The f16 delta is storage-only: `$if DTYPE == "half"` enables f16, the K/V cache binding becomes `array<${buffer_gvec_type(DTYPE, 4)}>` (QK/AV) or `array<${buffer_scalar_type(DTYPE)}>` (fd_split/update_cache), widened to f32 on read (`vec4<f32>(...)` / `f32(...)`); `update_cache` writes `f16(...)`. Mirrors the codegen usage of the landed `rms_norm` (an existing kernel retired into a template) and `steel_f16`.
- `WebGPUGraph::build`: collect the sdpa K/V cache value ids (`args[3]`, `args[4]`), allocate them as a dedicated half-size f16 buffer at the constant-allocation site, plus a defensive guard that throws if a non-sdpa op would misread an f16 cache.
- `Sdpa.cpp` / `SdpaFdDecode.cpp`: select the generated `_half` shader when `kv_f16()` via a local `const char*` (mirrors the steel-f16 gate — no function-signature/ABI change).
- Analogous to Vulkan's whole-graph `force_fp16` fp32→fp16 storage downcast (`backends/vulkan/serialization/vulkan_graph_builder.py` `get_effective_dtype`); here scoped to the K/V cache and gated on the negotiated shader-f16 feature rather than applied graph-wide.
**Constraints:** the default build (flag OFF) is byte-identical — all f16 code is `#ifdef WGPU_BACKEND_KV_F16`-gated, the templated f32 headers regenerate byte-identical (drift-check verified), no ABI change, no KV-cache-layout migration; the opt-in build fail-closes to f32 on a non-shader-f16 device; f16 storage requires `head_dim % 4 == 0` (already required and guarded).
Co-authored-with: Claude Code.
Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport
Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)
Pull Request resolved: #20773 Tests for the opt-in f16 KV cache (stacked op diff). When built with `EXECUTORCH_WEBGPU_KV_F16` and run on a shader-f16 device, the SDPA op stores the K/V cache as f16, so the entire existing `sdpa_with_kv_cache` golden suite (config sweep + replay + dynamic decode) exercises the f16-KV path with no new goldens needed; this diff loosens the SDPA numeric tolerance to the f16 read-precision floor when — and only when — that path is active. Key changes: - `test_webgpu_native.cpp` `sdpa_within_tol` — under `#ifdef WGPU_BACKEND_KV_F16`, compare at abs `2e-3` / rel `1e-2` when the device negotiated shader-f16, else keep the strict f32 abs `1e-4` / rel `1e-3`. Every SDPA config/replay/decode golden then validates the f16-KV output through the existing exemplars. Constraints: the default (flag-OFF) test build is unchanged; on a non-shader-f16 device (including the CI software adapter) the f16 KV path stays inactive and the strict f32 tolerance applies, so there is no CI behavior change. The f16-KV numeric validation is therefore shader-f16-device-only (Canary/Metal), matching the op's opt-in gating; no CI script change (mirrors the steel-f16 tests). Co-authored-with: Claude Code. ghstack-source-id: 401515179 @exported-using-ghexport Differential Revision: [D110919973](https://our.internmc.facebook.com/intern/diff/D110919973/)
Pull Request resolved: #20797 Problem: `WebGPUGraph::create_scratch_buffer` keeps every fused-op scratch `WGPUBuffer` alive in `scratch_buffers_` until graph teardown. The SDPA prefill workspace (`attn_weights` + `attn_weights_softmax`, each `Hq*S*Cmax*4` bytes) and the FlashDecoding partials (`part_o`/`part_ml`) are allocated fresh in every decoder layer's lowering, so a 16-layer Llama holds ~16x copies of scratch that is only live during a single op — a large slice of the device compute/scratch footprint. Solution: add an acquire/release reuse pool for SINGLE-OP-LIFETIME scratch. `acquire_scratch()` returns a free slot (best-fit, size in `[n, 2n]`) or creates one; the caller releases it at op-lowering scope exit via the RAII `ScopedScratch` guard. Across the layers the SDPA/FD scratch now recycles a small constant of buffers instead of Nx. Whole-graph-lifetime scratch stays on `create_scratch_buffer` unchanged. Correctness: reuse is bit-identical because WebGPU/Dawn auto-inserts RAW hazard barriers between dispatches that share a storage buffer regardless of pass structure — the same guarantee the shipped `mem_obj_id` tensor aliasing already relies on. A still-`in_use` slot is never handed to a co-live requester (co-live safety). `WEBGPU_NO_SCRATCH_POOL` falls back to the per-call path for A/B. ghstack-source-id: 401515183 @exported-using-ghexport Differential Revision: [D110948389](https://our.internmc.facebook.com/intern/diff/D110948389/)
Pull Request resolved: #20798 **+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 `PWDQ` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_wgsl.h`) — no standalone shader file. Selected at runtime when the device negotiated shader-f16 (`ctx->shader_f16_supported`) 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 generated shader variant 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. ghstack-source-id: 401564878 @exported-using-ghexport Differential Revision: [D111163510](https://our.internmc.facebook.com/intern/diff/D111163510/)
…eel GEMM Pull Request resolved: #20799 Locks the `group_size % BK` gate that selects the packed-word-dequant (`pwdq`) f16 steel GEMM. `pwdq` is bit-exact to the `half` kernel, so the existing `steel_f16`/`steel_f16_edge` configs (gs=32) already exercise it; these add the two group sizes the gate keys on but those omit. Key changes: - `test_quantized_linear.py` / `test_webgpu_native.cpp` — add `pwdq_gs64` (gs=64, `% BK == 0` → stays on `pwdq`) and `pwdq_gs8` (gs=8 `< BK=16` → falls back to the per-nibble `half` kernel, whose per-K scale read is correct there). Both goldened at the f16 tol (2.3e-4) in the `WGPU_BACKEND_STEEL_F16` build against the fp64 dequant-matmul truth. Co-authored-with: Claude Code. ghstack-source-id: 401564881 @exported-using-ghexport Differential Revision: [D111163551](https://our.internmc.facebook.com/intern/diff/D111163551/)
…GEMM Pull Request resolved: #20800 **+46-56% end-to-end prefill tok/s over the shipped f16 `steel` q4gsw GEMM (Apple M4 Pro / Chrome Canary), behind an opt-in runtime spec (default OFF); perplexity held (13.32 -> 13.37, +0.05).** **Problem:** the f16 `steel` prefill GEMM (and its packed-word-dequant variant `pwdq`) accumulates its 4x4 register tile in f32. WebLLM/MLC accumulate in f16, which halves the accumulator footprint and raises occupancy — the largest remaining prefill gap vs MLC on Apple. **Solution:** an f16-accumulate variant of the `pwdq` kernel — identical staging (dequant-once + hoisted scale) and 64x64/256-thread/BK=16 geometry, but the 4x4 accumulator is kept in f16 with `fma()` (mirrors MLC's `array<f16,16>` reduction) and cast to f32 in the epilogue for the f32 output/bias. Before (`pwdq`): f16 multiply, f32 accumulate. After (`pwdqf16acc`): f16 multiply, f16 accumulate, f32 epilogue. **Implementation:** - New `ACC=half` fork in the shared `q4gsw_linear_gemm_steel.wgsl` template (a `shader_variants` entry in `q4gsw_linear_gemm_steel.yaml` generates `q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h`) — no standalone shader file. - Opt-in via the `enable_f16_accumulate_gemm` runtime spec (a load-time `BackendOption` read in `WebGPUBackend::init`, threaded through `WebGPUGraph::build(..., f16_accumulate_gemm)` -> `graph.f16_accumulate_gemm()`), default OFF — no CMake option or compile flag. - When the spec is set, overrides the f32-accumulate steel kernels for M>1 prefill whenever the device negotiated shader-f16 and `group_size % BK == 0`; else the f32-accumulate `pwdq` / `half` / f32 kernels run (fail-closed). **Constraints:** LOSSY — f16 accumulation over the full K (up to 8192) is not bit-exact, so it ships on a perplexity bar, not a bit-exact gate (as MLC does): measured Llama-3.2-1B int4 perplexity 13.32 -> 13.37 (+0.05) on the real prefill path. Default-OFF keeps upstream builds on the strict f32-accumulate golden; the runtime spec is opt-in for latency-sensitive deployments. Co-authored-with: Claude Code. ghstack-source-id: 401564900 @exported-using-ghexport Differential Revision: [D111163606](https://our.internmc.facebook.com/intern/diff/D111163606/)
Pull Request resolved: #20801 Adds golden coverage for the f16-accumulate (`pwdqf16acc`) steel GEMM, which runs under `WGPU_BACKEND_STEEL_F16ACC` when `group_size % BK == 0`. Key changes: - `test_quantized_linear.py` / `test_webgpu_native.cpp` — add `pwdqf16acc` (96x2048x256, K=2048) and `pwdqf16acc_down` (128x8192x2048, deep-K worst case) under `#ifdef WGPU_BACKEND_STEEL_F16ACC`, goldened against the fp64 dequant-matmul truth. f16 accumulation error grows with K, so the tolerances are wider than the f16-multiply `steel_f16` (2.3e-4) and the deep-K shape gets the loosest gate; perplexity (the kernel diff) is the primary quality bar and this catches gross bit/index bugs. Co-authored-with: Claude Code. ghstack-source-id: 401515200 @exported-using-ghexport Differential Revision: [D111163651](https://our.internmc.facebook.com/intern/diff/D111163651/)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #20730 by @JCNTH
^ Please use this as the source of truth for the PR details, comments, and reviews
ghstack PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/4/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/4/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/JCNTH/3/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/4/orig
@diff-train-skip-merge