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

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

@@ -37,13 +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,
-    EmbeddingParams params,
+    const EmbeddingLayout& layout,
     uint32_t wg_size,
     size_t dispatch_idx,
     WGPUBuffer params_buf) {
@@ -52,17 +76,17 @@ void resize_embedding_q4gsw(
   if (ni == 0) {
     throw std::runtime_error("WebGPU embedding_q4gsw: zero indices");
   }
-  const uint64_t total_blocks = ni * params.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>(params.embed_dim));
+  od.push_back(static_cast<int64_t>(layout.embed_dim));
   g.set_cur_dims(out_id, od);
-  params.num_indices = static_cast<uint32_t>(ni);
-  params.total_blocks = static_cast<uint32_t>(total_blocks);
-  wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &params, sizeof(params));
+  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(
           g.device(),
@@ -166,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);
@@ -223,10 +247,10 @@ void embedding_q4gsw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   WGPUBuffer params_buf = uniform_buffer;
   graph.add_tensor_resize_hook(
       indices_id,
-      [indices_id, out_id, params, wg_size, dispatch_idx, params_buf](
+      [indices_id, out_id, layout, wg_size, dispatch_idx, params_buf](
           WebGPUGraph& g) {
         resize_embedding_q4gsw(
-            g, indices_id, out_id, params, 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;
@@ -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,12 +1690,14 @@ 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");
+// K3: linear/nonlinear-packed reuse must preserve nibble order across resizes.
+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);
+    }
   }
 }