[ExecuTorch][WebGPU] f16-accumulate (pwdqf16acc) steel q4gsw prefill GEMM

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -89,13 +89,22 @@ Result<DelegateHandle*> WebGPUBackend::init(
       enable_f16_kv_cache = spec.get();
     }
   }
+  bool enable_f16_accumulate_gemm = false;
+  {
+    Result<bool> spec =
+        context.get_runtime_spec<bool>("enable_f16_accumulate_gemm");
+    if (spec.ok()) {
+      enable_f16_accumulate_gemm = spec.get();
+    }
+  }
 
   try {
     graph->build(
         flatbuffer_data,
         constant_data,
         context.get_named_data_map(),
-        enable_f16_kv_cache);
+        enable_f16_kv_cache,
+        enable_f16_accumulate_gemm);
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU graph build failed: %s", e.what());
     graph->~WebGPUGraph();

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -359,7 +359,8 @@ void WebGPUGraph::build(
     const void* flatbuffer_data,
     const uint8_t* constant_data,
     const executorch::runtime::NamedDataMap* named_data_map,
-    bool f16_kv_cache) {
+    bool f16_kv_cache,
+    bool f16_accumulate_gemm) {
   if (!device_) {
     auto* ctx = get_default_webgpu_context();
     if (ctx) {
@@ -385,6 +386,10 @@ void WebGPUGraph::build(
   const WebGPUContext* kv_ctx = get_default_webgpu_context();
   kv_f16_ = f16_kv_cache && (kv_ctx != nullptr && kv_ctx->shader_f16_supported);
 
+  // f16-accumulate q4gsw steel prefill GEMM (runtime opt-in). QuantizedLinear
+  // additionally gates the kernel on the negotiated shader-f16 feature.
+  f16_accumulate_gemm_ = f16_accumulate_gemm;
+
   // Phase 1: Create all values
   const auto* values = graph->values();
   const int num_vals = values ? values->size() : 0;

backends/webgpu/runtime/WebGPUGraph.h

@@ -105,7 +105,8 @@ class WebGPUGraph {
       const void* flatbuffer_data,
       const uint8_t* constant_data,
       const executorch::runtime::NamedDataMap* named_data_map = nullptr,
-      bool f16_kv_cache = false);
+      bool f16_kv_cache = false,
+      bool f16_accumulate_gemm = false);
 
   // Copy input tensor data from host pointers into GPU buffers.
   void copy_inputs(const std::vector<InputData>& inputs);
@@ -350,9 +351,16 @@ class WebGPUGraph {
     return kv_f16_;
   }
 
+  // True when the q4gsw steel prefill GEMM uses the lossy f16-accumulate kernel
+  // (runtime opt-in; perplexity-gated, not bit-exact).
+  bool f16_accumulate_gemm() const {
+    return f16_accumulate_gemm_;
+  }
+
  private:
   bool kv_f16_ = false;
   std::unordered_set<int> kv_cache_ids_;
+  bool f16_accumulate_gemm_ = false;
 
  private:
   WGPUInstance instance_ = nullptr;

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

@@ -12,6 +12,7 @@
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h>
 #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_half_pwdq_f16acc_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h>
@@ -280,6 +281,17 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
           : kQ4gswLinearGemmSteelHalfWGSL;
     }
   }
+  // f16-accumulate: pwdq staging with an f16 register accumulator.
+  // Lossy (f16 accumulate over K) -> opt-in via the enable_f16_accumulate_gemm
+  // runtime spec (default off), gated on the negotiated shader-f16 feature and
+  // group_size % BK == 0 (same hoisted-scale requirement as pwdq). Overrides
+  // the f32-accumulate steel kernels.
+  if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) {
+    const WebGPUContext* ctx = get_default_webgpu_context();
+    if (ctx != nullptr && ctx->shader_f16_supported) {
+      shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL;
+    }
+  }
   const uint32_t workgroup_count = compute_q4gsw_workgroup_count(
       device,
       use_gemv,

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.yaml

@@ -16,3 +16,7 @@ q4gsw_linear_gemm_steel:
       DTYPE: half
       PWDQ: true
       ACC: float
+    - NAME: q4gsw_linear_gemm_steel_half_pwdq_f16acc
+      DTYPE: half
+      PWDQ: true
+      ACC: half

backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h

@@ -0,0 +1,141 @@
+/*
+ * 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: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce
+inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = 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.
+// 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<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<f16, 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.0h; }
+  }
+  // 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<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] = fma(a[m], bvec[n], acc[m][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 = f32(acc[m][n]);
+        if (params.has_bias != 0u) { v = v + t_bias[c]; }
+        t_out[r * params.N + c] = v;
+      }
+    }
+  }
+}
+)";
+
+inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeX =
+    16;
+inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeY =
+    16;
+inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqF16accWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu