[ExecuTorch][WebGPU] Add opt-in native f16 KV cache (−280 MiB, default-OFF)

**Add an opt-in native f16 KV cache for WebGPU SDPA — −280 MiB device memory, decode-neutral, default-OFF (byte-identical when off).**

**Problem:** The SDPA K/V cache is stored fp32, so at long context it dominates WebGPU device memory. On a device that negotiated the shader-f16 feature the cache can be stored as native f16 (half the bytes) while attention still accumulates in f32, at no measured decode-speed cost.

**Solution:** Behind a new opt-in build flag `EXECUTORCH_WEBGPU_KV_F16` (defines `WGPU_BACKEND_KV_F16`, default OFF), and only when the device supports shader-f16 (fail-closed to f32 otherwise):

- **Before:** the K/V caches are allocated fp32 (`numel*4` bytes) and read by the f32 SDPA kernels.
- **After:** the K/V caches are allocated on a dedicated f16 buffer (`numel*2` bytes, zero-init); SDPA and FlashDecoding select the generated f16 (`_half`) kernel variants that bind the cache as f16 and widen to f32 on read — compute and accumulation are unchanged.

**Implementation:**

- Retires the 4 SDPA/decode kernels (`sdpa_compute_attn_weights`, `sdpa_compute_out`, `sdpa_fd_split`, `update_cache`) into `DTYPE`-templated variants via the landed WGSL shader-variant codegen: each `<kernel>.wgsl` + a `<kernel>.yaml` (`DTYPE: float/half`) generates BOTH the existing f32 header (regenerated byte-identical) and a new `_half` header — the f16 variant is generated, not a hand-duplicated file. The f16 delta is storage-only: `$if DTYPE == "half"` enables f16, the K/V cache binding becomes `array<${buffer_gvec_type(DTYPE, 4)}>` (QK/AV) or `array<${buffer_scalar_type(DTYPE)}>` (fd_split/update_cache), widened to f32 on read (`vec4<f32>(...)` / `f32(...)`); `update_cache` writes `f16(...)`. Mirrors the codegen usage of the landed `rms_norm` (an existing kernel retired into a template) and `steel_f16`.
- `WebGPUGraph::build`: collect the sdpa K/V cache value ids (`args[3]`, `args[4]`), allocate them as a dedicated half-size f16 buffer at the constant-allocation site, plus a defensive guard that throws if a non-sdpa op would misread an f16 cache.
- `Sdpa.cpp` / `SdpaFdDecode.cpp`: select the generated `_half` shader when `kv_f16()` via a local `const char*` (mirrors the steel-f16 gate — no function-signature/ABI change).
- Analogous to Vulkan's whole-graph `force_fp16` fp32→fp16 storage downcast (`backends/vulkan/serialization/vulkan_graph_builder.py` `get_effective_dtype`); here scoped to the K/V cache and gated on the negotiated shader-f16 feature rather than applied graph-wide.

**Constraints:** the default build (flag OFF) is byte-identical — all f16 code is `#ifdef WGPU_BACKEND_KV_F16`-gated, the templated f32 headers regenerate byte-identical (drift-check verified), no ABI change, no KV-cache-layout migration; the opt-in build fail-closes to f32 on a non-shader-f16 device; f16 storage requires `head_dim % 4 == 0` (already required and guarded).

Co-authored-with: Claude Code.

Pull Request resolved: #20772
ghstack-source-id: 401515180
@exported-using-ghexport

Differential Revision: [D110919974](https://our.internmc.facebook.com/intern/diff/D110919974/)

Author: JCNTH

backends/webgpu/runtime/WebGPUBackend.cpp

@@ -79,8 +79,23 @@ 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();
+    }
+  }
+
   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);
   } catch (const std::exception& e) {
     ET_LOG(Error, "WebGPU graph build failed: %s", e.what());
     graph->~WebGPUGraph();

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -310,7 +310,8 @@ 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) {
   if (!device_) {
     auto* ctx = get_default_webgpu_context();
     if (ctx) {
@@ -331,6 +332,11 @@ 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);
+
   // Phase 1: Create all values
   const auto* values = graph->values();
   const int num_vals = values ? values->size() : 0;
@@ -358,6 +364,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 +391,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 +443,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,8 @@ 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);
 
   // Copy input tensor data from host pointers into GPU buffers.
   void copy_inputs(const std::vector<InputData>& inputs);
@@ -314,6 +315,16 @@ 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_;
+  }
+
+ private:
+  bool kv_f16_ = false;
+  std::unordered_set<int> kv_cache_ids_;
+
  private:
   WGPUInstance instance_ = nullptr;
   WGPUDevice device_ = nullptr;

backends/webgpu/runtime/ops/sdpa/Sdpa.cpp

@@ -9,10 +9,13 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h>
+#include <executorch/backends/webgpu/runtime/ops/update_cache/update_cache_half_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/update_cache/update_cache_wgsl.h>
 
 #include <webgpu/webgpu.h>
@@ -255,9 +258,13 @@ static WGPUBuffer record_update_cache_dispatch(
   WGPUBuffer ubuf = graph.make_uniform_buffer(&uc, sizeof(uc));
   BufferBinding bindings[2] = {
       {cache.buffer, cache.nbytes}, {src.buffer, src.nbytes}};
+  const char* uc_src = kUpdateCacheWGSL;
+  if (graph.kv_f16()) {
+    uc_src = kUpdateCacheHalfWGSL;
+  }
   build_dispatch(
       graph,
-      kUpdateCacheWGSL,
+      uc_src,
       bindings,
       2,
       ubuf,
@@ -494,9 +501,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         {attn_weights, aw_bytes},
         {q.buffer, q.nbytes},
         {k_cache.buffer, k_cache.nbytes}};
+    const char* qk_src = kSdpaComputeAttnWeightsWGSL;
+    if (graph.kv_f16()) {
+      qk_src = kSdpaComputeAttnWeightsHalfWGSL;
+    }
     build_dispatch(
         graph,
-        kSdpaComputeAttnWeightsWGSL,
+        qk_src,
         bindings,
         3,
         ubuf,
@@ -547,9 +558,13 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         {out.buffer, out.nbytes},
         {attn_weights_softmax, aw_bytes},
         {v_cache.buffer, v_cache.nbytes}};
+    const char* av_src = kSdpaComputeOutWGSL;
+    if (graph.kv_f16()) {
+      av_src = kSdpaComputeOutHalfWGSL;
+    }
     build_dispatch(
         graph,
-        kSdpaComputeOutWGSL,
+        av_src,
         bindings,
         3,
         ubuf,

backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl

@@ -1,6 +1,8 @@
+$if DTYPE == "half":
+  enable f16;
 @group(0) @binding(0) var<storage, read_write> t_attn_weights: array<f32>;
 @group(0) @binding(1) var<storage, read> t_q: array<vec4<f32>>;
-@group(0) @binding(2) var<storage, read> t_k_cache: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_k_cache: array<${buffer_gvec_type(DTYPE, 4)}>;
 
 struct Params {
   S: u32,
@@ -36,7 +38,10 @@ fn load_k_vec4(c: u32, kvh: u32, d4: u32) -> vec4<f32> {
     return vec4<f32>(0.0, 0.0, 0.0, 0.0);
   }
   let base = c * params.Hkv * params.D + kvh * params.D + d4;
-  return t_k_cache[base / 4u];
+  $if DTYPE == "half":
+    return vec4<f32>(t_k_cache[base / 4u]);
+  $else:
+    return t_k_cache[base / 4u];
 }
 
 fn store_qk(s: u32, c: u32, h: u32, raw: f32) {

backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.yaml

@@ -0,0 +1,11 @@
+sdpa_compute_attn_weights:
+  parameter_names_with_default_values:
+    DTYPE: float
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: float
+        SUFFIX: ""
+      - VALUE: half
+        SUFFIX: half
+  shader_variants:
+    - NAME: sdpa_compute_attn_weights

backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_half_wgsl.h

@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from sdpa_compute_attn_weights.wgsl - DO NOT EDIT.
+// wgsl-sha256: c8795d66b9b51516795fb0113b21fd55086b4a54d9a4b81c7f394ffd96d117b3
+inline constexpr const char* kSdpaComputeAttnWeightsHalfWGSL = R"(
+enable f16;
+@group(0) @binding(0) var<storage, read_write> t_attn_weights: array<f32>;
+@group(0) @binding(1) var<storage, read> t_q: array<vec4<f32>>;
+@group(0) @binding(2) var<storage, read> t_k_cache: array<vec4<f16>>;
+
+struct Params {
+  S: u32,
+  Hq: u32,
+  Hkv: u32,
+  D: u32,
+  context_len: u32,
+  input_pos: u32,
+  g: u32,
+  scale: f32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+// WGSL forbids literal -inf; large finite negative is a WGSL-safe stand-in.
+const NEG_INF: f32 = -1.0e30;
+
+override wg_size: u32 = 64;
+
+const TM: u32 = 4u;
+const TN: u32 = 4u;
+
+// D is a multiple of 4 (host-guarded), so a d4 chunk is fully in-bounds — no per-lane check.
+fn load_q_vec4(s: u32, h: u32, d4: u32) -> vec4<f32> {
+  if (s >= params.S) {
+    return vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+  let base = s * params.Hq * params.D + h * params.D + d4;
+  return t_q[base / 4u];
+}
+
+fn load_k_vec4(c: u32, kvh: u32, d4: u32) -> vec4<f32> {
+  if (c >= params.context_len) {
+    return vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  }
+  let base = c * params.Hkv * params.D + kvh * params.D + d4;
+  return vec4<f32>(t_k_cache[base / 4u]);
+}
+
+fn store_qk(s: u32, c: u32, h: u32, raw: f32) {
+  if (s >= params.S || c >= params.context_len) {
+    return;
+  }
+  var val = raw * params.scale;
+  // Causal mask: position c may not attend beyond s + input_pos.
+  if (c > s + params.input_pos) {
+    val = NEG_INF;
+  }
+  let idx = h * params.S * params.context_len + s * params.context_len + c;
+  t_attn_weights[idx] = val;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  let nrt = (params.S + TM - 1u) / TM;
+  let nct = (params.context_len + TN - 1u) / TN;
+  let tiles = nrt * nct;
+  let total = tiles * params.Hq;
+  // 2D dispatch fold: recover the linear tile index across x/y.
+  let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+  if (idx >= total) {
+    return;
+  }
+
+  let h = idx / tiles;
+  let rem = idx % tiles;
+  let row_tile = rem / nct;
+  let col_tile = rem % nct;
+  let kvh = h / params.g;
+  let s0 = row_tile * TM;
+  let c0 = col_tile * TN;
+
+  var acc: array<vec4<f32>, 4>;
+  acc[0] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  acc[1] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  acc[2] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+  acc[3] = vec4<f32>(0.0, 0.0, 0.0, 0.0);
+
+  // Skip fully-masked causal tiles; mirrors Vulkan attn_weights_tiled.glsl.
+  let skip_tile = c0 > s0 + (TM - 1u) + params.input_pos;
+  var d4: u32 = 0u;
+  loop {
+    if (d4 >= params.D || skip_tile) {
+      break;
+    }
+    var q: array<vec4<f32>, TM>;
+    var k: array<vec4<f32>, TN>;
+    for (var i: u32 = 0u; i < TM; i = i + 1u) {
+      q[i] = load_q_vec4(s0 + i, h, d4);
+    }
+    for (var j: u32 = 0u; j < TN; j = j + 1u) {
+      k[j] = load_k_vec4(c0 + j, kvh, d4);
+    }
+    for (var i: u32 = 0u; i < TM; i = i + 1u) {
+      acc[i] += vec4<f32>(
+          dot(q[i], k[0]),
+          dot(q[i], k[1]),
+          dot(q[i], k[2]),
+          dot(q[i], k[3]));
+    }
+    d4 = d4 + 4u;
+  }
+
+  var m: u32 = 0u;
+  loop {
+    if (m >= TM) {
+      break;
+    }
+    let av = acc[m];
+    store_qk(s0 + m, c0 + 0u, h, av.x);
+    store_qk(s0 + m, c0 + 1u, h, av.y);
+    store_qk(s0 + m, c0 + 2u, h, av.z);
+    store_qk(s0 + m, c0 + 3u, h, av.w);
+    m = m + 1u;
+  }
+}
+)";
+
+inline constexpr uint32_t kSdpaComputeAttnWeightsHalfWorkgroupSizeX = 64;
+inline constexpr uint32_t kSdpaComputeAttnWeightsHalfWorkgroupSizeY = 1;
+inline constexpr uint32_t kSdpaComputeAttnWeightsHalfWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu