Skip to content

Commit 2b362c8

Browse files
JCNTHJCNTH
committed
[ExecuTorch][WebGPU] et_vk.apply_rotary_emb test suite (export + native golden)
Pull Request resolved: #20290 Splits the `et_vk.apply_rotary_emb` tests into their own diff (op below, tests above), matching the `sdpa`/`update_cache`/`linear_q4gsw` convention, and brings them to the same rigor: a multi-shape config sweep run on-device (prefill + decode) and a library dual-oracle at both shapes. ghstack-source-id: 395549287 @exported-using-ghexport Differential Revision: [D108668384](https://our.internmc.facebook.com/intern/diff/D108668384/)
1 parent 72b1f79 commit 2b362c8

4 files changed

Lines changed: 303 additions & 0 deletions

File tree

‎backends/webgpu/scripts/test_webgpu_native_ci.sh‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,12 @@ EMBEDDING_GOLDEN="/tmp/webgpu_embedding_q4gsw_golden.bin"
5151
EMBEDDING_LLAMA1B_MODEL="/tmp/webgpu_embedding_q4gsw_llama1b.pte"
5252
EMBEDDING_LLAMA1B_INDICES="/tmp/webgpu_embedding_q4gsw_llama1b_indices.bin"
5353
EMBEDDING_LLAMA1B_GOLDEN="/tmp/webgpu_embedding_q4gsw_llama1b_golden.bin"
54+
ROPE_MODEL="/tmp/webgpu_rope.pte"
55+
ROPE_XQ_GOLDEN="/tmp/webgpu_rope_xq_golden.bin"
56+
ROPE_XK_GOLDEN="/tmp/webgpu_rope_xk_golden.bin"
57+
ROPE_DECODE_MODEL="/tmp/webgpu_rope_decode.pte"
58+
ROPE_DECODE_XQ_GOLDEN="/tmp/webgpu_rope_decode_xq_golden.bin"
59+
ROPE_DECODE_XK_GOLDEN="/tmp/webgpu_rope_decode_xk_golden.bin"
5460

5561
$PYTHON_EXECUTABLE -c "
5662
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
6369
export_embedding_q4gsw_model('${EMBEDDING_LLAMA1B_MODEL}', '${EMBEDDING_LLAMA1B_GOLDEN}', '${EMBEDDING_LLAMA1B_INDICES}', 'llama1b')
6470
" || echo "WARN: embedding_q4gsw export failed; embedding configs will FAIL in webgpu_native_test"
6571

72+
$PYTHON_EXECUTABLE -c "
73+
from executorch.backends.webgpu.test.ops.rope.test_rope import export_rope_model
74+
export_rope_model('${ROPE_MODEL}', '${ROPE_XQ_GOLDEN}', '${ROPE_XK_GOLDEN}')
75+
export_rope_model('${ROPE_DECODE_MODEL}', '${ROPE_DECODE_XQ_GOLDEN}', '${ROPE_DECODE_XK_GOLDEN}', 'decode')
76+
" || echo "WARN: rope export failed; apply_rotary_emb configs will FAIL in webgpu_native_test"
77+
6678
$PYTHON_EXECUTABLE -c "
6779
from executorch.backends.webgpu.test.ops.dispatch_order.test_dispatch_order import export_dispatch_order_cases
6880
export_dispatch_order_cases('${DISPATCH_ORDER_DIR}')
@@ -154,6 +166,12 @@ if [[ -x "${BIN_DIR}/webgpu_native_test" ]] &&
154166
WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_MODEL="${EMBEDDING_LLAMA1B_MODEL}" \
155167
WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_INDICES="${EMBEDDING_LLAMA1B_INDICES}" \
156168
WEBGPU_TEST_EMBEDDING_Q4GSW_LLAMA1B_GOLDEN="${EMBEDDING_LLAMA1B_GOLDEN}" \
169+
WEBGPU_TEST_ROPE_MODEL="${ROPE_MODEL}" \
170+
WEBGPU_TEST_ROPE_XQ_GOLDEN="${ROPE_XQ_GOLDEN}" \
171+
WEBGPU_TEST_ROPE_XK_GOLDEN="${ROPE_XK_GOLDEN}" \
172+
WEBGPU_TEST_ROPE_DECODE_MODEL="${ROPE_DECODE_MODEL}" \
173+
WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN="${ROPE_DECODE_XQ_GOLDEN}" \
174+
WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN="${ROPE_DECODE_XK_GOLDEN}" \
157175
"${BIN_DIR}/webgpu_native_test"
158176
else
159177
echo "(skipping webgpu_native_test: executorch wheel absent — exports did not run)"

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
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.

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

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
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()

‎backends/webgpu/test/test_webgpu_native.cpp‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,6 +425,117 @@ static bool quant_within_tol(
425425
return ok;
426426
}
427427

428+
static bool test_rope(
429+
const std::string& model_path,
430+
const std::string& xq_golden_path,
431+
const std::string& xk_golden_path,
432+
int S,
433+
int NH,
434+
int NKV,
435+
int HD,
436+
const char* label) {
437+
// Llama interleaved RoPE vs torch goldens; shapes/ramps per test_rope.py.
438+
const int xq_numel = S * NH * HD;
439+
const int xk_numel = S * NKV * HD;
440+
const int freqs_numel = S * (HD / 2);
441+
printf(
442+
"\n--- Test: apply_rotary_emb (%s: S=%d,NH=%d,NKV=%d,HD=%d) ---\n",
443+
label,
444+
S,
445+
NH,
446+
NKV,
447+
HD);
448+
449+
Module module(model_path);
450+
auto err = module.load_forward();
451+
if (err != Error::Ok) {
452+
printf("FAIL: could not load forward method (error %d)\n", (int)err);
453+
return false;
454+
}
455+
printf("Model loaded: %s\n", model_path.c_str());
456+
457+
// ((i % mod) - off) / 16: exact in fp32, matches test_rope.py::_ramp.
458+
auto ramp = [](int i, int mod, int off) {
459+
return static_cast<float>((i % mod) - off) / 16.0f;
460+
};
461+
std::vector<float> xq(xq_numel), xk(xk_numel), fc(freqs_numel),
462+
fs(freqs_numel);
463+
for (int i = 0; i < xq_numel; i++) {
464+
xq[i] = ramp(i, 17, 8);
465+
}
466+
for (int i = 0; i < xk_numel; i++) {
467+
xk[i] = ramp(i, 13, 6);
468+
}
469+
for (int i = 0; i < freqs_numel; i++) {
470+
fc[i] = ramp(i, 11, 5);
471+
fs[i] = ramp(i, 7, 3);
472+
}
473+
474+
auto xqt = make_tensor_ptr({1, S, NH, HD}, std::vector<float>(xq));
475+
auto xkt = make_tensor_ptr({1, S, NKV, HD}, std::vector<float>(xk));
476+
auto fct = make_tensor_ptr({S, HD / 2}, std::vector<float>(fc));
477+
auto fst = make_tensor_ptr({S, HD / 2}, std::vector<float>(fs));
478+
479+
auto result =
480+
module.forward({EValue(xqt), EValue(xkt), EValue(fct), EValue(fst)});
481+
if (!result.ok()) {
482+
printf("FAIL: forward failed (error %d)\n", (int)result.error());
483+
return false;
484+
}
485+
const auto& outputs = result.get();
486+
487+
// Outputs in graph order [0]=xq_out, [1]=xk_out (positional; the numel check
488+
// below guards a swap, since NH != NKV under GQA).
489+
if (outputs.size() < 2 || !outputs[0].isTensor() || !outputs[1].isTensor()) {
490+
printf("FAIL: expected 2 tensor outputs, got %zu\n", outputs.size());
491+
return false;
492+
}
493+
const auto& xq_t = outputs[0].toTensor();
494+
const auto& xk_t = outputs[1].toTensor();
495+
if (xq_t.numel() != xq_numel || xk_t.numel() != xk_numel) {
496+
printf(
497+
"FAIL: output shapes [%zu,%zu] != expected [%d,%d]\n",
498+
(size_t)xq_t.numel(),
499+
(size_t)xk_t.numel(),
500+
xq_numel,
501+
xk_numel);
502+
return false;
503+
}
504+
const float* xq_out = xq_t.const_data_ptr<float>();
505+
const float* xk_out = xk_t.const_data_ptr<float>();
506+
507+
std::vector<float> gq = load_golden(xq_golden_path, xq_numel);
508+
std::vector<float> gk = load_golden(xk_golden_path, xk_numel);
509+
if (gq.empty() || gk.empty()) {
510+
printf(
511+
"FAIL: could not load goldens %s / %s\n",
512+
xq_golden_path.c_str(),
513+
xk_golden_path.c_str());
514+
return false;
515+
}
516+
517+
// Per-element abs-OR-rel on xq and xk (shared helper, defined above).
518+
float maq = 0.0f, mrq = 0.0f, mak = 0.0f, mrk = 0.0f;
519+
const bool pass_q =
520+
quant_within_tol(xq_out, gq.data(), xq_numel, 1e-3f, 1e-3f, &maq, &mrq);
521+
const bool pass_k =
522+
quant_within_tol(xk_out, gk.data(), xk_numel, 1e-3f, 1e-3f, &mak, &mrk);
523+
const float max_abs_err = std::max(maq, mak);
524+
const float max_rel_err = std::max(mrq, mrk);
525+
526+
printf(
527+
"Max abs error: %e Max rel error: %e (checked %d elements)\n",
528+
max_abs_err,
529+
max_rel_err,
530+
xq_numel + xk_numel);
531+
if (!(pass_q && pass_k)) {
532+
printf("FAIL: apply_rotary_emb exceeds tolerance 1e-3 (abs AND rel)\n");
533+
return false;
534+
}
535+
printf("PASS: apply_rotary_emb test\n");
536+
return true;
537+
}
538+
428539
// Reconstruct _ramp_input bit-for-bit, run the op, compare to the fp64 golden.
429540
static bool test_q4gsw_config(
430541
const Q4gswConfig& cfg,
@@ -1472,6 +1583,37 @@ int main(int argc, char** argv) {
14721583
2048},
14731584
};
14741585

1586+
// apply_rotary_emb on-GPU configs: multi + decode (env-gated,
1587+
// run-if-present).
1588+
struct RopeConfig {
1589+
const char* name;
1590+
const char* model_env;
1591+
const char* xq_env;
1592+
const char* xk_env;
1593+
int S;
1594+
int NH;
1595+
int NKV;
1596+
int HD;
1597+
};
1598+
const RopeConfig rope_configs[] = {
1599+
{"multi",
1600+
"WEBGPU_TEST_ROPE_MODEL",
1601+
"WEBGPU_TEST_ROPE_XQ_GOLDEN",
1602+
"WEBGPU_TEST_ROPE_XK_GOLDEN",
1603+
5,
1604+
8,
1605+
2,
1606+
64},
1607+
{"decode",
1608+
"WEBGPU_TEST_ROPE_DECODE_MODEL",
1609+
"WEBGPU_TEST_ROPE_DECODE_XQ_GOLDEN",
1610+
"WEBGPU_TEST_ROPE_DECODE_XK_GOLDEN",
1611+
1,
1612+
32,
1613+
8,
1614+
64},
1615+
};
1616+
14751617
// SDPA sweep: configs self-discover their sdpa_<name>.pte/.golden.bin under
14761618
// this directory (default "" = the embedded-file root / cwd). Set
14771619
// WEBGPU_TEST_SDPA_DIR to point at the exported .pte directory (e.g. /tmp/).
@@ -1528,6 +1670,15 @@ int main(int argc, char** argv) {
15281670
}
15291671
}
15301672

1673+
for (const auto& c : rope_configs) {
1674+
const char* m = std::getenv(c.model_env);
1675+
const char* xq = std::getenv(c.xq_env);
1676+
const char* xk = std::getenv(c.xk_env);
1677+
if (m && xq && xk && *m && *xq && *xk) {
1678+
ok = test_rope(m, xq, xk, c.S, c.NH, c.NKV, c.HD, c.name) && ok;
1679+
}
1680+
}
1681+
15311682
bool sdpa_ran = false;
15321683
bool sdpa_ok = test_sdpa_sweep(sdpa_dir, &sdpa_ran);
15331684
if (sdpa_ran) {

0 commit comments

Comments
 (0)