diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 20f7024ca82..8e66d5c19e9 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 @@ -270,7 +271,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; } } const uint32_t workgroup_count = compute_q4gsw_workgroup_count( 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.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h new file mode 100644 index 00000000000..46057de2340 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -0,0 +1,139 @@ +/* + * 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: 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; +@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.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) + + 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] = 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 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