From e7bb34ea8e6753bd255ba6cf8ef1f3962b254716 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Sat, 4 Jul 2026 18:19:56 -0700 Subject: [PATCH 1/2] Update [ghstack-poisoned] --- .../ops/quantized_linear/QuantizedLinear.cpp | 74 ++++++++++- .../q4gsw_linear_gemm_steel.wgsl | 98 ++++++++++++++ .../q4gsw_linear_gemm_steel_wgsl.h | 122 ++++++++++++++++++ 3 files changed, 288 insertions(+), 6 deletions(-) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 130d81e600f..4021b3fce4e 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -52,6 +53,36 @@ constexpr int64_t kQ4gswShmemTileN = 32; constexpr uint32_t kQ4gswShmemMinDim = 4096u; constexpr uint32_t kQ4gswShmemNMinDim = 2048u; +// steel GEMM: 64x64 tile, 256 threads (16x16), fixed wg (no override). +constexpr uint32_t kQ4gswSteelTile = 64u; +constexpr uint32_t kQ4gswSteelBK = 16u; +constexpr uint32_t kQ4gswSteelInvocations = 256u; + +// steel needs 256-thread workgroups; fail-closed (query ok AND >=256). +bool steel_supported(WGPUDevice device) { + WGPULimits limits = {}; + return wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= kQ4gswSteelInvocations; +} + +// Not grid-strided: 0 (fall back) when K%BK != 0 or over the 1D dispatch limit. +uint32_t +steel_workgroup_count(WGPUDevice device, uint32_t m, uint32_t n, uint32_t K) { + if (K % kQ4gswSteelBK != 0u) { + return 0u; + } + const uint64_t total = + static_cast((m + kQ4gswSteelTile - 1u) / kQ4gswSteelTile) * + static_cast((n + kQ4gswSteelTile - 1u) / kQ4gswSteelTile); + WGPULimits limits = {}; + const uint32_t max_count = + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupsPerDimension > 0 + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; + return (total == 0u || total > max_count) ? 0u : static_cast(total); +} + // Workgroup count for a linear_q4gsw dispatch (bicol GEMV / shmem GEMM / tiled // GEMM), with the range/limit guards shared by the build-time path and the // resize hook. use_gemv/use_shmem_gemm are the build-time routing decision (the @@ -59,6 +90,7 @@ constexpr uint32_t kQ4gswShmemNMinDim = 2048u; uint32_t compute_q4gsw_workgroup_count( WGPUDevice device, bool use_gemv, + bool use_steel, bool use_shmem_gemm, uint32_t m, uint32_t n, @@ -80,6 +112,23 @@ uint32_t compute_q4gsw_workgroup_count( } return wgc; } + if (use_steel) { + // steel: one workgroup per 64x64 tile, no grid-stride -> throw over limit. + const int64_t total_wgs = utils::div_up(m, kQ4gswSteelTile) * + utils::div_up(n, kQ4gswSteelTile); + WGPULimits limits = {}; + const uint32_t max_wgs = + wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupsPerDimension > 0 + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; + if (total_wgs > static_cast(max_wgs)) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + + ": steel GEMM tile count exceeds the 1D dispatch limit"); + } + return static_cast(total_wgs); + } if (use_shmem_gemm) { // shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit. const int64_t total_wgs = utils::div_up(m, kQ4gswShmemTileM) * @@ -186,17 +235,28 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { "WebGPU linear_q4gsw: scales dims too small for K/N"); } - // M==1 -> bicol GEMV; M>1 -> shmem GEMM (large K/N) else tiled GEMM. + // M==1 -> bicol GEMV; M>1 -> steel GEMM (preferred) else shmem else tiled. const uint32_t wg_size = utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u); - const bool use_shmem_gemm = - !use_gemv && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); + // steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible. + const bool use_steel = !use_gemv && steel_supported(device) && + steel_workgroup_count(device, M, N, K) > 0u; + const bool use_shmem_gemm = !use_gemv && !use_steel && + (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL + : use_steel ? kQ4gswLinearGemmSteelWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL : kQ4gswLinearWGSL; const uint32_t workgroup_count = compute_q4gsw_workgroup_count( - device, use_gemv, use_shmem_gemm, M, N, wg_size, "linear_q4gsw"); + device, + use_gemv, + use_steel, + use_shmem_gemm, + M, + N, + wg_size, + "linear_q4gsw"); // Optional bias: real buffer if present, else a dummy for the fixed layout. uint32_t has_bias = 0; @@ -276,8 +336,8 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { pipeline_desc.layout = pipeline_layout; pipeline_desc.compute.module = shader; pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - // Only the tiled GEMM has a wg_size override; GEMV + shmem are fixed 64. - const bool fixed_wg = use_gemv || use_shmem_gemm; + // Only tiled GEMM overrides wg_size; GEMV/shmem (64) + steel (256) are fixed. + const bool fixed_wg = use_gemv || use_steel || use_shmem_gemm; pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u; pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant; WGPUComputePipeline pipeline = @@ -328,6 +388,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { has_bias, wg_size, use_gemv, + use_steel, use_shmem_gemm, dispatch_idx, uniform_buffer](WebGPUGraph& g) { @@ -356,6 +417,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const uint32_t wgc = compute_q4gsw_workgroup_count( g.device(), use_gemv, + use_steel, use_shmem_gemm, m, N, 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 new file mode 100644 index 00000000000..fafd9f9aeec --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl @@ -0,0 +1,98 @@ +@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. +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) + // 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 { + 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] = f32(t_input[base]); + As[ar * BK + ac + 1u] = f32(t_input[base + 1u]); + As[ar * BK + ac + 2u] = f32(t_input[base + 2u]); + As[ar * BK + ac + 3u] = f32(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; + } + // 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: f32 = 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; } + dqv = f32(i32(nib) - 8) * t_scales[scale_row + n]; + } + Bs[br * BN + bc + j] = dqv; + } + 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] + 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; + } + } + } +} 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 new file mode 100644 index 00000000000..9af92b4b448 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h @@ -0,0 +1,122 @@ +/* + * 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: 48918226043fc57781a449c4d42a4bb223f3816ce10a735991b7c10986d994fd +inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( +@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. +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) + // 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 { + 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] = f32(t_input[base]); + As[ar * BK + ac + 1u] = f32(t_input[base + 1u]); + As[ar * BK + ac + 2u] = f32(t_input[base + 2u]); + As[ar * BK + ac + 3u] = f32(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; + } + // 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: f32 = 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; } + dqv = f32(i32(nib) - 8) * t_scales[scale_row + n]; + } + Bs[br * BN + bc + j] = dqv; + } + 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] + 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 kQ4gswLinearGemmSteelWorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelWorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From 33b49b3d6c829c6bc680456880a16ca6b7352dd1 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Mon, 6 Jul 2026 13:47:45 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- .../q4gsw_linear_gemm_steel.wgsl | 18 +++++++++--------- .../q4gsw_linear_gemm_steel.yaml | 9 +++++++++ 2 files changed, 18 insertions(+), 9 deletions(-) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.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 fafd9f9aeec..abbb14f394d 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 @@ -18,8 +18,8 @@ struct Params { // "steel" prefill GEMM (M>1): 64x64 tile, 256 threads; K%16==0 host-guarded. const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; -var As: array; // BM*BK -var Bs: array; // BK*BN +var As: array<${buffer_scalar_type(DTYPE)}, 1024>; // BM*BK +var Bs: array<${buffer_scalar_type(DTYPE)}, 1024>; // BK*BN @compute @workgroup_size(16, 16) fn main(@builtin(workgroup_id) wid: vec3, @builtin(local_invocation_id) lid: vec3) { @@ -47,10 +47,10 @@ fn main(@builtin(workgroup_id) wid: vec3, let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - As[ar * BK + ac + 0u] = f32(t_input[base]); - As[ar * BK + ac + 1u] = f32(t_input[base + 1u]); - As[ar * BK + ac + 2u] = f32(t_input[base + 2u]); - As[ar * BK + ac + 3u] = f32(t_input[base + 3u]); + As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(t_input[base]); + As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 1u]); + 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; @@ -60,7 +60,7 @@ fn main(@builtin(workgroup_id) wid: vec3, 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: f32 = 0.0; + 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]; @@ -73,8 +73,8 @@ fn main(@builtin(workgroup_id) wid: vec3, } workgroupBarrier(); for (var k: u32 = 0u; k < BK; k = k + 1u) { - var a: array; - var bvec: array; + var a: array<${buffer_scalar_type(DTYPE)}, 4>; + var bvec: array<${buffer_scalar_type(DTYPE)}, 4>; 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) { 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 new file mode 100644 index 00000000000..862098e5260 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml @@ -0,0 +1,9 @@ +q4gsw_linear_gemm_steel: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: float + SUFFIX: "" + shader_variants: + - NAME: q4gsw_linear_gemm_steel