[ExecuTorch][WebGPU] et_vk.apply_rotary_emb test suite (export + native golden)

backends/webgpu/scripts/test_webgpu_native_ci.sh

@@ -51,6 +51,12 @@ EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
 EMBEDDING_LLAMA1B_MODEL="/tmp/webgpu_embedding_q4gsw_llama1b.pte"
 EMBEDDING_LLAMA1B_INDICES="/tmp/webgpu_embedding_q4gsw_llama1b_indices.bin"
 EMBEDDING_LLAMA1B_GOLDEN="/tmp/webgpu_embedding_q4gsw_llama1b_golden.bin"
+ROPE_MODEL="/tmp/webgpu_rope.pte"
+ROPE_XQ_GOLDEN="/tmp/webgpu_rope_xq_golden.bin"
+ROPE_XK_GOLDEN="/tmp/webgpu_rope_xk_golden.bin"
+ROPE_DECODE_MODEL="/tmp/webgpu_rope_decode.pte"
+ROPE_DECODE_XQ_GOLDEN="/tmp/webgpu_rope_decode_xq_golden.bin"
+ROPE_DECODE_XK_GOLDEN="/tmp/webgpu_rope_decode_xk_golden.bin"
 
 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import export_all_quantized_linear_models
@@ -63,6 +69,12 @@ export_embedding_q4gsw_model('${EMBEDDING_MODEL}', '${EMBEDDING_GOLDEN}', '${EMB
 export_embedding_q4gsw_model('${EMBEDDING_LLAMA1B_MODEL}', '${EMBEDDING_LLAMA1B_GOLDEN}', '${EMBEDDING_LLAMA1B_INDICES}', 'llama1b')
 " || echo "WARN: embedding_q4gsw export failed; embedding configs will FAIL in webgpu_native_test"
 
+$PYTHON_EXECUTABLE -c "
+from executorch.backends.webgpu.test.ops.rope.test_rope import export_rope_model
+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')
+" || echo "WARN: rope export failed; apply_rotary_emb configs will FAIL in webgpu_native_test"
+
 $PYTHON_EXECUTABLE -c "
 from executorch.backends.webgpu.test.ops.dispatch_order.test_dispatch_order import export_dispatch_order_cases
 export_dispatch_order_cases('${DISPATCH_ORDER_DIR}')
@@ -154,6 +166,12 @@ if [[ -x "${BIN_DIR}/webgpu_native_test" ]] &&
       WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \
       WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \
       WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \
+      WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \
+      WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \
+      WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \
+      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}" \
       "${BIN_DIR}/webgpu_native_test"
 else
   echo "(skipping webgpu_native_test: executorch wheel absent — exports did not run)"

backends/webgpu/test/ops/rope/__init__.py

@@ -0,0 +1,5 @@
+# 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.

backends/webgpu/test/ops/rope/test_rope.py

@@ -0,0 +1,129 @@
+# 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.
+
+"""Interleaved rotary positional embedding (`et_vk.apply_rotary_emb`) export +
+goldens for the WebGPU backend.
+
+Exports the Llama interleaved RoPE (use_hf_rope=False) with freqs_cos/freqs_sin
+as runtime forward inputs (no constant prepack), which fuses under
+VulkanPartitioner into `et_vk.apply_rotary_emb.default` (two outputs xq_out,
+xk_out serialized as a ValueList). Inputs are deterministic /16 ramps so the
+native binary reconstructs them bit-for-bit; the two torch-computed goldens are
+written for the native binary to compare (it has no ATen).
+
+Two shapes are exercised: a multi-token prefill shape and a single-token (S=1)
+decode shape at the Llama-3.2-1B head config (GQA 32:8), so the seq=1 / batch
+decompositions and the position->freqs indexing are covered at decode too.
+"""
+
+import unittest
+from collections import namedtuple
+
+import executorch.backends.vulkan.custom_ops_lib  # noqa: F401
+
+import torch
+from executorch.backends.vulkan import VulkanPartitioner
+from executorch.examples.models.llama.rope import apply_rotary_emb, RotaryEmbedding
+from executorch.exir import to_edge_transform_and_lower
+
+# B batch, S tokens, NH query heads, NKV kv heads (NH != NKV so the two outputs
+# are distinguishable by numel), HD head dim (even; HD/2 rotation pairs).
+Shape = namedtuple("Shape", ["name", "b", "s", "nh", "nkv", "hd"])
+SHAPES = [
+    Shape("multi", 1, 5, 8, 2, 64),
+    # Single-token decode at Llama-3.2-1B head config (GQA 32:8, head_dim 64).
+    Shape("decode", 1, 1, 32, 8, 64),
+]
+
+
+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)
+    return ((idx % mod) - off).to(torch.float32) / 16.0
+
+
+def _inputs(
+    shape: Shape,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    xq = _ramp(shape.b * shape.s * shape.nh * shape.hd, 17, 8).reshape(
+        shape.b, shape.s, shape.nh, shape.hd
+    )
+    xk = _ramp(shape.b * shape.s * shape.nkv * shape.hd, 13, 6).reshape(
+        shape.b, shape.s, shape.nkv, shape.hd
+    )
+    freqs_cos = _ramp(shape.s * (shape.hd // 2), 11, 5).reshape(shape.s, shape.hd // 2)
+    freqs_sin = _ramp(shape.s * (shape.hd // 2), 7, 3).reshape(shape.s, shape.hd // 2)
+    return xq, xk, freqs_cos, freqs_sin
+
+
+def _golden(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cos: torch.Tensor,
+    freqs_sin: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    # Reference = the registered et_vk op the kernel implements.
+    return torch.ops.et_vk.apply_rotary_emb.default(xq, xk, freqs_cos, freqs_sin)
+
+
+def _export(inputs):
+    # Export the real Llama RoPE module (not a hand-written copy) so the test
+    # exercises the same pattern the partitioner matches in production models.
+    ep = torch.export.export(RotaryEmbedding().eval(), inputs)
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+class TestRope(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 fusion)"
+                )
+
+    def test_golden_matches_eager(self) -> None:
+        # The et_vk golden must equal the real Llama apply_rotary_emb, so a buggy
+        # golden can't fake-pass the native kernel. Run at both shapes so the S=1
+        # decode position->freqs indexing is covered.
+        for shape in SHAPES:
+            with self.subTest(shape=shape.name):
+                xq, xk, fc, fs = _inputs(shape)
+                gq, gk = _golden(xq, xk, fc, fs)
+                eq, ek = apply_rotary_emb(xq, xk, fc, fs)
+                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 export_rope_model(
+    pte_path: str, xq_golden_path: str, xk_golden_path: str, shape_name: str = "multi"
+) -> None:
+    """Write the apply_rotary_emb .pte + the xq_out and xk_out torch goldens
+    (raw LE fp32). Inputs are /16 ramps reconstructed in the native test.
+    `shape_name` selects an entry from SHAPES (default the multi-token shape)."""
+    shape = next(s for s in SHAPES if s.name == shape_name)
+    xq, xk, fc, fs = _inputs(shape)
+    gq, gk = _golden(xq, xk, fc, fs)
+    et = _export((xq, xk, fc, fs))
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    gq.detach().numpy().astype("<f4").tofile(xq_golden_path)
+    gk.detach().numpy().astype("<f4").tofile(xk_golden_path)
+    print(
+        f"Exported {pte_path} (shape={shape_name}); xq_out golden {xq_golden_path} "
+        f"({gq.numel()} floats); xk_out golden {xk_golden_path} ({gk.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()

backends/webgpu/test/test_webgpu_native.cpp

@@ -425,6 +425,117 @@ static bool quant_within_tol(
   return ok;
 }
 
+static bool test_rope(
+    const std::string& model_path,
+    const std::string& xq_golden_path,
+    const std::string& xk_golden_path,
+    int S,
+    int NH,
+    int NKV,
+    int HD,
+    const char* label) {
+  // Llama interleaved RoPE vs torch goldens; shapes/ramps per test_rope.py.
+  const int xq_numel = S * NH * HD;
+  const int xk_numel = S * NKV * HD;
+  const int freqs_numel = S * (HD / 2);
+  printf(
+      "\n--- Test: apply_rotary_emb (%s: S=%d,NH=%d,NKV=%d,HD=%d) ---\n",
+      label,
+      S,
+      NH,
+      NKV,
+      HD);
+
+  Module module(model_path);
+  auto err = module.load_forward();
+  if (err != Error::Ok) {
+    printf("FAIL: could not load forward method (error %d)\n", (int)err);
+    return false;
+  }
+  printf("Model loaded: %s\n", model_path.c_str());
+
+  // ((i % mod) - off) / 16: exact in fp32, matches test_rope.py::_ramp.
+  auto ramp = [](int i, int mod, int off) {
+    return static_cast<float>((i % mod) - off) / 16.0f;
+  };
+  std::vector<float> xq(xq_numel), xk(xk_numel), fc(freqs_numel),
+      fs(freqs_numel);
+  for (int i = 0; i < xq_numel; i++) {
+    xq[i] = ramp(i, 17, 8);
+  }
+  for (int i = 0; i < xk_numel; i++) {
+    xk[i] = ramp(i, 13, 6);
+  }
+  for (int i = 0; i < freqs_numel; i++) {
+    fc[i] = ramp(i, 11, 5);
+    fs[i] = ramp(i, 7, 3);
+  }
+
+  auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector<float>(xq));
+  auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector<float>(xk));
+  auto fct = make_tensor_ptr({S, HD / 2}, std::vector<float>(fc));
+  auto fst = make_tensor_ptr({S, HD / 2}, std::vector<float>(fs));
+
+  auto result =
+      module.forward({EValue(xqt), EValue(xkt), EValue(fct), EValue(fst)});
+  if (!result.ok()) {
+    printf("FAIL: forward failed (error %d)\n", (int)result.error());
+    return false;
+  }
+  const auto& outputs = result.get();
+
+  // Outputs in graph order [0]=xq_out, [1]=xk_out (positional; the numel check
+  // below guards a swap, since NH != NKV under GQA).
+  if (outputs.size() < 2 || !outputs[0].isTensor() || !outputs[1].isTensor()) {
+    printf("FAIL: expected 2 tensor outputs, got %zu\n", outputs.size());
+    return false;
+  }
+  const auto& xq_t = outputs[0].toTensor();
+  const auto& xk_t = outputs[1].toTensor();
+  if (xq_t.numel() != xq_numel || xk_t.numel() != xk_numel) {
+    printf(
+        "FAIL: output shapes [%zu,%zu] != expected [%d,%d]\n",
+        (size_t)xq_t.numel(),
+        (size_t)xk_t.numel(),
+        xq_numel,
+        xk_numel);
+    return false;
+  }
+  const float* xq_out = xq_t.const_data_ptr<float>();
+  const float* xk_out = xk_t.const_data_ptr<float>();
+
+  std::vector<float> gq = load_golden(xq_golden_path, xq_numel);
+  std::vector<float> gk = load_golden(xk_golden_path, xk_numel);
+  if (gq.empty() || gk.empty()) {
+    printf(
+        "FAIL: could not load goldens %s / %s\n",
+        xq_golden_path.c_str(),
+        xk_golden_path.c_str());
+    return false;
+  }
+
+  // Per-element abs-OR-rel on xq and xk (shared helper, defined above).
+  float maq = 0.0f, mrq = 0.0f, mak = 0.0f, mrk = 0.0f;
+  const bool pass_q =
+      quant_within_tol(xq_out, gq.data(), xq_numel, 1e-3f, 1e-3f, &maq, &mrq);
+  const bool pass_k =
+      quant_within_tol(xk_out, gk.data(), xk_numel, 1e-3f, 1e-3f, &mak, &mrk);
+  const float max_abs_err = std::max(maq, mak);
+  const float max_rel_err = std::max(mrq, mrk);
+
+  printf(
+      "Max abs error: %e   Max rel error: %e (checked %d elements)\n",
+      max_abs_err,
+      max_rel_err,
+      xq_numel + xk_numel);
+  if (!(pass_q && pass_k)) {
+    printf("FAIL: apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)\n");
+    return false;
+  }
+  printf("PASS: apply_rotary_emb test\n");
+  return true;
+}
+
 // Reconstruct _ramp_input bit-for-bit, run the op, compare to the fp64 golden.
 static bool test_q4gsw_config(
     const Q4gswConfig& cfg,
@@ -1472,6 +1583,37 @@ int main(int argc, char** argv) {
        2048},
   };
 
+  // apply_rotary_emb on-GPU configs: multi + decode (env-gated,
+  // run-if-present).
+  struct RopeConfig {
+    const char* name;
+    const char* model_env;
+    const char* xq_env;
+    const char* xk_env;
+    int S;
+    int NH;
+    int NKV;
+    int HD;
+  };
+  const RopeConfig rope_configs[] = {
+      {"multi",
+       "WEBGPU_TEST_ROPE_MODEL",
+       "WEBGPU_TEST_ROPE_XQ_GOLDEN",
+       "WEBGPU_TEST_ROPE_XK_GOLDEN",
+       5,
+       8,
+       2,
+       64},
+      {"decode",
+       "WEBGPU_TEST_ROPE_DECODE_MODEL",
+       "WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN",
+       "WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN",
+       1,
+       32,
+       8,
+       64},
+  };
+
   // SDPA sweep: configs self-discover their sdpa_<name>.pte/.golden.bin under
   // this directory (default "" = the embedded-file root / cwd). Set
   // WEBGPU_TEST_SDPA_DIR to point at the exported .pte directory (e.g. /tmp/).
@@ -1528,6 +1670,15 @@ int main(int argc, char** argv) {
     }
   }
 
+  for (const auto& c : rope_configs) {
+    const char* m = std::getenv(c.model_env);
+    const char* xq = std::getenv(c.xq_env);
+    const char* xk = std::getenv(c.xk_env);
+    if (m && xq && xk && *m && *xq && *xk) {
+      ok = test_rope(m, xq, xk, c.S, c.NH, c.NKV, c.HD, c.name) && ok;
+    }
+  }
+
   bool sdpa_ran = false;
   bool sdpa_ok = test_sdpa_sweep(sdpa_dir, &sdpa_ran);
   if (sdpa_ran) {