From bb3fe70658023d6308e154704ede8745314fccae Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Mon, 15 Jun 2026 14:52:56 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- .../webgpu/scripts/test_webgpu_native_ci.sh | 18 +++ backends/webgpu/test/ops/rope/__init__.py | 5 + backends/webgpu/test/ops/rope/test_rope.py | 129 +++++++++++++++ backends/webgpu/test/test_webgpu_native.cpp | 152 ++++++++++++++++++ 4 files changed, 304 insertions(+) create mode 100644 backends/webgpu/test/ops/rope/__init__.py create mode 100644 backends/webgpu/test/ops/rope/test_rope.py diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 8603d50a753..4aa2fcc1121 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -53,6 +53,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.add.test_add import export_add_model, export_chained_add_model @@ -71,6 +77,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; webgpu_native_test embedding cases self-skip" +$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; webgpu_native_test apply_rotary_emb case self-skips" + $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.rms_norm.test_rms_norm import export_rms_norm_cases export_rms_norm_cases('${RMS_NORM_DIR}') @@ -167,6 +179,12 @@ if [[ -x "${BIN_DIR}/webgpu_native_test" && -f "${PTE_MODEL}" ]]; then 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: no exported .pte — needs the executorch python wheel)" diff --git a/backends/webgpu/test/ops/rope/__init__.py b/backends/webgpu/test/ops/rope/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/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. diff --git a/backends/webgpu/test/ops/rope/test_rope.py b/backends/webgpu/test/ops/rope/test_rope.py new file mode 100644 index 00000000000..136beb9d8c8 --- /dev/null +++ b/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("((i % mod) - off) / 16.0f; + }; + std::vector 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(xq)); + auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector(xk)); + auto fct = make_tensor_ptr({S, HD / 2}, std::vector(fc)); + auto fst = make_tensor_ptr({S, HD / 2}, std::vector(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 [xq_out, xk_out]; select by shape (handles 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(); + const float* xk_out = xk_t.const_data_ptr(); + + std::vector gq = load_golden(xq_golden_path, xq_numel); + std::vector 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; + } + + float max_abs_err = 0.0f, max_rel_err = 0.0f; + auto accum = [&](const float* out, const std::vector& g, int n) { + for (int i = 0; i < n; i++) { + const float ae = std::abs(out[i] - g[i]); + max_abs_err = std::max(max_abs_err, ae); + max_rel_err = std::max(max_rel_err, ae / std::max(std::abs(g[i]), 1e-6f)); + } + }; + accum(xq_out, gq, xq_numel); + accum(xk_out, gk, xk_numel); + + printf( + "Max abs error: %e Max rel error: %e (checked %d elements)\n", + max_abs_err, + max_rel_err, + xq_numel + xk_numel); + if (max_abs_err > 1e-3f || max_rel_err > 1e-3f) { + printf("FAIL: apply_rotary_emb exceeds tolerance 1e-3\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, @@ -1577,6 +1689,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_.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/). @@ -1639,6 +1782,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) {