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

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/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_;
 

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/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_;
 

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/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_;
 

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/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_;
 

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/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_;
 

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/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_;
 

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/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_;
 

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/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_;