[ExecuTorch][WebGPU] Preserve Linear and Q4 embedding params across dynamic resize

backends/webgpu/runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp

@@ -37,17 +37,37 @@ static_assert(
     sizeof(EmbeddingParams) == 32,
     "EmbeddingParams must be 32 bytes");
 
+struct EmbeddingLayout {
+  uint32_t embed_dim;
+  uint32_t blocks_per_row;
+  uint32_t group_size;
+  uint32_t groups_per_row;
+  uint32_t bytes_per_row;
+  bool is_linear_weight;
+};
+
+EmbeddingParams make_embedding_params(
+    const EmbeddingLayout& layout,
+    uint32_t num_indices,
+    uint32_t total_blocks) {
+  return {
+      layout.embed_dim,
+      layout.blocks_per_row,
+      num_indices,
+      layout.group_size,
+      layout.groups_per_row,
+      layout.bytes_per_row,
+      total_blocks,
+      layout.is_linear_weight ? 1u : 0u};
+}
+
 // Resize hook body: recompute counts/dispatch; out = indices dims +
 // [embed_dim].
 void resize_embedding_q4gsw(
     WebGPUGraph& g,
     int indices_id,
     int out_id,
-    uint32_t embed_dim,
-    uint32_t blocks_per_row,
-    uint32_t gs_u,
-    uint32_t groups_per_row,
-    uint32_t bytes_per_row,
+    const EmbeddingLayout& layout,
     uint32_t wg_size,
     size_t dispatch_idx,
     WGPUBuffer params_buf) {
@@ -56,22 +76,16 @@ void resize_embedding_q4gsw(
   if (ni == 0) {
     throw std::runtime_error("WebGPU embedding_q4gsw: zero indices");
   }
-  const uint64_t total_blocks = ni * blocks_per_row;
+  const uint64_t total_blocks = ni * layout.blocks_per_row;
   if (total_blocks > UINT32_MAX) {
     throw std::runtime_error(
         "WebGPU embedding_q4gsw: total_blocks exceeds uint32");
   }
   std::vector<int64_t> od = id;
-  od.push_back(static_cast<int64_t>(embed_dim));
+  od.push_back(static_cast<int64_t>(layout.embed_dim));
   g.set_cur_dims(out_id, od);
-  EmbeddingParams p = {};
-  p.embed_dim = embed_dim;
-  p.blocks_per_row = blocks_per_row;
-  p.num_indices = static_cast<uint32_t>(ni);
-  p.group_size = gs_u;
-  p.groups_per_row = groups_per_row;
-  p.bytes_per_row = bytes_per_row;
-  p.total_blocks = static_cast<uint32_t>(total_blocks);
+  EmbeddingParams p = make_embedding_params(
+      layout, static_cast<uint32_t>(ni), static_cast<uint32_t>(total_blocks));
   wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
   g.dispatch_at(dispatch_idx).workgroup_count_x =
       utils::compute_1d_workgroup_count(
@@ -176,15 +190,15 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
       device, static_cast<uint32_t>(total_blocks), wg_size, "embedding_q4gsw");
 
-  EmbeddingParams params = {};
-  params.embed_dim = embed_dim;
-  params.blocks_per_row = blocks_per_row;
-  params.num_indices = num_indices; // std140 layout only; shader derives it
-  params.group_size = static_cast<uint32_t>(group_size);
-  params.groups_per_row = groups_per_row;
-  params.bytes_per_row = bytes_per_row;
-  params.total_blocks = static_cast<uint32_t>(total_blocks);
-  params.is_linear_weight = is_linear ? 1u : 0u;
+  const EmbeddingLayout layout = {
+      embed_dim,
+      blocks_per_row,
+      static_cast<uint32_t>(group_size),
+      groups_per_row,
+      bytes_per_row,
+      is_linear};
+  EmbeddingParams params = make_embedding_params(
+      layout, num_indices, static_cast<uint32_t>(total_blocks));
 
   WGPUBufferDescriptor uniform_desc = {};
   uniform_desc.size = sizeof(EmbeddingParams);
@@ -230,32 +244,13 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
       {bundle.pipeline, bundle.bind_group, workgroup_count, "embedding_q4gsw"});
 
   // Dynamic shapes: recompute counts/dispatch; out = indices + [embed_dim].
-  const uint32_t gs_u = static_cast<uint32_t>(group_size);
   WGPUBuffer params_buf = uniform_buffer;
   graph.add_tensor_resize_hook(
       indices_id,
-      [indices_id,
-       out_id,
-       embed_dim,
-       blocks_per_row,
-       gs_u,
-       groups_per_row,
-       bytes_per_row,
-       wg_size,
-       dispatch_idx,
-       params_buf](WebGPUGraph& g) {
+      [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf](
+          WebGPUGraph& g) {
         resize_embedding_q4gsw(
-            g,
-            indices_id,
-            out_id,
-            embed_dim,
-            blocks_per_row,
-            gs_u,
-            groups_per_row,
-            bytes_per_row,
-            wg_size,
-            dispatch_idx,
-            params_buf);
+            g, indices_id, out_id, layout, wg_size, dispatch_idx, params_buf);
       });
 
   // Graph owns it so the resize hook can rewrite it; freed in the dtor.

backends/webgpu/runtime/ops/linear/Linear.cpp

@@ -33,6 +33,11 @@ static_assert(sizeof(LinearParams) == 16, "LinearParams must be 16 bytes");
 
 constexpr uint32_t kTile = 32u;
 
+LinearParams
+make_linear_params(uint32_t M, uint32_t N, uint32_t K, bool has_bias) {
+  return {M, N, K, has_bias ? 1u : 0u};
+}
+
 // aten.linear (+ optional bias); shared-memory tiled GEMM.
 void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   // args: [input, weight, bias?, out]; out is last. bias (arg 2) is a tensor
@@ -82,11 +87,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
     }
   }
 
-  LinearParams params = {};
-  params.M = M;
-  params.N = N;
-  params.K = K;
-  params.has_bias = has_bias ? 1u : 0u;
+  LinearParams params = make_linear_params(M, N, K, has_bias);
 
   // Bias binding (binding 4); a 4-byte dummy satisfies it when None
   // (WGSL-gated).
@@ -138,7 +139,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   WGPUBuffer params_buf = uniform_buffer;
   graph.add_tensor_resize_hook(
       in_id,
-      [in_id, out_id, M, N, K, dispatch_x, dispatch_idx, params_buf](
+      [in_id, out_id, M, N, K, has_bias, dispatch_x, dispatch_idx, params_buf](
           WebGPUGraph& g) {
         const auto& d = g.cur_dims(in_id);
         const uint64_t numel = utils::numel_of(d);
@@ -152,10 +153,7 @@ void linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
           throw std::runtime_error(
               "WebGPU linear: live M is 0 or exceeds the build-time max");
         }
-        LinearParams p = {};
-        p.M = m;
-        p.N = N;
-        p.K = K;
+        LinearParams p = make_linear_params(m, N, K, has_bias);
         wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
         g.dispatch_at(dispatch_idx).workgroup_count_x = dispatch_x;
         g.dispatch_at(dispatch_idx).workgroup_count_y =

backends/webgpu/test/native/test_dynamic_shape.cpp

@@ -132,6 +132,7 @@ constexpr int kLinK = 64;
 constexpr int kLinAltK = 72;
 constexpr int kLinN = 128;
 constexpr int kLinNShmem = 2048;
+constexpr int kFp32LinearN = 32;
 // Run <prefix> at [m_rows, kLinK] on an already-loaded module (so it can be
 // reused across M without a fresh load), and compare to the golden.
 void run_linear(
@@ -232,6 +233,14 @@ void check_linear_tiled(int m_rows) {
   run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK);
 }
 
+void check_fp32_linear_reused(const char* prefix, int k) {
+  Module module(g_dir + "/" + prefix + ".pte");
+  ASSERT_EQ(module.load_forward(), Error::Ok) << "load " << prefix << ".pte";
+  for (int m_rows : {128, 32, 1, 128}) {
+    run_linear(module, m_rows, prefix, kFp32LinearN, k, 1e-3f);
+  }
+}
+
 constexpr int kQkvNq = 2048;
 constexpr int kQkvNk = 512;
 constexpr int kQkvNv = 512;
@@ -738,31 +747,31 @@ void run_combined_routes(Module& m, 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) {
@@ -935,6 +944,23 @@ TEST(DynamicShape, RmsMul) {
   }
 }
 
+// I0: dynamic fp32 linear preserves bias across repeated resizes.
+TEST(DynamicShape, Fp32LinearVec4BiasedReusedGraph) {
+  check_fp32_linear_reused("dyn_linear_fp32_vec4_bias", 64);
+}
+
+TEST(DynamicShape, Fp32LinearVec4UnbiasedReusedGraph) {
+  check_fp32_linear_reused("dyn_linear_fp32_vec4_no_bias", 64);
+}
+
+TEST(DynamicShape, Fp32LinearTiledBiasedReusedGraph) {
+  check_fp32_linear_reused("dyn_linear_fp32_tiled_bias", 63);
+}
+
+TEST(DynamicShape, Fp32LinearTiledUnbiasedReusedGraph) {
+  check_fp32_linear_reused("dyn_linear_fp32_tiled_no_bias", 63);
+}
+
 // I: dynamic 4-bit quantized linear (prefill GEMM) at several M.
 TEST(DynamicShape, QuantizedLinear) {
   for (int m_rows : {128, 32, 1}) {
@@ -1664,6 +1690,16 @@ TEST(DynamicShape, EmbeddingReusedGraph) {
   }
 }
 
+TEST(DynamicShape, EmbeddingLayoutsReusedGraph) {
+  for (const char* prefix : {"emb_dyn_linear", "emb_dyn_nonlinear"}) {
+    Module m(g_dir + "/" + prefix + ".pte");
+    ASSERT_EQ(m.load_forward(), Error::Ok) << "load " << prefix << ".pte";
+    for (int n : {16, 8, 1, 16}) {
+      run_embedding(m, n, prefix);
+    }
+  }
+}
+
 // L: dynamic RoPE (two outputs) at several seq-len S.
 TEST(DynamicShape, Rope) {
   for (int s : {16, 8, 1}) {