[ExecuTorch][WebGPU] Add HuggingFace rotate-half RoPE

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<u32>) {
-  let pair = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  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<u32>) {
       ((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;
 }

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<storage, read_write> t_out: array<f32>;
 @group(0) @binding(1) var<storage, read> t_in: array<f32>;
@@ -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<u32>) {
-  let pair = gid.x;
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  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<u32>) {
       ((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;
 }
 )";
 

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}" \

backends/webgpu/test/native/test_compute_dispatch.cpp

@@ -11,6 +11,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
 #include <executorch/backends/webgpu/runtime/WebGPUShaderRegistry.h>
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/relu/relu_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sigmoid/sigmoid_wgsl.h>
 
@@ -20,6 +21,8 @@
 #include <cstdio>
 #include <limits>
 #include <stdexcept>
+#include <string>
+#include <vector>
 
 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<uint32_t> xq_dims;
+  std::vector<uint32_t> xk_dims;
+  std::vector<uint32_t> cos_dims;
+  std::vector<uint32_t> sin_dims;
+  std::vector<uint32_t> xq_out_dims;
+  std::vector<uint32_t> 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<vk::VkValue>> values;
+  auto add_tensor = [&](vk::VkDataType dtype,
+                        const std::vector<uint32_t>& 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<int32_t> output_value_ids = {4, 5};
+  values.push_back(vk::CreateVkValue(
+      fbb,
+      vk::GraphTypes::ValueList,
+      vk::CreateValueListDirect(fbb, &output_value_ids).Union()));
+
+  std::vector<int32_t> args = {0, 1, 2, 3, 6};
+  std::vector<::flatbuffers::Offset<vk::OperatorCall>> chain;
+  chain.push_back(vk::CreateOperatorCallDirect(
+      fbb, 0, "et_vk.apply_rotary_emb.default", &args));
+  std::vector<uint32_t> input_ids = {0, 1, 2, 3};
+  std::vector<uint32_t> 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<uint32_t> xq = {1, 2, 2, 4};
+  const std::vector<uint32_t> xk = {1, 2, 1, 4};
+  const std::vector<uint32_t> 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;