diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 130d81e600f..97327c18394 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,62 @@ 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; + +// Max workgroups per 1D dispatch dimension: the device limit, or 65535 when the +// query fails / reports 0. +uint32_t max_workgroups_per_dim(WGPUDevice device) { + WGPULimits limits = {}; + return (wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeWorkgroupsPerDimension > 0) + ? limits.maxComputeWorkgroupsPerDimension + : 65535u; +} + +// One workgroup per (tile_m x tile_n) tile, no grid-stride: throw when the tile +// count would exceed the 1D dispatch limit. Shared by the steel + shmem GEMM +// routes; `kind` names the route in the error message. +uint32_t tiled_wg_count( + WGPUDevice device, + uint32_t m, + uint32_t n, + int64_t tile_m, + int64_t tile_n, + const char* op_name, + const char* kind) { + const int64_t total_wgs = + utils::div_up(m, tile_m) * utils::div_up(n, tile_n); + if (total_wgs > static_cast(max_workgroups_per_dim(device))) { + throw std::runtime_error( + std::string("WebGPU ") + op_name + ": " + kind + + " tile count exceeds the 1D dispatch limit"); + } + return static_cast(total_wgs); +} + +// 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); + const uint32_t max_count = max_workgroups_per_dim(device); + 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 +116,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,22 +138,24 @@ uint32_t compute_q4gsw_workgroup_count( } return wgc; } + if (use_steel) { + // steel: one workgroup per 64x64 tile. Over-limit THROWS here -- unlike the + // build-time steel_workgroup_count, which returns 0 so the caller falls + // back to shmem/tiled. The routed kernel is baked into the pipeline at + // build, so the resize path cannot switch kernels for a larger live M. + return tiled_wg_count( + device, m, n, kQ4gswSteelTile, kQ4gswSteelTile, op_name, "steel GEMM"); + } 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) * - utils::div_up(n, kQ4gswShmemTileN); - 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 + - ": shmem GEMM tile count exceeds the 1D dispatch limit"); - } - return static_cast(total_wgs); + // shmem GEMM: one workgroup per tile. + return tiled_wg_count( + device, + m, + n, + kQ4gswShmemTileM, + kQ4gswShmemTileN, + op_name, + "shmem GEMM"); } const int64_t total_tiles = utils::div_up(m, kQ4gswTileM) * utils::div_up(n, kQ4gswTileN); @@ -186,17 +246,32 @@ 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; + // shmem GEMM is now a FALLBACK, not dead: steel shadows it whenever eligible, + // so shmem only wins when steel is ineligible (K % 16 != 0, or a + // <256-invocation device such as SwiftShader) and the shape still hits the + // large K/N thresholds; otherwise the register-tiled path handles it. + 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 +351,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 +403,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 +432,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..abbb14f394d --- /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<${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) { + 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] = ${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; + } + // 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; } + 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<${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) { + 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.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 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