From 527a196ee324eebc82337a428e06923fc41f6aef Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 9 Jul 2026 13:30:05 -0700 Subject: [PATCH 1/2] [ExecuTorch][WebGPU] f16-accumulate (pwdqf16acc) steel q4gsw prefill GEMM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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` 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/) --- backends/webgpu/runtime/WebGPUBackend.cpp | 11 +- backends/webgpu/runtime/WebGPUGraph.cpp | 7 +- backends/webgpu/runtime/WebGPUGraph.h | 10 +- .../ops/quantized_linear/QuantizedLinear.cpp | 12 ++ .../q4gsw_linear_gemm_steel.yaml | 4 + ..._linear_gemm_steel_half_pwdq_f16acc_wgsl.h | 141 ++++++++++++++++++ 6 files changed, 182 insertions(+), 3 deletions(-) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUBackend.cpp b/backends/webgpu/runtime/WebGPUBackend.cpp index 0b0b5254908..e1497bc4ed8 100644 --- a/backends/webgpu/runtime/WebGPUBackend.cpp +++ b/backends/webgpu/runtime/WebGPUBackend.cpp @@ -89,13 +89,22 @@ Result WebGPUBackend::init( enable_f16_kv_cache = spec.get(); } } + bool enable_f16_accumulate_gemm = false; + { + Result spec = + context.get_runtime_spec("enable_f16_accumulate_gemm"); + if (spec.ok()) { + enable_f16_accumulate_gemm = spec.get(); + } + } try { graph->build( flatbuffer_data, constant_data, context.get_named_data_map(), - enable_f16_kv_cache); + enable_f16_kv_cache, + enable_f16_accumulate_gemm); } catch (const std::exception& e) { ET_LOG(Error, "WebGPU graph build failed: %s", e.what()); graph->~WebGPUGraph(); diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 8fa4f5e35fd..d9ab6467882 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -359,7 +359,8 @@ void WebGPUGraph::build( const void* flatbuffer_data, const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map, - bool f16_kv_cache) { + bool f16_kv_cache, + bool f16_accumulate_gemm) { if (!device_) { auto* ctx = get_default_webgpu_context(); if (ctx) { @@ -385,6 +386,10 @@ void WebGPUGraph::build( const WebGPUContext* kv_ctx = get_default_webgpu_context(); kv_f16_ = f16_kv_cache && (kv_ctx != nullptr && kv_ctx->shader_f16_supported); + // f16-accumulate q4gsw steel prefill GEMM (runtime opt-in). QuantizedLinear + // additionally gates the kernel on the negotiated shader-f16 feature. + f16_accumulate_gemm_ = f16_accumulate_gemm; + // Phase 1: Create all values const auto* values = graph->values(); const int num_vals = values ? values->size() : 0; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index db4f8e499de..66f0e401de5 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -105,7 +105,8 @@ class WebGPUGraph { const void* flatbuffer_data, const uint8_t* constant_data, const executorch::runtime::NamedDataMap* named_data_map = nullptr, - bool f16_kv_cache = false); + bool f16_kv_cache = false, + bool f16_accumulate_gemm = false); // Copy input tensor data from host pointers into GPU buffers. void copy_inputs(const std::vector& inputs); @@ -350,9 +351,16 @@ class WebGPUGraph { return kv_f16_; } + // True when the q4gsw steel prefill GEMM uses the lossy f16-accumulate kernel + // (runtime opt-in; perplexity-gated, not bit-exact). + bool f16_accumulate_gemm() const { + return f16_accumulate_gemm_; + } + private: bool kv_f16_ = false; std::unordered_set kv_cache_ids_; + bool f16_accumulate_gemm_ = false; private: WGPUInstance instance_ = nullptr; diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 8e66d5c19e9..b0728764310 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -280,6 +281,17 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { : kQ4gswLinearGemmSteelHalfWGSL; } } + // f16-accumulate: pwdq staging with an f16 register accumulator. + // Lossy (f16 accumulate over K) -> opt-in via the enable_f16_accumulate_gemm + // runtime spec (default off), gated on the negotiated shader-f16 feature and + // group_size % BK == 0 (same hoisted-scale requirement as pwdq). Overrides + // the f32-accumulate steel kernels. + if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) { + const WebGPUContext* ctx = get_default_webgpu_context(); + if (ctx != nullptr && ctx->shader_f16_supported) { + shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + } + } const uint32_t workgroup_count = compute_q4gsw_workgroup_count( device, use_gemv, diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml index 554fab0f878..5a2cae5e499 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml @@ -16,3 +16,7 @@ q4gsw_linear_gemm_steel: DTYPE: half PWDQ: true ACC: float + - NAME: q4gsw_linear_gemm_steel_half_pwdq_f16acc + DTYPE: half + PWDQ: true + ACC: half diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h new file mode 100644 index 00000000000..efefd7edce1 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -0,0 +1,141 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. +// wgsl-sha256: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce +inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"( +enable f16; +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_input: array; +@group(0) @binding(2) var t_weight: array; +@group(0) @binding(3) var t_scales: array; +@group(0) @binding(4) var t_bias: array; + +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(5) var params: Params; + +// "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. +// The "steel" name + register-tiled dequant-to-shared GEMM structure are +// inspired by MLX's steel GEMM kernels (github.com/ml-explore/mlx, +// mlx/backend/metal/kernels/steel). One template, four variants: +// DTYPE=float f32 storage/multiply, per-nibble weight staging. +// DTYPE=half f16 storage/multiply, per-nibble weight staging. +// PWDQ (half only) packed-word dequant: load each u32 weight word ONCE, +// unpack all 16 nibbles of a column + hoist the per-column scale to one read +// (the per-nibble path re-reads each word ~8x). Requires K%BK==0 (steel +// route guarantees it) and group_size%BK==0 (hoisted scale across the tile). +// ACC=half (PWDQ only) f16 accumulate with fma(), cast to f32 in the epilogue +// -- LOSSY, perplexity-gated, opt-in via a runtime spec. ACC=float is f32 +// accumulate -- BIT-EXACT to the per-nibble half kernel. +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; +var As: array; // BM*BK +var Bs: array; // BK*BN +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let nbN = (params.N + BN - 1u) / BN; + let bx = wid.x % nbN; // decode 2D tile id from 1D dispatch + let by = wid.x / nbN; + let row0 = by * BM; + let col0 = bx * BN; + let tid = lid.y * 16u + lid.x; + var acc: array, 4>; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; } + } + // A staging coords: 256 threads load 64x16 = 1024 f32 -> 4 rows each (4 contiguous K). + let ar = tid / 4u; // 0..63 (row in tile) + let ac = (tid % 4u) * 4u; // 0,4,8,12 (K offset, 4 contiguous) + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + // stage activations (edge-masked on M; K is a multiple of BK for our shapes) + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + As[ar * BK + ac + 0u] = f16(t_input[base]); + As[ar * BK + ac + 1u] = f16(t_input[base + 1u]); + As[ar * BK + ac + 2u] = f16(t_input[base + 2u]); + As[ar * BK + ac + 3u] = f16(t_input[base + 3u]); + } else { + As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; + As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + } + // Packed-word dequant: threads [0,BN) each stage one full BK-column of Bs. + if (tid < BN) { + let c = tid; // Bs column within this tile + let n = col0 + c; // global output column + if (n < params.N) { + // Scale is constant across the BK tile (group_size % BK == 0 for all real + // group sizes; K%BK==0 on the steel route), so hoist it to one read. + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + // Column n's 16-nibble K-slice for this tile = two consecutive words. + // K_packed multiple of 8 => base_word stays inside column n's own region. + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); // word0 holds K-slice [0,8) + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < BK; k = k + 1u) { + var a: array; + var bvec: array; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + } + } + workgroupBarrier(); + k0 = k0 + BK; + } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { + let r = row0 + lid.y * 4u + m; + let c = col0 + lid.x * 4u + n; + if (r < params.M && c < params.N) { + var v = f32(acc[m][n]); + if (params.has_bias != 0u) { v = v + t_bias[c]; } + t_out[r * params.N + c] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeX = + 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY = + 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From e78229d2638679876df57a7acdfd678b80fe5910 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 9 Jul 2026 13:30:06 -0700 Subject: [PATCH 2/2] [ExecuTorch][WebGPU] Test coverage for the f16-accumulate steel GEMM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/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/) --- .../webgpu/test/ops/test_quantized_linear.py | 6 +++++ backends/webgpu/test/test_webgpu_native.cpp | 23 ++++++++++++++++++- 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/backends/webgpu/test/ops/test_quantized_linear.py b/backends/webgpu/test/ops/test_quantized_linear.py index ce14e714ada..72945c37d6d 100644 --- a/backends/webgpu/test/ops/test_quantized_linear.py +++ b/backends/webgpu/test/ops/test_quantized_linear.py @@ -77,6 +77,12 @@ class Q4gswConfig: # there). Same fp64 golden regardless of which kernel runs. Q4gswConfig("pwdq_gs64", 96, 2048, 256, group_size=64), # pwdq, non-32 group Q4gswConfig("pwdq_gs8", 96, 2048, 256, group_size=8), # steel_half fallback + # pwdqf16acc (f16-accumulate) runs when the enable_f16_accumulate_gemm runtime + # spec is set and gs % BK == 0 (perplexity-gated; see the kernel diff). Same + # .pte as the f32 configs -- only the accumulator dtype differs -- goldened at a + # looser f16-accumulate tol in the native test; deep-K stresses the worst case. + Q4gswConfig("pwdqf16acc", 96, 2048, 256), # f16-accumulate steel (runtime) + Q4gswConfig("pwdqf16acc_down", 128, 8192, 2048), # deep-K f16-accum worst case Q4gswConfig("gate_proj_pf", 128, 2048, 8192), # gate/up prefill (shmem via N) Q4gswConfig("down_proj_pf", 128, 8192, 2048), # down prefill (shmem via K) Q4gswConfig("shmem_edge", 130, 4096, 2056), # partial 32-tile bounds diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index bae9d8b782b..fbdfbd09076 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -11,6 +11,8 @@ #include #include #include +#include +#include #include @@ -300,6 +302,14 @@ const Q4gswConfig kQ4gswConfigs[] = { // gs=8 (< BK=16) falls back to the per-nibble steel_half kernel. {"pwdq_gs64", 96, 2048, 256, 2.3e-4f, 1e-3f, true, false}, {"pwdq_gs8", 96, 2048, 256, 2.3e-4f, 1e-3f, true, false}, + // f16-ACCUMULATE steel (pwdqf16acc): lossy, so a wider gate than the + // f16-multiply steel_f16 (2.3e-4). f16 accumulation error grows with K, so + // the deep-K down shape (K=8192) gets the loosest tol. Perplexity is the + // primary quality gate (see the kernel diff); this catches gross bit/index + // bugs. gs=32 (% BK == 0) selects pwdqf16acc; the sweep loads these rows + // with the enable_f16_accumulate_gemm runtime spec set. + {"pwdqf16acc", 96, 2048, 256, 2e-2f, 3e-2f, true, false}, + {"pwdqf16acc_down", 128, 8192, 2048, 5e-2f, 8e-2f, true, false}, {"gate_proj_pf", 128, 2048, 8192, 1e-4f, 1e-3f, true, false}, // shmem via N {"down_proj_pf", 128, 8192, 2048, 1e-3f, 1e-2f, true, false}, // shmem via K {"shmem_edge", 130, 4096, 2056, 1e-4f, 1e-3f, true, false}, // partial tiles @@ -563,7 +573,18 @@ void test_q4gsw_config( cfg.n); Module module(pte); - ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << pte; + // pwdqf16acc rows exercise the lossy f16-accumulate kernel, a runtime opt-in + // (default off); enable it via the backend option keyed by the registered id. + if (std::string(cfg.name).rfind("pwdqf16acc", 0) == 0) { + BackendOptions<1> opts; + opts.set_option("enable_f16_accumulate_gemm", true); + LoadBackendOptionsMap map; + ASSERT_EQ(map.set_options("VulkanBackend", opts.view()), Error::Ok); + ASSERT_EQ(module.load_forward(nullptr, nullptr, &map), Error::Ok) + << "could not load " << pte; + } else { + ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load " << pte; + } const int in_numel = cfg.m * cfg.k; const int out_numel = cfg.m * cfg.n;