From 42c7fe2f4b13ce08bd30cad412533e8b6f142abe Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 8 Jul 2026 20:44:37 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- .../ops/quantized_linear/QuantizedLinear.cpp | 9 +- .../q4gsw_linear_gemm_steel_half_pwdq.wgsl | 115 +++++++++++++++ .../q4gsw_linear_gemm_steel_half_pwdq.yaml | 9 ++ .../q4gsw_linear_gemm_steel_half_pwdq_wgsl.h | 136 ++++++++++++++++++ 4 files changed, 268 insertions(+), 1 deletion(-) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 91b834f50fe..aac7f259dac 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -15,6 +15,7 @@ #include #ifdef WGPU_BACKEND_STEEL_F16 #include +#include #include #endif @@ -273,7 +274,13 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_steel) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { - shader_src = kQ4gswLinearGemmSteelHalfWGSL; + // Packed-word dequant: bit-exact to the steel `half` kernel but loads + // each u32 weight word once + hoists the per-column scale (half re-reads + // them ~8x/~16x). Needs group_size % BK == 0 so the hoisted scale is + // constant across the BK tile; else the per-nibble `half` kernel. + shader_src = (gs % kQ4gswSteelBK == 0u) + ? kQ4gswLinearGemmSteelHalfPwdqWGSL + : kQ4gswLinearGemmSteelHalfWGSL; } } #endif diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl new file mode 100644 index 00000000000..824c71bc383 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl @@ -0,0 +1,115 @@ +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" if ACC == "half" else "0.0"}; } + } + 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) { + $if ACC == "half": + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + $else: + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + f32(a[m] * bvec[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 ACC == "half" else "acc[m][n]"}; + if (params.has_bias != 0u) { v = v + t_bias[c]; } + t_out[r * params.N + c] = v; + } + } + } +} 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 new file mode 100644 index 00000000000..a4ce843d0c9 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml @@ -0,0 +1,9 @@ +q4gsw_linear_gemm_steel_half_pwdq: + parameter_names_with_default_values: + ACC: float + generate_variant_forall: + ACC: + - VALUE: float + SUFFIX: "" + shader_variants: + - NAME: q4gsw_linear_gemm_steel_half_pwdq diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h new file mode 100644 index 00000000000..17b27e2c80d --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_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: d4e3ce11f873b8cec80b1a9457915d0e50a334e286f1d75753f9f3b023728ac4 +inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = 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.0; } + } + 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] = acc[m][n] + f32(a[m] * bvec[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 = acc[m][n]; + if (params.has_bias != 0u) { v = v + t_bias[c]; } + t_out[r * params.N + c] = v; + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 06fe79fd682c73c048afd9f1d79ca71b29264df6 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Thu, 9 Jul 2026 13:30:04 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- .../q4gsw_linear_gemm_steel.wgsl | 93 +++++++++----- .../q4gsw_linear_gemm_steel.yaml | 19 ++- .../q4gsw_linear_gemm_steel_half_pwdq.wgsl | 115 ------------------ .../q4gsw_linear_gemm_steel_half_pwdq.yaml | 9 -- .../q4gsw_linear_gemm_steel_half_pwdq_wgsl.h | 39 +++--- .../q4gsw_linear_gemm_steel_half_wgsl.h | 13 +- .../q4gsw_linear_gemm_steel_wgsl.h | 13 +- 7 files changed, 121 insertions(+), 180 deletions(-) delete mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl delete mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl index 788c9ffd941..4d7ab0b1d1e 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl @@ -21,7 +21,16 @@ struct 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). +// 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<${buffer_scalar_type(DTYPE)}, 1024>; // BM*BK var Bs: array<${buffer_scalar_type(DTYPE)}, 1024>; // BK*BN @@ -34,16 +43,17 @@ fn main(@builtin(workgroup_id) wid: vec3, let row0 = by * BM; let col0 = bx * BN; let tid = lid.y * 16u + lid.x; - var acc: array, 4>; + 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.0; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = ${"0.0h" if ACC == "half" else "0.0"}; } } // 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) - // B staging coords: 256 threads load 16x64 = 1024 dequant weights -> 4 cols each. - let br = tid / 16u; // 0..15 (K within BK) - let bc = (tid % 16u) * 4u; // 0,4,..60 (N offset, 4 contiguous) + $if not PWDQ: + // B staging coords: 256 threads load 16x64 = 1024 dequant weights -> 4 cols each. + let br = tid / 16u; // 0..15 (K within BK) + let bc = (tid % 16u) * 4u; // 0,4,..60 (N offset, 4 contiguous) var k0: u32 = 0u; loop { @@ -57,28 +67,53 @@ fn main(@builtin(workgroup_id) wid: vec3, As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 2u]); As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 3u]); } else { - As[ar * BK + ac + 0u] = 0.0; As[ar * BK + ac + 1u] = 0.0; - As[ar * BK + ac + 2u] = 0.0; As[ar * BK + ac + 3u] = 0.0; + As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"}; + As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "0.0"}; } - // stage DEQUANTIZED weights into Bs[k][n]: 4 contiguous N per thread. - let kk = k0 + br; // K index for this shmem row - let scale_row = (kk / params.group_size) * params.padded_N; - for (var j: u32 = 0u; j < 4u; j = j + 1u) { - let n = col0 + bc + j; - var dqv: ${buffer_scalar_type(DTYPE)} = 0.0; - if (n < params.N) { - let byte_idx = n * params.K_packed + (kk >> 1u); - let word = t_weight[byte_idx >> 2u]; - let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; - var nib: u32; - if ((kk & 1u) == 0u) { nib = b & 0x0Fu; } else { nib = (b >> 4u) & 0x0Fu; } - $if DTYPE == "half": - dqv = f16(i32(nib) - 8) * f16(t_scales[scale_row + n]); - $else: - dqv = f32(i32(nib) - 8) * t_scales[scale_row + n]; + $if PWDQ: + // 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; } + } + } + $else: + // stage DEQUANTIZED weights into Bs[k][n]: 4 contiguous N per thread. + let kk = k0 + br; // K index for this shmem row + let scale_row = (kk / params.group_size) * params.padded_N; + for (var j: u32 = 0u; j < 4u; j = j + 1u) { + let n = col0 + bc + j; + var dqv: ${buffer_scalar_type(DTYPE)} = 0.0; + if (n < params.N) { + let byte_idx = n * params.K_packed + (kk >> 1u); + let word = t_weight[byte_idx >> 2u]; + let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu; + var nib: u32; + if ((kk & 1u) == 0u) { nib = b & 0x0Fu; } else { nib = (b >> 4u) & 0x0Fu; } + $if DTYPE == "half": + dqv = f16(i32(nib) - 8) * f16(t_scales[scale_row + n]); + $else: + dqv = f32(i32(nib) - 8) * t_scales[scale_row + n]; + } + Bs[br * BN + bc + j] = dqv; } - Bs[br * BN + bc + j] = dqv; - } workgroupBarrier(); for (var k: u32 = 0u; k < BK; k = k + 1u) { var a: array<${buffer_scalar_type(DTYPE)}, 4>; @@ -86,7 +121,9 @@ fn main(@builtin(workgroup_id) wid: vec3, 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) { - $if DTYPE == "half": + $if ACC == "half": + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + $elif DTYPE == "half": for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + f32(a[m] * bvec[n]); } $else: for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + a[m] * bvec[n]; } @@ -100,7 +137,7 @@ fn main(@builtin(workgroup_id) wid: vec3, let r = row0 + lid.y * 4u + m; let c = col0 + lid.x * 4u + n; if (r < params.M && c < params.N) { - var v = acc[m][n]; + var v = ${"f32(acc[m][n])" if ACC == "half" else "acc[m][n]"}; if (params.has_bias != 0u) { v = v + t_bias[c]; } t_out[r * params.N + c] = v; } 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 1505d8b9924..554fab0f878 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 @@ -1,11 +1,18 @@ q4gsw_linear_gemm_steel: parameter_names_with_default_values: DTYPE: float - generate_variant_forall: - DTYPE: - - VALUE: float - SUFFIX: "" - - VALUE: half - SUFFIX: half + PWDQ: false + ACC: float shader_variants: - NAME: q4gsw_linear_gemm_steel + DTYPE: float + PWDQ: false + ACC: float + - NAME: q4gsw_linear_gemm_steel_half + DTYPE: half + PWDQ: false + ACC: float + - NAME: q4gsw_linear_gemm_steel_half_pwdq + DTYPE: half + PWDQ: true + ACC: float diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl deleted file mode 100644 index 3195adc7640..00000000000 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.wgsl +++ /dev/null @@ -1,115 +0,0 @@ -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(), cast to f32 in the -// epilogue -- LOSSY, perplexity-gated, opt-in via a runtime spec. -// 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" if ACC == "half" else "0.0"}; } - } - 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) { - $if ACC == "half": - for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } - $else: - for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + f32(a[m] * bvec[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 ACC == "half" else "acc[m][n]"}; - if (params.has_bias != 0u) { v = v + t_bias[c]; } - t_out[r * params.N + c] = v; - } - } - } -} 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 deleted file mode 100644 index a4ce843d0c9..00000000000 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq.yaml +++ /dev/null @@ -1,9 +0,0 @@ -q4gsw_linear_gemm_steel_half_pwdq: - parameter_names_with_default_values: - ACC: float - generate_variant_forall: - ACC: - - VALUE: float - SUFFIX: "" - shader_variants: - - NAME: q4gsw_linear_gemm_steel_half_pwdq diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h index a1470246b15..46057de2340 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -12,11 +12,10 @@ namespace executorch::backends::webgpu { -// @generated from q4gsw_linear_gemm_steel_half_pwdq.wgsl - DO NOT EDIT. -// wgsl-sha256: 9ba15b2fe5dd05fe19d94725f83b3efeed3ebac521fafc217d823e2e9a1c6f8e +// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. +// wgsl-sha256: 1f916bcd30dbbbcc7eca37e795ecc26e3c72e645ccd2c361fa0ac4e66f1a174a inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = 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; @@ -35,20 +34,22 @@ struct Params { } @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(), cast to f32 in the -// epilogue -- LOSSY, perplexity-gated, opt-in via a runtime spec. -// 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). +// "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, staged as f16 (multiply operand only) -var Bs: array; // BK*BN, dequantized straight to f16 +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) { @@ -62,12 +63,14 @@ fn main(@builtin(workgroup_id) wid: vec3, for (var m: u32 = 0u; m < 4u; m = m + 1u) { for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0; } } - let ar = tid / 4u; - let ac = (tid % 4u) * 4u; + // 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; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h index 7e9363e3b36..ae03f63fb5c 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: e3c21e7db7c18f6e085de71e283988f0bd3b2543807ddc17774a1c607e69c766 +// wgsl-sha256: 00cdd2f2fb98a5c7343d16fdf7e59f1b840e180cec3f82bf9b569513c0a45396 inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; @@ -37,7 +37,16 @@ struct 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). +// 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 diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h index 9e73a0c9a66..71b0f45bdf7 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: 43536b16026d3b62d77087b86289885606c04834d9783c3c871512d6789ee6f6 +// wgsl-sha256: dd771b9ab096410f3ad0d9259bef7816e41330a434325ea28baa5abbfb2841d2 inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_input: array; @@ -36,7 +36,16 @@ struct 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). +// 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