From 582b4f934f73c69f2def99630f214a277b0294dd Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:11:27 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../runtime/ops/rope/RotaryEmbedding.cpp | 360 ++++++++++++++++++ .../runtime/ops/rope/rotary_embedding_hf.wgsl | 49 +++ .../ops/rope/rotary_embedding_hf_wgsl.h | 73 ++++ backends/webgpu/test/ops/test_rope_hf.py | 176 +++++++++ backends/webgpu/test/test_build_webgpu.sh | 9 + backends/webgpu/test/test_webgpu_native.cpp | 115 ++++++ 6 files changed, 782 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 create mode 100644 backends/webgpu/test/ops/test_rope_hf.py diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 76eafb0e738..3a206fc9611 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 @@ -378,10 +379,369 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { // pipeline_q/pipeline_k owned by their dispatches; graph dtor frees. } +// Mirrors Vulkan's full-dimension HuggingFace rotate-half RoPE. +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"); + +WGPUBuffer add_rope_hf_dispatch( + WebGPUGraph& graph, + WGPUDevice device, + WGPUComputePipeline pipeline, + WGPUBindGroupLayout bgl, + 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)); + + WGPUBindGroupEntry bg_entries[5] = {}; + bg_entries[0].binding = 0; + bg_entries[0].buffer = out.buffer; + bg_entries[0].size = out.nbytes; + bg_entries[1].binding = 1; + bg_entries[1].buffer = x.buffer; + bg_entries[1].size = x.nbytes; + bg_entries[2].binding = 2; + bg_entries[2].buffer = freqs_cos.buffer; + bg_entries[2].size = freqs_cos.nbytes; + bg_entries[3].binding = 3; + bg_entries[3].buffer = freqs_sin.buffer; + bg_entries[3].size = freqs_sin.nbytes; + bg_entries[4].binding = 4; + bg_entries[4].buffer = uniform_buffer; + bg_entries[4].size = sizeof(RotaryHfParams); + + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 5; + bg_desc.entries = bg_entries; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); + + graph.add_dispatch( + {pipeline, bind_group, workgroup_count, "apply_rotary_emb_hf"}); + graph.own_uniform_buffer(uniform_buffer); + return uniform_buffer; +} + +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]); + + 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"); + } + 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; + + // Decode uses a SymInt position; static graphs use an Int. + 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); + } 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"); + } + if (static_cast(start_pos) + seq > max_seq) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos + seq exceeds freqs max_seq"); + } + + 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); + 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"); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kRotaryEmbeddingHfWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUBindGroupLayoutEntry entries[5] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_Storage; + for (uint32_t i = 1; i <= 3; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + } + entries[4].binding = 4; + entries[4].visibility = WGPUShaderStage_Compute; + entries[4].buffer.type = WGPUBufferBindingType_Uniform; + + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 5; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device, &pl_desc); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(wg_size); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_size_constant; + WGPUComputePipeline pipeline_q = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + WGPUComputePipeline pipeline_k = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + + WGPUBuffer q_ubuf = add_rope_hf_dispatch( + graph, + device, + pipeline_q, + bgl, + xq, + xq_out, + freqs_cos, + freqs_sin, + n_heads_q, + seq, + head_dim, + half_dim, + rotary_dim, + static_cast(start_pos), + xq_wgc); + const size_t q_idx = graph.num_dispatches() - 1; + WGPUBuffer k_ubuf = add_rope_hf_dispatch( + graph, + device, + pipeline_k, + bgl, + xk, + xk_out, + freqs_cos, + freqs_sin, + n_heads_k, + seq, + head_dim, + half_dim, + rotary_dim, + static_cast(start_pos), + xk_wgc); + const size_t k_idx = graph.num_dispatches() - 1; + + 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_hf_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, + max_seq, + wg_size, + q_idx, + k_idx, + q_ubuf, + k_ubuf](WebGPUGraph& g) { + const auto& qd = g.cur_dims(xq_id); + const auto& kd = g.cur_dims(xk_id); + 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); + 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); + } + if (static_cast(start_pos) + s > max_seq) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos + seq exceeds freqs max_seq"); + } + 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, + "rope_hf_q(resize)"); + g.dispatch_at(k_idx).workgroup_count_x = utils::compute_1d_workgroup_count( + g.device(), + static_cast(kn / 2u), + wg_size, + "rope_hf_k(resize)"); + g.set_cur_dims(xq_out_id, qd); + g.set_cur_dims(xk_out_id, kd); + }; + graph.add_tensor_resize_hook(xq_id, rope_hf_hook); + graph.add_tensor_resize_hook(xk_id, rope_hf_hook); + if (dynamic_pos) { + graph.add_resize_hook(start_pos_id, rope_hf_hook); + } + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); +} + } // 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 diff --git a/backends/webgpu/test/ops/test_rope_hf.py b/backends/webgpu/test/ops/test_rope_hf.py new file mode 100644 index 00000000000..f10ebbebb65 --- /dev/null +++ b/backends/webgpu/test/ops/test_rope_hf.py @@ -0,0 +1,176 @@ +# 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. + +"""Dynamic HuggingFace RoPE export and native-test goldens.""" + +import os +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + +import torch +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.examples.models.llama.rope import ( + hf_apply_rotary_emb, + hf_precompute_freqs_cis, +) +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes + +BATCH = 1 +SEQ = 1 +N_HEADS_Q = 16 +N_HEADS_K = 8 +HEAD_DIM = 128 +MAX_SEQ = 16 +POSITIONS = (0, 7, 15) + + +class DynamicHfRope(torch.nn.Module): + def forward( + self, + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + input_pos: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + start_pos = input_pos[0].item() + torch._check_is_size(start_pos) + torch._check(start_pos < freqs_cos.shape[0]) + return hf_apply_rotary_emb( + xq, + xk, + freqs_cos.narrow(0, start_pos, xq.shape[1]), + freqs_sin.narrow(0, start_pos, xq.shape[1]), + ) + + +def _ramp(numel: int, mod: int, off: int) -> torch.Tensor: + idx = torch.arange(numel, dtype=torch.int64) + return ((idx % mod) - off).to(torch.float32) / 16.0 + + +def _inputs() -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + xq = _ramp(BATCH * SEQ * N_HEADS_Q * HEAD_DIM, 17, 8).reshape( + BATCH, SEQ, N_HEADS_Q, HEAD_DIM + ) + xk = _ramp(BATCH * SEQ * N_HEADS_K * HEAD_DIM, 13, 6).reshape( + BATCH, SEQ, N_HEADS_K, HEAD_DIM + ) + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + HEAD_DIM, + MAX_SEQ, + theta=10000.0, + ) + input_pos = torch.tensor([0], dtype=torch.long) + return xq, xk, freqs_cos, freqs_sin, input_pos + + +def _golden( + xq: torch.Tensor, + xk: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + position: int, +) -> tuple[torch.Tensor, torch.Tensor]: + return hf_apply_rotary_emb( + xq, + xk, + freqs_cos[position : position + SEQ], + freqs_sin[position : position + SEQ], + ) + + +def _export_program(): + inputs = _inputs() + with torch._dynamo.config.patch(capture_scalar_outputs=True): + ep = torch.export.export(DynamicHfRope().eval(), inputs) + + symints = [ + node + for node in ep.graph_module.graph.nodes + if isinstance(node.meta.get("val"), torch.SymInt) + ] + if not symints: + raise AssertionError("input_pos did not lower to a SymInt") + + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + portable = get_non_lowered_nodes(graph) + if len(delegates) != 1: + raise AssertionError(f"expected one delegate, got {len(delegates)}") + if portable: + raise AssertionError(f"unexpected non-lowered nodes: {portable}") + + et = edge.to_executorch() + delegate_ids = [ + delegate.id + for plan in et.executorch_program.execution_plan + for delegate in plan.delegates + ] + if delegate_ids != ["VulkanBackend"]: + raise AssertionError(f"unexpected delegates: {delegate_ids}") + return et + + +def export_rope_hf_dynamic(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + xq, xk, freqs_cos, freqs_sin, _ = _inputs() + et = _export_program() + with open(os.path.join(out_dir, "rope_hf_dynamic.pte"), "wb") as output: + output.write(et.buffer) + + for name, tensor in ( + ("xq", xq), + ("xk", xk), + ("freqs_cos", freqs_cos), + ("freqs_sin", freqs_sin), + ): + tensor.detach().numpy().astype(" None: + self.assertIsNotNone(_export_program()) + + def test_position_goldens_match_custom_op(self) -> None: + xq, xk, freqs_cos, freqs_sin, _ = _inputs() + self.assertNotEqual(xq.shape[2], xk.shape[2]) + position_outputs = [] + for position in POSITIONS: + with self.subTest(position=position): + expected_q, expected_k = _golden(xq, xk, freqs_cos, freqs_sin, position) + position_outputs.append(expected_q) + actual_q, actual_k = torch.ops.et_vk.apply_rotary_emb_hf.default( + xq, xk, freqs_cos, freqs_sin, position + ) + torch.testing.assert_close(actual_q, expected_q) + torch.testing.assert_close(actual_k, expected_k) + self.assertFalse(torch.allclose(position_outputs[0], position_outputs[1])) + self.assertFalse(torch.allclose(position_outputs[1], position_outputs[2])) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/webgpu/test/test_build_webgpu.sh b/backends/webgpu/test/test_build_webgpu.sh index 1c79d17cf06..5be78ec4916 100755 --- a/backends/webgpu/test/test_build_webgpu.sh +++ b/backends/webgpu/test/test_build_webgpu.sh @@ -27,12 +27,14 @@ $PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/test_wgsl_codegen.py" -v echo "=== Step 1: Run Python export tests ===" $PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/ops/test_add.py" -v $PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/ops/test_rms_norm.py" -v +$PYTHON_EXECUTABLE -m pytest "${SCRIPT_DIR}/ops/test_rope_hf.py" -v # ── Step 2: Export .pte model ───────────────────────────────────────────────── echo "=== Step 2: Export test models ===" DISPATCH_ORDER_DIR="/tmp/dispatch_order" PTE_UPDATE_CACHE_MODEL="/tmp/webgpu_update_cache_test.pte" +ROPE_HF_DIR="/tmp/webgpu_rope_hf" cd "${EXECUTORCH_ROOT}" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_dispatch_order import export_dispatch_order_cases @@ -46,6 +48,12 @@ from executorch.backends.webgpu.test.ops.test_update_cache import export_update_ export_update_cache_model('${PTE_UPDATE_CACHE_MODEL}') " || { echo "WARN: update_cache export failed; skipping update_cache native test"; UPDATE_CACHE_OK=0; } +echo "=== Export dynamic HuggingFace RoPE model and goldens ===" +$PYTHON_EXECUTABLE -c " +from executorch.backends.webgpu.test.ops.test_rope_hf import export_rope_hf_dynamic +export_rope_hf_dynamic('${ROPE_HF_DIR}') +" + echo "=== Export SDPA sweep models (sdpa_.pte + .golden.bin to /tmp) ===" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import export_all_sdpa_models @@ -108,6 +116,7 @@ else fi env \ ${UPDATE_CACHE_ENV_VAR} \ + WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}" \ WEBGPU_TEST_SDPA_DIR=/tmp/ \ "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_native_test" diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index bae1e4e8863..394abada288 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -764,6 +764,109 @@ void test_rope( << "apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)"; } +void test_rope_hf_dynamic(const std::string& dir) { + constexpr int S = 1; + constexpr int NH = 16; + constexpr int NKV = 8; + constexpr int HD = 128; + constexpr int MAXS = 16; + constexpr int positions[] = {0, 7, 15}; + constexpr int xq_numel = S * NH * HD; + constexpr int xk_numel = S * NKV * HD; + constexpr int freqs_numel = MAXS * HD; + + Module module(dir + "rope_hf_dynamic.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load HF RoPE dynamic model"; + + std::vector xq = load_golden(dir + "rope_hf_dynamic.xq.bin", xq_numel); + std::vector xk = load_golden(dir + "rope_hf_dynamic.xk.bin", xk_numel); + std::vector freqs_cos = + load_golden(dir + "rope_hf_dynamic.freqs_cos.bin", freqs_numel); + std::vector freqs_sin = + load_golden(dir + "rope_hf_dynamic.freqs_sin.bin", freqs_numel); + ASSERT_FALSE( + xq.empty() || xk.empty() || freqs_cos.empty() || freqs_sin.empty()) + << "could not load HF RoPE input binaries from " << dir; + + auto has_shape = [](const executorch::aten::Tensor& tensor, + const std::vector& expected) { + if (tensor.dim() != static_cast(expected.size())) { + return false; + } + for (size_t i = 0; i < expected.size(); i++) { + if (tensor.size(static_cast(i)) != expected[i]) { + return false; + } + } + return true; + }; + + for (const int position : positions) { + auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::vector(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::vector(freqs_sin)); + auto post = make_tensor_ptr( + {1}, std::vector{static_cast(position)}); + auto result = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + ASSERT_TRUE(result.ok()) + << "HF RoPE forward failed at position " << position << " (error " + << static_cast(result.error()) << ")"; + const auto& outputs = result.get(); + ASSERT_TRUE( + outputs.size() == 2 && outputs[0].isTensor() && outputs[1].isTensor()) + << "expected exactly two HF RoPE tensor outputs"; + const auto& xq_out = outputs[0].toTensor(); + const auto& xk_out = outputs[1].toTensor(); + ASSERT_TRUE(has_shape(xq_out, {1, S, NH, HD})) + << "HF RoPE query output has the wrong shape at position " << position; + ASSERT_TRUE(has_shape(xk_out, {1, S, NKV, HD})) + << "HF RoPE key output has the wrong shape at position " << position; + + const std::string prefix = + dir + "rope_hf_dynamic.pos" + std::to_string(position); + const std::vector golden_q = + load_golden(prefix + ".xq.golden.bin", xq_numel); + const std::vector golden_k = + load_golden(prefix + ".xk.golden.bin", xk_numel); + ASSERT_FALSE(golden_q.empty() || golden_k.empty()) + << "could not load HF RoPE goldens for position " << position; + + float q_abs = 0.0f, q_rel = 0.0f, k_abs = 0.0f, k_rel = 0.0f; + const bool q_ok = quant_within_tol( + xq_out.const_data_ptr(), + golden_q.data(), + xq_numel, + 1e-4f, + 1e-3f, + &q_abs, + &q_rel); + const bool k_ok = quant_within_tol( + xk_out.const_data_ptr(), + golden_k.data(), + xk_numel, + 1e-4f, + 1e-3f, + &k_abs, + &k_rel); + EXPECT_TRUE(q_ok && k_ok) + << "HF RoPE mismatch at position " << position << ": q abs=" << q_abs + << " rel=" << q_rel << ", k abs=" << k_abs << " rel=" << k_rel; + } + + auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({MAXS, HD}, std::move(freqs_cos)); + auto fst = make_tensor_ptr({MAXS, HD}, std::move(freqs_sin)); + auto post = make_tensor_ptr({1}, std::vector{MAXS}); + auto out_of_range = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + EXPECT_FALSE(out_of_range.ok()) + << "HF RoPE accepted start_pos + seq beyond the frequency table"; +} + void test_prepack( const std::string& model_path, const std::string& golden_path, @@ -1799,6 +1902,18 @@ TEST(WebGPUNative, Rope) { } } +TEST(WebGPUNative, RopeHfDynamic) { + const char* env = std::getenv("WEBGPU_TEST_ROPE_HF_DIR"); + if (env == nullptr || *env == '\0') { + GTEST_SKIP() << "WEBGPU_TEST_ROPE_HF_DIR not set"; + } + std::string dir = env; + if (dir.back() != '/') { + dir += '/'; + } + test_rope_hf_dynamic(dir); +} + TEST(WebGPUNative, Prepack) { if (g_prepack_model_path.empty() || g_prepack_golden_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_PREPACK_MODEL/GOLDEN not set";