From 3649dd468715a3b8b2cff94b31acbac01adbabd9 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 7 Jul 2026 10:59:05 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 13 ++ backends/webgpu/runtime/WebGPUGraph.cpp | 58 ++++++ backends/webgpu/runtime/WebGPUGraph.h | 12 ++ backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 29 ++- .../sdpa/sdpa_compute_attn_weights_f16.wgsl | 123 +++++++++++++ .../sdpa/sdpa_compute_attn_weights_f16_wgsl.h | 147 ++++++++++++++++ .../ops/sdpa/sdpa_compute_out_f16.wgsl | 141 +++++++++++++++ .../ops/sdpa/sdpa_compute_out_f16_wgsl.h | 165 ++++++++++++++++++ .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 11 +- .../ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl | 134 ++++++++++++++ .../sdpa_fd_decode/sdpa_fd_split_f16_wgsl.h | 158 +++++++++++++++++ .../ops/update_cache/update_cache_f16.wgsl | 27 +++ .../ops/update_cache/update_cache_f16_wgsl.h | 51 ++++++ 13 files changed, 1065 insertions(+), 4 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16_wgsl.h create mode 100644 backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl create mode 100644 backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index be8e0593be7..ec1db8d495e 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/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() diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index eb10fb44e28..b8223fcf2f3 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/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(a->Get(3))); + kv_cache_ids_.insert(static_cast(a->Get(4))); + } +#endif for (unsigned j = 0; j < a->size(); j++) { int id = static_cast(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(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 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 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(); diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 0ebcd8071f9..4debd3a3b69 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/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 kv_cache_ids_; +#endif + private: WGPUInstance instance_ = nullptr; WGPUDevice device_ = nullptr; diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 17918863a6e..8f79f979707 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -14,6 +14,11 @@ #include #include #include +#ifdef WGPU_BACKEND_KV_F16 +#include +#include +#include +#endif #include @@ -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& 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& 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, diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16.wgsl new file mode 100644 index 00000000000..2c950bc6912 --- /dev/null +++ b/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 t_attn_weights: array; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; + +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 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 { + if (s >= params.S) { + return vec4(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 { + if (c >= params.context_len) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let base = c * params.Hkv * params.D + kvh * params.D + d4; + return vec4(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, + @builtin(num_workgroups) num_workgroups: vec3) { + 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, 4>; + acc[0] = vec4(0.0, 0.0, 0.0, 0.0); + acc[1] = vec4(0.0, 0.0, 0.0, 0.0); + acc[2] = vec4(0.0, 0.0, 0.0, 0.0); + acc[3] = vec4(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, TM>; + var k: array, 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( + 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; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h new file mode 100644 index 00000000000..e4a895a5f52 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h @@ -0,0 +1,147 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +// @generated from sdpa_compute_attn_weights_f16.wgsl - DO NOT EDIT. +// wgsl-sha256: 64d7a574615bfccc3338d44e4be33d5052a64e31910eb7a07c37ab525cf948be +inline constexpr const char* kSdpaComputeAttnWeightsF16WGSL = R"( +enable f16; + +// f16 K-cache variant of sdpa_compute_attn_weights.wgsl (f32 compute). +@group(0) @binding(0) var t_attn_weights: array; +@group(0) @binding(1) var t_q: array>; +@group(0) @binding(2) var t_k_cache: array>; + +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 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 { + if (s >= params.S) { + return vec4(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 { + if (c >= params.context_len) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let base = c * params.Hkv * params.D + kvh * params.D + d4; + return vec4(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, + @builtin(num_workgroups) num_workgroups: vec3) { + 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, 4>; + acc[0] = vec4(0.0, 0.0, 0.0, 0.0); + acc[1] = vec4(0.0, 0.0, 0.0, 0.0); + acc[2] = vec4(0.0, 0.0, 0.0, 0.0); + acc[3] = vec4(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, TM>; + var k: array, 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( + 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 kSdpaComputeAttnWeightsF16WorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaComputeAttnWeightsF16WorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaComputeAttnWeightsF16WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl new file mode 100644 index 00000000000..3c9b621f719 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl @@ -0,0 +1,141 @@ +enable f16; + +// f16 V-cache variant of sdpa_compute_out.wgsl (f32 compute). +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_attn_weights_softmax: array; +@group(0) @binding(2) var t_v_cache: array>; + +struct Params { + S: u32, + Hq: u32, + Hkv: u32, + D: u32, + context_len: u32, + g: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64; + +const TM: u32 = 4u; +const TN: u32 = 4u; + +// Checked loaders mask context lanes past context_len (D%4==0, host-guarded). +fn load_a_vec4(s: u32, h: u32, c4: u32) -> vec4 { + var r = vec4(0.0, 0.0, 0.0, 0.0); + if (s >= params.S) { + return r; + } + let base = h * params.S * params.context_len + s * params.context_len; + if (c4 + 0u < params.context_len) { r.x = t_attn_weights_softmax[base + c4 + 0u]; } + if (c4 + 1u < params.context_len) { r.y = t_attn_weights_softmax[base + c4 + 1u]; } + if (c4 + 2u < params.context_len) { r.z = t_attn_weights_softmax[base + c4 + 2u]; } + if (c4 + 3u < params.context_len) { r.w = t_attn_weights_softmax[base + c4 + 3u]; } + return r; +} + +fn load_v_d4(c: u32, kvh: u32, d0: u32) -> vec4 { + if (c >= params.context_len) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let base = c * params.Hkv * params.D + kvh * params.D + d0; + return vec4(t_v_cache[base / 4u]); +} + +// Branch-free loaders for the aligned body: caller guarantees c4..c4+3 < context_len. +fn load_a_vec4_nc(s: u32, h: u32, c4: u32) -> vec4 { + if (s >= params.S) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let base = h * params.S * params.context_len + s * params.context_len + c4; + return vec4(t_attn_weights_softmax[base], t_attn_weights_softmax[base + 1u], t_attn_weights_softmax[base + 2u], t_attn_weights_softmax[base + 3u]); +} + +fn load_v_d4_nc(c: u32, kvh: u32, d0: u32) -> vec4 { + let base = c * params.Hkv * params.D + kvh * params.D + d0; + return vec4(t_v_cache[base / 4u]); +} + +fn store_out_vec4(s: u32, d0: u32, h: u32, val: vec4) { + if (s >= params.S) { + return; + } + let idx = s * params.Hq * params.D + h * params.D + d0; + t_out[idx / 4u] = val; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let nrt = (params.S + TM - 1u) / TM; + let nct = (params.D + 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 d0 = col_tile * TN; + + var acc: array, 4>; + acc[0] = vec4(0.0, 0.0, 0.0, 0.0); + acc[1] = vec4(0.0, 0.0, 0.0, 0.0); + acc[2] = vec4(0.0, 0.0, 0.0, 0.0); + acc[3] = vec4(0.0, 0.0, 0.0, 0.0); + + // Branch-free aligned body + checked tail; mirrors Vulkan out_tiled.glsl. + let ctx_aligned = params.context_len - (params.context_len & 3u); + var c4: u32 = 0u; + loop { + if (c4 >= ctx_aligned) { + break; + } + let a0 = load_a_vec4_nc(s0 + 0u, h, c4); + let a1 = load_a_vec4_nc(s0 + 1u, h, c4); + let a2 = load_a_vec4_nc(s0 + 2u, h, c4); + let a3 = load_a_vec4_nc(s0 + 3u, h, c4); + let v0 = load_v_d4_nc(c4 + 0u, kvh, d0); + let v1 = load_v_d4_nc(c4 + 1u, kvh, d0); + let v2 = load_v_d4_nc(c4 + 2u, kvh, d0); + let v3 = load_v_d4_nc(c4 + 3u, kvh, d0); + acc[0] += a0.x * v0 + a0.y * v1 + a0.z * v2 + a0.w * v3; + acc[1] += a1.x * v0 + a1.y * v1 + a1.z * v2 + a1.w * v3; + acc[2] += a2.x * v0 + a2.y * v1 + a2.z * v2 + a2.w * v3; + acc[3] += a3.x * v0 + a3.y * v1 + a3.z * v2 + a3.w * v3; + c4 = c4 + 4u; + } + if (c4 < params.context_len) { + let a0 = load_a_vec4(s0 + 0u, h, c4); + let a1 = load_a_vec4(s0 + 1u, h, c4); + let a2 = load_a_vec4(s0 + 2u, h, c4); + let a3 = load_a_vec4(s0 + 3u, h, c4); + let v0 = load_v_d4(c4 + 0u, kvh, d0); + let v1 = load_v_d4(c4 + 1u, kvh, d0); + let v2 = load_v_d4(c4 + 2u, kvh, d0); + let v3 = load_v_d4(c4 + 3u, kvh, d0); + acc[0] += a0.x * v0 + a0.y * v1 + a0.z * v2 + a0.w * v3; + acc[1] += a1.x * v0 + a1.y * v1 + a1.z * v2 + a1.w * v3; + acc[2] += a2.x * v0 + a2.y * v1 + a2.z * v2 + a2.w * v3; + acc[3] += a3.x * v0 + a3.y * v1 + a3.z * v2 + a3.w * v3; + } + + var m: u32 = 0u; + loop { + if (m >= TM) { + break; + } + store_out_vec4(s0 + m, d0, h, acc[m]); + m = m + 1u; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h new file mode 100644 index 00000000000..488003f540e --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h @@ -0,0 +1,165 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +// @generated from sdpa_compute_out_f16.wgsl - DO NOT EDIT. +// wgsl-sha256: 64fa4352362fd055d8dde0d4fe5cde022ea1e655187fce1e1a6a0731a6c2372f +inline constexpr const char* kSdpaComputeOutF16WGSL = R"( +enable f16; + +// f16 V-cache variant of sdpa_compute_out.wgsl (f32 compute). +@group(0) @binding(0) var t_out: array>; +@group(0) @binding(1) var t_attn_weights_softmax: array; +@group(0) @binding(2) var t_v_cache: array>; + +struct Params { + S: u32, + Hq: u32, + Hkv: u32, + D: u32, + context_len: u32, + g: u32, + _pad0: u32, + _pad1: u32, +} +@group(0) @binding(3) var params: Params; + +override wg_size: u32 = 64; + +const TM: u32 = 4u; +const TN: u32 = 4u; + +// Checked loaders mask context lanes past context_len (D%4==0, host-guarded). +fn load_a_vec4(s: u32, h: u32, c4: u32) -> vec4 { + var r = vec4(0.0, 0.0, 0.0, 0.0); + if (s >= params.S) { + return r; + } + let base = h * params.S * params.context_len + s * params.context_len; + if (c4 + 0u < params.context_len) { r.x = t_attn_weights_softmax[base + c4 + 0u]; } + if (c4 + 1u < params.context_len) { r.y = t_attn_weights_softmax[base + c4 + 1u]; } + if (c4 + 2u < params.context_len) { r.z = t_attn_weights_softmax[base + c4 + 2u]; } + if (c4 + 3u < params.context_len) { r.w = t_attn_weights_softmax[base + c4 + 3u]; } + return r; +} + +fn load_v_d4(c: u32, kvh: u32, d0: u32) -> vec4 { + if (c >= params.context_len) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let base = c * params.Hkv * params.D + kvh * params.D + d0; + return vec4(t_v_cache[base / 4u]); +} + +// Branch-free loaders for the aligned body: caller guarantees c4..c4+3 < context_len. +fn load_a_vec4_nc(s: u32, h: u32, c4: u32) -> vec4 { + if (s >= params.S) { + return vec4(0.0, 0.0, 0.0, 0.0); + } + let base = h * params.S * params.context_len + s * params.context_len + c4; + return vec4(t_attn_weights_softmax[base], t_attn_weights_softmax[base + 1u], t_attn_weights_softmax[base + 2u], t_attn_weights_softmax[base + 3u]); +} + +fn load_v_d4_nc(c: u32, kvh: u32, d0: u32) -> vec4 { + let base = c * params.Hkv * params.D + kvh * params.D + d0; + return vec4(t_v_cache[base / 4u]); +} + +fn store_out_vec4(s: u32, d0: u32, h: u32, val: vec4) { + if (s >= params.S) { + return; + } + let idx = s * params.Hq * params.D + h * params.D + d0; + t_out[idx / 4u] = val; +} + +@compute @workgroup_size(wg_size, 1, 1) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let nrt = (params.S + TM - 1u) / TM; + let nct = (params.D + 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 d0 = col_tile * TN; + + var acc: array, 4>; + acc[0] = vec4(0.0, 0.0, 0.0, 0.0); + acc[1] = vec4(0.0, 0.0, 0.0, 0.0); + acc[2] = vec4(0.0, 0.0, 0.0, 0.0); + acc[3] = vec4(0.0, 0.0, 0.0, 0.0); + + // Branch-free aligned body + checked tail; mirrors Vulkan out_tiled.glsl. + let ctx_aligned = params.context_len - (params.context_len & 3u); + var c4: u32 = 0u; + loop { + if (c4 >= ctx_aligned) { + break; + } + let a0 = load_a_vec4_nc(s0 + 0u, h, c4); + let a1 = load_a_vec4_nc(s0 + 1u, h, c4); + let a2 = load_a_vec4_nc(s0 + 2u, h, c4); + let a3 = load_a_vec4_nc(s0 + 3u, h, c4); + let v0 = load_v_d4_nc(c4 + 0u, kvh, d0); + let v1 = load_v_d4_nc(c4 + 1u, kvh, d0); + let v2 = load_v_d4_nc(c4 + 2u, kvh, d0); + let v3 = load_v_d4_nc(c4 + 3u, kvh, d0); + acc[0] += a0.x * v0 + a0.y * v1 + a0.z * v2 + a0.w * v3; + acc[1] += a1.x * v0 + a1.y * v1 + a1.z * v2 + a1.w * v3; + acc[2] += a2.x * v0 + a2.y * v1 + a2.z * v2 + a2.w * v3; + acc[3] += a3.x * v0 + a3.y * v1 + a3.z * v2 + a3.w * v3; + c4 = c4 + 4u; + } + if (c4 < params.context_len) { + let a0 = load_a_vec4(s0 + 0u, h, c4); + let a1 = load_a_vec4(s0 + 1u, h, c4); + let a2 = load_a_vec4(s0 + 2u, h, c4); + let a3 = load_a_vec4(s0 + 3u, h, c4); + let v0 = load_v_d4(c4 + 0u, kvh, d0); + let v1 = load_v_d4(c4 + 1u, kvh, d0); + let v2 = load_v_d4(c4 + 2u, kvh, d0); + let v3 = load_v_d4(c4 + 3u, kvh, d0); + acc[0] += a0.x * v0 + a0.y * v1 + a0.z * v2 + a0.w * v3; + acc[1] += a1.x * v0 + a1.y * v1 + a1.z * v2 + a1.w * v3; + acc[2] += a2.x * v0 + a2.y * v1 + a2.z * v2 + a2.w * v3; + acc[3] += a3.x * v0 + a3.y * v1 + a3.z * v2 + a3.w * v3; + } + + var m: u32 = 0u; + loop { + if (m >= TM) { + break; + } + store_out_vec4(s0 + m, d0, h, acc[m]); + m = m + 1u; + } +} +)"; + +inline constexpr uint32_t kSdpaComputeOutF16WorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaComputeOutF16WorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaComputeOutF16WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index 70108beb892..9a604c90aca 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -13,6 +13,9 @@ #include #include #include +#ifdef WGPU_BACKEND_KV_F16 +#include +#endif #include @@ -218,9 +221,15 @@ void sdpa_fd_decode_dispatch( static_cast(split_threads), kSdpaFdSplitWorkgroupSizeX, "fd_split"); + const char* split_shader = kSdpaFdSplitWGSL; +#ifdef WGPU_BACKEND_KV_F16 + if (graph.kv_f16()) { + split_shader = kSdpaFdSplitF16WGSL; + } +#endif build_dispatch( graph, - kSdpaFdSplitWGSL, + split_shader, split_bindings, 5, 2, diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl new file mode 100644 index 00000000000..236cb96764b --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl @@ -0,0 +1,134 @@ +enable f16; + +// f16 K/V-cache variant of sdpa_fd_split.wgsl (f32 compute + softmax stats). +@group(0) @binding(0) var t_part_o: array; +@group(0) @binding(1) var t_part_ml: array; +@group(0) @binding(2) var t_q: array; +@group(0) @binding(3) var t_k_cache: array; +@group(0) @binding(4) var t_v_cache: array; + +struct Params { + _pad0: u32, + Hkv: u32, + D: u32, + context_len: u32, + g: u32, + num_splits: u32, + split_len: u32, + scale: f32, +} +@group(0) @binding(5) var params: Params; + +const WG_SIZE: u32 = 64u; +const MAX_SPLITS: u32 = 128u; +const MAX_D_PER_LANE: u32 = 2u; +const NEG_INF: f32 = -1.0e30; + +// sh_s: block scores then softmax weights; sh_red: max/sum reduction scratch. +var sh_s: array; +var sh_red: array; + +// FlashDecoding pass 1: per-(head,split) unnormalized softmax partial. +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let h = wid.x / params.num_splits; + let split_i = wid.x % params.num_splits; + let t = lid.x; + let D = params.D; + let D4 = D / 4u; // D is a multiple of 4 (guarded host-side); vec4 QK dot + let ctx = params.context_len; + let kv = h / params.g; + let q_base = h * D; + let kv_row_stride = params.Hkv * D; + + let c0 = split_i * params.split_len; + var c1 = c0 + params.split_len; + if (c1 > ctx) { c1 = ctx; } + + var m: f32 = NEG_INF; + var l: f32 = 0.0; + var o_acc: array; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { o_acc[nd] = 0.0; } + + // Stream the split in blocks of WG_SIZE KV positions. + var block: u32 = c0; + loop { + if (block >= c1) { break; } + var n: u32 = c1 - block; + if (n > WG_SIZE) { n = WG_SIZE; } + + // Phase 1: lane t computes the full QK dot for position block+t (vec4), one + // K row read once. Out-of-block lanes hold NEG_INF (safe for the max). + var s: f32 = NEG_INF; + if (t < n) { + let kvbase = (block + t) * kv_row_stride + kv * D; + var acc4 = vec4(0.0, 0.0, 0.0, 0.0); + for (var i4: u32 = 0u; i4 < D4; i4 = i4 + 1u) { + let qi = q_base + i4 * 4u; + let ki = kvbase + i4 * 4u; + let qv = vec4(t_q[qi], t_q[qi + 1u], t_q[qi + 2u], t_q[qi + 3u]); + let kvv = vec4( + f32(t_k_cache[ki]), f32(t_k_cache[ki + 1u]), + f32(t_k_cache[ki + 2u]), f32(t_k_cache[ki + 3u])); + acc4 = acc4 + qv * kvv; + } + s = (acc4.x + acc4.y + acc4.z + acc4.w) * params.scale; + } + sh_s[t] = s; + + // Phase 2a: block max via tree reduction (sh_red written from register s). + sh_red[t] = s; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { sh_red[t] = max(sh_red[t], sh_red[t + stride]); } + workgroupBarrier(); + } + let m_new = max(m, sh_red[0]); + let rescale = exp(m - m_new); + + // Phase 2b: each lane exponentiates ITS position once -> p (reuse sh_s), + // and reduce the block sum of p. + var p_t: f32 = 0.0; + if (t < n) { p_t = exp(sh_s[t] - m_new); } + workgroupBarrier(); // all reads of sh_s (the scores) done before overwrite + sh_s[t] = p_t; + sh_red[t] = p_t; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { sh_red[t] = sh_red[t] + sh_red[t + stride]; } + workgroupBarrier(); + } + l = rescale * l + sh_red[0]; + + // Phase 2c: each lane accumulates V for its own output dims over the block, + // reading the shared per-position weights (no exp in this loop). + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + var acc: f32 = rescale * o_acc[nd]; + for (var j: u32 = 0u; j < n; j = j + 1u) { + let vbase = (block + j) * kv_row_stride + kv * D; + acc = acc + sh_s[j] * f32(t_v_cache[vbase + d]); + } + o_acc[nd] = acc; + } + } + m = m_new; + workgroupBarrier(); // before the next block overwrites sh_s / sh_red + block = block + WG_SIZE; + } + + let part = h * MAX_SPLITS + split_i; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + t_part_o[part * D + d] = o_acc[nd]; + } + } + if (t == 0u) { + t_part_ml[part * 2u + 0u] = m; + t_part_ml[part * 2u + 1u] = l; + } +} diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16_wgsl.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16_wgsl.h new file mode 100644 index 00000000000..71dfcfb8af0 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16_wgsl.h @@ -0,0 +1,158 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +// @generated from sdpa_fd_split_f16.wgsl - DO NOT EDIT. +// wgsl-sha256: 3396d4f0d990d59ded70734cb92feb6c776635695b5d5c21b8cdfb7e55c8ed6d +inline constexpr const char* kSdpaFdSplitF16WGSL = R"( +enable f16; + +// f16 K/V-cache variant of sdpa_fd_split.wgsl (f32 compute + softmax stats). +@group(0) @binding(0) var t_part_o: array; +@group(0) @binding(1) var t_part_ml: array; +@group(0) @binding(2) var t_q: array; +@group(0) @binding(3) var t_k_cache: array; +@group(0) @binding(4) var t_v_cache: array; + +struct Params { + _pad0: u32, + Hkv: u32, + D: u32, + context_len: u32, + g: u32, + num_splits: u32, + split_len: u32, + scale: f32, +} +@group(0) @binding(5) var params: Params; + +const WG_SIZE: u32 = 64u; +const MAX_SPLITS: u32 = 128u; +const MAX_D_PER_LANE: u32 = 2u; +const NEG_INF: f32 = -1.0e30; + +// sh_s: block scores then softmax weights; sh_red: max/sum reduction scratch. +var sh_s: array; +var sh_red: array; + +// FlashDecoding pass 1: per-(head,split) unnormalized softmax partial. +@compute @workgroup_size(64, 1, 1) +fn main( + @builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let h = wid.x / params.num_splits; + let split_i = wid.x % params.num_splits; + let t = lid.x; + let D = params.D; + let D4 = D / 4u; // D is a multiple of 4 (guarded host-side); vec4 QK dot + let ctx = params.context_len; + let kv = h / params.g; + let q_base = h * D; + let kv_row_stride = params.Hkv * D; + + let c0 = split_i * params.split_len; + var c1 = c0 + params.split_len; + if (c1 > ctx) { c1 = ctx; } + + var m: f32 = NEG_INF; + var l: f32 = 0.0; + var o_acc: array; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { o_acc[nd] = 0.0; } + + // Stream the split in blocks of WG_SIZE KV positions. + var block: u32 = c0; + loop { + if (block >= c1) { break; } + var n: u32 = c1 - block; + if (n > WG_SIZE) { n = WG_SIZE; } + + // Phase 1: lane t computes the full QK dot for position block+t (vec4), one + // K row read once. Out-of-block lanes hold NEG_INF (safe for the max). + var s: f32 = NEG_INF; + if (t < n) { + let kvbase = (block + t) * kv_row_stride + kv * D; + var acc4 = vec4(0.0, 0.0, 0.0, 0.0); + for (var i4: u32 = 0u; i4 < D4; i4 = i4 + 1u) { + let qi = q_base + i4 * 4u; + let ki = kvbase + i4 * 4u; + let qv = vec4(t_q[qi], t_q[qi + 1u], t_q[qi + 2u], t_q[qi + 3u]); + let kvv = vec4( + f32(t_k_cache[ki]), f32(t_k_cache[ki + 1u]), + f32(t_k_cache[ki + 2u]), f32(t_k_cache[ki + 3u])); + acc4 = acc4 + qv * kvv; + } + s = (acc4.x + acc4.y + acc4.z + acc4.w) * params.scale; + } + sh_s[t] = s; + + // Phase 2a: block max via tree reduction (sh_red written from register s). + sh_red[t] = s; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { sh_red[t] = max(sh_red[t], sh_red[t + stride]); } + workgroupBarrier(); + } + let m_new = max(m, sh_red[0]); + let rescale = exp(m - m_new); + + // Phase 2b: each lane exponentiates ITS position once -> p (reuse sh_s), + // and reduce the block sum of p. + var p_t: f32 = 0.0; + if (t < n) { p_t = exp(sh_s[t] - m_new); } + workgroupBarrier(); // all reads of sh_s (the scores) done before overwrite + sh_s[t] = p_t; + sh_red[t] = p_t; + workgroupBarrier(); + for (var stride: u32 = WG_SIZE / 2u; stride > 0u; stride = stride >> 1u) { + if (t < stride) { sh_red[t] = sh_red[t] + sh_red[t + stride]; } + workgroupBarrier(); + } + l = rescale * l + sh_red[0]; + + // Phase 2c: each lane accumulates V for its own output dims over the block, + // reading the shared per-position weights (no exp in this loop). + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + var acc: f32 = rescale * o_acc[nd]; + for (var j: u32 = 0u; j < n; j = j + 1u) { + let vbase = (block + j) * kv_row_stride + kv * D; + acc = acc + sh_s[j] * f32(t_v_cache[vbase + d]); + } + o_acc[nd] = acc; + } + } + m = m_new; + workgroupBarrier(); // before the next block overwrites sh_s / sh_red + block = block + WG_SIZE; + } + + let part = h * MAX_SPLITS + split_i; + for (var nd: u32 = 0u; nd < MAX_D_PER_LANE; nd = nd + 1u) { + let d = t + nd * WG_SIZE; + if (d < D) { + t_part_o[part * D + d] = o_acc[nd]; + } + } + if (t == 0u) { + t_part_ml[part * 2u + 0u] = m; + t_part_ml[part * 2u + 1u] = l; + } +} +)"; + +inline constexpr uint32_t kSdpaFdSplitF16WorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaFdSplitF16WorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaFdSplitF16WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl b/backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl new file mode 100644 index 00000000000..7d2c0ceb5c1 --- /dev/null +++ b/backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl @@ -0,0 +1,27 @@ +enable f16; + +// f16 KV-cache variant of update_cache.wgsl: cache stored as f16, source f32. +@group(0) @binding(0) var t_cache: array; +@group(0) @binding(1) var t_value: array; + +struct Params { + numel: u32, + dst_offset: u32, + cache_numel: u32, + _pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.numel) { + return; + } + if (params.dst_offset + i >= params.cache_numel) { + return; + } + t_cache[params.dst_offset + i] = f16(t_value[i]); +} diff --git a/backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h b/backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h new file mode 100644 index 00000000000..3a838f5dc2c --- /dev/null +++ b/backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h @@ -0,0 +1,51 @@ +/* + * 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 + +namespace executorch::backends::webgpu { + +// @generated from update_cache_f16.wgsl - DO NOT EDIT. +// wgsl-sha256: e115bcc12d0b0d09e9f72e1f63fee21c0ce56f787dfbd728b9f4adebb1bf5630 +inline constexpr const char* kUpdateCacheF16WGSL = R"( +enable f16; + +// f16 KV-cache variant of update_cache.wgsl: cache stored as f16, source f32. +@group(0) @binding(0) var t_cache: array; +@group(0) @binding(1) var t_value: array; + +struct Params { + numel: u32, + dst_offset: u32, + cache_numel: u32, + _pad0: u32, +} +@group(0) @binding(2) var params: Params; + +override wg_size: u32 = 256; + +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let i = gid.x; + if (i >= params.numel) { + return; + } + if (params.dst_offset + i >= params.cache_numel) { + return; + } + t_cache[params.dst_offset + i] = f16(t_value[i]); +} +)"; + +inline constexpr uint32_t kUpdateCacheF16WorkgroupSizeX = 256; +inline constexpr uint32_t kUpdateCacheF16WorkgroupSizeY = 1; +inline constexpr uint32_t kUpdateCacheF16WorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu