[ExecuTorch][WebGPU] Packed-word-dequant f16 steel q4gsw prefill GEMM

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_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>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
@@ -270,7 +271,13 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   if (use_steel) {
     const WebGPUContext* ctx = get_default_webgpu_context();
     if (ctx != nullptr && ctx->shader_f16_supported) {
-      shader_src = kQ4gswLinearGemmSteelHalfWGSL;
+      // Packed-word dequant: bit-exact to the steel `half` kernel but loads
+      // each u32 weight word once + hoists the per-column scale (half re-reads
+      // them ~8x/~16x). Needs group_size % BK == 0 so the hoisted scale is
+      // constant across the BK tile; else the per-nibble `half` kernel.
+      shader_src = (gs % kQ4gswSteelBK == 0u)
+          ? kQ4gswLinearGemmSteelHalfPwdqWGSL
+          : kQ4gswLinearGemmSteelHalfWGSL;
     }
   }
   const uint32_t workgroup_count = compute_q4gsw_workgroup_count(

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

@@ -21,7 +21,16 @@ struct 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).
+// 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<${buffer_scalar_type(DTYPE)}, 1024>;   // BM*BK
 var<workgroup> Bs: array<${buffer_scalar_type(DTYPE)}, 1024>;   // BK*BN
@@ -34,16 +43,17 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
   let row0 = by * BM;
   let col0 = bx * BN;
   let tid = lid.y * 16u + lid.x;
-  var acc: array<array<f32, 4>, 4>;
+  var acc: array<array<${buffer_scalar_type(ACC)}, 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; }
+    for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = ${"0.0h" if ACC == "half" else "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)
+  $if not PWDQ:
+    // 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 {
@@ -57,36 +67,63 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
       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;
+      As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"};
+      As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "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; }
-        $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];
+    $if PWDQ:
+      // 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; }
+        }
+      }
+    $else:
+      // 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; }
+          $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;
       }
-      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) {
-        $if DTYPE == "half":
+        $if ACC == "half":
+          for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); }
+        $elif 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]; }
@@ -100,7 +137,7 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
       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];
+        var v = ${"f32(acc[m][n])" if ACC == "half" else "acc[m][n]"};
         if (params.has_bias != 0u) { v = v + t_bias[c]; }
         t_out[r * params.N + c] = v;
       }

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

@@ -1,11 +1,18 @@
 q4gsw_linear_gemm_steel:
   parameter_names_with_default_values:
     DTYPE: float
-  generate_variant_forall:
-    DTYPE:
-      - VALUE: float
-        SUFFIX: ""
-      - VALUE: half
-        SUFFIX: half
+    PWDQ: false
+    ACC: float
   shader_variants:
     - NAME: q4gsw_linear_gemm_steel
+      DTYPE: float
+      PWDQ: false
+      ACC: float
+    - NAME: q4gsw_linear_gemm_steel_half
+      DTYPE: half
+      PWDQ: false
+      ACC: float
+    - NAME: q4gsw_linear_gemm_steel_half_pwdq
+      DTYPE: half
+      PWDQ: true
+      ACC: float

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

@@ -0,0 +1,139 @@
+/*
+ * 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: 1f916bcd30dbbbcc7eca37e795ecc26e3c72e645ccd2c361fa0ac4e66f1a174a
+inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = 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<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)
+
+  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] = 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 kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeX = 16;
+inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeY = 16;
+inline constexpr uint32_t kQ4gswLinearGemmSteelHalfPwdqWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu

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

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: e3c21e7db7c18f6e085de71e283988f0bd3b2543807ddc17774a1c607e69c766
+// wgsl-sha256: 00cdd2f2fb98a5c7343d16fdf7e59f1b840e180cec3f82bf9b569513c0a45396
 inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"(
 enable f16;
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@@ -37,7 +37,16 @@ struct 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).
+// 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

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

@@ -13,7 +13,7 @@
 namespace executorch::backends::webgpu {
 
 // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
-// wgsl-sha256: 43536b16026d3b62d77087b86289885606c04834d9783c3c871512d6789ee6f6
+// wgsl-sha256: dd771b9ab096410f3ad0d9259bef7816e41330a434325ea28baa5abbfb2841d2
 inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"(
 @group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_input: array<f32>;
@@ -36,7 +36,16 @@ struct 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).
+// 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<f32, 1024>;   // BM*BK
 var<workgroup> Bs: array<f32, 1024>;   // BK*BN