diff --git a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp index 012875b21e7..29f00c5823f 100644 --- a/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp +++ b/backends/webgpu/runtime/ops/rope/RotaryEmbedding.cpp @@ -7,23 +7,22 @@ */ #include +#include #include #include -#include -#include #include #include -#include -#include #include -#include namespace executorch::backends::webgpu { namespace { +constexpr const char* kRotaryShader = "rotary_embedding"; +constexpr const char* kRotaryHfShader = "rotary_embedding_hf"; + // Uniform layout matching the WGSL Params struct (16-byte aligned, 32 bytes). struct RotaryParams { uint32_t n_heads; @@ -37,184 +36,109 @@ struct RotaryParams { }; static_assert(sizeof(RotaryParams) == 32, "RotaryParams must be 32 bytes"); -// A rope dispatch: its param-uniform (rewritten on resize) and its index in the -// graph's dispatch list (so a resize hook can update the workgroup count). -struct RopeDispatch { - WGPUBuffer uniform; - size_t dispatch_index; +enum class RopeGridPolicy { + OneDimensional, + FoldedTwoDimensional, +}; + +struct RopeGridContext { + int tensor_id; + uint32_t workgroup_size; + RopeGridPolicy policy; + const char* op_name; }; -// Rotate one (x->out) with the shared shader; freqs shared between xq and xk. -RopeDispatch add_rope_dispatch( +WebGPUDispatchGrid pick_rope_grid( + const WebGPUGraph& graph, + const RopeGridContext& context) { + const uint32_t num_pairs = static_cast( + utils::numel_of(graph.cur_dims(context.tensor_id)) / 2u); + if (context.policy == RopeGridPolicy::OneDimensional) { + return { + utils::compute_1d_workgroup_count( + graph.device(), num_pairs, context.workgroup_size, context.op_name), + 1u}; + } + const utils::WgCount grid = utils::compute_2d_workgroup_count( + graph.device(), num_pairs, context.workgroup_size, context.op_name); + return {grid.x, grid.y}; +} + +void preflight_rope_grids( + const WebGPUGraph& graph, + const RopeGridContext& q_context, + const RopeGridContext& k_context) { + (void)pick_rope_grid(graph, q_context); + (void)pick_rope_grid(graph, k_context); +} + +template +WGPUBuffer add_rope_dispatch( WebGPUGraph& graph, - WGPUDevice device, - std::optional& shared_resources, - uint32_t wg_size, + const char* shader_name, + const char* kernel_name, 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 workgroup_count) { - const uint32_t half_dim = head_dim / 2u; - // out.dims == in.dims (asserted in impl), so this matches the caller's wgc. - const uint32_t num_pairs = - static_cast(utils::numel_of(out.dims) / 2u); - - RotaryParams params = {}; - params.n_heads = n_heads; - params.seq = seq; - params.head_dim = head_dim; - params.half_dim = half_dim; - params.num_pairs = num_pairs; - - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(RotaryParams); - 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(RotaryParams)); - std::memcpy(mapped, ¶ms, sizeof(RotaryParams)); - wgpuBufferUnmap(uniform_buffer); - graph.add_uniform_buffer_bytes(sizeof(RotaryParams)); - - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = static_cast(wg_size); - - const std::vector bindings = { - {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(RotaryParams)}, - }; - utils::ComputePipelineBundle bundle = shared_resources.has_value() - ? utils::make_compute_pipeline( - device, *shared_resources, bindings, &wg_size_constant, 1) - : utils::make_compute_pipeline( - device, kRotaryEmbeddingWGSL, bindings, &wg_size_constant, 1); - - const size_t dispatch_index = graph.add_dispatch( - {bundle.pipeline, - bundle.bind_group, - workgroup_count, - "apply_rotary_emb"}); - if (!shared_resources.has_value()) { - shared_resources.emplace(std::move(bundle)); - } - - // Graph owns it so a resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); - return {uniform_buffer, dispatch_index}; + const Params& params, + int trigger_tensor_id, + const RopeGridContext& grid_context, + uint32_t wg_size) { + WGPUBuffer uniform_buffer = graph.create_params_buffer(params); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = shader_name; + descriptor.kernel_name = kernel_name; + descriptor.bindings = { + {out.buffer, 0u, out.nbytes}, + {x.buffer, 0u, x.nbytes}, + {freqs_cos.buffer, 0u, freqs_cos.nbytes}, + {freqs_sin.buffer, 0u, freqs_sin.nbytes}, + {uniform_buffer, 0u, sizeof(Params)}}; + descriptor.constants = {{"wg_size", static_cast(wg_size)}}; + graph.add_dynamic_compute_dispatch( + descriptor, trigger_tensor_id, pick_rope_grid, grid_context); + return uniform_buffer; } -// Resize hook body: recompute S/num_pairs + both dispatches; out follows xq/xk. -void resize_rope( - WebGPUGraph& g, - int xq_id, - int xk_id, - int xq_out_id, - int xk_out_id, - uint32_t n_heads_q, - uint32_t n_heads_k, - uint32_t head_dim, - uint32_t half_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(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); - // pk = pq (seq=s); require k's seq == s, not silently q's. - if (static_cast(kd[kd.size() - 3]) != s) { - throw std::runtime_error( - "apply_rotary_emb(resize): q and k seq lengths differ"); - } - // freqs stay max-allocated; shader indexes by position (S = prefix). - RotaryParams 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); - RotaryParams 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(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(resize)"); - g.set_cur_dims(xq_out_id, qd); - g.set_cur_dims(xk_out_id, kd); -} - -// args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. -void apply_rotary_emb_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 std::vector& out_list = graph.get_value_list(args.at(4)); - if (out_list.size() != 2) { - throw std::runtime_error( - "WebGPU apply_rotary_emb: 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]); +struct RotaryGeometry { + uint32_t head_dim; + uint32_t seq; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t half_dim; + uint64_t xq_numel; + uint64_t xk_numel; +}; - // Vulkan shape contract: xq/xk (B,S,n_heads,head_dim), freqs (S,head_dim/2). +RotaryGeometry validate_rope_inputs( + const WebGPUTensor& xq, + const WebGPUTensor& xk, + const WebGPUTensor& freqs_cos, + const WebGPUTensor& freqs_sin, + const WebGPUTensor& xq_out, + const WebGPUTensor& xk_out) { if (xq.dims.size() < 3 || xk.dims.size() < 3 || freqs_cos.dims.size() < 2) { throw std::runtime_error("WebGPU apply_rotary_emb: 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]); + RotaryGeometry geometry = {}; + geometry.head_dim = static_cast(xq.dims.back()); + geometry.seq = static_cast(xq.dims[xq.dims.size() - 3]); + geometry.n_heads_q = static_cast(xq.dims[xq.dims.size() - 2]); + geometry.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 half_dim = static_cast(freqs_cos.dims.back()); + geometry.half_dim = static_cast(freqs_cos.dims.back()); - if (head_dim == 0 || head_dim % 2 != 0) { + if (geometry.head_dim == 0 || geometry.head_dim % 2 != 0) { throw std::runtime_error( "WebGPU apply_rotary_emb: head_dim must be a nonzero multiple of 2"); } - if (static_cast(xk.dims.back()) != head_dim || seq_k != seq) { + if (static_cast(xk.dims.back()) != geometry.head_dim || + seq_k != geometry.seq) { throw std::runtime_error( "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"); } - if (half_dim * 2u != head_dim) { + if (geometry.half_dim * 2u != geometry.head_dim) { throw std::runtime_error( "WebGPU apply_rotary_emb: head_dim != 2 * freqs_cos last dim"); } @@ -222,127 +146,163 @@ void apply_rotary_emb_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error( "WebGPU apply_rotary_emb: freqs_cos and freqs_sin shapes differ"); } - 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: null buffer binding"); } - // 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); + geometry.xq_numel = utils::numel_of(xq.dims); + geometry.xk_numel = utils::numel_of(xk.dims); const uint64_t freqs_numel = utils::numel_of(freqs_cos.dims); - if (freqs_numel != static_cast(seq) * half_dim || - xq.nbytes != xq_numel * sizeof(float) || - xk.nbytes != xk_numel * sizeof(float) || + if (freqs_numel != static_cast(geometry.seq) * geometry.half_dim || + xq.nbytes != geometry.xq_numel * sizeof(float) || + xk.nbytes != geometry.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)) { + xq_out.nbytes != geometry.xq_numel * sizeof(float) || + xk_out.nbytes != geometry.xk_numel * sizeof(float)) { throw std::runtime_error( "WebGPU apply_rotary_emb: dtype/byte-size mismatch (all fp32) or " "freqs shape != [seq, head_dim/2]"); } + if (geometry.xq_numel > UINT32_MAX || geometry.xk_numel > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb: element index exceeds uint32 range"); + } + return geometry; +} - if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { +struct RotaryResizeContext { + int xq_id; + int xk_id; + int xq_out_id; + int xk_out_id; + uint32_t n_heads_q; + uint32_t n_heads_k; + uint32_t head_dim; + uint32_t half_dim; + WGPUBuffer q_uniform; + WGPUBuffer k_uniform; +}; + +// Resize hook body: update parameters and outputs; the graph owns grid refresh. +void resize_rope(WebGPUGraph& graph, const RotaryResizeContext& context) { + const auto& q_dims = graph.cur_dims(context.xq_id); + const auto& k_dims = graph.cur_dims(context.xk_id); + if (q_dims.size() < 3 || k_dims.size() < 3) { + throw std::runtime_error("apply_rotary_emb(resize): q/k rank must be >= 3"); + } + const uint32_t seq = static_cast(q_dims[q_dims.size() - 3]); + const uint64_t q_numel = utils::numel_of(q_dims); + const uint64_t k_numel = utils::numel_of(k_dims); + // pk = pq (seq=s); require k's seq == s, not silently q's. + if (static_cast(k_dims[k_dims.size() - 3]) != seq) { throw std::runtime_error( - "WebGPU apply_rotary_emb: pair count exceeds uint32 dispatch range"); + "apply_rotary_emb(resize): q and k seq lengths differ"); } + // freqs stay max-allocated; shader indexes by position (S = prefix). + RotaryParams q_params = {}; + q_params.n_heads = context.n_heads_q; + q_params.seq = seq; + q_params.head_dim = context.head_dim; + q_params.half_dim = context.half_dim; + q_params.num_pairs = static_cast(q_numel / 2u); + RotaryParams k_params = q_params; + k_params.n_heads = context.n_heads_k; + k_params.num_pairs = static_cast(k_numel / 2u); + wgpuQueueWriteBuffer( + graph.queue(), context.q_uniform, 0, &q_params, sizeof(q_params)); + wgpuQueueWriteBuffer( + graph.queue(), context.k_uniform, 0, &k_params, sizeof(k_params)); + graph.set_cur_dims(context.xq_out_id, q_dims); + graph.set_cur_dims(context.xk_out_id, k_dims); +} - const uint32_t wg_size = - utils::clamp_workgroup_size(device, kRotaryEmbeddingWorkgroupSizeX); - // 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"); - const uint32_t xk_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xk_numel / 2u), - wg_size, - "apply_rotary_emb"); +// args: [xq, xk, freqs_cos, freqs_sin, out_list(ValueList[xq_out, xk_out])]. +void apply_rotary_emb_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 std::vector& out_list = graph.get_value_list(args.at(4)); + if (out_list.size() != 2) { + throw std::runtime_error( + "WebGPU apply_rotary_emb: expected an output ValueList of size 2"); + } - std::optional shared_resources; - RopeDispatch q_disp = add_rope_dispatch( + 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]); + + const RotaryGeometry geometry = + validate_rope_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); + + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), get_webgpu_shader_info(kRotaryShader).workgroup_size_x); + const RopeGridContext q_grid = { + xq_id, wg_size, RopeGridPolicy::OneDimensional, "apply_rotary_emb"}; + const RopeGridContext k_grid = { + xk_id, wg_size, RopeGridPolicy::OneDimensional, "apply_rotary_emb"}; + preflight_rope_grids(graph, q_grid, k_grid); + + RotaryParams q_params = {}; + q_params.n_heads = geometry.n_heads_q; + q_params.seq = geometry.seq; + q_params.head_dim = geometry.head_dim; + q_params.half_dim = geometry.half_dim; + q_params.num_pairs = static_cast(geometry.xq_numel / 2u); + RotaryParams k_params = q_params; + k_params.n_heads = geometry.n_heads_k; + k_params.num_pairs = static_cast(geometry.xk_numel / 2u); + const WGPUBuffer q_uniform = add_rope_dispatch( graph, - device, - shared_resources, - wg_size, + kRotaryShader, + "apply_rotary_emb", xq, xq_out, freqs_cos, freqs_sin, - n_heads_q, - seq, - head_dim, - xq_wgc); - RopeDispatch k_disp = add_rope_dispatch( + q_params, + xq_id, + q_grid, + wg_size); + const WGPUBuffer k_uniform = add_rope_dispatch( graph, - device, - shared_resources, - wg_size, + kRotaryShader, + "apply_rotary_emb", xk, xk_out, freqs_cos, freqs_sin, - n_heads_k, - seq, - head_dim, - 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; - - // Dynamic shapes: recompute S/num_pairs + both dispatches; out follows xq/xk. - const int xq_out_id = out_list[0]; - const int xk_out_id = out_list[1]; + k_params, + xk_id, + k_grid, + wg_size); + // Register on both xq and xk so the recompute fires whichever is marked dirty // (q and k co-resize on S; resize_rope is idempotent, so a double-fire when // both are dirty is harmless). - auto rope_hook = [xq_id, - xk_id, - xq_out_id, - xk_out_id, - n_heads_q, - n_heads_k, - head_dim, - half_dim, - wg_size, - q_idx, - k_idx, - q_ubuf, - k_ubuf](WebGPUGraph& g) { - resize_rope( - g, - xq_id, - xk_id, - xq_out_id, - xk_out_id, - n_heads_q, - n_heads_k, - head_dim, - half_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); + const RotaryResizeContext resize_context = { + xq_id, + xk_id, + out_list[0], + out_list[1], + geometry.n_heads_q, + geometry.n_heads_k, + geometry.head_dim, + geometry.half_dim, + q_uniform, + k_uniform}; + graph.add_tensor_resize_hook(xq_id, resize_rope, resize_context); + graph.add_tensor_resize_hook(xk_id, resize_rope, resize_context); } -// 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). +// Mirrors Vulkan's full-dimension HuggingFace rotate-half RoPE. struct RotaryHfParams { uint32_t n_heads; uint32_t seq; @@ -355,187 +315,65 @@ struct RotaryHfParams { }; static_assert(sizeof(RotaryHfParams) == 32, "RotaryHfParams must be 32 bytes"); -RopeDispatch add_rope_hf_dispatch( - WebGPUGraph& graph, - 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; - - WGPUBuffer uniform_buffer = - utils::make_uniform(graph.device(), ¶ms, sizeof(RotaryHfParams)); - 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( - graph.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); -} - -// Validated HF-rope shape, derived from the input tensors. -struct RopeHfShape { +struct RotaryHfGeometry { uint32_t head_dim; uint32_t seq; uint32_t n_heads_q; uint32_t n_heads_k; + uint32_t max_seq; uint32_t rotary_dim; uint32_t half_dim; uint64_t xq_numel; uint64_t xk_numel; }; -// Derive and validate the HF-rope input shapes; throws on any malformed input. -RopeHfShape validate_rope_hf_inputs( - const WebGPUTensor& xq, +RotaryHfGeometry validate_rope_hf_inputs( + const WebGPUTensor& x, const WebGPUTensor& xk, const WebGPUTensor& freqs_cos, const WebGPUTensor& freqs_sin, - const WebGPUTensor& xq_out, + const WebGPUTensor& x_out, const WebGPUTensor& xk_out) { - // 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) { + if (x.dims.size() < 3 || xk.dims.size() != x.dims.size() || + 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) { + if (x_out.dims != x.dims || xk_out.dims != xk.dims) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: output shapes must match q/k inputs"); + } + for (size_t i = 0; i + 3 < x.dims.size(); i++) { + if (x.dims[i] != xk.dims[i]) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: q/k batch dimensions differ"); + } + } + const auto positive_u32 = [](int64_t value, const char* label) { + if (value <= 0 || static_cast(value) > UINT32_MAX) { + throw std::runtime_error( + std::string("WebGPU apply_rotary_emb_hf: invalid ") + label); + } + return static_cast(value); + }; + RotaryHfGeometry geometry = {}; + geometry.head_dim = positive_u32(x.dims.back(), "head_dim"); + geometry.seq = positive_u32(x.dims[x.dims.size() - 3], "sequence length"); + geometry.n_heads_q = + positive_u32(x.dims[x.dims.size() - 2], "query head count"); + geometry.n_heads_k = + positive_u32(xk.dims[xk.dims.size() - 2], "key head count"); + geometry.max_seq = positive_u32(freqs_cos.dims[0], "frequency row count"); + geometry.rotary_dim = positive_u32(freqs_cos.dims[1], "rotary_dim"); + if (geometry.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) { + if (xk.dims.back() != static_cast(geometry.head_dim) || + xk.dims[xk.dims.size() - 3] != static_cast(geometry.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) { + if (geometry.rotary_dim != geometry.head_dim) { throw std::runtime_error( "WebGPU apply_rotary_emb_hf: partial rotary (rotary_dim != head_dim) " "not supported"); @@ -544,51 +382,134 @@ RopeHfShape validate_rope_hf_inputs( throw std::runtime_error( "WebGPU apply_rotary_emb_hf: freqs_cos and freqs_sin shapes differ"); } - if (max_seq < seq) { + if (geometry.max_seq < geometry.seq) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: freqs max_seq < seq"); } - - if (xq.buffer == nullptr || xk.buffer == nullptr || + if (x.buffer == nullptr || xk.buffer == nullptr || freqs_cos.buffer == nullptr || freqs_sin.buffer == nullptr || - xq_out.buffer == nullptr || xk_out.buffer == nullptr) { + x_out.buffer == nullptr || xk_out.buffer == nullptr) { throw std::runtime_error("WebGPU apply_rotary_emb_hf: null buffer binding"); } + const WebGPUTensor* tensors[] = { + &x, &xk, &freqs_cos, &freqs_sin, &x_out, &xk_out}; + for (const WebGPUTensor* tensor : tensors) { + if (tensor->is_int || tensor->elem_size != sizeof(float)) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: all tensors must be fp32"); + } + } - // 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); + geometry.half_dim = geometry.rotary_dim / 2u; + geometry.xq_numel = utils::numel_of(x.dims); + geometry.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) || + if (freqs_numel != + static_cast(geometry.max_seq) * geometry.rotary_dim || + x.nbytes != geometry.xq_numel * sizeof(float) || + xk.nbytes != geometry.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)) { + x_out.nbytes != geometry.xq_numel * sizeof(float) || + xk_out.nbytes != geometry.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 (geometry.xq_numel == 0 || geometry.xk_numel == 0 || + geometry.xq_numel > UINT32_MAX || geometry.xk_numel > UINT32_MAX) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: element index exceeds uint32 range"); + } + return geometry; +} + +struct RotaryHfResizeContext { + 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 max_seq; + WGPUBuffer q_uniform; + WGPUBuffer k_uniform; +}; - if (xq_numel / 2u > UINT32_MAX || xk_numel / 2u > UINT32_MAX) { +void resize_rope_hf(WebGPUGraph& graph, const RotaryHfResizeContext& context) { + const auto& q_dims = graph.cur_dims(context.xq_id); + const auto& k_dims = graph.cur_dims(context.xk_id); + if (q_dims.size() < 3 || k_dims.size() != q_dims.size()) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k rank must be >= 3"); + } + const int64_t seq_value = q_dims[q_dims.size() - 3]; + if (seq_value <= 0 || static_cast(seq_value) > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): invalid sequence length"); + } + const uint32_t seq = static_cast(seq_value); + if (k_dims[k_dims.size() - 3] != seq_value) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q and k seq lengths differ"); + } + if (q_dims.back() != static_cast(context.head_dim) || + k_dims.back() != static_cast(context.head_dim) || + q_dims[q_dims.size() - 2] != static_cast(context.n_heads_q) || + k_dims[k_dims.size() - 2] != static_cast(context.n_heads_k)) { throw std::runtime_error( - "WebGPU apply_rotary_emb_hf: pair count exceeds uint32 dispatch range"); - } - - return { - head_dim, - seq, - n_heads_q, - n_heads_k, - rotary_dim, - rotary_dim / 2u, - xq_numel, - xk_numel}; + "apply_rotary_emb_hf(resize): q/k head geometry changed"); + } + for (size_t i = 0; i + 3 < q_dims.size(); i++) { + if (q_dims[i] != k_dims[i]) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): q/k batch dimensions differ"); + } + } + const uint64_t q_numel = utils::numel_of(q_dims); + const uint64_t k_numel = utils::numel_of(k_dims); + if (q_numel == 0 || k_numel == 0 || q_numel > UINT32_MAX || + k_numel > UINT32_MAX) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): element index exceeds uint32 range"); + } + uint32_t start_pos = context.baked_start_pos; + if (context.dynamic_pos) { + const int64_t pos = graph.read_symint(context.start_pos_id); + if (pos < 0 || static_cast(pos) > UINT32_MAX) { + 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) + seq > context.max_seq) { + throw std::runtime_error( + "apply_rotary_emb_hf(resize): start_pos + seq exceeds freqs max_seq"); + } + RotaryHfParams q_params = {}; + q_params.n_heads = context.n_heads_q; + q_params.seq = seq; + q_params.head_dim = context.head_dim; + q_params.half_dim = context.half_dim; + q_params.num_pairs = static_cast(q_numel / 2u); + q_params.rotary_dim = context.rotary_dim; + q_params.start_pos = start_pos; + RotaryHfParams k_params = q_params; + k_params.n_heads = context.n_heads_k; + k_params.num_pairs = static_cast(k_numel / 2u); + wgpuQueueWriteBuffer( + graph.queue(), context.q_uniform, 0, &q_params, sizeof(q_params)); + wgpuQueueWriteBuffer( + graph.queue(), context.k_uniform, 0, &k_params, sizeof(k_params)); + graph.set_cur_dims(context.xq_out_id, q_dims); + graph.set_cur_dims(context.xk_out_id, k_dims); } -// 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) { @@ -604,8 +525,6 @@ void apply_rotary_emb_hf_impl( "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); @@ -613,24 +532,15 @@ void apply_rotary_emb_hf_impl( const auto& xq_out = graph.get_tensor(out_list[0]); const auto& xk_out = graph.get_tensor(out_list[1]); - const RopeHfShape shp = + const RotaryHfGeometry geometry = validate_rope_hf_inputs(xq, xk, freqs_cos, freqs_sin, xq_out, xk_out); - const uint32_t head_dim = shp.head_dim; - const uint32_t seq = shp.seq; - const uint32_t n_heads_q = shp.n_heads_q; - const uint32_t n_heads_k = shp.n_heads_k; - const uint32_t rotary_dim = shp.rotary_dim; - const uint32_t half_dim = shp.half_dim; - const uint64_t xq_numel = shp.xq_numel; - const uint64_t xk_numel = shp.xk_numel; - - // start_pos: build-time Int (baked) OR runtime SymInt (dynamic decode); - // mirrors sdpa's input_pos handling. + + // 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); // build placeholder (e.g. 0) + 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 { @@ -641,99 +551,82 @@ void apply_rotary_emb_hf_impl( throw std::runtime_error( "WebGPU apply_rotary_emb_hf: start_pos must be non-negative"); } + if (static_cast(start_pos) + geometry.seq > geometry.max_seq) { + throw std::runtime_error( + "WebGPU apply_rotary_emb_hf: start_pos + seq exceeds freqs max_seq"); + } - 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), + const uint32_t wg_size = utils::clamp_workgroup_size( + graph.device(), get_webgpu_shader_info(kRotaryHfShader).workgroup_size_x); + const RopeGridContext q_grid = { + xq_id, wg_size, - "apply_rotary_emb_hf"); - const uint32_t xk_wgc = utils::compute_1d_workgroup_count( - device, - static_cast(xk_numel / 2u), + RopeGridPolicy::FoldedTwoDimensional, + "apply_rotary_emb_hf"}; + const RopeGridContext k_grid = { + xk_id, wg_size, - "apply_rotary_emb_hf"); - - RopeDispatch q_disp = add_rope_hf_dispatch( + RopeGridPolicy::FoldedTwoDimensional, + "apply_rotary_emb_hf"}; + preflight_rope_grids(graph, q_grid, k_grid); + + RotaryHfParams q_params = {}; + q_params.n_heads = geometry.n_heads_q; + q_params.seq = geometry.seq; + q_params.head_dim = geometry.head_dim; + q_params.half_dim = geometry.half_dim; + q_params.num_pairs = static_cast(geometry.xq_numel / 2u); + q_params.rotary_dim = geometry.rotary_dim; + q_params.start_pos = static_cast(start_pos); + RotaryHfParams k_params = q_params; + k_params.n_heads = geometry.n_heads_k; + k_params.num_pairs = static_cast(geometry.xk_numel / 2u); + + const WGPUBuffer q_uniform = add_rope_dispatch( graph, - wg_size, + kRotaryHfShader, + "apply_rotary_emb_hf", 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( + q_params, + xq_id, + q_grid, + wg_size); + const WGPUBuffer k_uniform = add_rope_dispatch( graph, - wg_size, + kRotaryHfShader, + "apply_rotary_emb_hf", xk, xk_out, freqs_cos, freqs_sin, - n_heads_k, - seq, - head_dim, - half_dim, - rotary_dim, + k_params, + xk_id, + k_grid, + wg_size); + + const RotaryHfResizeContext resize_context = { + xq_id, + xk_id, + out_list[0], + out_list[1], + start_pos_id, + dynamic_pos, 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. + geometry.n_heads_q, + geometry.n_heads_k, + geometry.head_dim, + geometry.half_dim, + geometry.rotary_dim, + geometry.max_seq, + q_uniform, + k_uniform}; + graph.add_tensor_resize_hook(xq_id, resize_rope_hf, resize_context); + graph.add_tensor_resize_hook(xk_id, resize_rope_hf, resize_context); if (dynamic_pos) { - graph.add_resize_hook(start_pos_id, rope_hook); + graph.add_resize_hook(start_pos_id, resize_rope_hf, resize_context); } } diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl index 14a6853afa3..859b24ca6fc 100644 --- a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf.wgsl @@ -18,11 +18,13 @@ struct 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. +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table indexed at row +// (start_pos + s); each output half uses its corresponding frequency column. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let pair = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let pair = gid.x + gid.y * (num_workgroups.x * wg_size); if (pair >= params.num_pairs) { return; } @@ -38,12 +40,16 @@ fn main(@builtin(global_invocation_id) gid: vec3) { ((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 freqs_base = (s + params.start_pos) * params.rotary_dim; + let freqs_a_idx = freqs_base + pair_i; + let freqs_b_idx = freqs_a_idx + half_dim; - let c = t_freqs_cos[freqs_idx]; - let si = t_freqs_sin[freqs_idx]; + let c_a = t_freqs_cos[freqs_a_idx]; + let si_a = t_freqs_sin[freqs_a_idx]; + let c_b = t_freqs_cos[freqs_b_idx]; + let si_b = t_freqs_sin[freqs_b_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; + t_out[a_idx] = x_a * c_a - x_b * si_a; + t_out[b_idx] = x_b * c_b + x_a * si_b; } diff --git a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h index 191ec710e66..21242fb480a 100644 --- a/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h +++ b/backends/webgpu/runtime/ops/rope/rotary_embedding_hf_wgsl.h @@ -13,7 +13,7 @@ namespace executorch::backends::webgpu { // @generated from rotary_embedding_hf.wgsl - DO NOT EDIT. -// wgsl-sha256: 5ba8d45925f00f12af17bf3092a1af9513a9e501c5c35e6b0d48cfb3dac7b5d6 +// wgsl-sha256: 4f081ed4c8165f021cbb722d379e437f30b8dfb08bf03bfcbaa406ed7799c7b6 inline constexpr const char* kRotaryEmbeddingHfWGSL = R"( @group(0) @binding(0) var t_out: array; @group(0) @binding(1) var t_in: array; @@ -35,11 +35,13 @@ struct 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. +// xq/xk shader. freqs is the FULL [max_seq, rotary_dim] table indexed at row +// (start_pos + s); each output half uses its corresponding frequency column. @compute @workgroup_size(wg_size, 1, 1) -fn main(@builtin(global_invocation_id) gid: vec3) { - let pair = gid.x; +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) num_workgroups: vec3) { + let pair = gid.x + gid.y * (num_workgroups.x * wg_size); if (pair >= params.num_pairs) { return; } @@ -55,14 +57,18 @@ fn main(@builtin(global_invocation_id) gid: vec3) { ((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 freqs_base = (s + params.start_pos) * params.rotary_dim; + let freqs_a_idx = freqs_base + pair_i; + let freqs_b_idx = freqs_a_idx + half_dim; - let c = t_freqs_cos[freqs_idx]; - let si = t_freqs_sin[freqs_idx]; + let c_a = t_freqs_cos[freqs_a_idx]; + let si_a = t_freqs_sin[freqs_a_idx]; + let c_b = t_freqs_cos[freqs_b_idx]; + let si_b = t_freqs_sin[freqs_b_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; + t_out[a_idx] = x_a * c_a - x_b * si_a; + t_out[b_idx] = x_b * c_b + x_a * si_b; } )"; diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 75af3ff84e6..0d5d7a9e799 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -65,6 +65,7 @@ DISPATCH_ORDER_DIR="/tmp/dispatch_order" UPDATE_CACHE_DIR="/tmp/update_cache" INDEX_DIR="/tmp/index" DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape" +ROPE_HF_DIR="/tmp/webgpu_rope_hf" SYMINT_BLOB="/tmp/sdpa_dyn_small.pte" OUTPUT_SUPPRESSION_DIR="/tmp/output_suppression" EMBEDDING_MODEL="/tmp/webgpu_embedding_q4gsw.pte" @@ -104,6 +105,11 @@ export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}') export_rope_model('${ROPE_DECODE_MODEL}', '${ROPE_DECODE_XQ_GOLDEN}', '${ROPE_DECODE_XK_GOLDEN}', 'decode') " +$PYTHON_EXECUTABLE -c " +from executorch.backends.webgpu.test.ops.test_rope_hf import export_rope_hf_dynamic +export_rope_hf_dynamic('${ROPE_HF_DIR}') +" + $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_prepack import export_prepack_model, export_prepack_two_const_model, export_prepack_tied_const_model export_prepack_model('${PREPACK_MODEL}', '${PREPACK_GOLDEN}') @@ -150,6 +156,7 @@ export_dynamic_decode('/tmp') export_incache_decode('/tmp') " +require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte" require_file "${SYMINT_BLOB}" require_file "${OUTPUT_SUPPRESSION_DIR}/input.bin" @@ -201,6 +208,7 @@ run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \ WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \ WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \ + WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}" \ WEBGPU_TEST_SYMINT_BLOB="${SYMINT_BLOB}" \ WEBGPU_TEST_PREPACK_MODEL="${PREPACK_MODEL}" \ WEBGPU_TEST_PREPACK_GOLDEN="${PREPACK_GOLDEN}" \ diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 0963168ccd5..feb6854c59d 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include @@ -20,6 +21,8 @@ #include #include #include +#include +#include namespace executorch::backends::webgpu { namespace { @@ -653,6 +656,140 @@ TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { expect_dispatch_grid(graph, 2u, 13u, 17u); } +struct InvalidRopeGraphCase { + const char* name; + std::vector xq_dims; + std::vector xk_dims; + std::vector cos_dims; + std::vector sin_dims; + std::vector xq_out_dims; + std::vector xk_out_dims; + vkgraph::VkDataType xq_dtype; + const char* expected_error; +}; + +void expect_invalid_rope_graph(const InvalidRopeGraphCase& test_case) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + }; + add_tensor(test_case.xq_dtype, test_case.xq_dims, 0); + add_tensor(vk::VkDataType::FLOAT32, test_case.xk_dims, 1); + add_tensor(vk::VkDataType::FLOAT32, test_case.cos_dims, 2); + add_tensor(vk::VkDataType::FLOAT32, test_case.sin_dims, 3); + add_tensor(vk::VkDataType::FLOAT32, test_case.xq_out_dims, 4); + add_tensor(vk::VkDataType::FLOAT32, test_case.xk_out_dims, 5); + std::vector output_value_ids = {4, 5}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_value_ids).Union())); + + std::vector args = {0, 1, 2, 3, 6}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb.default", &args)); + std::vector input_ids = {0, 1, 2, 3}; + std::vector output_ids = {4, 5}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + std::string error; + try { + graph.build(fbb.GetBufferPointer(), nullptr, nullptr); + } catch (const std::exception& exception) { + error = exception.what(); + } + EXPECT_FALSE(error.empty()) << test_case.name << " unexpectedly built"; + EXPECT_EQ(error, test_case.expected_error) + << test_case.name << " rejected for the wrong reason"; + const WebGPUMemoryStats stats = graph.memory_stats(); + EXPECT_EQ(stats.num_dispatches, 0) << test_case.name; + EXPECT_EQ(stats.uniform_buffer_bytes, 0u) << test_case.name; + EXPECT_EQ(stats.num_cached_shaders, 0) << test_case.name; + EXPECT_EQ(stats.num_cached_pipelines, 0) << test_case.name; +} + +TEST(WebGPURopeValidation, RejectsMalformedGraphsBeforeDispatchAllocation) { + ASSERT_TRUE( + webgpu_operator_registry().has_op("et_vk.apply_rotary_emb.default")); + const std::vector xq = {1, 2, 2, 4}; + const std::vector xk = {1, 2, 1, 4}; + const std::vector freqs = {2, 2}; + const InvalidRopeGraphCase cases[] = { + {"query rank", + {2, 4}, + xk, + freqs, + freqs, + {2, 4}, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: malformed dims"}, + {"sequence mismatch", + xq, + {1, 3, 1, 4}, + freqs, + freqs, + xq, + {1, 3, 1, 4}, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"}, + {"head dimension mismatch", + xq, + {1, 2, 1, 6}, + freqs, + freqs, + xq, + {1, 2, 1, 6}, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: xq/xk head_dim and seq must match"}, + {"frequency width mismatch", + xq, + xk, + {2, 3}, + {2, 3}, + xq, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: head_dim != 2 * freqs_cos last dim"}, + {"cosine/sine shape mismatch", + xq, + xk, + freqs, + {2, 1}, + xq, + xk, + vkgraph::VkDataType::FLOAT32, + "WebGPU apply_rotary_emb: freqs_cos and freqs_sin shapes differ"}, + {"query byte size mismatch", + xq, + xk, + freqs, + freqs, + xq, + xk, + vkgraph::VkDataType::INT64, + "WebGPU apply_rotary_emb: dtype/byte-size mismatch (all fp32) or " + "freqs shape != [seq, head_dim/2]"}, + }; + for (const InvalidRopeGraphCase& test_case : cases) { + SCOPED_TRACE(test_case.name); + expect_invalid_rope_graph(test_case); + } +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/ops/test_rope_hf.py b/backends/webgpu/test/ops/test_rope_hf.py index 6e26fe53b17..f1929ac8d3b 100644 --- a/backends/webgpu/test/ops/test_rope_hf.py +++ b/backends/webgpu/test/ops/test_rope_hf.py @@ -19,6 +19,7 @@ compare (it has no ATen). """ +import os import unittest from collections import namedtuple @@ -26,8 +27,12 @@ import torch from executorch.backends.vulkan import VulkanPartitioner -from executorch.examples.models.llama.rope import hf_apply_rotary_emb +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 # B batch, S tokens, NH query heads, NKV kv heads (NH != NKV so the two outputs # are distinguishable by numel), HD head dim (even; full rotary, rotary_dim==HD). @@ -39,6 +44,20 @@ Shape("decode", 1, 1, 16, 8, 128), ] +DYNAMIC_BATCH = 1 +DYNAMIC_SEQ = 1 +DYNAMIC_N_HEADS_Q = 16 +DYNAMIC_N_HEADS_K = 8 +DYNAMIC_HEAD_DIM = 128 +DYNAMIC_MAX_SEQ = 16 +DYNAMIC_POSITIONS = (0, 7, 15) +DYNAMIC_SEQUENCE_CASES = ( + (DYNAMIC_MAX_SEQ, 0), + (5, 7), + (1, DYNAMIC_MAX_SEQ - 1), + (DYNAMIC_MAX_SEQ, 0), +) + class HfRope(torch.nn.Module): # unsqueeze_dim=1: freqs [S, HD] -> [S, 1, HD] broadcasts over (B, NH) of the @@ -47,6 +66,26 @@ def forward(self, xq, xk, freqs_cos, freqs_sin): return hf_apply_rotary_emb(xq, xk, freqs_cos, freqs_sin, unsqueeze_dim=1) +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 + xq.shape[1] <= 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: # ((i % mod) - off) / 16: exact in fp32, matches test_webgpu_native.cpp. idx = torch.arange(numel, dtype=torch.int64) @@ -68,6 +107,32 @@ def _inputs( return xq, xk, freqs_cos, freqs_sin +def _dynamic_inputs(seq: int = DYNAMIC_SEQ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, +]: + xq = _ramp( + DYNAMIC_BATCH * seq * DYNAMIC_N_HEADS_Q * DYNAMIC_HEAD_DIM, + 17, + 8, + ).reshape(DYNAMIC_BATCH, seq, DYNAMIC_N_HEADS_Q, DYNAMIC_HEAD_DIM) + xk = _ramp( + DYNAMIC_BATCH * seq * DYNAMIC_N_HEADS_K * DYNAMIC_HEAD_DIM, + 13, + 6, + ).reshape(DYNAMIC_BATCH, seq, DYNAMIC_N_HEADS_K, DYNAMIC_HEAD_DIM) + freqs_cos, freqs_sin = hf_precompute_freqs_cis( + DYNAMIC_HEAD_DIM, + DYNAMIC_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, @@ -78,27 +143,123 @@ def _golden( return torch.ops.et_vk.apply_rotary_emb_hf.default(xq, xk, freqs_cos, freqs_sin, 0) -def _export(inputs): +def _dynamic_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 + xq.shape[1]], + freqs_sin[position : position + xq.shape[1]], + ) + + +def _assert_fully_delegated(edge) -> None: + 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}") + + +def _lower(inputs): ep = torch.export.export(HfRope().eval(), inputs) - return to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _export(inputs): + return _lower(inputs).to_executorch() + + +def _lower_dynamic_program(): + inputs = _dynamic_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()]) + _assert_fully_delegated(edge) + return edge + + +def _lower_dynamic_sequence_program(): + inputs = _dynamic_inputs(DYNAMIC_MAX_SEQ) + s_dim = torch.export.Dim("rope_hf_s", min=1, max=DYNAMIC_MAX_SEQ) + dynamic_shapes = ({1: s_dim}, {1: s_dim}, None, None, None) + with torch._dynamo.config.patch(capture_scalar_outputs=True): + ep = torch.export.export( + DynamicHfRope().eval(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + + scalar_symints = [ + node + for node in ep.graph_module.graph.nodes + if isinstance(node.meta.get("val"), torch.SymInt) + ] + if not scalar_symints: + raise AssertionError("input_pos did not lower to a SymInt") + xq_placeholder = next( + node + for node in ep.graph_module.graph.nodes + if node.op == "placeholder" and node.target == "xq" + ) + if not isinstance(xq_placeholder.meta["val"].shape[1], torch.SymInt): + raise AssertionError("query sequence dimension did not remain symbolic") + + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + _assert_fully_delegated(edge) + return edge + + +def _export_dynamic_program(): + edge = _lower_dynamic_program() + + 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_dynamic_sequence_program(): + edge = _lower_dynamic_sequence_program() + 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 class TestRopeHf(unittest.TestCase): def test_export_delegates(self) -> None: for shape in SHAPES: with self.subTest(shape=shape.name): - et = _export(_inputs(shape)) - found = any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ) - self.assertTrue( - found, - "Expected a VulkanBackend delegate (apply_rotary_emb_hf " "fusion)", - ) + self.assertIsNotNone(_lower(_inputs(shape))) def test_golden_matches_eager(self) -> None: # The et_vk golden must equal the real HF rotate-half apply_rotary_emb, @@ -112,6 +273,44 @@ def test_golden_matches_eager(self) -> None: torch.testing.assert_close(gq, eq, atol=1e-5, rtol=1e-5) torch.testing.assert_close(gk, ek, atol=1e-5, rtol=1e-5) + def test_dynamic_export_is_fully_delegated(self) -> None: + self.assertIsNotNone(_lower_dynamic_program()) + + def test_dynamic_position_goldens_match_custom_op(self) -> None: + xq, xk, freqs_cos, freqs_sin, _ = _dynamic_inputs() + self.assertNotEqual(xq.shape[2], xk.shape[2]) + position_outputs = [] + for position in DYNAMIC_POSITIONS: + with self.subTest(position=position): + expected_q, expected_k = _dynamic_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])) + + def test_dynamic_sequence_export_is_fully_delegated(self) -> None: + self.assertIsNotNone(_lower_dynamic_sequence_program()) + + def test_dynamic_sequence_goldens_match_custom_op(self) -> None: + _, _, freqs_cos, freqs_sin, _ = _dynamic_inputs(DYNAMIC_MAX_SEQ) + for seq, position in dict.fromkeys(DYNAMIC_SEQUENCE_CASES): + with self.subTest(seq=seq, position=position): + xq, xk, _, _, _ = _dynamic_inputs(seq) + expected_q, expected_k = _dynamic_golden( + xq, xk, freqs_cos, freqs_sin, position + ) + 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) + def export_rope_hf_model( pte_path: str, xq_golden_path: str, xk_golden_path: str, shape_name: str = "multi" @@ -132,5 +331,58 @@ def export_rope_hf_model( ) +def export_rope_hf_dynamic(out_dir: str) -> None: + os.makedirs(out_dir, exist_ok=True) + xq, xk, freqs_cos, freqs_sin, _ = _dynamic_inputs() + et = _export_dynamic_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: + os.makedirs(out_dir, exist_ok=True) + _, _, freqs_cos, freqs_sin, _ = _dynamic_inputs(DYNAMIC_MAX_SEQ) + et = _export_dynamic_sequence_program() + prefix = "rope_hf_dynamic_sequence" + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as output: + output.write(et.buffer) + + for name, tensor in (("freqs_cos", freqs_cos), ("freqs_sin", freqs_sin)): + tensor.detach().numpy().astype(".pte + .golden.bin to /tmp) ===" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import export_all_sdpa_models @@ -96,6 +105,7 @@ cmake \ "${EXECUTORCH_ROOT}" cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_native_test -j${NPROC} +cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_compute_dispatch_test -j${NPROC} cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_dispatch_order_test -j${NPROC} cmake --build "${NATIVE_BUILD_DIR}" --target webgpu_scratch_buffer_test -j${NPROC} @@ -108,9 +118,11 @@ 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" +"${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_compute_dispatch_test" "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" "${NATIVE_BUILD_DIR}/backends/webgpu/webgpu_scratch_buffer_test" diff --git a/backends/webgpu/test/test_native_ci_contract.py b/backends/webgpu/test/test_native_ci_contract.py index 011354a0e0f..43b8ec6f56e 100644 --- a/backends/webgpu/test/test_native_ci_contract.py +++ b/backends/webgpu/test/test_native_ci_contract.py @@ -64,3 +64,12 @@ def test_requires_symint_and_suppression_fixtures(self) -> None: "${OUTPUT_SUPPRESSION_DIR}/input.bin", ): self.assertIn(f'require_file "{fixture}"', script) + + def test_requires_dynamic_rope_fixture(self) -> None: + script = ( + pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" + ).read_text() + + self.assertIn("export_rope_hf_dynamic('${ROPE_HF_DIR}')", script) + self.assertIn('WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}"', script) + self.assertIn('require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte"', script) diff --git a/backends/webgpu/test/test_webgpu_native.cpp b/backends/webgpu/test/test_webgpu_native.cpp index f64cccec550..a57c41e53a8 100644 --- a/backends/webgpu/test/test_webgpu_native.cpp +++ b/backends/webgpu/test/test_webgpu_native.cpp @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -513,6 +514,152 @@ std::vector run_rms_norm_at_wg( return result; } +struct RotaryHfProbeParams { + 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 _pad; +}; + +std::vector run_rope_hf_2d_probe(const WebGPUContext& ctx) { + constexpr uint32_t kWorkgroupSize = 2; + constexpr uint32_t kWorkgroupsX = 2; + constexpr uint32_t kWorkgroupsY = 2; + constexpr uint32_t kNumPairs = kWorkgroupSize * kWorkgroupsX * kWorkgroupsY; + constexpr uint32_t kHeadDim = kNumPairs * 2; + + std::vector input(kHeadDim); + std::vector output(kHeadDim, 0.0f); + std::vector freqs_cos(kHeadDim, 1.0f); + std::vector freqs_sin(kHeadDim, 0.0f); + for (uint32_t i = 0; i < kHeadDim; i++) { + input[i] = static_cast(i + 1u); + if (i >= kNumPairs) { + freqs_cos[i] = 2.0f; + } + } + + WGPUDevice device = ctx.device; + 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); + + WGPUConstantEntry wg_const = {}; + wg_const.key = {"wg_size", WGPU_STRLEN}; + wg_const.value = static_cast(kWorkgroupSize); + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = 1; + pipeline_desc.compute.constants = &wg_const; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + WGPUBindGroupLayout layout = + wgpuComputePipelineGetBindGroupLayout(pipeline, 0); + + auto make_buffer = + [device](const void* data, uint64_t size, WGPUBufferUsage usage) { + WGPUBufferDescriptor desc = {}; + desc.size = size; + desc.usage = usage; + desc.mappedAtCreation = true; + WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &desc); + std::memcpy(wgpuBufferGetMappedRange(buffer, 0, size), data, size); + wgpuBufferUnmap(buffer); + return buffer; + }; + + const uint64_t data_bytes = kHeadDim * sizeof(float); + WGPUBuffer out_buffer = make_buffer( + output.data(), + data_bytes, + WGPUBufferUsage_Storage | WGPUBufferUsage_CopySrc); + WGPUBuffer in_buffer = + make_buffer(input.data(), data_bytes, WGPUBufferUsage_Storage); + WGPUBuffer cos_buffer = + make_buffer(freqs_cos.data(), data_bytes, WGPUBufferUsage_Storage); + WGPUBuffer sin_buffer = + make_buffer(freqs_sin.data(), data_bytes, WGPUBufferUsage_Storage); + const RotaryHfProbeParams params = { + 1u, 1u, kHeadDim, kNumPairs, kNumPairs, kHeadDim, 0u, 0u}; + WGPUBuffer params_buffer = + make_buffer(¶ms, sizeof(params), WGPUBufferUsage_Uniform); + + WGPUBindGroupEntry entries[5] = {}; + const WGPUBuffer buffers[] = { + out_buffer, in_buffer, cos_buffer, sin_buffer, params_buffer}; + const uint64_t sizes[] = { + data_bytes, data_bytes, data_bytes, data_bytes, sizeof(params)}; + for (uint32_t i = 0; i < 5; i++) { + entries[i].binding = i; + entries[i].buffer = buffers[i]; + entries[i].size = sizes[i]; + } + WGPUBindGroupDescriptor bind_group_desc = {}; + bind_group_desc.layout = layout; + bind_group_desc.entryCount = 5; + bind_group_desc.entries = entries; + WGPUBindGroup bind_group = + wgpuDeviceCreateBindGroup(device, &bind_group_desc); + + WGPUBufferDescriptor staging_desc = {}; + staging_desc.size = data_bytes; + staging_desc.usage = WGPUBufferUsage_MapRead | WGPUBufferUsage_CopyDst; + WGPUBuffer staging = wgpuDeviceCreateBuffer(device, &staging_desc); + + WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device, nullptr); + WGPUComputePassDescriptor pass_desc = {}; + WGPUComputePassEncoder pass = + wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); + wgpuComputePassEncoderSetPipeline(pass, pipeline); + wgpuComputePassEncoderSetBindGroup(pass, 0, bind_group, 0, nullptr); + wgpuComputePassEncoderDispatchWorkgroups(pass, kWorkgroupsX, kWorkgroupsY, 1); + wgpuComputePassEncoderEnd(pass); + wgpuComputePassEncoderRelease(pass); + wgpuCommandEncoderCopyBufferToBuffer( + encoder, out_buffer, 0, staging, 0, data_bytes); + WGPUCommandBuffer command = wgpuCommandEncoderFinish(encoder, nullptr); + wgpuQueueSubmit(ctx.queue, 1, &command); + wgpuCommandBufferRelease(command); + wgpuCommandEncoderRelease(encoder); + + WgMapData callback = {}; + WGPUBufferMapCallbackInfo callback_info = {}; + callback_info.mode = WGPUCallbackMode_WaitAnyOnly; + callback_info.callback = wg_map_cb; + callback_info.userdata1 = &callback; + WGPUFuture future = wgpuBufferMapAsync( + staging, WGPUMapMode_Read, 0, data_bytes, callback_info); + const WGPUWaitStatus wait = webgpu_wait(ctx.instance, future); + if (wait == WGPUWaitStatus_Success && + callback.status == WGPUMapAsyncStatus_Success) { + const void* mapped = wgpuBufferGetConstMappedRange(staging, 0, data_bytes); + std::memcpy(output.data(), mapped, data_bytes); + wgpuBufferUnmap(staging); + } else { + output.clear(); + } + + wgpuBufferRelease(staging); + wgpuBindGroupRelease(bind_group); + wgpuBufferRelease(params_buffer); + wgpuBufferRelease(sin_buffer); + wgpuBufferRelease(cos_buffer); + wgpuBufferRelease(in_buffer); + wgpuBufferRelease(out_buffer); + wgpuBindGroupLayoutRelease(layout); + wgpuComputePipelineRelease(pipeline); + wgpuShaderModuleRelease(shader); + return output; +} + // linear_q4gsw sweep config; mirrors CONFIGS in test_quantized_linear.py. struct Q4gswConfig { const char* name; @@ -767,6 +914,220 @@ void test_rope( << "apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)"; } +bool 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; +} + +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; + + 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 overflow_post = + make_tensor_ptr({1}, std::vector{INT64_C(1) << 32}); + auto overflow = module.forward({ + EValue(xqt), + EValue(xkt), + EValue(fct), + EValue(fst), + EValue(overflow_post), + }); + EXPECT_FALSE(overflow.ok()) + << "HF RoPE accepted a start_pos that aliases to zero when narrowed"; + + 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"; + + auto negative_post = make_tensor_ptr({1}, std::vector{-1}); + auto negative = module.forward({ + EValue(xqt), + EValue(xkt), + EValue(fct), + EValue(fst), + EValue(negative_post), + }); + EXPECT_FALSE(negative.ok()) << "HF RoPE accepted a negative start_pos"; +} + +void test_rope_hf_dynamic_sequence_reused_graph(const std::string& dir) { + constexpr int NH = 16; + constexpr int NKV = 8; + constexpr int HD = 128; + constexpr int MAXS = 16; + struct Case { + int seq; + int position; + }; + constexpr Case cases[] = {{16, 0}, {5, 7}, {1, 15}, {16, 0}}; + + const std::string prefix = dir + "rope_hf_dynamic_sequence"; + Module module(prefix + ".pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load HF RoPE dynamic-sequence model"; + + const int freqs_numel = MAXS * HD; + const std::vector freqs_cos = + load_golden(prefix + ".freqs_cos.bin", freqs_numel); + const std::vector freqs_sin = + load_golden(prefix + ".freqs_sin.bin", freqs_numel); + ASSERT_FALSE(freqs_cos.empty() || freqs_sin.empty()) + << "could not load HF RoPE dynamic-sequence frequencies"; + + for (const Case& c : cases) { + const int xq_numel = c.seq * NH * HD; + const int xk_numel = c.seq * NKV * HD; + const std::string case_prefix = prefix + ".S" + std::to_string(c.seq) + + ".pos" + std::to_string(c.position); + const std::vector xq = + load_golden(case_prefix + ".xq.bin", xq_numel); + const std::vector xk = + load_golden(case_prefix + ".xk.bin", xk_numel); + const std::vector golden_q = + load_golden(case_prefix + ".xq.golden.bin", xq_numel); + const std::vector golden_k = + load_golden(case_prefix + ".xk.golden.bin", xk_numel); + ASSERT_FALSE( + xq.empty() || xk.empty() || golden_q.empty() || golden_k.empty()) + << "could not load HF RoPE dynamic-sequence case " << case_prefix; + + auto xqt = make_tensor_ptr({1, c.seq, NH, HD}, std::vector(xq)); + auto xkt = make_tensor_ptr({1, c.seq, 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(c.position)}); + auto result = module.forward( + {EValue(xqt), EValue(xkt), EValue(fct), EValue(fst), EValue(post)}); + ASSERT_TRUE(result.ok()) + << "HF RoPE dynamic-sequence forward failed for " << case_prefix + << " (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 dynamic-sequence tensor outputs"; + const auto& xq_out = outputs[0].toTensor(); + const auto& xk_out = outputs[1].toTensor(); + ASSERT_TRUE(has_shape(xq_out, {1, c.seq, NH, HD})) + << "HF RoPE query output has the wrong shape for " << case_prefix; + ASSERT_TRUE(has_shape(xk_out, {1, c.seq, NKV, HD})) + << "HF RoPE key output has the wrong shape for " << case_prefix; + + 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 dynamic-sequence mismatch for " << case_prefix + << ": q abs=" << q_abs << " rel=" << q_rel << ", k abs=" << k_abs + << " rel=" << k_rel; + } +} + void test_prepack( const std::string& model_path, const std::string& golden_path, @@ -2125,6 +2486,82 @@ static bool test_select_double_scalars() { return ok; } +void expect_rope_hf_resize_numel_overflow(uint32_t q_heads, uint32_t k_heads) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + + const std::vector q_dims = {1u, 1u, q_heads, 2u}; + const std::vector k_dims = {1u, 1u, k_heads, 2u}; + const std::vector freqs_dims = {2u, 2u}; + std::vector<::flatbuffers::Offset> values; + const auto add_tensor = [&](const std::vector& dims, int mem_id) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + /*mem_obj_id=*/mem_id) + .Union())); + }; + add_tensor(q_dims, 0); + add_tensor(k_dims, 1); + add_tensor(freqs_dims, 2); + add_tensor(freqs_dims, 3); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, 0).Union())); + add_tensor(q_dims, 4); + add_tensor(k_dims, 5); + const std::vector output_items = {5, 6}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &output_items).Union())); + + const std::vector args = {0, 1, 2, 3, 4, 7}; + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, 0, "et_vk.apply_rotary_emb_hf.default", &args)); + const std::vector input_ids = {0, 1, 2, 3}; + const std::vector output_ids = {5, 6}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + WebGPUGraph graph; + ASSERT_NO_THROW(graph.build(fbb.GetBufferPointer(), nullptr, nullptr)); + ASSERT_EQ(graph.num_dispatches(), 2u); + const uint32_t q_x = graph.dispatch_at(0).workgroup_count_x; + const uint32_t q_y = graph.dispatch_at(0).workgroup_count_y; + const uint32_t k_x = graph.dispatch_at(1).workgroup_count_x; + const uint32_t k_y = graph.dispatch_at(1).workgroup_count_y; + + constexpr int64_t kLargeBatch = INT64_C(1) << 30; + const std::vector q_live = { + kLargeBatch, 1, static_cast(q_heads), 2}; + const std::vector k_live = { + kLargeBatch, 1, static_cast(k_heads), 2}; + graph.get_tensor(0).dims = q_live; + graph.get_tensor(1).dims = k_live; + ASSERT_NO_THROW(graph.resize_input(0, q_live)); + ASSERT_NO_THROW(graph.resize_input(1, k_live)); + + try { + graph.propagate_resize(); + FAIL() << "accepted q/k element count outside uint32 range"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ( + error.what(), + "apply_rotary_emb_hf(resize): element index exceeds uint32 range"); + } + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_x, q_x); + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_y, q_y); + EXPECT_EQ(graph.dispatch_at(1).workgroup_count_x, k_x); + EXPECT_EQ(graph.dispatch_at(1).workgroup_count_y, k_y); +} + // apply_rotary_emb on-GPU configs: multi + decode (env-gated, run-if-present). struct RopeConfig { const char* name; @@ -2298,6 +2735,48 @@ 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, RopeHfDynamicSequenceReusedGraph) { + 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_sequence_reused_graph(dir); +} + +TEST(WebGPUNative, RopeHfUsesFull2DGridStride) { + const WebGPUContext* ctx = get_default_webgpu_context(); + ASSERT_NE(ctx, nullptr); + const std::vector output = run_rope_hf_2d_probe(*ctx); + ASSERT_EQ(output.size(), 16u) << "HF RoPE probe output map failed"; + for (size_t i = 0; i < output.size(); i++) { + const float expected = + static_cast(i + 1u) * (i < output.size() / 2u ? 1.0f : 2.0f); + EXPECT_EQ(output[i], expected) + << "HF RoPE 2D grid or second-half frequency mismatch at element " << i; + } +} + +TEST(WebGPUNative, RopeHfResizeRejectsQOrKNumelOverflow) { + expect_rope_hf_resize_numel_overflow(/*q_heads=*/2, /*k_heads=*/1); + expect_rope_hf_resize_numel_overflow(/*q_heads=*/1, /*k_heads=*/2); +} + TEST(WebGPUNative, Prepack) { if (g_prepack_model_path.empty() || g_prepack_golden_path.empty()) { GTEST_SKIP() << "WEBGPU_TEST_PREPACK_MODEL/GOLDEN not set"; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index fe8340d6528..0cd110d0ff7 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -178,6 +178,28 @@ def test_committed_headers_match_generator(self) -> None: got, want, f"{header.name} stale; run scripts/gen_wgsl_headers.py" ) + def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: + shader = ( + g.BACKEND_ROOT / "runtime" / "ops" / "rope" / "rotary_embedding_hf.wgsl" + ).read_text() + self.assertIn("@builtin(num_workgroups) num_workgroups", shader) + self.assertIn( + "gid.x + gid.y * (num_workgroups.x * wg_size)", + shader, + ) + self.assertIn("let freqs_b_idx = freqs_a_idx + half_dim;", shader) + self.assertIn("t_out[b_idx] = x_b * c_b + x_a * si_b;", shader) + + wg_size = 2 + workgroups_x = 2 + indices = [ + group_x * wg_size + lane + group_y * (workgroups_x * wg_size) + for group_y in range(2) + for group_x in range(workgroups_x) + for lane in range(wg_size) + ] + self.assertEqual(indices, list(range(8))) + def test_parse_workgroup_allows_space(self) -> None: # @workgroup_size (64) — the spec-legal spaced form must still parse. self.assertEqual(