From 276826804be71fe5afd00b9dcbb7e7e6c8543330 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 7 Jul 2026 11:02:21 -0700 Subject: [PATCH 1/2] 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 From 6e03bac0b92be84b04889eb32f45e30a39efe961 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 7 Jul 2026 11:29:22 -0700 Subject: [PATCH 2/2] Update [ghstack-poisoned] --- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 12 +- .../ops/sdpa/sdpa_compute_attn_weights.wgsl | 9 +- .../ops/sdpa/sdpa_compute_attn_weights.yaml | 11 ++ .../sdpa/sdpa_compute_attn_weights_f16.wgsl | 123 --------------- ... => sdpa_compute_attn_weights_half_wgsl.h} | 14 +- .../runtime/ops/sdpa/sdpa_compute_out.wgsl | 14 +- .../runtime/ops/sdpa/sdpa_compute_out.yaml | 11 ++ .../ops/sdpa/sdpa_compute_out_f16.wgsl | 141 ------------------ ...16_wgsl.h => sdpa_compute_out_half_wgsl.h} | 14 +- .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 4 +- .../ops/sdpa_fd_decode/sdpa_fd_split.wgsl | 22 ++- .../ops/sdpa_fd_decode/sdpa_fd_split.yaml | 11 ++ .../ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl | 134 ----------------- ...t_f16_wgsl.h => sdpa_fd_split_half_wgsl.h} | 14 +- .../ops/update_cache/update_cache.wgsl | 9 +- .../ops/update_cache/update_cache.yaml | 11 ++ .../ops/update_cache/update_cache_f16.wgsl | 27 ---- ...he_f16_wgsl.h => update_cache_half_wgsl.h} | 14 +- 18 files changed, 117 insertions(+), 478 deletions(-) create mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.yaml delete mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16.wgsl rename backends/webgpu/runtime/ops/sdpa/{sdpa_compute_attn_weights_f16_wgsl.h => sdpa_compute_attn_weights_half_wgsl.h} (87%) create mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.yaml delete mode 100644 backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl rename backends/webgpu/runtime/ops/sdpa/{sdpa_compute_out_f16_wgsl.h => sdpa_compute_out_half_wgsl.h} (91%) create mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.yaml delete mode 100644 backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl rename backends/webgpu/runtime/ops/sdpa_fd_decode/{sdpa_fd_split_f16_wgsl.h => sdpa_fd_split_half_wgsl.h} (91%) create mode 100644 backends/webgpu/runtime/ops/update_cache/update_cache.yaml delete mode 100644 backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl rename backends/webgpu/runtime/ops/update_cache/{update_cache_f16_wgsl.h => update_cache_half_wgsl.h} (67%) diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index 8f79f979707..30e385822fe 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -15,9 +15,9 @@ #include #include #ifdef WGPU_BACKEND_KV_F16 -#include -#include -#include +#include +#include +#include #endif #include @@ -263,7 +263,7 @@ static WGPUBuffer record_update_cache_dispatch( const char* uc_src = kUpdateCacheWGSL; #ifdef WGPU_BACKEND_KV_F16 if (graph.kv_f16()) { - uc_src = kUpdateCacheF16WGSL; + uc_src = kUpdateCacheHalfWGSL; } #endif build_dispatch( @@ -508,7 +508,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const char* qk_src = kSdpaComputeAttnWeightsWGSL; #ifdef WGPU_BACKEND_KV_F16 if (graph.kv_f16()) { - qk_src = kSdpaComputeAttnWeightsF16WGSL; + qk_src = kSdpaComputeAttnWeightsHalfWGSL; } #endif build_dispatch( @@ -567,7 +567,7 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const char* av_src = kSdpaComputeOutWGSL; #ifdef WGPU_BACKEND_KV_F16 if (graph.kv_f16()) { - av_src = kSdpaComputeOutF16WGSL; + av_src = kSdpaComputeOutHalfWGSL; } #endif build_dispatch( diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl index 014f0039048..097d87ecce0 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.wgsl @@ -1,6 +1,8 @@ +$if DTYPE == "half": + enable f16; @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>; +@group(0) @binding(2) var 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 { return vec4(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(t_k_cache[base / 4u]); + $else: + return t_k_cache[base / 4u]; } fn store_qk(s: u32, c: u32, h: u32, raw: f32) { diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.yaml b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights.yaml new file mode 100644 index 00000000000..d031e2864f4 --- /dev/null +++ b/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 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 deleted file mode 100644 index 2c950bc6912..00000000000 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16.wgsl +++ /dev/null @@ -1,123 +0,0 @@ -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_half_wgsl.h similarity index 87% rename from backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h rename to backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_half_wgsl.h index e4a895a5f52..dc6f5858c7d 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_f16_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_attn_weights_half_wgsl.h @@ -12,12 +12,10 @@ namespace executorch::backends::webgpu { -// @generated from sdpa_compute_attn_weights_f16.wgsl - DO NOT EDIT. -// wgsl-sha256: 64d7a574615bfccc3338d44e4be33d5052a64e31910eb7a07c37ab525cf948be -inline constexpr const char* kSdpaComputeAttnWeightsF16WGSL = R"( +// @generated from sdpa_compute_attn_weights.wgsl - DO NOT EDIT. +// wgsl-sha256: c8795d66b9b51516795fb0113b21fd55086b4a54d9a4b81c7f394ffd96d117b3 +inline constexpr const char* kSdpaComputeAttnWeightsHalfWGSL = 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>; @@ -140,8 +138,8 @@ fn main( } )"; -inline constexpr uint32_t kSdpaComputeAttnWeightsF16WorkgroupSizeX = 64; -inline constexpr uint32_t kSdpaComputeAttnWeightsF16WorkgroupSizeY = 1; -inline constexpr uint32_t kSdpaComputeAttnWeightsF16WorkgroupSizeZ = 1; +inline constexpr uint32_t kSdpaComputeAttnWeightsHalfWorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaComputeAttnWeightsHalfWorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaComputeAttnWeightsHalfWorkgroupSizeZ = 1; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl index 713345c0afa..53067447832 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.wgsl @@ -1,6 +1,8 @@ +$if DTYPE == "half": + enable f16; @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>; +@group(0) @binding(2) var t_v_cache: array<${buffer_gvec_type(DTYPE, 4)}>; struct Params { S: u32, @@ -38,7 +40,10 @@ fn load_v_d4(c: u32, kvh: u32, d0: u32) -> vec4 { return vec4(0.0, 0.0, 0.0, 0.0); } let base = c * params.Hkv * params.D + kvh * params.D + d0; - return t_v_cache[base / 4u]; + $if DTYPE == "half": + return vec4(t_v_cache[base / 4u]); + $else: + return t_v_cache[base / 4u]; } // Branch-free loaders for the aligned body: caller guarantees c4..c4+3 < context_len. @@ -52,7 +57,10 @@ fn load_a_vec4_nc(s: u32, h: u32, c4: u32) -> vec4 { fn load_v_d4_nc(c: u32, kvh: u32, d0: u32) -> vec4 { let base = c * params.Hkv * params.D + kvh * params.D + d0; - return t_v_cache[base / 4u]; + $if DTYPE == "half": + return vec4(t_v_cache[base / 4u]); + $else: + return t_v_cache[base / 4u]; } fn store_out_vec4(s: u32, d0: u32, h: u32, val: vec4) { diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.yaml b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.yaml new file mode 100644 index 00000000000..14b418f8a60 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out.yaml @@ -0,0 +1,11 @@ +sdpa_compute_out: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: float + SUFFIX: "" + - VALUE: half + SUFFIX: half + shader_variants: + - NAME: sdpa_compute_out diff --git a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl deleted file mode 100644 index 3c9b621f719..00000000000 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16.wgsl +++ /dev/null @@ -1,141 +0,0 @@ -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_half_wgsl.h similarity index 91% rename from backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h rename to backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_half_wgsl.h index 488003f540e..228921af58b 100644 --- a/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_f16_wgsl.h +++ b/backends/webgpu/runtime/ops/sdpa/sdpa_compute_out_half_wgsl.h @@ -12,12 +12,10 @@ namespace executorch::backends::webgpu { -// @generated from sdpa_compute_out_f16.wgsl - DO NOT EDIT. -// wgsl-sha256: 64fa4352362fd055d8dde0d4fe5cde022ea1e655187fce1e1a6a0731a6c2372f -inline constexpr const char* kSdpaComputeOutF16WGSL = R"( +// @generated from sdpa_compute_out.wgsl - DO NOT EDIT. +// wgsl-sha256: ed9709c966538edf2cbc6be97c284b89a9d921b6a4dbf115c6cbd76af301a1be +inline constexpr const char* kSdpaComputeOutHalfWGSL = 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>; @@ -158,8 +156,8 @@ fn main( } )"; -inline constexpr uint32_t kSdpaComputeOutF16WorkgroupSizeX = 64; -inline constexpr uint32_t kSdpaComputeOutF16WorkgroupSizeY = 1; -inline constexpr uint32_t kSdpaComputeOutF16WorkgroupSizeZ = 1; +inline constexpr uint32_t kSdpaComputeOutHalfWorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaComputeOutHalfWorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaComputeOutHalfWorkgroupSizeZ = 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 9a604c90aca..f6f39cc4bc8 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -14,7 +14,7 @@ #include #include #ifdef WGPU_BACKEND_KV_F16 -#include +#include #endif #include @@ -224,7 +224,7 @@ void sdpa_fd_decode_dispatch( const char* split_shader = kSdpaFdSplitWGSL; #ifdef WGPU_BACKEND_KV_F16 if (graph.kv_f16()) { - split_shader = kSdpaFdSplitF16WGSL; + split_shader = kSdpaFdSplitHalfWGSL; } #endif build_dispatch( diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl index c14c6bd07bd..da67489cc4d 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.wgsl @@ -1,8 +1,10 @@ +$if DTYPE == "half": + enable f16; @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; +@group(0) @binding(3) var t_k_cache: array<${buffer_scalar_type(DTYPE)}>; +@group(0) @binding(4) var t_v_cache: array<${buffer_scalar_type(DTYPE)}>; struct Params { _pad0: u32, @@ -66,9 +68,14 @@ fn main( 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( - t_k_cache[ki], t_k_cache[ki + 1u], - t_k_cache[ki + 2u], t_k_cache[ki + 3u]); + $if DTYPE == "half": + 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])); + $else: + let kvv = vec4( + t_k_cache[ki], t_k_cache[ki + 1u], + t_k_cache[ki + 2u], t_k_cache[ki + 3u]); acc4 = acc4 + qv * kvv; } s = (acc4.x + acc4.y + acc4.z + acc4.w) * params.scale; @@ -107,7 +114,10 @@ fn main( 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] * t_v_cache[vbase + d]; + $if DTYPE == "half": + acc = acc + sh_s[j] * f32(t_v_cache[vbase + d]); + $else: + acc = acc + sh_s[j] * t_v_cache[vbase + d]; } o_acc[nd] = acc; } diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.yaml b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.yaml new file mode 100644 index 00000000000..283e6516996 --- /dev/null +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split.yaml @@ -0,0 +1,11 @@ +sdpa_fd_split: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: float + SUFFIX: "" + - VALUE: half + SUFFIX: half + shader_variants: + - NAME: sdpa_fd_split 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 deleted file mode 100644 index 236cb96764b..00000000000 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16.wgsl +++ /dev/null @@ -1,134 +0,0 @@ -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_half_wgsl.h similarity index 91% rename from backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_f16_wgsl.h rename to backends/webgpu/runtime/ops/sdpa_fd_decode/sdpa_fd_split_half_wgsl.h index 71dfcfb8af0..bc69a444edb 100644 --- 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_half_wgsl.h @@ -12,12 +12,10 @@ namespace executorch::backends::webgpu { -// @generated from sdpa_fd_split_f16.wgsl - DO NOT EDIT. -// wgsl-sha256: 3396d4f0d990d59ded70734cb92feb6c776635695b5d5c21b8cdfb7e55c8ed6d -inline constexpr const char* kSdpaFdSplitF16WGSL = R"( +// @generated from sdpa_fd_split.wgsl - DO NOT EDIT. +// wgsl-sha256: 147fc5775b76f3626dc934b2df5e34148bef12c345789332f098d13143bb646e +inline constexpr const char* kSdpaFdSplitHalfWGSL = 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; @@ -151,8 +149,8 @@ fn main( } )"; -inline constexpr uint32_t kSdpaFdSplitF16WorkgroupSizeX = 64; -inline constexpr uint32_t kSdpaFdSplitF16WorkgroupSizeY = 1; -inline constexpr uint32_t kSdpaFdSplitF16WorkgroupSizeZ = 1; +inline constexpr uint32_t kSdpaFdSplitHalfWorkgroupSizeX = 64; +inline constexpr uint32_t kSdpaFdSplitHalfWorkgroupSizeY = 1; +inline constexpr uint32_t kSdpaFdSplitHalfWorkgroupSizeZ = 1; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/update_cache/update_cache.wgsl b/backends/webgpu/runtime/ops/update_cache/update_cache.wgsl index 62f882ad547..cdda4aac9cd 100644 --- a/backends/webgpu/runtime/ops/update_cache/update_cache.wgsl +++ b/backends/webgpu/runtime/ops/update_cache/update_cache.wgsl @@ -1,4 +1,6 @@ -@group(0) @binding(0) var t_cache: array; +$if DTYPE == "half": + enable f16; +@group(0) @binding(0) var t_cache: array<${buffer_scalar_type(DTYPE)}>; @group(0) @binding(1) var t_value: array; struct Params { @@ -20,5 +22,8 @@ fn main(@builtin(global_invocation_id) gid: vec3) { if (params.dst_offset + i >= params.cache_numel) { return; } - t_cache[params.dst_offset + i] = t_value[i]; + $if DTYPE == "half": + t_cache[params.dst_offset + i] = f16(t_value[i]); + $else: + t_cache[params.dst_offset + i] = t_value[i]; } diff --git a/backends/webgpu/runtime/ops/update_cache/update_cache.yaml b/backends/webgpu/runtime/ops/update_cache/update_cache.yaml new file mode 100644 index 00000000000..4df425c763e --- /dev/null +++ b/backends/webgpu/runtime/ops/update_cache/update_cache.yaml @@ -0,0 +1,11 @@ +update_cache: + parameter_names_with_default_values: + DTYPE: float + generate_variant_forall: + DTYPE: + - VALUE: float + SUFFIX: "" + - VALUE: half + SUFFIX: half + shader_variants: + - NAME: update_cache diff --git a/backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl b/backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl deleted file mode 100644 index 7d2c0ceb5c1..00000000000 --- a/backends/webgpu/runtime/ops/update_cache/update_cache_f16.wgsl +++ /dev/null @@ -1,27 +0,0 @@ -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_half_wgsl.h similarity index 67% rename from backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h rename to backends/webgpu/runtime/ops/update_cache/update_cache_half_wgsl.h index 3a838f5dc2c..d31728178a8 100644 --- a/backends/webgpu/runtime/ops/update_cache/update_cache_f16_wgsl.h +++ b/backends/webgpu/runtime/ops/update_cache/update_cache_half_wgsl.h @@ -12,12 +12,10 @@ namespace executorch::backends::webgpu { -// @generated from update_cache_f16.wgsl - DO NOT EDIT. -// wgsl-sha256: e115bcc12d0b0d09e9f72e1f63fee21c0ce56f787dfbd728b9f4adebb1bf5630 -inline constexpr const char* kUpdateCacheF16WGSL = R"( +// @generated from update_cache.wgsl - DO NOT EDIT. +// wgsl-sha256: 390daabe0d4545311dd5c6768d427fc9d133125bb1143869184c7d7631a88954 +inline constexpr const char* kUpdateCacheHalfWGSL = 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; @@ -44,8 +42,8 @@ fn main(@builtin(global_invocation_id) gid: vec3) { } )"; -inline constexpr uint32_t kUpdateCacheF16WorkgroupSizeX = 256; -inline constexpr uint32_t kUpdateCacheF16WorkgroupSizeY = 1; -inline constexpr uint32_t kUpdateCacheF16WorkgroupSizeZ = 1; +inline constexpr uint32_t kUpdateCacheHalfWorkgroupSizeX = 256; +inline constexpr uint32_t kUpdateCacheHalfWorkgroupSizeY = 1; +inline constexpr uint32_t kUpdateCacheHalfWorkgroupSizeZ = 1; } // namespace executorch::backends::webgpu