diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index f41adbf09ec..3752e4ee53f 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -193,31 +193,31 @@ void check_sdpa(int s) { constexpr int kEmbDim = 64; // Run emb_dyn at N tokens on an already-loaded module (so it can be reused // across N), and compare to the golden. -void run_embedding(Module& m, int n) { - const std::string b = g_dir + "/emb_dyn.S" + std::to_string(n) + "."; +void run_embedding(Module& m, int n, const char* prefix = "emb_dyn") { + const std::string b = g_dir + "/" + prefix + ".S" + std::to_string(n) + "."; std::ifstream f(b + "idx.bin", std::ios::binary | std::ios::ate); - ASSERT_TRUE(f.good()) << "missing emb_dyn.S" << n; + ASSERT_TRUE(f.good()) << "missing " << prefix << ".S" << n; const std::streamsize nb = f.tellg(); - ASSERT_GE(nb, 0) << "missing emb_dyn.S" << n; + ASSERT_GE(nb, 0) << "missing " << prefix << ".S" << n; f.seekg(0); std::vector idx(static_cast(nb) / sizeof(int64_t)); f.read(reinterpret_cast(idx.data()), nb); ASSERT_EQ(idx.size(), static_cast(n)) - << "wrong emb_dyn idx size S" << n; + << "wrong " << prefix << " idx size S" << n; auto golden = read_bin(b + "golden.bin"); auto t = make_tensor_ptr({n}, std::move(idx)); // int64 (Long) host input auto r = m.forward({EValue(t)}); ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) - << "emb N=" << n + << prefix << " N=" << n << " forward failed (err=" << (r.ok() ? 0 : (int)r.error()) << ")"; const auto& out = r.get()[0].toTensor(); const size_t numel = static_cast(n) * kEmbDim; ASSERT_EQ(static_cast(out.numel()), numel) - << "emb N=" << n << " output numel mismatch"; + << prefix << " N=" << n << " output numel mismatch"; std::vector got( out.const_data_ptr(), out.const_data_ptr() + numel); const float e = max_err(got, golden); - EXPECT_LT(e, 5e-3f) << "emb_dyn N=" << n << " max_err=" << e; + EXPECT_LT(e, 5e-3f) << prefix << " N=" << n << " max_err=" << e; } void check_embedding(int n) { @@ -459,6 +459,15 @@ TEST(DynamicShape, EmbeddingReusedGraph) { } } +// K3: linear-packed reuse must preserve nibble order across resizes. +TEST(DynamicShape, LinearPackedEmbeddingReusedGraph) { + Module m(g_dir + "/emb_dyn_linear.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load emb_dyn_linear.pte"; + for (int n : {16, 8, 1, 16}) { + run_embedding(m, n, "emb_dyn_linear"); + } +} + // L: dynamic RoPE (two outputs) at several seq-len S. TEST(DynamicShape, Rope) { for (int s : {16, 8, 1}) { diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index 969c5e06343..19cb41451e0 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -322,8 +322,53 @@ def cfg(s: int) -> "SdpaConfig": EMB_MAXN = 16 +class _LinearPackedEmbedding(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + packed = torch.arange(EMB_VOCAB * (EMB_DIM // 2), dtype=torch.int64).reshape( + EMB_VOCAB, EMB_DIM // 2 + ) + self.register_buffer("weight", (packed % 256).to(torch.uint8)) + self.register_buffer("scales", torch.ones(EMB_VOCAB, EMB_DIM // EMB_GROUP)) + + def forward(self, indices: torch.Tensor) -> torch.Tensor: + return torch.ops.et_vk.embedding_q4gsw.default( + self.weight, self.scales, EMB_GROUP, indices, True + ) + + +def _write_embedding_goldens( + out_dir: str, + prefix: str, + weight: torch.Tensor, + scales: torch.Tensor, + group_size: int, + is_linear: bool, +) -> None: + for n in [EMB_MAXN, 8, 1]: + idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB + golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, is_linear + ) + if is_linear: + nonlinear_golden = torch.ops.et_vk.embedding_q4gsw.default( + weight, scales, group_size, idx, False + ) + if torch.equal(golden, nonlinear_golden): + raise RuntimeError( + "emb_dyn_linear fixture does not distinguish nibble packing" + ) + idx.detach().numpy().astype(" None: - from executorch.backends.webgpu.test.ops.embedding_q4gsw.test_embedding_q4gsw import ( + from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, Shape, @@ -346,18 +391,23 @@ def _export_dynamic_embedding(out_dir: str) -> None: f.write(et.buffer) print("Exported emb_dyn.pte") weight, scales, group_size = _quant_params(qm) - for n in [EMB_MAXN, 8, 1]: - idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB - g = torch.ops.et_vk.embedding_q4gsw.default( - weight, scales, group_size, idx, False - ) - idx.detach().numpy().astype("