[ExecuTorch][WebGPU] Test coverage for the f16 KV cache

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

@@ -90,6 +90,49 @@ WGPUBuffer WebGPUGraph::create_scratch_buffer(size_t nbytes) {
   return buffer;
 }
 
+WGPUBuffer WebGPUGraph::acquire_scratch(size_t nbytes) {
+  nbytes = nbytes > 0 ? nbytes : 4;
+  // Best-fit reuse: smallest free slot with size in [nbytes, 2*nbytes] -- the
+  // 2x cap stops a large Cmax-sized buffer from backing a tiny request. Never
+  // reuse an in_use slot (co-live safety).
+  ScratchSlot* best = nullptr;
+  for (auto& s : scratch_pool_) {
+    // s.size - nbytes (safe: s.size >= nbytes) avoids overflowing 2 * nbytes.
+    if (!s.in_use && s.size >= nbytes && s.size - nbytes <= nbytes) {
+      if (best == nullptr || s.size < best->size) {
+        best = &s;
+      }
+    }
+  }
+  if (best != nullptr) {
+    best->in_use = true;
+    return best->buffer;
+  }
+  // None reusable -> create a new slot (freed in the dtor, like
+  // scratch_buffers_).
+  WGPUBufferDescriptor buf_desc = {};
+  buf_desc.size = nbytes;
+  buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |
+      WGPUBufferUsage_CopySrc;
+  buf_desc.mappedAtCreation = false;
+  WGPUBuffer buffer = wgpuDeviceCreateBuffer(device_, &buf_desc);
+  scratch_pool_.push_back({buffer, nbytes, true});
+  return buffer;
+}
+
+void WebGPUGraph::release_scratch(WGPUBuffer buffer) {
+  if (!buffer) {
+    return;
+  }
+  for (auto& s : scratch_pool_) {
+    if (s.buffer == buffer) {
+      s.in_use = false;
+      return;
+    }
+  }
+  // Not a pooled buffer -> no-op; the dtor frees it via scratch_buffers_.
+}
+
 WGPUBuffer WebGPUGraph::make_uniform_buffer(const void* data, size_t size) {
   WGPUBufferDescriptor desc = {};
   desc.size = size;
@@ -267,6 +310,11 @@ WebGPUGraph::~WebGPUGraph() {
       wgpuBufferRelease(buf);
     }
   }
+  for (auto& s : scratch_pool_) {
+    if (s.buffer) {
+      wgpuBufferRelease(s.buffer);
+    }
+  }
   for (auto& buf : owned_uniform_buffers_) {
     if (buf) {
       wgpuBufferRelease(buf);
@@ -311,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) {
@@ -337,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);
@@ -268,6 +269,35 @@ class WebGPUGraph {
   // Graph-owned scratch storage buffer for fused-op intermediates (e.g. SDPA).
   WGPUBuffer create_scratch_buffer(size_t nbytes);
 
+  // Reusable scratch pool for SINGLE-OP-LIFETIME fused-op scratch (SDPA
+  // attn_weights/softmax, FlashDecoding partials). acquire_scratch() reuses a
+  // free slot (best-fit, size in [n,2n]) or creates one; the caller RELEASES it
+  // at op-lowering scope exit (use ScopedScratch), so N layers' scratch reuses
+  // a small constant of buffers instead of N x held to graph teardown.
+  // Correctness: WebGPU/Dawn auto-inserts RAW hazard barriers between
+  // dispatches on a shared storage buffer regardless of pass structure -- the
+  // SAME guarantee mem_obj_id aliasing already relies on -- so reuse is
+  // bit-identical. Never hand a still-in_use slot to a co-live requester.
+  WGPUBuffer acquire_scratch(size_t nbytes);
+  void release_scratch(WGPUBuffer buffer);
+  // RAII: releases an acquired scratch slot when the op-lowering scope exits
+  // (leak-safe vs early returns).
+  struct ScopedScratch {
+    WebGPUGraph* g = nullptr;
+    WGPUBuffer buf = nullptr;
+    ScopedScratch(WebGPUGraph* graph, WGPUBuffer b) : g(graph), buf(b) {}
+    ~ScopedScratch() {
+      if (g && buf) {
+        g->release_scratch(buf);
+      }
+    }
+    ScopedScratch(const ScopedScratch&) = delete;
+    ScopedScratch& operator=(const ScopedScratch&) = delete;
+    operator WGPUBuffer() const {
+      return buf;
+    }
+  };
+
   // Create a mapped-at-creation uniform buffer from `size` bytes and track it
   // in the memory stats. Shared helper for ops needing a uniform Params buffer.
   WGPUBuffer make_uniform_buffer(const void* data, size_t size);
@@ -321,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;
@@ -377,6 +414,16 @@ class WebGPUGraph {
   // Long-lived scratch storage buffers for fused ops (e.g. SDPA temporaries).
   std::vector<WGPUBuffer> scratch_buffers_;
 
+  // Reusable scratch pool: single-op-lifetime buffers recycled across ops
+  // (acquire_scratch/release_scratch). Each slot is freed in the dtor. See
+  // acquire_scratch() for the reuse policy.
+  struct ScratchSlot {
+    WGPUBuffer buffer = nullptr;
+    size_t size = 0;
+    bool in_use = false;
+  };
+  std::vector<ScratchSlot> scratch_pool_;
+
   // Uniform buffers owned for the graph's lifetime; released in the dtor.
   std::vector<WGPUBuffer> owned_uniform_buffers_;
 

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

@@ -12,6 +12,8 @@
 #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>
 #include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
@@ -270,7 +272,24 @@ 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;
+    }
+  }
+  // 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(

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,22 @@
 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
+    - NAME: q4gsw_linear_gemm_steel_half_pwdq_f16acc
+      DTYPE: half
+      PWDQ: true
+      ACC: half