diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index adbf4301413..be8e0593be7 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -104,6 +104,17 @@ if(EXECUTORCH_BUILD_WEBGPU_PROFILING) ) endif() +# Opt-in f16-multiply steel q4gsw prefill GEMM (QuantizedLinear.cpp); OFF so +# default builds keep the f32 kernel + strict golden. Runtime-gated on the +# negotiated shader-f16 feature (fail-closed). Mirrors the profiling gate above. +option(EXECUTORCH_WEBGPU_STEEL_F16 + "Enable the f16-multiply steel q4gsw prefill GEMM (needs shader-f16)" + OFF +) +if(EXECUTORCH_WEBGPU_STEEL_F16) + target_compile_definitions(webgpu_backend PRIVATE WGPU_BACKEND_STEEL_F16) +endif() + # Link with --whole-archive for static registration of backend + ops executorch_target_link_options_shared_lib(webgpu_backend) @@ -147,6 +158,9 @@ function(add_webgpu_native_test test_name test_src) ${test_name} PRIVATE WGPU_BACKEND_ENABLE_PROFILING ) endif() + if(EXECUTORCH_WEBGPU_STEEL_F16) + target_compile_definitions(${test_name} PRIVATE WGPU_BACKEND_STEEL_F16) + endif() set_property(TARGET ${test_name} PROPERTY CXX_STANDARD 17) endfunction() diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 97327c18394..91b834f50fe 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -13,6 +13,10 @@ #include #include #include +#ifdef WGPU_BACKEND_STEEL_F16 +#include +#include +#endif #include @@ -263,6 +267,16 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { : use_steel ? kQ4gswLinearGemmSteelWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL : kQ4gswLinearWGSL; +#ifdef WGPU_BACKEND_STEEL_F16 + // Opt-in f16-multiply steel: only when the device negotiated shader-f16; + // else the f32 steel kernel runs (fail-closed). Same bindings and tile. + if (use_steel) { + const WebGPUContext* ctx = get_default_webgpu_context(); + if (ctx != nullptr && ctx->shader_f16_supported) { + shader_src = kQ4gswLinearGemmSteelHalfWGSL; + } + } +#endif const uint32_t workgroup_count = compute_q4gsw_workgroup_count( device, use_gemv, 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 abbb14f394d..cc117434b73 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 @@ -1,3 +1,5 @@ +$if DTYPE == "half": + 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; @@ -67,7 +69,10 @@ fn main(@builtin(workgroup_id) wid: vec3, 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]; + $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; } @@ -78,7 +83,10 @@ 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) { - for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = acc[m][n] + a[m] * bvec[n]; } + $if 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]; } } } workgroupBarrier(); 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 862098e5260..1505d8b9924 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 @@ -5,5 +5,7 @@ q4gsw_linear_gemm_steel: DTYPE: - VALUE: float SUFFIX: "" + - VALUE: half + SUFFIX: half shader_variants: - NAME: q4gsw_linear_gemm_steel 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 new file mode 100644 index 00000000000..ead42673541 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h @@ -0,0 +1,123 @@ +/* + * 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: 22617bb5a5d2bf629a76833a5b55007173b1301a430e014337e9f4f04cf57baa +inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = 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. +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] = 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.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: f16 = 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 = f16(i32(nib) - 8) * f16(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] + 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 kQ4gswLinearGemmSteelHalfWorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfWorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswLinearGemmSteelHalfWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu