|
| 1 | +# Copyright (c) Meta Platforms, Inc. and affiliates. |
| 2 | +# All rights reserved. |
| 3 | +# |
| 4 | +# This source code is licensed under the BSD-style license found in the |
| 5 | +# LICENSE file in the root directory of this source tree. |
| 6 | + |
| 7 | +"""Interleaved rotary positional embedding (`et_vk.apply_rotary_emb`) export + |
| 8 | +goldens for the WebGPU backend. |
| 9 | +
|
| 10 | +Exports the Llama interleaved RoPE (use_hf_rope=False) with freqs_cos/freqs_sin |
| 11 | +as runtime forward inputs (no constant prepack), which fuses under |
| 12 | +VulkanPartitioner into `et_vk.apply_rotary_emb.default` (two outputs xq_out, |
| 13 | +xk_out serialized as a ValueList). Inputs are deterministic /16 ramps so the |
| 14 | +native binary reconstructs them bit-for-bit; the two torch-computed goldens are |
| 15 | +written for the native binary to compare (it has no ATen). |
| 16 | +
|
| 17 | +Two shapes are exercised: a multi-token prefill shape and a single-token (S=1) |
| 18 | +decode shape at the Llama-3.2-1B head config (GQA 32:8), so the seq=1 / batch |
| 19 | +decompositions and the position->freqs indexing are covered at decode too. |
| 20 | +""" |
| 21 | + |
| 22 | +import unittest |
| 23 | +from collections import namedtuple |
| 24 | + |
| 25 | +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 |
| 26 | + |
| 27 | +import torch |
| 28 | +from executorch.backends.vulkan import VulkanPartitioner |
| 29 | +from executorch.examples.models.llama.rope import apply_rotary_emb, RotaryEmbedding |
| 30 | +from executorch.exir import to_edge_transform_and_lower |
| 31 | + |
| 32 | +# B batch, S tokens, NH query heads, NKV kv heads (NH != NKV so the two outputs |
| 33 | +# are distinguishable by numel), HD head dim (even; HD/2 rotation pairs). |
| 34 | +Shape = namedtuple("Shape", ["name", "b", "s", "nh", "nkv", "hd"]) |
| 35 | +SHAPES = [ |
| 36 | + Shape("multi", 1, 5, 8, 2, 64), |
| 37 | + # Single-token decode at Llama-3.2-1B head config (GQA 32:8, head_dim 64). |
| 38 | + Shape("decode", 1, 1, 32, 8, 64), |
| 39 | +] |
| 40 | + |
| 41 | + |
| 42 | +def _ramp(numel: int, mod: int, off: int) -> torch.Tensor: |
| 43 | + # ((i % mod) - off) / 16: exact in fp32, matches test_webgpu_native.cpp. |
| 44 | + idx = torch.arange(numel, dtype=torch.int64) |
| 45 | + return ((idx % mod) - off).to(torch.float32) / 16.0 |
| 46 | + |
| 47 | + |
| 48 | +def _inputs( |
| 49 | + shape: Shape, |
| 50 | +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: |
| 51 | + xq = _ramp(shape.b * shape.s * shape.nh * shape.hd, 17, 8).reshape( |
| 52 | + shape.b, shape.s, shape.nh, shape.hd |
| 53 | + ) |
| 54 | + xk = _ramp(shape.b * shape.s * shape.nkv * shape.hd, 13, 6).reshape( |
| 55 | + shape.b, shape.s, shape.nkv, shape.hd |
| 56 | + ) |
| 57 | + freqs_cos = _ramp(shape.s * (shape.hd // 2), 11, 5).reshape(shape.s, shape.hd // 2) |
| 58 | + freqs_sin = _ramp(shape.s * (shape.hd // 2), 7, 3).reshape(shape.s, shape.hd // 2) |
| 59 | + return xq, xk, freqs_cos, freqs_sin |
| 60 | + |
| 61 | + |
| 62 | +def _golden( |
| 63 | + xq: torch.Tensor, |
| 64 | + xk: torch.Tensor, |
| 65 | + freqs_cos: torch.Tensor, |
| 66 | + freqs_sin: torch.Tensor, |
| 67 | +) -> tuple[torch.Tensor, torch.Tensor]: |
| 68 | + # Reference = the registered et_vk op the kernel implements. |
| 69 | + return torch.ops.et_vk.apply_rotary_emb.default(xq, xk, freqs_cos, freqs_sin) |
| 70 | + |
| 71 | + |
| 72 | +def _export(inputs): |
| 73 | + # Export the real Llama RoPE module (not a hand-written copy) so the test |
| 74 | + # exercises the same pattern the partitioner matches in production models. |
| 75 | + ep = torch.export.export(RotaryEmbedding().eval(), inputs) |
| 76 | + return to_edge_transform_and_lower( |
| 77 | + ep, partitioner=[VulkanPartitioner()] |
| 78 | + ).to_executorch() |
| 79 | + |
| 80 | + |
| 81 | +class TestRope(unittest.TestCase): |
| 82 | + def test_export_delegates(self) -> None: |
| 83 | + for shape in SHAPES: |
| 84 | + with self.subTest(shape=shape.name): |
| 85 | + et = _export(_inputs(shape)) |
| 86 | + found = any( |
| 87 | + d.id == "VulkanBackend" |
| 88 | + for plan in et.executorch_program.execution_plan |
| 89 | + for d in plan.delegates |
| 90 | + ) |
| 91 | + self.assertTrue( |
| 92 | + found, "Expected a VulkanBackend delegate (apply_rotary_emb fusion)" |
| 93 | + ) |
| 94 | + |
| 95 | + def test_golden_matches_eager(self) -> None: |
| 96 | + # The et_vk golden must equal the real Llama apply_rotary_emb, so a buggy |
| 97 | + # golden can't fake-pass the native kernel. Run at both shapes so the S=1 |
| 98 | + # decode position->freqs indexing is covered. |
| 99 | + for shape in SHAPES: |
| 100 | + with self.subTest(shape=shape.name): |
| 101 | + xq, xk, fc, fs = _inputs(shape) |
| 102 | + gq, gk = _golden(xq, xk, fc, fs) |
| 103 | + eq, ek = apply_rotary_emb(xq, xk, fc, fs) |
| 104 | + torch.testing.assert_close(gq, eq, atol=1e-5, rtol=1e-5) |
| 105 | + torch.testing.assert_close(gk, ek, atol=1e-5, rtol=1e-5) |
| 106 | + |
| 107 | + |
| 108 | +def export_rope_model( |
| 109 | + pte_path: str, xq_golden_path: str, xk_golden_path: str, shape_name: str = "multi" |
| 110 | +) -> None: |
| 111 | + """Write the apply_rotary_emb .pte + the xq_out and xk_out torch goldens |
| 112 | + (raw LE fp32). Inputs are /16 ramps reconstructed in the native test. |
| 113 | + `shape_name` selects an entry from SHAPES (default the multi-token shape).""" |
| 114 | + shape = next(s for s in SHAPES if s.name == shape_name) |
| 115 | + xq, xk, fc, fs = _inputs(shape) |
| 116 | + gq, gk = _golden(xq, xk, fc, fs) |
| 117 | + et = _export((xq, xk, fc, fs)) |
| 118 | + with open(pte_path, "wb") as f: |
| 119 | + f.write(et.buffer) |
| 120 | + gq.detach().numpy().astype("<f4").tofile(xq_golden_path) |
| 121 | + gk.detach().numpy().astype("<f4").tofile(xk_golden_path) |
| 122 | + print( |
| 123 | + f"Exported {pte_path} (shape={shape_name}); xq_out golden {xq_golden_path} " |
| 124 | + f"({gq.numel()} floats); xk_out golden {xk_golden_path} ({gk.numel()} floats)" |
| 125 | + ) |
| 126 | + |
| 127 | + |
| 128 | +if __name__ == "__main__": |
| 129 | + unittest.main() |
0 commit comments