[ExecuTorch][WebGPU] Add 256-thread "steel" q4gsw prefill GEMM

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

@@ -11,6 +11,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_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
 
 #include <webgpu/webgpu.h>
@@ -52,13 +53,70 @@ 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<int64_t>(m, tile_m) * utils::div_up<int64_t>(n, tile_n);
+  if (total_wgs > static_cast<int64_t>(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<uint32_t>(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<uint64_t>((m + kQ4gswSteelTile - 1u) / kQ4gswSteelTile) *
+      static_cast<uint64_t>((n + kQ4gswSteelTile - 1u) / kQ4gswSteelTile);
+  const uint32_t max_count = max_workgroups_per_dim(device);
+  return (total == 0u || total > max_count) ? 0u : static_cast<uint32_t>(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
 // shader/pipeline is fixed at build); the resize hook re-runs this with live m.
 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<int64_t>(m, kQ4gswShmemTileM) *
-        utils::div_up<int64_t>(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<int64_t>(max_wgs)) {
-      throw std::runtime_error(
-          std::string("WebGPU ") + op_name +
-          ": shmem GEMM tile count exceeds the 1D dispatch limit");
-    }
-    return static_cast<uint32_t>(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<int64_t>(m, kQ4gswTileM) *
       utils::div_up<int64_t>(n, kQ4gswTileN);
@@ -186,17 +246,32 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& 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<int>& 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<int>& 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<int>& args) {
         const uint32_t wgc = compute_q4gsw_workgroup_count(
             g.device(),
             use_gemv,
+            use_steel,
             use_shmem_gemm,
             m,
             N,

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

@@ -0,0 +1,98 @@
+@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<${buffer_scalar_type(DTYPE)}, 1024>;   // BM*BK
+var<workgroup> Bs: array<${buffer_scalar_type(DTYPE)}, 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] = ${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;
+      }
+    }
+  }
+}

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