From d29f10124b83d6d1885a8913ec62eaf00736c14f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 8 Jul 2026 20:44:46 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 16 +++ .../ops/quantized_linear/QuantizedLinear.cpp | 18 +++ .../q4gsw_linear_gemm_steel_half_pwdq.yaml | 2 + ..._linear_gemm_steel_half_pwdq_f16acc_wgsl.h | 136 ++++++++++++++++++ 4 files changed, 172 insertions(+) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index ec1db8d495e..3f5b2be843c 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -115,6 +115,19 @@ if(EXECUTORCH_WEBGPU_STEEL_F16) target_compile_definitions(webgpu_backend PRIVATE WGPU_BACKEND_STEEL_F16) endif() +# Opt-in f16-ACCUMULATE steel q4gsw prefill GEMM (QuantizedLinear.cpp); OFF so +# default builds keep f32 accumulation + the strict golden. LOSSY (f16 +# accumulate over K) -- ships on a perplexity bar, not a bit-exact gate. +# Runtime-gated on the negotiated shader-f16 feature (fail-closed). Mirrors the +# steel-f16 gate above. +option(EXECUTORCH_WEBGPU_STEEL_F16ACC + "Enable the f16-accumulate steel q4gsw prefill GEMM (needs shader-f16)" + OFF +) +if(EXECUTORCH_WEBGPU_STEEL_F16ACC) + target_compile_definitions(webgpu_backend PRIVATE WGPU_BACKEND_STEEL_F16ACC) +endif() + # Opt-in native f16 KV cache (Sdpa.cpp/WebGPUGraph.cpp); OFF so default builds # keep the f32 cache + a byte-identical graph. Runtime-gated on the negotiated # shader-f16 feature (fail-closed). Mirrors the steel-f16 gate above. @@ -171,6 +184,9 @@ function(add_webgpu_native_test test_name test_src) if(EXECUTORCH_WEBGPU_STEEL_F16) target_compile_definitions(${test_name} PRIVATE WGPU_BACKEND_STEEL_F16) endif() + if(EXECUTORCH_WEBGPU_STEEL_F16ACC) + target_compile_definitions(${test_name} PRIVATE WGPU_BACKEND_STEEL_F16ACC) + endif() if(EXECUTORCH_WEBGPU_KV_F16) target_compile_definitions(${test_name} PRIVATE WGPU_BACKEND_KV_F16) endif() diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index aac7f259dac..13f6c69e071 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -18,6 +18,10 @@ #include #include #endif +#ifdef WGPU_BACKEND_STEEL_F16ACC +#include +#include +#endif #include @@ -283,6 +287,20 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { : kQ4gswLinearGemmSteelHalfWGSL; } } +#endif +#ifdef WGPU_BACKEND_STEEL_F16ACC + // f16-accumulate: pwdq staging with an MLC-style f16 register accumulator. + // Lossy (f16 accumulate over K) -> ships on a perplexity bar (Llama-3.2-1B + // int4: 13.32 -> 13.37, +0.05), behind its own build option (default off). + // Overrides the f32-accumulate steel kernels when the device negotiated + // shader-f16 and group_size % BK == 0 (same hoisted-scale requirement as + // pwdq). + if (use_steel && (gs % kQ4gswSteelBK == 0u)) { + const WebGPUContext* ctx = get_default_webgpu_context(); + if (ctx != nullptr && ctx->shader_f16_supported) { + shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + } + } #endif const uint32_t workgroup_count = compute_q4gsw_workgroup_count( device, diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml index a4ce843d0c9..d6723b63970 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml @@ -5,5 +5,7 @@ q4gsw_linear_gemm_steel_half_pwdq: ACC: - VALUE: float SUFFIX: "" + - VALUE: half + SUFFIX: f16acc shader_variants: - NAME: q4gsw_linear_gemm_steel_half_pwdq 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..b1ed1266849 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -0,0 +1,136 @@ +/* + * 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_half_pwdq.wgsl - DO NOT EDIT. +// wgsl-sha256: e6608df2f08b4a3a2528466874eda7ecd614e0df8bb410563bd2af3b4206f431 +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; + +// Packed-word-dequant f16 "steel" GEMM (the `half` variant of +// q4gsw_linear_gemm_steel). Loads each u32 weight word ONCE, unpacks all 16 +// nibbles of one BN column, and hoists the per-column scale to one read (the +// per-nibble steel `half` re-reads each ~8x/~16x). 64x64 tile / 256-thread / +// BK=16 geometry. Two ACC variants from this one template: +// ACC=float ("pwdq"): f32 accumulate -- BIT-EXACT to the steel `half` kernel. +// ACC=half ("pwdqf16acc"): f16 accumulate with fma() (MLC-style), cast to f32 +// in the epilogue -- LOSSY, perplexity-gated, opt-in via STEEL_F16ACC. +// Requires K%BK==0 (steel route guarantees it, so K_packed=K/2 is a multiple of 8 +// and every column is u32-aligned) and group_size%BK==0 (hoisted scale constant +// across the BK tile). +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; +var As: array; // BM*BK, staged as f16 (multiply operand only) +var Bs: array; // BK*BN, dequantized straight to f16 +@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; } + } + let ar = tid / 4u; + let ac = (tid % 4u) * 4u; + + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + 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