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

backends/webgpu/CMakeLists.txt

@@ -115,6 +115,16 @@ if(EXECUTORCH_WEBGPU_STEEL_F16)
   target_compile_definitions(webgpu_backend PRIVATE WGPU_BACKEND_STEEL_F16)
 endif()
 
+# Opt-in native f16 KV cache (Sdpa.cpp/WebGPUGraph.cpp); OFF so default builds
+# keep the f32 cache + a byte-identical graph. Runtime-gated on the negotiated
+# shader-f16 feature (fail-closed). Mirrors the steel-f16 gate above.
+option(EXECUTORCH_WEBGPU_KV_F16 "Enable native f16 KV cache (needs shader-f16)"
+       OFF
+)
+if(EXECUTORCH_WEBGPU_KV_F16)
+  target_compile_definitions(webgpu_backend PRIVATE WGPU_BACKEND_KV_F16)
+endif()
+
 # Link with --whole-archive for static registration of backend + ops
 executorch_target_link_options_shared_lib(webgpu_backend)
 
@@ -161,6 +171,9 @@ function(add_webgpu_native_test test_name test_src)
   if(EXECUTORCH_WEBGPU_STEEL_F16)
     target_compile_definitions(${test_name} PRIVATE WGPU_BACKEND_STEEL_F16)
   endif()
+  if(EXECUTORCH_WEBGPU_KV_F16)
+    target_compile_definitions(${test_name} PRIVATE WGPU_BACKEND_KV_F16)
+  endif()
   set_property(TARGET ${test_name} PROPERTY CXX_STANDARD 17)
 endfunction()
 

backends/webgpu/runtime/WebGPUGraph.cpp

@@ -331,6 +331,12 @@ void WebGPUGraph::build(
   constant_data_ = constant_data;
   named_data_map_ = named_data_map;
 
+#ifdef WGPU_BACKEND_KV_F16
+  // f16 KV cache (opt-in): store K/V caches as f16 iff shader-f16 negotiated.
+  const WebGPUContext* kv_ctx = get_default_webgpu_context();
+  kv_f16_ = (kv_ctx != nullptr && kv_ctx->shader_f16_supported);
+#endif
+
   // Phase 1: Create all values
   const auto* values = graph->values();
   const int num_vals = values ? values->size() : 0;
@@ -358,6 +364,14 @@ void WebGPUGraph::build(
       if (!a) {
         continue;
       }
+#ifdef WGPU_BACKEND_KV_F16
+      // f16 KV: tag sdpa K/V cache values (args[3],[4]) for half-size alloc.
+      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)));
+      }
+#endif
       for (unsigned j = 0; j < a->size(); j++) {
         int id = static_cast<int>(a->Get(j));
         if (is_prepack && j == 0) {
@@ -378,6 +392,30 @@ void WebGPUGraph::build(
     }
   }
 
+#ifdef WGPU_BACKEND_KV_F16
+  // f16 KV defensive guard: fail loud if a non-sdpa op reads an f16 cache.
+  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");
+        }
+      }
+    }
+  }
+#endif
+
   for (int i = 0; i < num_vals; i++) {
     const auto* val = values->Get(i);
     if (!val || val->value_type() == vkgraph::GraphTypes::NONE) {
@@ -407,6 +445,26 @@ void WebGPUGraph::build(
         tensor.cur_dims = tensor.dims;
         tensor.cur_nbytes = tensor.nbytes;
 
+#ifdef WGPU_BACKEND_KV_F16
+        // f16 KV cache: dedicated half-size array<f16> buffer, zero-init.
+        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);
+          std::vector<uint8_t> zeros(tensor.nbytes, 0);
+          wgpuQueueWriteBuffer(
+              queue_, tensor.buffer, 0, zeros.data(), tensor.nbytes);
+          break;
+        }
+#endif
+
         int constant_id = vk_tensor->constant_id();
         int mem_obj_id = vk_tensor->mem_obj_id();
 

backends/webgpu/runtime/WebGPUGraph.h

@@ -314,6 +314,18 @@ class WebGPUGraph {
     return value_types_[id];
   }
 
+#ifdef WGPU_BACKEND_KV_F16
+ public:
+  // True when the sdpa K/V cache is stored f16-packed (opt-in build).
+  bool kv_f16() const {
+    return kv_f16_;
+  }
+
+ private:
+  bool kv_f16_ = false;
+  std::unordered_set<int> kv_cache_ids_;
+#endif
+
  private:
   WGPUInstance instance_ = nullptr;
   WGPUDevice device_ = nullptr;

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

@@ -14,6 +14,11 @@
 #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_wgsl.h>
+#ifdef WGPU_BACKEND_KV_F16
+#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h>
+#endif
 
 #include <webgpu/webgpu.h>
 
@@ -255,9 +260,15 @@ 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;
+#ifdef WGPU_BACKEND_KV_F16
+  if (graph.kv_f16()) {
+    uc_src = kUpdateCacheF16WGSL;
+  }
+#endif
   build_dispatch(
       graph,
-      kUpdateCacheWGSL,
+      uc_src,
       bindings,
       2,
       ubuf,
@@ -494,9 +505,15 @@ 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;
+#ifdef WGPU_BACKEND_KV_F16
+    if (graph.kv_f16()) {
+      qk_src = kSdpaComputeAttnWeightsF16WGSL;
+    }
+#endif
     build_dispatch(
         graph,
-        kSdpaComputeAttnWeightsWGSL,
+        qk_src,
         bindings,
         3,
         ubuf,
@@ -547,9 +564,15 @@ 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;
+#ifdef WGPU_BACKEND_KV_F16
+    if (graph.kv_f16()) {
+      av_src = kSdpaComputeOutF16WGSL;
+    }
+#endif
     build_dispatch(
         graph,
-        kSdpaComputeOutWGSL,
+        av_src,
         bindings,
         3,
         ubuf,

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

@@ -0,0 +1,123 @@
+enable f16;
+
+// f16 K-cache variant of sdpa_compute_attn_weights.wgsl (f32 compute).
+@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;
+  }
+}