[ExecuTorch][WebGPU] WGSL shader-variant codegen (vec1/vec4, fp16-ready) + rms_norm dedup (#20812)

Author: pytorchbot

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -79,8 +79,32 @@ Result<DelegateHandle*> WebGPUBackend::init(
     return Error::DelegateInvalidCompatibility;
   }
 
+  // Load-time backend option (BackendOption / LoadBackendOptionsMap), keyed by
+  // the registered backend name; default false. Mirrors the CoreML/XNNPACK
+  // runtime-spec pattern -- no compile flag and no .pte re-export needed.
+  bool enable_f16_kv_cache = false;
+  {
+    Result<bool> spec = context.get_runtime_spec<bool>("enable_f16_kv_cache");
+    if (spec.ok()) {
+      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());
+    graph->build(
+        flatbuffer_data,
+        constant_data,
+        context.get_named_data_map(),
+        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/WebGPUDevice.cpp

@@ -13,9 +13,7 @@
 #include <cstdlib>
 #include <memory>
 #include <stdexcept>
-#ifdef WGPU_BACKEND_ENABLE_PROFILING
 #include <vector>
-#endif // WGPU_BACKEND_ENABLE_PROFILING
 
 namespace executorch {
 namespace backends {
@@ -143,16 +141,28 @@ WebGPUContext create_webgpu_context() {
     device_desc.requiredLimits = &supported_limits;
   }
 
+  // Request optional features when the adapter advertises them. A feature the
+  // adapter lacks is skipped (its fast path stays disabled). A feature the
+  // adapter advertises becomes required of the device, so if the device were
+  // to reject it, context creation fails below rather than falling back. The
+  // vector must outlive wgpuAdapterRequestDevice below (device_desc points
+  // into it).
+  std::vector<WGPUFeatureName> required_features;
+  if (wgpuAdapterHasFeature(ctx.adapter, WGPUFeatureName_ShaderF16)) {
+    required_features.push_back(WGPUFeatureName_ShaderF16);
+    ctx.shader_f16_supported = true;
+  }
 #ifdef WGPU_BACKEND_ENABLE_PROFILING
   // Bench: enable TimestampQuery if available; fail-open (skip timing if not).
-  std::vector<WGPUFeatureName> required_features;
   if (wgpuAdapterHasFeature(ctx.adapter, WGPUFeatureName_TimestampQuery)) {
     required_features.push_back(WGPUFeatureName_TimestampQuery);
-    device_desc.requiredFeatureCount = required_features.size();
-    device_desc.requiredFeatures = required_features.data();
     ctx.timestamp_supported = true;
   }
 #endif // WGPU_BACKEND_ENABLE_PROFILING
+  if (!required_features.empty()) {
+    device_desc.requiredFeatureCount = required_features.size();
+    device_desc.requiredFeatures = required_features.data();
+  }
 
   device_desc.uncapturedErrorCallbackInfo.callback = on_device_error;
 

backends/webgpu/runtime/WebGPUDevice.h

@@ -25,6 +25,9 @@ struct WebGPUContext {
   WGPUAdapter adapter = nullptr;
   WGPUDevice device = nullptr;
   WGPUQueue queue = nullptr;
+  // True if the device was created with the ShaderF16 feature; reserved for a
+  // future fp16 storage/compute path (fp32 is used when false or unset).
+  bool shader_f16_supported = false;
 #ifdef WGPU_BACKEND_ENABLE_PROFILING
   // True if the device was created with the TimestampQuery feature (bench).
   bool timestamp_supported = false;

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);
@@ -310,7 +358,9 @@ WebGPUGraph::~WebGPUGraph() {
 void WebGPUGraph::build(
     const void* flatbuffer_data,
     const uint8_t* constant_data,
-    const executorch::runtime::NamedDataMap* named_data_map) {
+    const executorch::runtime::NamedDataMap* named_data_map,
+    bool f16_kv_cache,
+    bool f16_accumulate_gemm) {
   if (!device_) {
     auto* ctx = get_default_webgpu_context();
     if (ctx) {
@@ -331,6 +381,15 @@ void WebGPUGraph::build(
   constant_data_ = constant_data;
   named_data_map_ = named_data_map;
 
+  // f16 KV cache (runtime opt-in): store K/V caches as f16 iff the opt-in is
+  // set AND the device negotiated shader-f16 (fail-closed).
+  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;
@@ -358,6 +417,13 @@ void WebGPUGraph::build(
       if (!a) {
         continue;
       }
+      // f16 KV: tag sdpa K/V cache values (args[3],[4]) for half-size alloc.
+      // Inert unless kv_f16_ (runtime opt-in) is set.
+      if (kv_f16_ && a->size() > 4 &&
+          oc->name()->str() == "sdpa_with_kv_cache.default") {
+        kv_cache_ids_.insert(static_cast<int>(a->Get(3)));
+        kv_cache_ids_.insert(static_cast<int>(a->Get(4)));
+      }
       for (unsigned j = 0; j < a->size(); j++) {
         int id = static_cast<int>(a->Get(j));
         if (is_prepack && j == 0) {
@@ -378,6 +444,29 @@ void WebGPUGraph::build(
     }
   }
 
+  // f16 KV defensive guard: fail loud if a non-sdpa op reads an f16 cache.
+  // Inert unless kv_f16_ (runtime opt-in) is set.
+  if (kv_f16_ && !kv_cache_ids_.empty() && chain_prescan) {
+    for (unsigned ci = 0; ci < chain_prescan->size(); ci++) {
+      const auto* oc = chain_prescan->Get(ci);
+      const std::string nm = oc->name()->str();
+      if (nm == "sdpa_with_kv_cache.default" || nm == kPrepackOpName) {
+        continue;
+      }
+      const auto* a = oc->args();
+      if (!a) {
+        continue;
+      }
+      for (unsigned j = 0; j < a->size(); j++) {
+        if (kv_cache_ids_.count(static_cast<int>(a->Get(j))) != 0) {
+          throw std::runtime_error(
+              "WebGPU f16 KV: cache tensor consumed by non-sdpa op '" + nm +
+              "' would misread the f16 buffer");
+        }
+      }
+    }
+  }
+
   for (int i = 0; i < num_vals; i++) {
     const auto* val = values->Get(i);
     if (!val || val->value_type() == vkgraph::GraphTypes::NONE) {
@@ -407,6 +496,23 @@ void WebGPUGraph::build(
         tensor.cur_dims = tensor.dims;
         tensor.cur_nbytes = tensor.nbytes;
 
+        // f16 KV cache: dedicated half-size array<f16> buffer. WebGPU
+        // zero-initializes freshly-created buffers, so no explicit clear is
+        // needed. Inert unless kv_f16_ (runtime opt-in) is set.
+        if (kv_f16_ && kv_cache_ids_.count(i) != 0) {
+          tensor.elem_size = 2;
+          tensor.nbytes = numel * 2;
+          tensor.cur_nbytes = tensor.nbytes;
+          tensor_mem_obj_ids_[i] = -1;
+          WGPUBufferDescriptor buf_desc = {};
+          buf_desc.size = std::max(tensor.nbytes, size_t(4));
+          buf_desc.usage = WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |
+              WGPUBufferUsage_CopySrc;
+          buf_desc.mappedAtCreation = false;
+          tensor.buffer = wgpuDeviceCreateBuffer(device_, &buf_desc);
+          break;
+        }
+
         int constant_id = vk_tensor->constant_id();
         int mem_obj_id = vk_tensor->mem_obj_id();
 

backends/webgpu/runtime/WebGPUGraph.h

@@ -104,7 +104,9 @@ class WebGPUGraph {
   void build(
       const void* flatbuffer_data,
       const uint8_t* constant_data,
-      const executorch::runtime::NamedDataMap* named_data_map = nullptr);
+      const executorch::runtime::NamedDataMap* named_data_map = nullptr,
+      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);
@@ -267,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);
@@ -314,6 +345,23 @@ class WebGPUGraph {
     return value_types_[id];
   }
 
+ public:
+  // True when the sdpa K/V cache is stored f16-packed (runtime opt-in).
+  bool kv_f16() const {
+    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;
   WGPUDevice device_ = nullptr;
@@ -366,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_;