From 91348a24ec9f05271cdb0ad3b580dd4ae681dd0d Mon Sep 17 00:00:00 2001
From: Julian Ng-Thow-Hing <juliannth@meta.com>
Date: Mon, 27 Jul 2026 11:15:46 -0700
Subject: [PATCH] Update

[ghstack-poisoned]
---
 .../webgpu/test/native/test_dynamic_shape.cpp | 21 ++++++---
 .../test_dynamic_shape_export.py              | 46 ++++++++++++++++++-
 2 files changed, 60 insertions(+), 7 deletions(-)

diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp
index f41adbf09ec..1487348f897 100644
--- a/backends/webgpu/test/native/test_dynamic_shape.cpp
+++ b/backends/webgpu/test/native/test_dynamic_shape.cpp
@@ -193,17 +193,17 @@ 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)});
@@ -217,7 +217,7 @@ void run_embedding(Module& m, int n) {
   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}) {
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..e85b3076fdc 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,23 @@ 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 _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,
@@ -359,6 +374,35 @@ def _export_dynamic_embedding(out_dir: str) -> None:
         )
         print(f"  golden emb_dyn N={n} (shape {tuple(g.shape)})")
 
+    linear_model = _LinearPackedEmbedding().eval()
+    _export(
+        linear_model,
+        (idx_max,),
+        ({0: n_dim},),
+        os.path.join(out_dir, "emb_dyn_linear.pte"),
+    )
+    for n in [EMB_MAXN, 8, 1]:
+        idx = (torch.arange(n, dtype=torch.long) * 7) % EMB_VOCAB
+        linear_golden = linear_model(idx)
+        nonlinear_golden = torch.ops.et_vk.embedding_q4gsw.default(
+            linear_model.weight,
+            linear_model.scales,
+            EMB_GROUP,
+            idx,
+            False,
+        )
+        if torch.equal(linear_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"emb_dyn_linear.S{n}.idx.bin")
+        )
+        linear_golden.detach().numpy().astype("<f4").tofile(
+            os.path.join(out_dir, f"emb_dyn_linear.S{n}.golden.bin")
+        )
+        print(f"  golden emb_dyn_linear N={n}")
+
 
 # Dynamic RoPE: xq/xk + freqs all share a dynamic seq-len S.
 ROPE_NH = 8