[ExecuTorch][WebGPU] Add f16-multiply variant of the steel q4gsw prefill GEMM

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()
 

backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp

@@ -13,6 +13,10 @@
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
+#ifdef WGPU_BACKEND_STEEL_F16
+#include <executorch/backends/webgpu/runtime/WebGPUDevice.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h>
+#endif
 
 #include <webgpu/webgpu.h>
 
@@ -263,6 +267,16 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& 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,

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<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
 @group(0) @binding(2) var<storage, read> t_weight: array<u32>;
@@ -67,7 +69,10 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
         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<u32>,
       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();

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

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 <cstdint>
+
+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<storage, read_write> t_out: array<f32>;
+@group(0) @binding(1) var<storage, read> t_input: array<f32>;
+@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
+@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
+@group(0) @binding(4) var<storage, read> t_bias: array<f32>;
+
+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<uniform> 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<workgroup> As: array<f16, 1024>;   // BM*BK
+var<workgroup> Bs: array<f16, 1024>;   // BK*BN
+@compute @workgroup_size(16, 16)
+fn main(@builtin(workgroup_id) wid: vec3<u32>,
+        @builtin(local_invocation_id) lid: vec3<u32>) {
+  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<array<f32, 4>, 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<f16, 4>;
+      var bvec: array<f16, 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] + 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