[executorch][webgpu] Test linear embedding dynamic reuse

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<int64_t> idx(static_cast<size_t>(nb) / sizeof(int64_t));
   f.read(reinterpret_cast<char*>(idx.data()), nb);
   ASSERT_EQ(idx.size(), static_cast<size_t>(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<size_t>(n) * kEmbDim;
   ASSERT_EQ(static_cast<size_t>(out.numel()), numel)
-      << "emb N=" << n << " output numel mismatch";
+      << prefix << " N=" << n << " output numel mismatch";
   std::vector<float> got(
       out.const_data_ptr<float>(), out.const_data_ptr<float>() + 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}) {

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("<i8").tofile(
+            os.path.join(out_dir, f"{prefix}.S{n}.idx.bin")
+        )
+        golden.detach().numpy().astype("<f4").tofile(
+            os.path.join(out_dir, f"{prefix}.S{n}.golden.bin")
+        )
+        print(f"  golden {prefix} N={n} (shape {tuple(golden.shape)})")
+
+
 def _export_dynamic_embedding(out_dir: str) -> 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("<i8").tofile(
-            os.path.join(out_dir, f"emb_dyn.S{n}.idx.bin")
-        )
-        g.detach().numpy().astype("<f4").tofile(
-            os.path.join(out_dir, f"emb_dyn.S{n}.golden.bin")
-        )
-        print(f"  golden emb_dyn N={n} (shape {tuple(g.shape)})")
+    _write_embedding_goldens(out_dir, "emb_dyn", weight, scales, group_size, False)
+
+    linear_model = _LinearPackedEmbedding().eval()
+    _export(
+        linear_model,
+        (idx_max,),
+        ({0: n_dim},),
+        os.path.join(out_dir, "emb_dyn_linear.pte"),
+    )
+    _write_embedding_goldens(
+        out_dir,
+        "emb_dyn_linear",
+        linear_model.weight,
+        linear_model.scales,
+        EMB_GROUP,
+        True,
+    )
 
 
 # Dynamic RoPE: xq/xk + freqs all share a dynamic seq-len S.