From 8874028928690c26bd034afecd29ebf67da5caed Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 10 Jul 2026 11:26:57 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../runtime/ops/rope/RotaryEmbedding.cpp | 375 ++++++++++++++++++ .../runtime/ops/rope/rotary_embedding_hf.wgsl | 49 +++ .../ops/rope/rotary_embedding_hf_wgsl.h | 73 ++++ 3 files changed, 497 insertions(+) create mode 100644 backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl create mode 100644 backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 03de78233b0..12fd2e50565 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include @@ -328,10 +329,384 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { graph.add_tensor_resize_hook(xk_id, rope_hook); } +// HuggingFace rotate-half RoPE (Qwen3 etc.). Structural sibling of the +// interleaved path above (same one-thread-per-pair scalar dispatch, wg_size, +// and resize hook); differs only in element pairing (i with i+half_dim vs +// even/odd), a full [max_seq, rotary_dim] freqs table, and a start_pos offset. +// Mirrors Vulkan's et_vk.apply_rotary_emb_hf +// (backends/vulkan/runtime/graph/ops/impl/RotaryEmbedding.cpp:211). + +// Uniform layout matching the HF WGSL Params struct (32 bytes). +struct RotaryHfParams { + uint32_t n_heads; + uint32_t seq; + uint32_t head_dim; + uint32_t half_dim; + uint32_t num_pairs; + uint32_t rotary_dim; + uint32_t start_pos; + uint32_t _pad0; +}; +static_assert(sizeof(RotaryHfParams) == 32, "RotaryHfParams must be 32 bytes"); + +RopeDispatch add_rope_hf_dispatch( + WebGPUGraph& graph, + WGPUDevice device, + uint32_t wg_size, + const WebGPUTensor& x, + const WebGPUTensor& out, + const WebGPUTensor& freqs_cos, + const WebGPUTensor& freqs_sin, + uint32_t n_heads, + uint32_t seq, + uint32_t head_dim, + uint32_t half_dim, + uint32_t rotary_dim, + uint32_t start_pos, + uint32_t workgroup_count) { + const uint32_t num_pairs = + static_cast(utils::numel_of(out.dims) / 2u); + + RotaryHfParams params = {}; + params.n_heads = n_heads; + params.seq = seq; + params.head_dim = head_dim; + params.half_dim = half_dim; + params.num_pairs = num_pairs; + params.rotary_dim = rotary_dim; + params.start_pos = start_pos; + + WGPUBufferDescriptor uniform_desc = {}; + uniform_desc.size = sizeof(RotaryHfParams); + uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + uniform_desc.mappedAtCreation = true; + WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); + void* mapped = + wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(RotaryHfParams)); + std::memcpy(mapped, ¶ms, sizeof(RotaryHfParams)); + wgpuBufferUnmap(uniform_buffer); + graph.add_uniform_buffer_bytes(sizeof(RotaryHfParams)); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + kRotaryEmbeddingHfWGSL, + { + {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, + {1, WGPUBufferBindingType_ReadOnlyStorage, x.buffer, x.nbytes}, + {2, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_cos.buffer, + freqs_cos.nbytes}, + {3, + WGPUBufferBindingType_ReadOnlyStorage, + freqs_sin.buffer, + freqs_sin.nbytes}, + {4, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(RotaryHfParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_index = graph.add_dispatch( + {bundle.pipeline, + bundle.bind_group, + workgroup_count, + "apply_rotary_emb_hf"}); + + graph.own_uniform_buffer(uniform_buffer); + return {uniform_buffer, dispatch_index}; +} + +// Resize hook body: recompute S/num_pairs + (dynamic) start_pos for both +// dispatches; out follows xq/xk. Fires on xq/xk seq resize and, when start_pos +// is a runtime SymInt (KV-cache decode), on each start_pos change; idempotent. +void resize_rope_hf( + WebGPUGraph& g, + int xq_id, + int xk_id, + int xq_out_id, + int xk_out_id, + int start_pos_id, + bool dynamic_pos, + uint32_t baked_start_pos, + uint32_t n_heads_q, + uint32_t n_heads_k, + uint32_t head_dim, + uint32_t half_dim, + uint32_t rotary_dim, + uint32_t wg_size, + size_t q_idx, + size_t k_idx, + WGPUBuffer q_ubuf, + WGPUBuffer k_ubuf) { + const auto& qd = g.cur_dims(xq_id); + const auto& kd = g.cur_dims(xk_id); + if (qd.size() < 3 || kd.size() < 3) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); + } + const uint32_t s = static_cast(qd[qd.size() - 3]); + const uint64_t qn = utils::numel_of(qd); + const uint64_t kn = utils::numel_of(kd); + if (static_cast(kd[kd.size() - 3]) != s) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q and k seq lengths differ"); + } + uint32_t start_pos = baked_start_pos; + if (dynamic_pos) { + const int32_t pos = g.read_symint(start_pos_id); + if (pos < 0) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos must be non-negative"); + } + start_pos = static_cast(pos); + } + RotaryHfParams pq = {}; + pq.n_heads = n_heads_q; + pq.seq = s; + pq.head_dim = head_dim; + pq.half_dim = half_dim; + pq.num_pairs = static_cast(qn / 2u); + pq.rotary_dim = rotary_dim; + pq.start_pos = start_pos; + RotaryHfParams pk = pq; + pk.n_heads = n_heads_k; + pk.num_pairs = static_cast(kn / 2u); + wgpuQueueWriteBuffer(g.queue(), q_ubuf, 0, &pq, sizeof(pq)); + wgpuQueueWriteBuffer(g.queue(), k_ubuf, 0, &pk, sizeof(pk)); + g.dispatch_at(q_idx).workgroup_count_x = utils::compute_1d_workgroup_count( + g.device(), + static_cast(qn / 2u), + wg_size, + "apply_rotary_emb_hf(resize)"); + g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( + g.device(), + static_cast(kn / 2u), + wg_size, + "apply_rotary_emb_hf(resize)"); + g.set_cur_dims(xq_out_id, qd); + g.set_cur_dims(xk_out_id, kd); +} + +// args: [xq, xk, freqs_cos, freqs_sin, start_pos, out_list(ValueList[xq_out, +// xk_out])]. freqs is the FULL [max_seq, rotary_dim] table (start_pos offsets +// into it), unlike the pre-sliced interleaved freqs. +void apply_rotary_emb_hf_impl( + WebGPUGraph& graph, + const std::vector& args) { + const int xq_id = args.at(0); + const int xk_id = args.at(1); + const int freqs_cos_id = args.at(2); + const int freqs_sin_id = args.at(3); + const int start_pos_id = args.at(4); + + const std::vector& out_list = graph.get_value_list(args.at(5)); + if (out_list.size() != 2) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: expected an output ValueList of size 2"); + } + + WGPUDevice device = graph.device(); + + const auto& xq = graph.get_tensor(xq_id); + const auto& xk = graph.get_tensor(xk_id); + const auto& freqs_cos = graph.get_tensor(freqs_cos_id); + const auto& freqs_sin = graph.get_tensor(freqs_sin_id); + const auto& xq_out = graph.get_tensor(out_list[0]); + const auto& xk_out = graph.get_tensor(out_list[1]); + + // Shape contract: xq/xk (B,S,n_heads,head_dim), freqs (max_seq, rotary_dim). + if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { + throw std::runtime_error("WebGPU apply_rotary_emb_hf: malformed dims"); + } + const uint32_t head_dim = static_cast(xq.dims.back()); + const uint32_t seq = static_cast(xq.dims[xq.dims.size() - 3]); + const uint32_t n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); + const uint32_t n_heads_k = static_cast(xk.dims[xk.dims.size() - 2]); + const uint32_t seq_k = static_cast(xk.dims[xk.dims.size() - 3]); + const uint32_t max_seq = + static_cast(freqs_cos.dims[freqs_cos.dims.size() - 2]); + const uint32_t rotary_dim = static_cast(freqs_cos.dims.back()); + + if (head_dim == 0 || head_dim % 2 != 0) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: head_dim must be a nonzero multiple of 2"); + } + if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: xq/xk head_dim and seq must match"); + } + // Full rotary only (rotary_dim == head_dim); partial-rotary passthrough is a + // documented follow-up (Qwen3 uses full RoPE). Throw rather than mis-rotate. + if (rotary_dim != head_dim) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: partial rotary (rotary_dim != head_dim) " + "not supported"); + } + if (freqs_cos.dims != freqs_sin.dims) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: freqs_cos and freqs_sin shapes differ"); + } + if (max_seq < seq) { + throw std::runtime_error("WebGPU apply_rotary_emb_hf: freqs max_seq < seq"); + } + + if (xq.buffer == nullptr || xk.buffer == nullptr || + freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || + xq_out.buffer == nullptr || xk_out.buffer == nullptr) { + throw std::runtime_error("WebGPU apply_rotary_emb_hf: null buffer binding"); + } + + const uint32_t half_dim = rotary_dim / 2u; + + // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); + // mirrors sdpa's input_pos handling. + int64_t start_pos = 0; + const auto start_pos_type = graph.get_value_type(start_pos_id); + const bool dynamic_pos = start_pos_type == WebGPUGraph::ValueType::SymInt; + if (dynamic_pos) { + start_pos = graph.read_symint(start_pos_id); // build placeholder (e.g. 0) + } else if (start_pos_type == WebGPUGraph::ValueType::Int) { + start_pos = graph.get_int(start_pos_id); + } else { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos must be Int or SymInt"); + } + if (start_pos < 0) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); + } + + // All tensors are fp32; output shapes equal their inputs. + const uint64_t xq_numel = utils::numel_of(xq.dims); + const uint64_t xk_numel = utils::numel_of(xk.dims); + const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); + if (freqs_numel != static_cast(max_seq) * rotary_dim || + xq.nbytes != xq_numel * sizeof(float) || + xk.nbytes != xk_numel * sizeof(float) || + freqs_cos.nbytes != freqs_numel * sizeof(float) || + freqs_sin.nbytes != freqs_numel * sizeof(float) || + xq_out.nbytes != xq_numel * sizeof(float) || + xk_out.nbytes != xk_numel * sizeof(float)) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: dtype/byte-size mismatch (all fp32) or " + "freqs shape != [max_seq, rotary_dim]"); + } + + if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: pair count exceeds uint32 dispatch range"); + } + + const uint32_t wg_size = + utils::clamp_workgroup_size(device, kRotaryEmbeddingHfWorkgroupSizeX); + // Validate both dispatches before any GPU-object alloc (no leak on throw). + const uint32_t xq_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(xq_numel / 2u), + wg_size, + "apply_rotary_emb_hf"); + const uint32_t xk_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(xk_numel / 2u), + wg_size, + "apply_rotary_emb_hf"); + + RopeDispatch q_disp = add_rope_hf_dispatch( + graph, + device, + wg_size, + xq, + xq_out, + freqs_cos, + freqs_sin, + n_heads_q, + seq, + head_dim, + half_dim, + rotary_dim, + static_cast(start_pos), + xq_wgc); + RopeDispatch k_disp = add_rope_hf_dispatch( + graph, + device, + wg_size, + xk, + xk_out, + freqs_cos, + freqs_sin, + n_heads_k, + seq, + head_dim, + half_dim, + rotary_dim, + static_cast(start_pos), + xk_wgc); + WGPUBuffer q_ubuf = q_disp.uniform; + WGPUBuffer k_ubuf = k_disp.uniform; + const size_t q_idx = q_disp.dispatch_index; + const size_t k_idx = k_disp.dispatch_index; + + const int xq_out_id = out_list[0]; + const int xk_out_id = out_list[1]; + const uint32_t baked_start_pos = static_cast(start_pos); + auto rope_hook = [xq_id, + xk_id, + xq_out_id, + xk_out_id, + start_pos_id, + dynamic_pos, + baked_start_pos, + n_heads_q, + n_heads_k, + head_dim, + half_dim, + rotary_dim, + wg_size, + q_idx, + k_idx, + q_ubuf, + k_ubuf](WebGPUGraph& g) { + resize_rope_hf( + g, + xq_id, + xk_id, + xq_out_id, + xk_out_id, + start_pos_id, + dynamic_pos, + baked_start_pos, + n_heads_q, + n_heads_k, + head_dim, + half_dim, + rotary_dim, + wg_size, + q_idx, + k_idx, + q_ubuf, + k_ubuf); + }; + graph.add_tensor_resize_hook(xq_id, rope_hook); + graph.add_tensor_resize_hook(xk_id, rope_hook); + // Dynamic decode: re-fire when the runtime start_pos SymInt changes. + if (dynamic_pos) { + graph.add_resize_hook(start_pos_id, rope_hook); + } +} + } // namespace WEBGPU_REGISTER_OPERATORS { WEBGPU_REGISTER_OP(et_vk.apply_rotary_emb.default, apply_rotary_emb_impl); + WEBGPU_REGISTER_OP( + et_vk.apply_rotary_emb_hf.default, apply_rotary_emb_hf_impl); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl new file mode 100644 index 00000000000..14a6853afa3 --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl @@ -0,0 +1,49 @@ +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_freqs_cos: array; +@group(0) @binding(3) var t_freqs_sin: array; + +struct Params { + n_heads: u32, + seq: u32, + head_dim: u32, + half_dim: u32, + num_pairs: u32, + rotary_dim: u32, + start_pos: u32, + _pad0: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +// One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated +// halves) indexed at row (start_pos + s); only the first-half column is read. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let pair = gid.x; + if (pair >= params.num_pairs) { + return; + } + let half_dim = params.half_dim; + let pair_i = pair % half_dim; + let t1 = pair / half_dim; + let head = t1 % params.n_heads; + let t2 = t1 / params.n_heads; + let s = t2 % params.seq; + let b = t2 / params.seq; + + let head_base = + ((b * params.seq + s) * params.n_heads + head) * params.head_dim; + let a_idx = head_base + pair_i; + let b_idx = head_base + pair_i + half_dim; + let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + + let c = t_freqs_cos[freqs_idx]; + let si = t_freqs_sin[freqs_idx]; + let x_a = t_in[a_idx]; + let x_b = t_in[b_idx]; + t_out[a_idx] = x_a * c - x_b * si; + t_out[b_idx] = x_b * c + x_a * si; +} diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h new file mode 100644 index 00000000000..191ec710e66 --- /dev/null +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h @@ -0,0 +1,73 @@ +/* + * 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 rotary_embedding_hf.wgsl - DO NOT EDIT. +// wgsl-sha256: 5ba8d45925f00f12af17bf3092a1af9513a9e501c5c35e6b0d48cfb3dac7b5d6 +inline constexpr const char* kRotaryEmbeddingHfWGSL = R"( +@group(0) @binding(0) var t_out: array; +@group(0) @binding(1) var t_in: array; +@group(0) @binding(2) var t_freqs_cos: array; +@group(0) @binding(3) var t_freqs_sin: array; + +struct Params { + n_heads: u32, + seq: u32, + head_dim: u32, + half_dim: u32, + num_pairs: u32, + rotary_dim: u32, + start_pos: u32, + _pad0: u32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64u; + +// One thread per (i, i+half_dim) pair; HuggingFace rotate-half RoPE, shared +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table (duplicated +// halves) indexed at row (start_pos + s); only the first-half column is read. +@compute @workgroup_size(wg_size, 1, 1) +fn main(@builtin(global_invocation_id) gid: vec3) { + let pair = gid.x; + if (pair >= params.num_pairs) { + return; + } + let half_dim = params.half_dim; + let pair_i = pair % half_dim; + let t1 = pair / half_dim; + let head = t1 % params.n_heads; + let t2 = t1 / params.n_heads; + let s = t2 % params.seq; + let b = t2 / params.seq; + + let head_base = + ((b * params.seq + s) * params.n_heads + head) * params.head_dim; + let a_idx = head_base + pair_i; + let b_idx = head_base + pair_i + half_dim; + let freqs_idx = (s + params.start_pos) * params.rotary_dim + pair_i; + + let c = t_freqs_cos[freqs_idx]; + let si = t_freqs_sin[freqs_idx]; + let x_a = t_in[a_idx]; + let x_b = t_in[b_idx]; + t_out[a_idx] = x_a * c - x_b * si; + t_out[b_idx] = x_b * c + x_a * si; +} +)"; + +inline constexpr uint32_t kRotaryEmbeddingHfWorkgroupSizeX = 64; +inline constexpr uint32_t kRotaryEmbeddingHfWorkgroupSizeY = 1; +inline constexpr uint32_t kRotaryEmbeddingHfWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu