diff --git a/backends/vulkan/custom_ops_lib.py b/backends/vulkan/custom_ops_lib.py
index 4364f67123d..a074597466a 100644
--- a/backends/vulkan/custom_ops_lib.py
+++ b/backends/vulkan/custom_ops_lib.py
@@ -289,6 +289,67 @@ def linear_dq8ca_q4gsw(
 lib.impl(name, linear_q4gsw, "CompositeExplicitAutograd")
 linear_qc4w_op = getattr(getattr(torch.ops, namespace), name)
 
+
+# Backward of linear_q4gsw wrt input (for on-device LoRA training through a frozen
+# 4-bit base): d_x = d_out @ dequant(W). Reference impl extracts dequant(W) via the
+# forward on an identity so it is layout-agnostic; the runtime dispatches this op to
+# the tiled q4gsw_backward WGSL kernel (contracts over N).
+def linear_q4gsw_backward_impl(
+    d_out: torch.Tensor,
+    weights: torch.Tensor,
+    weight_scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    in_features = int(weights.shape[1]) * 2
+    eye = torch.eye(in_features, dtype=d_out.dtype, device=d_out.device)
+    w_t = linear_q4gsw(eye, weights, weight_scales, group_size)  # [in, out]
+    return d_out @ w_t.t()  # [M, out] @ [out, in] = [M, in]
+
+
+def linear_q4gsw_backward_meta(
+    d_out: torch.Tensor,
+    weights: torch.Tensor,
+    weight_scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    return d_out.new_empty(d_out.shape[:-1] + (int(weights.shape[1]) * 2,))
+
+
+name = "linear_q4gsw_backward"
+lib.define(
+    f"{name}(Tensor d_out, Tensor weights, Tensor weight_scales, int group_size) -> Tensor"
+)
+lib.impl(name, linear_q4gsw_backward_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_q4gsw_backward_meta, "Meta")
+linear_q4gsw_backward_op = getattr(getattr(torch.ops, namespace), name)
+
+
+def linear_q4gsw_setup_context(ctx, inputs, output) -> None:
+    _x, weights, weight_scales, group_size, _bias = inputs
+    ctx.save_for_backward(weights, weight_scales)
+    ctx.group_size = group_size
+
+
+def linear_q4gsw_backward(ctx, grad_out):
+    weights, weight_scales = ctx.saved_tensors
+    d_x = torch.ops.et_vk.linear_q4gsw_backward(
+        grad_out, weights, weight_scales, ctx.group_size
+    )
+    return (
+        d_x,
+        None,
+        None,
+        None,
+        None,
+    )  # grads for (x, weights, scales, group_size, bias)
+
+
+torch.library.register_autograd(
+    f"{namespace}::linear_q4gsw",
+    linear_q4gsw_backward,
+    setup_context=linear_q4gsw_setup_context,
+)
+
 name = "linear_dq8ca_q4gsw"
 lib.define(
     f"""
@@ -1090,3 +1151,171 @@ def rms_norm_impl(
 lib.define(f"{name}(Tensor x, Tensor weight, float eps) -> Tensor")
 lib.impl(name, rms_norm_impl, "CompositeExplicitAutograd")
 rms_norm_op = getattr(getattr(torch.ops, namespace), name)
+
+
+########################
+## fused_ce (training) ##
+########################
+
+
+def fused_ce_impl(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    n_valid: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    mask = labels >= 0
+    safe = labels.clamp(min=0).long()
+    lse = torch.logsumexp(logits, dim=-1)
+    picked = logits.gather(-1, safe[:, None]).squeeze(-1)
+    loss = torch.where(mask, (lse - picked) / n_valid, torch.zeros_like(lse)).sum()
+    softmax = torch.softmax(logits, dim=-1)
+    onehot = torch.nn.functional.one_hot(safe, logits.shape[-1]).to(logits.dtype)
+    dlogits = torch.where(
+        mask[:, None], (softmax - onehot) / n_valid, torch.zeros_like(softmax)
+    )
+    return loss, dlogits
+
+
+def fused_ce_meta(
+    logits: torch.Tensor,
+    labels: torch.Tensor,
+    n_valid: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    return logits.new_empty([]), torch.empty_like(logits)
+
+
+def fused_ce_setup_context(ctx, inputs, output) -> None:
+    ctx.save_for_backward(output[1])
+
+
+def fused_ce_backward(ctx, grad_loss, grad_dlogits):
+    (dlogits,) = ctx.saved_tensors
+    return grad_loss * dlogits, None, None
+
+
+name = "fused_ce"
+lib.define(f"{name}(Tensor logits, Tensor labels, float n_valid) -> (Tensor, Tensor)")
+lib.impl(name, fused_ce_impl, "CompositeExplicitAutograd")
+lib.impl(name, fused_ce_meta, "Meta")
+torch.library.register_autograd(
+    f"{namespace}::{name}", fused_ce_backward, setup_context=fused_ce_setup_context
+)
+fused_ce_op = getattr(getattr(torch.ops, namespace), name)
+
+
+###########################
+## adamw_step (training) ##
+###########################
+
+
+def adamw_step_impl(
+    param: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    grad: torch.Tensor,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    bias_correction1: float,
+    bias_correction2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    param.mul_(1.0 - lr * weight_decay)
+    m.mul_(beta1).add_(grad, alpha=1.0 - beta1)
+    v.mul_(beta2).addcmul_(grad, grad, value=1.0 - beta2)
+    mhat = m / bias_correction1
+    denom = (v / bias_correction2).sqrt() + eps
+    param.addcdiv_(mhat, denom, value=-lr)
+    return param, m, v
+
+
+def adamw_step_meta(
+    param: torch.Tensor,
+    m: torch.Tensor,
+    v: torch.Tensor,
+    grad: torch.Tensor,
+    lr: float,
+    beta1: float,
+    beta2: float,
+    eps: float,
+    weight_decay: float,
+    bias_correction1: float,
+    bias_correction2: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    return param, m, v
+
+
+name = "adamw_step"
+lib.define(
+    f"{name}(Tensor(a!) param, Tensor(b!) m, Tensor(c!) v, Tensor grad, float lr, float beta1, float beta2, float eps, float weight_decay, float bias_correction1, float bias_correction2) -> (Tensor(a!), Tensor(b!), Tensor(c!))"
+)
+lib.impl(name, adamw_step_impl, "CompositeExplicitAutograd")
+lib.impl(name, adamw_step_meta, "Meta")
+adamw_step_op = getattr(getattr(torch.ops, namespace), name)
+
+
+# STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
+def linear_dW_impl(
+    d_out: torch.Tensor,
+    x: torch.Tensor,
+) -> torch.Tensor:
+    return d_out.reshape(-1, d_out.shape[-1]).t() @ x.reshape(-1, x.shape[-1])
+
+
+def linear_dW_meta(
+    d_out: torch.Tensor,
+    x: torch.Tensor,
+) -> torch.Tensor:
+    return d_out.new_empty((d_out.shape[-1], x.shape[-1]))
+
+
+name = "linear_dW"
+lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
+lib.impl(name, linear_dW_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_dW_meta, "Meta")
+linear_dW_op = getattr(getattr(torch.ops, namespace), name)
+
+
+##################
+## q4gsw_requant ##
+##################
+
+
+# STE re-quant of fp32 latent weights into the frozen-scale 4-bit codes.
+def q4gsw_requant_impl(
+    latent: torch.Tensor,
+    scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    n, k = latent.shape
+    group_idx = torch.arange(k, device=latent.device) // group_size
+    scale_full = scales.t()[:, group_idx]  # [N, K]: scales[k // group_size, n]
+    nonzero = scale_full != 0
+    safe = torch.where(nonzero, scale_full, torch.ones_like(scale_full))
+    q = torch.round(latent / safe)
+    q = torch.where(nonzero, q, torch.zeros_like(q))
+    codes = (torch.clamp(q, -8, 7).to(torch.int32) + 8) & 0xF  # [N, K] in 0..15
+    k_packed = (k + 1) // 2
+    packed = torch.zeros((n, k_packed), dtype=torch.uint8, device=latent.device)
+    packed[:, :] = codes[:, 0::2].to(torch.uint8)
+    if k > 1:
+        high = codes[:, 1::2].to(torch.uint8)
+        packed[:, : high.shape[1]] |= high << 4
+    return packed
+
+
+def q4gsw_requant_meta(
+    latent: torch.Tensor,
+    scales: torch.Tensor,
+    group_size: int,
+) -> torch.Tensor:
+    n, k = latent.shape
+    return latent.new_empty((n, (k + 1) // 2), dtype=torch.uint8)
+
+
+name = "q4gsw_requant"
+lib.define(f"{name}(Tensor latent, Tensor scales, int group_size) -> Tensor")
+lib.impl(name, q4gsw_requant_impl, "CompositeExplicitAutograd")
+lib.impl(name, q4gsw_requant_meta, "Meta")
+q4gsw_requant_op = getattr(getattr(torch.ops, namespace), name)
diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py
index 105bcea89a6..78270be4373 100644
--- a/backends/vulkan/op_registry.py
+++ b/backends/vulkan/op_registry.py
@@ -462,6 +462,15 @@ def register_quantizedlinear_cpp_ops():
     )
 
 
+@update_features(exir_ops.edge.et_vk.linear_q4gsw_backward.default)
+def register_linear_q4gsw_backward():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+        supports_prepacking=True,
+    )
+
+
 @update_features(exir_ops.edge.et_vk.linear_dq8ca_q4gsw.default)
 def register_linear_dq8ca_q4gsw():
     return OpFeatures(
@@ -1746,6 +1755,74 @@ def register_rms_norm():
     )
 
 
+# =============================================================================
+# FusedCe.cpp (training)
+# =============================================================================
+
+
+@update_features(exir_ops.edge.et_vk.fused_ce.default)
+def register_fused_ce():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=[utils.FP_T, utils.INT_T, utils.NONE_T],
+        outputs_dtypes=[utils.FP_T, utils.FP_T],
+    )
+
+
+@update_features(
+    [
+        exir_ops.edge.aten.ne.Scalar,
+        exir_ops.edge.aten.lt.Scalar,
+        exir_ops.edge.aten.le.Scalar,
+        exir_ops.edge.aten.ge.Scalar,
+        exir_ops.edge.aten.gt.Scalar,
+    ]
+)
+def register_compare_scalar_ops():
+    return OpFeatures(
+        inputs_storage=utils.ANY_STORAGE,
+        inputs_dtypes=utils.FP_INT_T,
+        outputs_dtypes=utils.BOOL_T,
+        supports_resize=True,
+        supports_highdim=True,
+    )
+
+
+@update_features(exir_ops.edge.aten.logical_not.default)
+def register_logical_not():
+    return OpFeatures(
+        inputs_storage=utils.ANY_STORAGE,
+        inputs_dtypes=utils.BOOL_T,
+        supports_resize=True,
+        supports_highdim=True,
+    )
+
+
+@update_features("et_vk::adamw_step")
+def register_adamw_step():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+    )
+
+
+@update_features(exir_ops.edge.et_vk.linear_dW.default)
+def register_linear_dW():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+        supports_prepacking=True,
+    )
+
+
+@update_features(exir_ops.edge.et_vk.q4gsw_requant.default)
+def register_q4gsw_requant():
+    return OpFeatures(
+        inputs_storage=utils.CONTIGUOUS_ANY,
+        inputs_dtypes=utils.FP_T,
+    )
+
+
 #######################
 ## Utility functions ##
 #######################
diff --git a/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl
new file mode 100644
index 00000000000..ccb7b34bcb7
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.glsl
@@ -0,0 +1,54 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type(STORAGE)}
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "rw", "t_param", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "rw", "t_m", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "rw", "t_v", DTYPE, STORAGE)}
+${layout_declare_tensor(B, "r", "t_grad", DTYPE, STORAGE)}
+
+layout(push_constant) uniform restrict Block {
+  int numel;
+  float lr;
+  float beta1;
+  float beta2;
+  float eps;
+  float weight_decay;
+  float bias_correction1;
+  float bias_correction2;
+};
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  const int i = int(gl_GlobalInvocationID.x);
+  if (i >= numel) {
+    return;
+  }
+  T g = t_grad[i];
+  T p = t_param[i];
+  p = p - lr * weight_decay * p;
+  T m = beta1 * t_m[i] + (1.0 - beta1) * g;
+  T v = beta2 * t_v[i] + (1.0 - beta2) * g * g;
+  t_m[i] = m;
+  t_v[i] = v;
+  T mhat = m / bias_correction1;
+  T vhat = v / bias_correction2;
+  t_param[i] = p - lr * mhat / (sqrt(vhat) + eps);
+}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml
new file mode 100644
index 00000000000..114a257cfad
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml
@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+adamw_step:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: adamw_step
diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl
new file mode 100644
index 00000000000..044301d5330
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.glsl
@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+#define PRECISION ${PRECISION}
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type(STORAGE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_dlogits", DTYPE, "buffer")}
+${layout_declare_tensor(B, "w", "t_loss_partial", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_logits", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_labels", "int", "buffer")}
+
+${layout_declare_ubo(B, "int", "vocab")}
+${layout_declare_ubo(B, "int", "n_rows")}
+${layout_declare_ubo(B, "float", "n_valid")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+#define NWORKERS 64
+
+shared float red_m[NWORKERS];
+shared float red_l[NWORKERS];
+
+// Fused cross-entropy: one workgroup per row cooperatively reduces the vocab
+// dimension with a single-pass online softmax (running max + rescaled running
+// sum), then writes per-row loss and dlogits. Rows with label < 0 are masked.
+void main() {
+  const uint row = gl_GlobalInvocationID.y;
+  if (int(row) >= n_rows) {
+    return;
+  }
+
+  const uint tid = gl_LocalInvocationID.x;
+  const uint V = uint(vocab);
+  const uint base = row * V;
+  const int lbl = t_labels[row];
+
+  if (lbl < 0) {
+    for (uint j = tid; j < V; j += NWORKERS) {
+      t_dlogits[base + j] = T(0);
+    }
+    if (tid == 0u) {
+      t_loss_partial[row] = T(0);
+    }
+    return;
+  }
+
+  // Single read pass: maintain a running max m and the sum l of exp(x - m),
+  // rescaling l whenever a larger value updates m. Finite -3.4e38 init.
+  float m = -3.4e38;
+  float l = 0.0;
+  for (uint j = tid; j < V; j += NWORKERS) {
+    float x = float(t_logits[base + j]);
+    if (x > m) {
+      l = l * exp(m - x) + 1.0;
+      m = x;
+    } else {
+      l = l + exp(x - m);
+    }
+  }
+  red_m[tid] = m;
+  red_l[tid] = l;
+  memoryBarrierShared();
+  barrier();
+
+  // Tree-combine the (m, l) pairs: m becomes the max, l is rescaled to it.
+  for (uint s = NWORKERS / 2u; s > 0u; s >>= 1u) {
+    if (tid < s) {
+      float ma = red_m[tid];
+      float la = red_l[tid];
+      float mb = red_m[tid + s];
+      float lb = red_l[tid + s];
+      float mm = max(ma, mb);
+      red_m[tid] = mm;
+      red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm);
+    }
+    memoryBarrierShared();
+    barrier();
+  }
+
+  const float row_max = red_m[0];
+  const float denom = red_l[0];
+  const float inv = 1.0 / denom;
+  const float scale = 1.0 / n_valid;
+
+  if (tid == 0u) {
+    const float lse = row_max + log(denom);
+    t_loss_partial[row] = T((lse - float(t_logits[base + uint(lbl)])) * scale);
+  }
+
+  for (uint j = tid; j < V; j += NWORKERS) {
+    float g = exp(float(t_logits[base + j]) - row_max) * inv * scale;
+    if (j == uint(lbl)) {
+      g = g - scale;
+    }
+    t_dlogits[base + j] = T(g);
+  }
+}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml
new file mode 100644
index 00000000000..8cf9cdac04e
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce.yaml
@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+fused_ce:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: fused_ce_buffer
diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl
new file mode 100644
index 00000000000..c61820e074c
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.glsl
@@ -0,0 +1,55 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+#define PRECISION ${PRECISION}
+#define T ${buffer_scalar_type(DTYPE)}
+
+${define_active_storage_type(STORAGE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_loss", DTYPE, "buffer")}
+${layout_declare_tensor(B, "r", "t_loss_partial", DTYPE, "buffer")}
+
+${layout_declare_ubo(B, "int", "n_rows")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+#define NWORKERS 64
+
+shared float red[NWORKERS];
+
+// Self-contained [N] -> [1] tree-sum of the per-row losses in one workgroup, so
+// fused_ce carries no cross-op reduce dependency.
+void main() {
+  const uint tid = gl_LocalInvocationID.x;
+
+  float s = 0.0;
+  for (uint j = tid; j < uint(n_rows); j += NWORKERS) {
+    s += float(t_loss_partial[j]);
+  }
+  red[tid] = s;
+  memoryBarrierShared();
+  barrier();
+
+  for (uint k = NWORKERS / 2u; k > 0u; k >>= 1u) {
+    if (tid < k) {
+      red[tid] += red[tid + k];
+    }
+    memoryBarrierShared();
+    barrier();
+  }
+
+  if (tid == 0u) {
+    t_loss[0] = T(red[0]);
+  }
+}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml
new file mode 100644
index 00000000000..cc1bd16cebb
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/fused_ce_sum.yaml
@@ -0,0 +1,15 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+fused_ce_sum:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: fused_ce_sum_buffer
diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl
new file mode 100644
index 00000000000..1a8ed94031b
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.glsl
@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define T ${texel_load_component_type(DTYPE, STORAGE)}
+
+#define TILE_N 4
+#define TILE_K 4
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_dw", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_dout", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_x", DTYPE, STORAGE, is_scalar_array=True)}
+
+${layout_declare_ubo(B, "ivec4", "dout_sizes")}
+${layout_declare_ubo(B, "ivec4", "x_sizes")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+void main() {
+  // dW[N, K] = sum_m d_out[m, N] * x[m, K]; contraction over the flattened M.
+  const int N = dout_sizes.x;
+  const int M = dout_sizes.y * dout_sizes.z * dout_sizes.w;
+  const int K = x_sizes.x;
+
+  const int nnt = (N + TILE_N - 1) / TILE_N;
+  const int nkt = (K + TILE_K - 1) / TILE_K;
+  const int tiles = nnt * nkt;
+
+  const int tile_idx = int(gl_GlobalInvocationID.x);
+  if (tile_idx >= tiles) {
+    return;
+  }
+
+  const int n0 = (tile_idx / nkt) * TILE_N;
+  const int k0 = (tile_idx % nkt) * TILE_K;
+
+  float acc[TILE_N * TILE_K];
+  for (int i = 0; i < TILE_N * TILE_K; ++i) {
+    acc[i] = 0.0;
+  }
+
+  for (int m = 0; m < M; ++m) {
+    float dout_reg[TILE_N];
+    for (int nl = 0; nl < TILE_N; ++nl) {
+      const int n_eff = min(n0 + nl, N - 1);
+      dout_reg[nl] = float(t_dout[m * N + n_eff]);
+    }
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int k_eff = min(k0 + kl, K - 1);
+      const float xv = float(t_x[m * K + k_eff]);
+      for (int nl = 0; nl < TILE_N; ++nl) {
+        acc[nl * TILE_K + kl] += dout_reg[nl] * xv;
+      }
+    }
+  }
+
+  for (int nl = 0; nl < TILE_N; ++nl) {
+    const int n = n0 + nl;
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int k = k0 + kl;
+      if (n < N && k < K) {
+        t_dw[n * K + k] = T(acc[nl * TILE_K + kl]);
+      }
+    }
+  }
+}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml
new file mode 100644
index 00000000000..a2d38101f3d
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/linear_dW.yaml
@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+linear_dW:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: linear_dW
diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl
new file mode 100644
index 00000000000..d89bb83feed
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.glsl
@@ -0,0 +1,104 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+#define T ${texel_load_component_type(DTYPE, STORAGE)}
+
+#define TILE_M 4
+#define TILE_K 4
+
+${define_required_extensions(STORAGE, DTYPE)}
+${define_required_extensions("buffer", DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_dx", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_dout", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_q4_weights", "int", "buffer", is_scalar_array=False, vec_size=4)}
+${layout_declare_tensor(B, "r", "t_scales", DTYPE, "buffer", is_scalar_array=True)}
+
+${layout_declare_ubo(B, "ivec4", "dout_sizes")}
+${layout_declare_ubo(B, "ivec4", "dx_sizes")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+${layout_declare_spec_const(C, "int", "group_size", "32")}
+
+// d_x[M, K] = d_out[M, N] @ dequant(W)[N, K], contracting over N.
+// dequant(W[n, k]) = (code - 8) * scale, with code read from the same W_4X8
+// block-packed weight the forward reads (mirrors q4gsw_linear_gemm__w_4x8.glsl).
+void main() {
+  const int N = dout_sizes.x;
+  const int M = dout_sizes.y * dout_sizes.z * dout_sizes.w;
+  const int K = dx_sizes.x;
+
+  const int nmt = (M + TILE_M - 1) / TILE_M;
+  const int nkt = (K + TILE_K - 1) / TILE_K;
+  const int tiles = nmt * nkt;
+
+  const int tile_idx = int(gl_GlobalInvocationID.x);
+  if (tile_idx >= tiles) {
+    return;
+  }
+
+  const int m0 = (tile_idx / nkt) * TILE_M;
+  const int k0 = (tile_idx % nkt) * TILE_K;
+
+  // K and N are multiples of 4 (prepack guarantees), so k0 is 4-aligned: the
+  // tile's 4 K lanes are byte b = kl of one k4 group and share one scale group.
+  const int k4 = k0 >> 2;
+  const int N4 = (N + 3) >> 2;
+  const int N4_padded = (N4 + 1) & ~1;
+  const int N8 = N4_padded >> 1;
+  const int group = k0 / group_size;
+
+  float acc[TILE_M * TILE_K];
+  for (int i = 0; i < TILE_M * TILE_K; ++i) {
+    acc[i] = 0.0;
+  }
+
+  for (int n = 0; n < N; ++n) {
+    float dout_reg[TILE_M];
+    for (int ml = 0; ml < TILE_M; ++ml) {
+      const int m_eff = min(m0 + ml, M - 1);
+      dout_reg[ml] = float(t_dout[m_eff * N + n]);
+    }
+
+    // W_4X8 address for column n: ivec4 at (k4, n8); component by (n4 parity,
+    // n-in-tile half); low/high nibble by n parity (even-N low, odd-N high).
+    const int n4 = n >> 2;
+    const int ni = n & 3;
+    const int n8 = n4 >> 1;
+    const int comp = (n4 & 1) * 2 + (ni >> 1);
+    const int nib_hi = (ni & 1) * 4;
+    const ivec4 w_block = t_q4_weights[k4 * N8 + n8];
+    const int w_int = w_block[comp];
+    const float scale = float(t_scales[group * N + n]);
+
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int code = int((uint(w_int) >> (8 * kl + nib_hi)) & 0xFu);
+      const float dq = float(code - 8) * scale;
+      for (int ml = 0; ml < TILE_M; ++ml) {
+        acc[ml * TILE_K + kl] += dout_reg[ml] * dq;
+      }
+    }
+  }
+
+  for (int ml = 0; ml < TILE_M; ++ml) {
+    const int m = m0 + ml;
+    for (int kl = 0; kl < TILE_K; ++kl) {
+      const int k = k0 + kl;
+      if (m < M && k < K) {
+        t_dx[m * K + k] = T(acc[ml * TILE_K + kl]);
+      }
+    }
+  }
+}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml
new file mode 100644
index 00000000000..b2bb8ea1173
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_backward.yaml
@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+q4gsw_backward:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: q4gsw_backward
diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl
new file mode 100644
index 00000000000..896da905648
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.glsl
@@ -0,0 +1,97 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#version 450 core
+
+#define PRECISION ${PRECISION}
+
+${define_required_extensions(STORAGE, DTYPE)}
+
+layout(std430) buffer;
+
+${layout_declare_tensor(B, "w", "t_packed", "int", "buffer", is_scalar_array=False, vec_size=4)}
+${layout_declare_tensor(B, "r", "t_latent", DTYPE, STORAGE, is_scalar_array=True)}
+${layout_declare_tensor(B, "r", "t_scales", DTYPE, "buffer", is_scalar_array=True)}
+
+${layout_declare_ubo(B, "ivec4", "latent_sizes")}
+
+layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
+
+${layout_declare_spec_const(C, "int", "group_size", "32")}
+
+// STE re-quant of fp32 latent [N, K] into the W_4X8 block-packed 4-bit codes
+// the forward reads (mirrors pack_q4_linear_weight__w_4x8.glsl). Each thread
+// writes one 4K x 8N block (an ivec4) at (k4, n8).
+uint quant_nibble(const int n, const int k, const int N, const int K) {
+  const float s = float(t_scales[(k / group_size) * N + n]);
+  // roundEven + clamp-in-float match torch.round (half-to-even) then clamp.
+  float qf = 0.0;
+  if (s != 0.0) {
+    qf = clamp(roundEven(float(t_latent[n * K + k]) / s), -8.0, 7.0);
+  }
+  return uint(int(qf) + 8) & 0xFu;
+}
+
+// Pack a 4K x 4N tile into the (packed_x, packed_y) int pair. Byte b holds one
+// (even-N low nibble, odd-N high nibble) pair at K = k4*4 + b; rows N0,N1 go to
+// packed_x, rows N2,N3 to packed_y.
+void pack_tile(
+    out uint packed_x,
+    out uint packed_y,
+    const int k4,
+    const int n4,
+    const int N,
+    const int K) {
+  packed_x = 0u;
+  packed_y = 0u;
+  for (int ni = 0; ni < 4; ++ni) {
+    const int n = n4 * 4 + ni;
+    for (int b = 0; b < 4; ++b) {
+      const uint code = quant_nibble(n, k4 * 4 + b, N, K);
+      const int shift = 8 * b + (ni & 1) * 4;
+      if (ni < 2) {
+        packed_x |= code << shift;
+      } else {
+        packed_y |= code << shift;
+      }
+    }
+  }
+}
+
+void main() {
+  const int k4 = int(gl_GlobalInvocationID.x);
+  const int n8 = int(gl_GlobalInvocationID.y);
+
+  const int K = latent_sizes.x;
+  const int N = latent_sizes.y;
+  const int K4 = K >> 2;
+  const int N4 = (N + 3) >> 2;
+  const int N8 = (N4 + 1) >> 1;
+
+  if (k4 >= K4 || n8 >= N8) {
+    return;
+  }
+
+  const int n4_a = 2 * n8;
+  const int n4_b = n4_a + 1;
+
+  uint packed_x_a = 0u;
+  uint packed_y_a = 0u;
+  // OOB upper tile (odd N4 boundary) is the bias-zero pattern, matching the
+  // forward's prepack padding so the whole block is always readable.
+  uint packed_x_b = 0x88888888u;
+  uint packed_y_b = 0x88888888u;
+
+  pack_tile(packed_x_a, packed_y_a, k4, n4_a, N, K);
+  if (n4_b < N4) {
+    pack_tile(packed_x_b, packed_y_b, k4, n4_b, N, K);
+  }
+
+  t_packed[k4 * N8 + n8] = ivec4(
+      int(packed_x_a), int(packed_y_a), int(packed_x_b), int(packed_y_b));
+}
diff --git a/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml
new file mode 100644
index 00000000000..275bec5bae2
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/glsl/q4gsw_requant.yaml
@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+q4gsw_requant:
+  parameter_names_with_default_values:
+    DTYPE: float
+    STORAGE: buffer
+  generate_variant_forall:
+    STORAGE:
+      - VALUE: buffer
+    DTYPE:
+      - VALUE: float
+  shader_variants:
+    - NAME: q4gsw_requant
diff --git a/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp
new file mode 100644
index 00000000000..489e95a773d
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/impl/AdamwStep.cpp
@@ -0,0 +1,126 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+#include <array>
+
+namespace vkcompute {
+
+void check_adamw_step_args(
+    ComputeGraph& graph,
+    const ValueRef param,
+    const ValueRef m,
+    const ValueRef v,
+    const ValueRef grad) {
+  VK_CHECK_COND(graph.dtype_of(param) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(m) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(v) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(grad) == vkapi::kFloat);
+
+  const int32_t numel = graph.numel_of(param);
+  VK_CHECK_COND(graph.numel_of(m) == numel);
+  VK_CHECK_COND(graph.numel_of(v) == numel);
+  VK_CHECK_COND(graph.numel_of(grad) == numel);
+}
+
+void add_adamw_step_node(
+    ComputeGraph& graph,
+    const ValueRef param,
+    const ValueRef m,
+    const ValueRef v,
+    const ValueRef grad,
+    const ValueRef lr,
+    const ValueRef beta1,
+    const ValueRef beta2,
+    const ValueRef eps,
+    const ValueRef weight_decay,
+    const ValueRef bias_correction1,
+    const ValueRef bias_correction2) {
+  check_adamw_step_args(graph, param, m, v, grad);
+
+  const float bc1 = graph.extract_scalar<float>(bias_correction1);
+  const float bc2 = graph.extract_scalar<float>(bias_correction2);
+  VK_CHECK_COND(bc1 != 0.0f);
+  VK_CHECK_COND(bc2 != 0.0f);
+
+  // Split into <=16-byte push-constant entries (PushConstantData.h limit).
+  const std::array<float, 4> pc0 = {
+      graph.extract_scalar<float>(lr),
+      graph.extract_scalar<float>(beta1),
+      graph.extract_scalar<float>(beta2),
+      graph.extract_scalar<float>(eps)};
+  const std::array<float, 3> pc1 = {
+      graph.extract_scalar<float>(weight_decay), bc1, bc2};
+
+  std::string kernel_name("adamw_step");
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(param));
+  add_dtype_suffix(kernel_name, graph.dtype_of(param));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      default_pick_global_wg_size,
+      default_pick_local_wg_size,
+      // Inputs and Outputs
+      {{param, vkapi::kReadWrite},
+       {m, vkapi::kReadWrite},
+       {v, vkapi::kReadWrite},
+       {grad, vkapi::kRead}},
+      // Shader params buffers
+      {},
+      // Push Constants
+      {graph.numel_pc_of(param),
+       PushConstantDataInfo(pc0.data(), sizeof(pc0)),
+       PushConstantDataInfo(pc1.data(), sizeof(pc1))},
+      // Specialization Constants
+      {},
+      // Resize Args
+      {},
+      // Resizing Logic
+      nullptr));
+}
+
+void adamw_step(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  int arg_idx = 0;
+  const ValueRef param = args[arg_idx++];
+  const ValueRef m = args[arg_idx++];
+  const ValueRef v = args[arg_idx++];
+  const ValueRef grad = args[arg_idx++];
+  const ValueRef lr = args[arg_idx++];
+  const ValueRef beta1 = args[arg_idx++];
+  const ValueRef beta2 = args[arg_idx++];
+  const ValueRef eps = args[arg_idx++];
+  const ValueRef weight_decay = args[arg_idx++];
+  const ValueRef bias_correction1 = args[arg_idx++];
+  const ValueRef bias_correction2 = args[arg_idx++];
+  add_adamw_step_node(
+      graph,
+      param,
+      m,
+      v,
+      grad,
+      lr,
+      beta1,
+      beta2,
+      eps,
+      weight_decay,
+      bias_correction1,
+      bias_correction2);
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.adamw_step.default, adamw_step);
+}
+
+} // namespace vkcompute
diff --git a/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp
new file mode 100644
index 00000000000..f66e0e4fc46
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/impl/FusedCe.cpp
@@ -0,0 +1,197 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+using namespace utils;
+
+utils::uvec3 fused_ce_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)shader;
+  (void)resize_args;
+  const ValueRef loss_partial = args.at(0).refs.at(1);
+  return {
+      1u, utils::safe_downcast<uint32_t>(graph->numel_of(loss_partial)), 1u};
+}
+
+utils::uvec3 fused_ce_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {64u, 1u, 1u};
+}
+
+utils::uvec3 fused_ce_sum_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)args;
+  (void)resize_args;
+  return {1u, 1u, 1u};
+}
+
+utils::uvec3 fused_ce_sum_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {64u, 1u, 1u};
+}
+
+void resize_fused_ce_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)resize_args;
+  const ValueRef dlogits = args.at(0).refs.at(0);
+  const ValueRef loss_partial = args.at(0).refs.at(1);
+  const ValueRef logits = args.at(1).refs.at(0);
+
+  const std::vector<int64_t> logits_sizes = graph->sizes_of(logits);
+  graph->virtual_resize(dlogits, logits_sizes);
+  graph->virtual_resize(loss_partial, {logits_sizes.at(0)});
+}
+
+void resize_fused_ce_sum_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)resize_args;
+  const ValueRef loss = args.at(0).refs.at(0);
+  // Rank-0 scalar, matching fused_ce_meta's logits.new_empty([]).
+  graph->virtual_resize(loss, {});
+}
+
+void fused_ce(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  int arg_idx = 0;
+  const ValueRef logits = args[arg_idx++];
+  const ValueRef labels = args[arg_idx++];
+  const ValueRef n_valid_ref = args[arg_idx++];
+  const ValueRef out_tuple_ref = args[arg_idx];
+
+  // Read loss/dlogits via short-lived get_value_list() temporaries: no get_*()
+  // pointer may be alive when loss_partial is added below, else adding Values
+  // can reallocate values_ and check_no_active_value_ptrs() throws.
+  const ValueRef loss = graph.get_value_list(out_tuple_ref)->at(0);
+  const ValueRef dlogits = graph.get_value_list(out_tuple_ref)->at(1);
+
+  VK_CHECK_COND(
+      graph.is_buffer_storage(logits) && graph.is_buffer_storage(dlogits),
+      "fused_ce: logits and dlogits must use buffer storage");
+  VK_CHECK_COND(
+      graph.dim_of(logits) == 2, "fused_ce: logits must be 2D [n_rows, vocab]");
+  VK_CHECK_COND(
+      graph.sizes_of(dlogits) == graph.sizes_of(logits),
+      "fused_ce: dlogits must match logits shape");
+  VK_CHECK_COND(
+      graph.dtype_of(labels) == vkapi::kInt, "fused_ce: labels must be int32");
+  VK_CHECK_COND(graph.dim_of(labels) == 1, "fused_ce: labels must be 1D [N]");
+  VK_CHECK_COND(
+      graph.size_at<int64_t>(0, labels) == graph.size_at<int64_t>(0, logits),
+      "fused_ce: labels length must equal number of rows");
+  VK_CHECK_COND(graph.numel_of(loss) == 1, "fused_ce: loss must be a scalar");
+
+  const int32_t n_rows = graph.size_at<int32_t>(0, logits);
+  const int32_t vocab = graph.size_at<int32_t>(1, logits);
+  const float n_valid = graph.extract_scalar<float>(n_valid_ref);
+
+  // One workgroup per row; the Vulkan spec guarantees maxComputeWorkGroupCount
+  // >= 65535 per dimension.
+  VK_CHECK_COND(
+      n_rows <= 65535, "fused_ce: n_rows exceeds max workgroup count");
+
+  // Per-row loss contributions, reduced to the scalar loss by the sum node.
+  TmpTensor loss_partial(
+      &graph, {n_rows}, graph.dtype_of(logits), utils::kBuffer);
+
+  std::string kernel_name = "fused_ce";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(dlogits));
+  add_dtype_suffix(kernel_name, graph.dtype_of(dlogits));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      fused_ce_global_wg_size,
+      fused_ce_local_wg_size,
+      // Inputs and Outputs
+      {{{dlogits, loss_partial}, vkapi::kWrite},
+       {{logits, labels}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.create_params_buffer(vocab),
+       graph.create_params_buffer(n_rows),
+       graph.create_params_buffer(n_valid)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_fused_ce_node));
+
+  std::string sum_kernel_name = "fused_ce_sum";
+  sum_kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(sum_kernel_name, graph.storage_type_of(loss));
+  add_dtype_suffix(sum_kernel_name, graph.dtype_of(loss));
+
+  // Separate dispatch node so the loss_partial write (above) is ordered before
+  // this read: the graph inserts a per-tensor pipeline barrier between them.
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(sum_kernel_name),
+      fused_ce_sum_global_wg_size,
+      fused_ce_sum_local_wg_size,
+      // Inputs and Outputs
+      {{loss, vkapi::kWrite}, {loss_partial, vkapi::kRead}},
+      // Shader params buffers
+      {graph.create_params_buffer(n_rows)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_fused_ce_sum_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.fused_ce.default, fused_ce);
+}
+
+} // namespace vkcompute
diff --git a/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp b/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp
new file mode 100644
index 00000000000..1ab5e6976ca
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/impl/LinearDw.cpp
@@ -0,0 +1,112 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+void resize_linear_dW_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)resize_args;
+  const ValueRef dW = args.at(0).refs.at(0);
+  const ValueRef d_out = args.at(1).refs.at(0);
+  const ValueRef x = args.at(1).refs.at(1);
+  const int64_t N = graph->size_at<int64_t>(-1, d_out);
+  const int64_t K = graph->size_at<int64_t>(-1, x);
+  graph->virtual_resize(dW, {N, K});
+}
+
+utils::uvec3 linear_dW_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)shader;
+  (void)resize_args;
+  const ValueRef dW = args.at(0).refs.at(0);
+  const uint32_t N = graph->size_at<uint32_t>(-2, dW);
+  const uint32_t K = graph->size_at<uint32_t>(-1, dW);
+  const uint32_t tiles = utils::div_up_4(N) * utils::div_up_4(K);
+  return {tiles, 1u, 1u};
+}
+
+utils::uvec3 linear_dW_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {64u, 1u, 1u};
+}
+
+void linear_dW(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  int32_t i = 0;
+  const ValueRef d_out = args.at(i++);
+  const ValueRef x = args.at(i++);
+  const ValueRef dW = args.at(i);
+
+  VK_CHECK_COND(graph.dtype_of(d_out) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(x) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(dW) == vkapi::kFloat);
+  VK_CHECK_COND(graph.is_buffer_storage(d_out));
+  VK_CHECK_COND(graph.is_buffer_storage(x));
+  VK_CHECK_COND(graph.is_buffer_storage(dW));
+
+  const int64_t N = graph.size_at<int64_t>(-1, d_out);
+  const int64_t K = graph.size_at<int64_t>(-1, x);
+  VK_CHECK_COND(N > 0 && K > 0);
+  VK_CHECK_COND(graph.numel_of(d_out) % N == 0);
+  VK_CHECK_COND(graph.numel_of(x) % K == 0);
+  VK_CHECK_COND(graph.numel_of(d_out) / N == graph.numel_of(x) / K);
+
+  const uint32_t tiles = utils::div_up_4(static_cast<uint32_t>(N)) *
+      utils::div_up_4(static_cast<uint32_t>(K));
+  VK_CHECK_COND(
+      (tiles + 63u) / 64u <= 65535u,
+      "linear_dW: tile count exceeds max workgroup count");
+
+  std::string kernel_name = "linear_dW";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(dW));
+  add_dtype_suffix(kernel_name, graph.dtype_of(dW));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      linear_dW_global_wg_size,
+      linear_dW_local_wg_size,
+      // Inputs and Outputs
+      {{dW, vkapi::kWrite}, {{d_out, x}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.sizes_ubo(d_out), graph.sizes_ubo(x)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {},
+      // Resize Args
+      {},
+      // Resizing Logic
+      resize_linear_dW_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.linear_dW.default, linear_dW);
+}
+
+} // namespace vkcompute
diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp
new file mode 100644
index 00000000000..493daa47126
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearBackward.cpp
@@ -0,0 +1,130 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Q4gswLinear.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+// Resize d_x to d_out.shape[:-1] + (K,), mirroring linear_q4gsw_backward_meta.
+// extra_args = { weight_data, d_out }.
+void resize_linear_q4gsw_backward_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  const ValueRef d_x = args.at(0).refs.at(0);
+  const ValueRef weight_data = extra_args.at(0);
+  const ValueRef d_out = extra_args.at(1);
+  const int64_t K = graph->sizes_of(weight_data).at(1) * 2;
+  std::vector<int64_t> new_sizes = graph->sizes_of(d_out);
+  new_sizes.back() = K;
+  graph->virtual_resize(d_x, new_sizes);
+}
+
+utils::uvec3 linear_q4gsw_backward_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)shader;
+  (void)resize_args;
+  const ValueRef d_x = args.at(0).refs.at(0);
+  const uint32_t K = graph->size_at<uint32_t>(-1, d_x);
+  const uint32_t M = utils::safe_downcast<uint32_t>(graph->numel_of(d_x) / K);
+  const uint32_t tiles = utils::div_up_4(M) * utils::div_up_4(K);
+  return {tiles, 1u, 1u};
+}
+
+utils::uvec3 linear_q4gsw_backward_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {64u, 1u, 1u};
+}
+
+void linear_q4gsw_backward(
+    ComputeGraph& graph,
+    const std::vector<ValueRef>& args) {
+  int32_t i = 0;
+  const ValueRef d_out = args.at(i++);
+  const ValueRef weight_data = args.at(i++);
+  const ValueRef weight_scales_data = args.at(i++);
+  const ValueRef group_size_ref = args.at(i++);
+  const ValueRef d_x = args.at(i++);
+
+  VK_CHECK_COND(graph.dtype_of(d_out) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(d_x) == vkapi::kFloat);
+  VK_CHECK_COND(graph.is_buffer_storage(d_out));
+  VK_CHECK_COND(graph.is_buffer_storage(d_x));
+
+  const vkapi::ScalarType in_dtype = graph.dtype_of(d_out);
+  const int64_t group_size_val = graph.extract_scalar<int64_t>(group_size_ref);
+  VK_CHECK_COND(group_size_val > 0 && group_size_val % 4 == 0);
+
+  const std::vector<int64_t> weight_sizes = graph.sizes_of(weight_data);
+  const int64_t N = weight_sizes.at(0);
+  const int64_t K = weight_sizes.at(1) * 2;
+  VK_CHECK_COND(N > 0 && K > 0);
+  VK_CHECK_COND(N % 4 == 0 && K % 4 == 0);
+  VK_CHECK_COND(K % group_size_val == 0);
+  VK_CHECK_COND(graph.size_at<int64_t>(-1, d_out) == N);
+
+  const ValueRef packed_weight = prepack_q4_w_4x8_nc_buffer(graph, weight_data);
+  const ValueRef packed_scales =
+      prepack_q4_scales(graph, weight_scales_data, in_dtype);
+
+  const uint32_t M = utils::safe_downcast<uint32_t>(graph.numel_of(d_out) / N);
+  const uint32_t tiles =
+      utils::div_up_4(M) * utils::div_up_4(static_cast<uint32_t>(K));
+  VK_CHECK_COND(
+      (tiles + 63u) / 64u <= 65535u,
+      "linear_q4gsw_backward: tile count exceeds max workgroup count");
+
+  std::string kernel_name = "q4gsw_backward";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(d_x));
+  add_dtype_suffix(kernel_name, graph.dtype_of(d_x));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      linear_q4gsw_backward_global_wg_size,
+      linear_q4gsw_backward_local_wg_size,
+      // Inputs and Outputs
+      {{d_x, vkapi::kWrite},
+       {{d_out, packed_weight, packed_scales}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.sizes_ubo(d_out), graph.sizes_ubo(d_x)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {static_cast<uint32_t>(group_size_val)},
+      // Resize Args
+      {weight_data, d_out},
+      // Resizing Logic
+      resize_linear_q4gsw_backward_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.linear_q4gsw_backward.default, linear_q4gsw_backward);
+}
+
+} // namespace vkcompute
diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp
new file mode 100644
index 00000000000..ab7336fcc33
--- /dev/null
+++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearRequant.cpp
@@ -0,0 +1,131 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/DynamicDispatchNode.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
+
+#include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
+
+namespace vkcompute {
+
+// Resize the packed output to the W_4X8 int buffer size K4 * N4_padded * 2,
+// matching prepack_q4_w_4x8_nc_buffer's layout. extra_args = { latent }.
+void resize_q4gsw_requant_node(
+    ComputeGraph* graph,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& extra_args) {
+  const ValueRef packed = args.at(0).refs.at(0);
+  const ValueRef latent = extra_args.at(0);
+  const std::vector<int64_t> latent_sizes = graph->sizes_of(latent);
+  const int64_t N = latent_sizes.at(0);
+  const int64_t K = latent_sizes.at(1);
+  const int64_t K4 = K / 4;
+  const int64_t N4 = N / 4;
+  const int64_t N4_padded = (N4 + 1) & ~int64_t{1};
+  graph->virtual_resize(packed, {K4 * N4_padded * 2});
+}
+
+utils::uvec3 q4gsw_requant_global_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)shader;
+  (void)resize_args;
+  const ValueRef latent = args.at(1).refs.at(0);
+  const std::vector<int64_t> latent_sizes = graph->sizes_of(latent);
+  const uint32_t N = utils::safe_downcast<uint32_t>(latent_sizes.at(0));
+  const uint32_t K = utils::safe_downcast<uint32_t>(latent_sizes.at(1));
+  const uint32_t K4 = K / 4u;
+  const uint32_t N4 = (N + 3u) / 4u;
+  const uint32_t N8 = (N4 + 1u) / 2u;
+  return {K4, N8, 1u};
+}
+
+utils::uvec3 q4gsw_requant_local_wg_size(
+    ComputeGraph* graph,
+    const vkapi::ShaderInfo& shader,
+    const utils::uvec3& global_workgroup_size,
+    const std::vector<ArgGroup>& args,
+    const std::vector<ValueRef>& resize_args) {
+  (void)graph;
+  (void)shader;
+  (void)global_workgroup_size;
+  (void)args;
+  (void)resize_args;
+  return {8u, 8u, 1u};
+}
+
+void q4gsw_requant(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  int32_t i = 0;
+  const ValueRef latent = args.at(i++);
+  const ValueRef scales = args.at(i++);
+  const ValueRef group_size_ref = args.at(i++);
+  const ValueRef packed = args.at(i++);
+
+  VK_CHECK_COND(graph.dtype_of(latent) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(scales) == vkapi::kFloat);
+  VK_CHECK_COND(graph.dtype_of(packed) == vkapi::kInt);
+  VK_CHECK_COND(graph.is_buffer_storage(latent));
+  VK_CHECK_COND(graph.is_buffer_storage(packed));
+
+  const int64_t group_size_val = graph.extract_scalar<int64_t>(group_size_ref);
+  VK_CHECK_COND(group_size_val > 0 && group_size_val % 4 == 0);
+
+  const std::vector<int64_t> latent_sizes = graph.sizes_of(latent);
+  VK_CHECK_COND(latent_sizes.size() == 2);
+  const int64_t N = latent_sizes.at(0);
+  const int64_t K = latent_sizes.at(1);
+  VK_CHECK_COND(N > 0 && K > 0);
+  VK_CHECK_COND(N % 4 == 0 && K % 4 == 0);
+  VK_CHECK_COND(K % group_size_val == 0);
+
+  const uint32_t K4 = utils::safe_downcast<uint32_t>(K / 4);
+  const uint32_t N4 = utils::safe_downcast<uint32_t>(N / 4);
+  const uint32_t N8 = (N4 + 1u) / 2u;
+  VK_CHECK_COND(
+      K4 <= 65535u && N8 <= 65535u,
+      "q4gsw_requant: dispatch grid exceeds max workgroup count");
+
+  // Scales are a frozen constant; materialize them to a GPU buffer once.
+  const ValueRef packed_scales =
+      prepack_standard(graph, scales, utils::kBuffer, utils::kWidthPacked);
+
+  std::string kernel_name = "q4gsw_requant";
+  kernel_name.reserve(kShaderNameReserve);
+  add_storage_type_suffix(kernel_name, graph.storage_type_of(latent));
+  add_dtype_suffix(kernel_name, graph.dtype_of(latent));
+
+  graph.execute_nodes().emplace_back(new DynamicDispatchNode(
+      graph,
+      VK_KERNEL_FROM_STR(kernel_name),
+      q4gsw_requant_global_wg_size,
+      q4gsw_requant_local_wg_size,
+      // Inputs and Outputs
+      {{packed, vkapi::kWrite}, {{latent, packed_scales}, vkapi::kRead}},
+      // Shader params buffers
+      {graph.sizes_ubo(latent)},
+      // Push Constants
+      {},
+      // Specialization Constants
+      {static_cast<uint32_t>(group_size_val)},
+      // Resize Args
+      {latent},
+      // Resizing Logic
+      resize_q4gsw_requant_node));
+}
+
+REGISTER_OPERATORS {
+  VK_REGISTER_OP(et_vk.q4gsw_requant.default, q4gsw_requant);
+}
+
+} // namespace vkcompute
diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt
index 5e8991f8e50..0f8456accf5 100644
--- a/backends/vulkan/test/op_tests/CMakeLists.txt
+++ b/backends/vulkan/test/op_tests/CMakeLists.txt
@@ -116,6 +116,23 @@ if(TARGET vulkan_backend AND LIB_TORCH)
     quantized_linear_test ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_test.cpp
     test_utils
   )
+  vulkan_op_test(
+    adamw_step_test ${CMAKE_CURRENT_SOURCE_DIR}/adamw_step_test.cpp test_utils
+  )
+  vulkan_op_test(
+    linear_dW_test ${CMAKE_CURRENT_SOURCE_DIR}/linear_dW_test.cpp test_utils
+  )
+  vulkan_op_test(
+    fused_ce_test ${CMAKE_CURRENT_SOURCE_DIR}/fused_ce_test.cpp test_utils
+  )
+  vulkan_op_test(
+    quantized_linear_backward_test
+    ${CMAKE_CURRENT_SOURCE_DIR}/quantized_linear_backward_test.cpp test_utils
+  )
+  vulkan_op_test(
+    q4gsw_requant_test ${CMAKE_CURRENT_SOURCE_DIR}/q4gsw_requant_test.cpp
+    test_utils
+  )
 
   # Only build generated op tests if a path to tags.yaml and
   # native_functions.yaml is provided. These files are required for codegen.
diff --git a/backends/vulkan/test/op_tests/adamw_step_test.cpp b/backends/vulkan/test/op_tests/adamw_step_test.cpp
new file mode 100644
index 00000000000..e2db8592712
--- /dev/null
+++ b/backends/vulkan/test/op_tests/adamw_step_test.cpp
@@ -0,0 +1,191 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include "test_utils.h"
+
+#include <cmath>
+#include <tuple>
+
+//
+// Reference implementation: an ATen transcription of the CPU-eager
+// adamw_step_impl golden in backends/vulkan/custom_ops_lib.py.
+//
+
+std::tuple<at::Tensor, at::Tensor, at::Tensor> adamw_reference_impl(
+    const at::Tensor& param_in,
+    const at::Tensor& m_in,
+    const at::Tensor& v_in,
+    const at::Tensor& grad,
+    const double lr,
+    const double beta1,
+    const double beta2,
+    const double eps,
+    const double weight_decay,
+    const double bias_correction1,
+    const double bias_correction2) {
+  at::Tensor param = param_in.clone();
+  at::Tensor m = m_in.clone();
+  at::Tensor v = v_in.clone();
+
+  param.mul_(1.0 - lr * weight_decay);
+  m.mul_(beta1).add_(grad, 1.0 - beta1);
+  v.mul_(beta2).addcmul_(grad, grad, 1.0 - beta2);
+  at::Tensor mhat = m / bias_correction1;
+  at::Tensor denom = (v / bias_correction2).sqrt() + eps;
+  param.addcdiv_(mhat, denom, -lr);
+
+  return std::make_tuple(param, m, v);
+}
+
+//
+// Test harness
+//
+
+void test_vulkan_adamw_step(
+    const std::vector<int64_t>& sizes,
+    const double lr,
+    const double beta1,
+    const double beta2,
+    const double eps,
+    const double weight_decay,
+    const int step) {
+  at::manual_seed(0);
+
+  const double bias_correction1 = 1.0 - std::pow(beta1, step);
+  const double bias_correction2 = 1.0 - std::pow(beta2, step);
+
+  const auto options = at::device(at::kCPU).dtype(at::kFloat);
+  at::Tensor param = at::randn(sizes, options);
+  at::Tensor m = at::randn(sizes, options);
+  at::Tensor v = at::rand(sizes, options); // second moment is non-negative
+  at::Tensor grad = at::randn(sizes, options);
+
+  at::Tensor ref_param;
+  at::Tensor ref_m;
+  at::Tensor ref_v;
+  std::tie(ref_param, ref_m, ref_v) = adamw_reference_impl(
+      param,
+      m,
+      v,
+      grad,
+      lr,
+      beta1,
+      beta2,
+      eps,
+      weight_decay,
+      bias_correction1,
+      bias_correction2);
+
+  using namespace vkcompute;
+
+  GraphConfig config;
+  ComputeGraph graph(config);
+
+  ValueRef r_param = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+  ValueRef r_m = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+  ValueRef r_v = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+  ValueRef r_grad = graph.add_tensor(sizes, vkapi::kFloat, utils::kBuffer);
+
+  ValueRef staging_param = graph.set_input_tensor(r_param);
+  ValueRef staging_m = graph.set_input_tensor(r_m);
+  ValueRef staging_v = graph.set_input_tensor(r_v);
+  ValueRef staging_grad = graph.set_input_tensor(r_grad);
+
+  VK_GET_OP_FN("et_vk.adamw_step.default")
+  (graph,
+   {
+       r_param,
+       r_m,
+       r_v,
+       r_grad,
+       graph.add_scalar<double>(lr),
+       graph.add_scalar<double>(beta1),
+       graph.add_scalar<double>(beta2),
+       graph.add_scalar<double>(eps),
+       graph.add_scalar<double>(weight_decay),
+       graph.add_scalar<double>(bias_correction1),
+       graph.add_scalar<double>(bias_correction2),
+   });
+
+  ValueRef staging_param_out = graph.set_output_tensor(r_param);
+  ValueRef staging_m_out = graph.set_output_tensor(r_m);
+  ValueRef staging_v_out = graph.set_output_tensor(r_v);
+
+  graph.prepare();
+  graph.prepack();
+
+  graph.maybe_cast_and_copy_into_staging(
+      staging_param, param.const_data_ptr(), param.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      staging_m, m.const_data_ptr(), m.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      staging_v, v.const_data_ptr(), v.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      staging_grad, grad.const_data_ptr(), grad.numel(), vkapi::kFloat);
+
+  graph.execute();
+
+  at::Tensor vk_param = at::empty_like(param);
+  at::Tensor vk_m = at::empty_like(m);
+  at::Tensor vk_v = at::empty_like(v);
+
+  graph.maybe_cast_and_copy_from_staging(
+      staging_param_out,
+      vk_param.mutable_data_ptr(),
+      vk_param.numel(),
+      vkapi::kFloat);
+  graph.maybe_cast_and_copy_from_staging(
+      staging_m_out, vk_m.mutable_data_ptr(), vk_m.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_from_staging(
+      staging_v_out, vk_v.mutable_data_ptr(), vk_v.numel(), vkapi::kFloat);
+
+  ASSERT_TRUE(at::allclose(ref_param, vk_param, /*rtol=*/1e-4, /*atol=*/1e-4));
+  ASSERT_TRUE(at::allclose(ref_m, vk_m, /*rtol=*/1e-4, /*atol=*/1e-4));
+  ASSERT_TRUE(at::allclose(ref_v, vk_v, /*rtol=*/1e-4, /*atol=*/1e-4));
+}
+
+TEST(VulkanAdamwStepTest, test_adamw_step_no_weight_decay) {
+  test_vulkan_adamw_step(
+      /*sizes=*/{4, 8},
+      /*lr=*/1e-3,
+      /*beta1=*/0.9,
+      /*beta2=*/0.999,
+      /*eps=*/1e-8,
+      /*weight_decay=*/0.0,
+      /*step=*/1);
+}
+
+TEST(VulkanAdamwStepTest, test_adamw_step_with_weight_decay) {
+  test_vulkan_adamw_step(
+      /*sizes=*/{3, 5, 7},
+      /*lr=*/1e-2,
+      /*beta1=*/0.9,
+      /*beta2=*/0.999,
+      /*eps=*/1e-8,
+      /*weight_decay=*/0.01,
+      /*step=*/10);
+}
+
+TEST(VulkanAdamwStepTest, test_adamw_step_large_1d) {
+  test_vulkan_adamw_step(
+      /*sizes=*/{1000},
+      /*lr=*/1e-3,
+      /*beta1=*/0.9,
+      /*beta2=*/0.999,
+      /*eps=*/1e-8,
+      /*weight_decay=*/0.05,
+      /*step=*/5);
+}
diff --git a/backends/vulkan/test/op_tests/fused_ce_test.cpp b/backends/vulkan/test/op_tests/fused_ce_test.cpp
new file mode 100644
index 00000000000..fa2a8ae22bc
--- /dev/null
+++ b/backends/vulkan/test/op_tests/fused_ce_test.cpp
@@ -0,0 +1,145 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include "test_utils.h"
+
+#include <iostream>
+#include <vector>
+
+//
+// Reference implementation: transcription of the CPU-eager fused_ce_impl
+// (backends/vulkan/custom_ops_lib.py), computed with ATen (library golden).
+//
+
+std::pair<at::Tensor, at::Tensor> fused_ce_reference(
+    const at::Tensor& logits,
+    const at::Tensor& labels,
+    double n_valid) {
+  at::Tensor mask = labels.ge(0);
+  at::Tensor safe = labels.clamp_min(0).to(at::kLong);
+  at::Tensor lse = at::logsumexp(logits, -1);
+  at::Tensor picked = logits.gather(-1, safe.unsqueeze(-1)).squeeze(-1);
+  at::Tensor loss =
+      at::where(mask, (lse - picked) / n_valid, at::zeros_like(lse)).sum();
+  at::Tensor softmax = at::softmax(logits, -1);
+  at::Tensor onehot = at::one_hot(safe, logits.size(-1)).to(logits.dtype());
+  at::Tensor dlogits = at::where(
+      mask.unsqueeze(-1),
+      (softmax - onehot) / n_valid,
+      at::zeros_like(softmax));
+  return {loss, dlogits};
+}
+
+void test_vulkan_fused_ce(
+    const int n_rows,
+    const int vocab,
+    const std::vector<int32_t>& labels_data,
+    const double n_valid) {
+  at::manual_seed(0);
+
+  at::Tensor logits =
+      at::rand({n_rows, vocab}, at::device(at::kCPU).dtype(at::kFloat)) * 4.0 -
+      2.0;
+  at::Tensor labels =
+      at::from_blob(
+          const_cast<int32_t*>(labels_data.data()), {n_rows}, at::kInt)
+          .clone();
+
+  auto ref = fused_ce_reference(logits, labels, n_valid);
+  at::Tensor ref_loss = ref.first;
+  at::Tensor ref_dlogits = ref.second;
+
+  using namespace vkcompute;
+
+  GraphConfig config;
+  ComputeGraph graph(config);
+
+  IOValueRef r_logits = graph.add_input_tensor(
+      logits.sizes().vec(), vkapi::kFloat, utils::kBuffer);
+  IOValueRef r_labels =
+      graph.add_input_tensor(labels.sizes().vec(), vkapi::kInt, utils::kBuffer);
+
+  const ValueRef r_n_valid = graph.add_scalar<double>(n_valid);
+
+  const ValueRef r_loss = graph.add_tensor({}, vkapi::kFloat, utils::kBuffer);
+  const ValueRef r_dlogits =
+      graph.add_tensor({n_rows, vocab}, vkapi::kFloat, utils::kBuffer);
+  const ValueRef r_out = graph.add_value_list({r_loss, r_dlogits});
+
+  VK_GET_OP_FN("et_vk.fused_ce.default")
+  (graph, {r_logits.value, r_labels.value, r_n_valid, r_out});
+
+  ValueRef staging_loss = graph.set_output_tensor(r_loss);
+  ValueRef staging_dlogits = graph.set_output_tensor(r_dlogits);
+
+  graph.prepare();
+  graph.prepack();
+  graph.propagate_resize();
+
+  graph.maybe_cast_and_copy_into_staging(
+      r_logits.staging, logits.const_data_ptr(), logits.numel(), vkapi::kFloat);
+  graph.maybe_cast_and_copy_into_staging(
+      r_labels.staging, labels.const_data_ptr(), labels.numel(), vkapi::kInt);
+
+  graph.execute();
+
+  at::Tensor vk_loss = at::zeros({}, at::device(at::kCPU).dtype(at::kFloat));
+  graph.maybe_cast_and_copy_from_staging(
+      staging_loss, vk_loss.mutable_data_ptr(), 1, vkapi::kFloat);
+
+  at::Tensor vk_dlogits =
+      at::zeros({n_rows, vocab}, at::device(at::kCPU).dtype(at::kFloat));
+  graph.maybe_cast_and_copy_from_staging(
+      staging_dlogits,
+      vk_dlogits.mutable_data_ptr(),
+      vk_dlogits.numel(),
+      vkapi::kFloat);
+
+  const double atol = 1e-4;
+  const double rtol = 1e-4;
+
+  const bool loss_ok = at::allclose(ref_loss, vk_loss, rtol, atol);
+  const bool dlogits_ok = at::allclose(ref_dlogits, vk_dlogits, rtol, atol);
+
+  if (!loss_ok || !dlogits_ok) {
+    std::cout << "fused_ce mismatch: n_rows=" << n_rows << " vocab=" << vocab
+              << " n_valid=" << n_valid << std::endl;
+    std::cout << "loss ref=" << ref_loss.item<float>()
+              << " vk=" << vk_loss.item<float>() << std::endl;
+    std::cout << "max dlogits diff="
+              << at::max(at::abs(ref_dlogits - vk_dlogits)).item<float>()
+              << std::endl;
+  }
+  ASSERT_TRUE(loss_ok);
+  ASSERT_TRUE(dlogits_ok);
+}
+
+TEST(VulkanFusedCeTest, all_valid_small) {
+  test_vulkan_fused_ce(
+      /*n_rows=*/4, /*vocab=*/8, /*labels=*/{0, 3, 7, 1}, /*n_valid=*/4.0);
+}
+
+TEST(VulkanFusedCeTest, masked_label) {
+  // One masked row (label < 0): contributes 0 to loss and 0 gradient.
+  test_vulkan_fused_ce(
+      /*n_rows=*/4, /*vocab=*/8, /*labels=*/{2, -1, 5, 0}, /*n_valid=*/3.0);
+}
+
+TEST(VulkanFusedCeTest, large_vocab_strided_reduce) {
+  // vocab >> NWORKERS exercises the strided per-row reduction.
+  test_vulkan_fused_ce(
+      /*n_rows=*/3, /*vocab=*/200, /*labels=*/{17, -1, 199}, /*n_valid=*/2.0);
+}
diff --git a/backends/vulkan/test/op_tests/linear_dW_test.cpp b/backends/vulkan/test/op_tests/linear_dW_test.cpp
new file mode 100644
index 00000000000..1d9f31328ed
--- /dev/null
+++ b/backends/vulkan/test/op_tests/linear_dW_test.cpp
@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include "test_utils.h"
+
+//
+// Reference Implementation
+//
+
+// Golden: dW[N, K] = d_out^T @ x, contracting over the flattened leading dims.
+// Mirrors the CPU-eager linear_dW_impl in custom_ops_lib.py.
+at::Tensor linear_dW_reference_impl(
+    const at::Tensor& d_out,
+    const at::Tensor& x) {
+  const int64_t N = d_out.size(-1);
+  const int64_t K = x.size(-1);
+  return d_out.reshape({-1, N}).t().matmul(x.reshape({-1, K})).contiguous();
+}
+
+//
+// Test function
+//
+
+void test_vulkan_linear_dW_impl(
+    const std::vector<int64_t>& d_out_sizes,
+    const std::vector<int64_t>& x_sizes,
+    const vkcompute::utils::StorageType storage = vkcompute::utils::kBuffer) {
+  at::Tensor d_out =
+      at::rand(d_out_sizes, at::device(at::kCPU).dtype(at::kFloat));
+  at::Tensor x = at::rand(x_sizes, at::device(at::kCPU).dtype(at::kFloat));
+
+  at::Tensor dW_ref = linear_dW_reference_impl(d_out, x);
+
+  // Build Vulkan graph
+  using namespace vkcompute;
+
+  GraphConfig config;
+  ComputeGraph graph(config);
+
+  IOValueRef r_d_out = graph.add_input_tensor(
+      d_out.sizes().vec(), from_at_scalartype(d_out.scalar_type()), storage);
+  IOValueRef r_x = graph.add_input_tensor(
+      x.sizes().vec(), from_at_scalartype(x.scalar_type()), storage);
+
+  const ValueRef r_dW = graph.add_tensor(
+      dW_ref.sizes().vec(), from_at_scalartype(dW_ref.scalar_type()), storage);
+
+  VK_GET_OP_FN("et_vk.linear_dW.default")
+  (graph, {r_d_out.value, r_x.value, r_dW});
+
+  ValueRef staging_out = graph.set_output_tensor(r_dW);
+
+  graph.prepare();
+  graph.prepack();
+  graph.propagate_resize();
+
+  graph.maybe_cast_and_copy_into_staging(
+      r_d_out.staging,
+      d_out.const_data_ptr(),
+      d_out.numel(),
+      from_at_scalartype(d_out.scalar_type()));
+  graph.maybe_cast_and_copy_into_staging(
+      r_x.staging,
+      x.const_data_ptr(),
+      x.numel(),
+      from_at_scalartype(x.scalar_type()));
+
+  graph.execute();
+
+  at::Tensor vk_dW = at::empty_like(dW_ref);
+  graph.maybe_cast_and_copy_from_staging(
+      staging_out,
+      vk_dW.mutable_data_ptr(),
+      vk_dW.numel(),
+      from_at_scalartype(vk_dW.scalar_type()));
+
+  ASSERT_TRUE(at::allclose(vk_dW, dW_ref, 1e-3, 1e-3));
+}
+
+// Tile-aligned 2D shapes (M, N, K all multiples of 4).
+TEST(VulkanLinearDwTest, test_tile_aligned) {
+  test_vulkan_linear_dW_impl(
+      /*d_out_sizes=*/{8, 16}, /*x_sizes=*/{8, 32});
+}
+
+// Non-tile-multiple shapes (M, N, K each not a multiple of 4) to exercise the
+// partial-tile min()-clamp paths in the shader.
+TEST(VulkanLinearDwTest, test_non_tile_multiple) {
+  test_vulkan_linear_dW_impl(
+      /*d_out_sizes=*/{5, 6}, /*x_sizes=*/{5, 10});
+}
+
+// Leading dims > 2D: M is the flattened product of all leading dims.
+TEST(VulkanLinearDwTest, test_leading_dims_flatten) {
+  test_vulkan_linear_dW_impl(
+      /*d_out_sizes=*/{2, 3, 16}, /*x_sizes=*/{2, 3, 32});
+}
diff --git a/backends/vulkan/test/op_tests/q4gsw_requant_test.cpp b/backends/vulkan/test/op_tests/q4gsw_requant_test.cpp
new file mode 100644
index 00000000000..a0233bc2706
--- /dev/null
+++ b/backends/vulkan/test/op_tests/q4gsw_requant_test.cpp
@@ -0,0 +1,195 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include "test_utils.h"
+
+//
+// Reference Implementation
+//
+
+// Pack [N, K] codes (0..15) into the W_4X8 block-packed int buffer the forward
+// reads. Mirrors pack_q4_linear_weight__w_4x8.glsl: one ivec4 per (k4, n8),
+// byte b holds an (even-N low nibble, odd-N high nibble) pair at K = k4*4 + b.
+std::vector<int32_t> pack_codes_w4x8(const at::Tensor& codes) {
+  const int64_t N = codes.size(0);
+  const int64_t K = codes.size(1);
+  const int64_t K4 = K / 4;
+  const int64_t N4 = N / 4;
+  const int64_t N4_padded = (N4 + 1) & ~int64_t{1};
+  const int64_t N8 = N4_padded / 2;
+  std::vector<int32_t> buf(K4 * N4_padded * 2, 0);
+  auto ca = codes.accessor<int, 2>();
+
+  auto pack_tile = [&](int64_t k4, int64_t n4, uint32_t& px, uint32_t& py) {
+    px = 0u;
+    py = 0u;
+    for (int ni = 0; ni < 4; ++ni) {
+      const int64_t n = n4 * 4 + ni;
+      for (int b = 0; b < 4; ++b) {
+        const uint32_t code = static_cast<uint32_t>(ca[n][k4 * 4 + b] & 0xF);
+        const int shift = 8 * b + (ni & 1) * 4;
+        if (ni < 2) {
+          px |= code << shift;
+        } else {
+          py |= code << shift;
+        }
+      }
+    }
+  };
+
+  for (int64_t k4 = 0; k4 < K4; ++k4) {
+    for (int64_t n8 = 0; n8 < N8; ++n8) {
+      const int64_t n4_a = 2 * n8;
+      const int64_t n4_b = n4_a + 1;
+      uint32_t px_a, py_a, px_b = 0x88888888u, py_b = 0x88888888u;
+      pack_tile(k4, n4_a, px_a, py_a);
+      if (n4_b < N4) {
+        pack_tile(k4, n4_b, px_b, py_b);
+      }
+      const int64_t base = (k4 * N8 + n8) * 4;
+      buf[base + 0] = static_cast<int32_t>(px_a);
+      buf[base + 1] = static_cast<int32_t>(py_a);
+      buf[base + 2] = static_cast<int32_t>(px_b);
+      buf[base + 3] = static_cast<int32_t>(py_b);
+    }
+  }
+  return buf;
+}
+
+//
+// Test function
+//
+
+void test_vulkan_q4gsw_requant_impl(
+    const int64_t N,
+    const int64_t K,
+    const int64_t group_size,
+    const bool with_zero_scale) {
+  const int64_t num_groups = K / group_size;
+
+  at::Tensor scales =
+      at::rand({num_groups, N}, at::device(at::kCPU).dtype(at::kFloat)) + 0.5;
+  if (with_zero_scale) {
+    scales.index_put_({0, 0}, 0.0);
+  }
+
+  const at::Tensor group_idx =
+      at::arange(K, at::device(at::kCPU).dtype(at::kLong))
+          .div(group_size, "floor");
+  const at::Tensor scale_full =
+      scales.t().contiguous().index_select(1, group_idx); // [N, K]
+
+  // Deterministic quotient targets, each >=0.2 from any .5 tie, so GPU fp32
+  // division (~2.5 ULP, not correctly rounded) and the CPU golden round
+  // identically. Covers round both directions and clamp past [-8, 7].
+  const std::vector<float> pattern = {
+      0.3f, -0.4f, 2.7f, -3.3f, 6.4f, -6.4f, 13.2f, -21.7f};
+  const at::Tensor pat =
+      at::tensor(pattern, at::device(at::kCPU).dtype(at::kFloat));
+  const at::Tensor q_idx =
+      at::arange(N * K, at::device(at::kCPU).dtype(at::kLong))
+          .remainder(static_cast<int64_t>(pattern.size()));
+  const at::Tensor target_q = pat.index_select(0, q_idx).reshape({N, K});
+  at::Tensor latent = target_q * scale_full;
+
+  // Golden codes, mirroring quant_nibble: q=0 where scale==0, else roundEven
+  // (matches at::round half-to-even); clamp to [-8, 7]; code = (q + 8) & 0xF.
+  const at::Tensor nonzero = scale_full != 0;
+  const at::Tensor safe =
+      at::where(nonzero, scale_full, at::ones_like(scale_full));
+  at::Tensor q = at::round(latent / safe);
+  q = at::where(nonzero, q, at::zeros_like(q));
+  q = at::clamp(q, -8, 7);
+  const at::Tensor golden_codes =
+      (q.to(at::kInt) + 8).bitwise_and(0xF); // [N, K] in 0..15
+
+  const std::vector<int32_t> expected = pack_codes_w4x8(golden_codes);
+
+  using namespace vkcompute;
+
+  GraphConfig config;
+  ComputeGraph graph(config);
+
+  IOValueRef r_latent = graph.add_input_tensor(
+      latent.sizes().vec(),
+      from_at_scalartype(latent.scalar_type()),
+      utils::kBuffer);
+  ValueRef r_scales = graph.add_tensorref(
+      scales.sizes().vec(),
+      from_at_scalartype(scales.scalar_type()),
+      scales.const_data_ptr());
+  const ValueRef r_group_size = graph.add_scalar<int64_t>(group_size);
+
+  const int64_t N4 = N / 4;
+  const int64_t N4_padded = (N4 + 1) & ~int64_t{1};
+  const ValueRef r_packed =
+      graph.add_tensor({(K / 4) * N4_padded * 2}, vkapi::kInt, utils::kBuffer);
+
+  VK_GET_OP_FN("et_vk.q4gsw_requant.default")
+  (graph, {r_latent.value, r_scales, r_group_size, r_packed});
+
+  ValueRef staging_out = graph.set_output_tensor(r_packed);
+
+  graph.prepare();
+  graph.prepack();
+  graph.propagate_resize();
+
+  graph.maybe_cast_and_copy_into_staging(
+      r_latent.staging,
+      latent.const_data_ptr(),
+      latent.numel(),
+      from_at_scalartype(latent.scalar_type()));
+
+  graph.execute();
+
+  at::Tensor vk_packed = at::empty(
+      {static_cast<int64_t>(expected.size())},
+      at::device(at::kCPU).dtype(at::kInt));
+  graph.maybe_cast_and_copy_from_staging(
+      staging_out,
+      vk_packed.mutable_data_ptr(),
+      vk_packed.numel(),
+      from_at_scalartype(vk_packed.scalar_type()));
+
+  auto va = vk_packed.accessor<int, 1>();
+  for (size_t i = 0; i < expected.size(); ++i) {
+    ASSERT_EQ(va[i], expected[i]) << "mismatch at packed int " << i;
+  }
+}
+
+// Tile-aligned single-group.
+TEST(VulkanQ4gswRequantTest, test_tile_aligned) {
+  test_vulkan_q4gsw_requant_impl(
+      /*N=*/16, /*K=*/32, /*group_size=*/32, /*with_zero_scale=*/false);
+}
+
+// Multiple quantization groups along K.
+TEST(VulkanQ4gswRequantTest, test_grouped) {
+  test_vulkan_q4gsw_requant_impl(
+      /*N=*/32, /*K=*/64, /*group_size=*/32, /*with_zero_scale=*/false);
+}
+
+// N not a multiple of 8 (odd N4 -> padded stride + bias-zero OOB tile).
+TEST(VulkanQ4gswRequantTest, test_odd_n4) {
+  test_vulkan_q4gsw_requant_impl(
+      /*N=*/12, /*K=*/16, /*group_size=*/16, /*with_zero_scale=*/false);
+}
+
+// A zero scale must produce the bias-zero code (8), not a divide-by-zero.
+TEST(VulkanQ4gswRequantTest, test_zero_scale) {
+  test_vulkan_q4gsw_requant_impl(
+      /*N=*/16, /*K=*/32, /*group_size=*/32, /*with_zero_scale=*/true);
+}
diff --git a/backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp b/backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp
new file mode 100644
index 00000000000..8135c339f99
--- /dev/null
+++ b/backends/vulkan/test/op_tests/quantized_linear_backward_test.cpp
@@ -0,0 +1,160 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <gtest/gtest.h>
+
+#include <ATen/ATen.h>
+
+#include <executorch/backends/vulkan/runtime/api/api.h>
+#include <executorch/backends/vulkan/runtime/graph/ComputeGraph.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/OperatorRegistry.h>
+
+#include "test_utils.h"
+
+//
+// Reference Implementation
+//
+
+// Pack unpacked [N, K] codes (0..15) into the flat [N, K/2] uint8 weight the
+// forward's prepack consumes: even-K in the low nibble, odd-K in the high.
+at::Tensor pack_codes_flat(const at::Tensor& codes) {
+  const int64_t N = codes.size(0);
+  const int64_t K = codes.size(1);
+  at::Tensor packed =
+      at::empty({N, K / 2}, at::device(at::kCPU).dtype(at::kByte));
+  auto ca = codes.accessor<int, 2>();
+  auto pa = packed.accessor<uint8_t, 2>();
+  for (int64_t n = 0; n < N; ++n) {
+    for (int64_t kb = 0; kb < K / 2; ++kb) {
+      const int lo = ca[n][2 * kb] & 0xF;
+      const int hi = ca[n][2 * kb + 1] & 0xF;
+      pa[n][kb] = static_cast<uint8_t>(lo | (hi << 4));
+    }
+  }
+  return packed;
+}
+
+// Golden d_x[M, K] = d_out[M, N] @ dequant(W)[N, K], with
+// dequant(W[n, k]) = (code(n, k) - 8) * scales[k / group_size, n].
+// Mirrors the CPU-eager linear_q4gsw_backward_impl in custom_ops_lib.py.
+at::Tensor linear_q4gsw_backward_reference_impl(
+    const at::Tensor& d_out,
+    const at::Tensor& codes,
+    const at::Tensor& scales,
+    const int64_t group_size) {
+  const int64_t N = codes.size(0);
+  const int64_t K = codes.size(1);
+  const at::Tensor group_idx =
+      at::arange(K, at::device(at::kCPU).dtype(at::kLong))
+          .div(group_size, "floor");
+  const at::Tensor scale_full =
+      scales.t().contiguous().index_select(1, group_idx); // [N, K]
+  const at::Tensor dequant_w =
+      (codes.to(at::kFloat) - 8.0) * scale_full; // [N, K]
+  const at::Tensor d_x_flat = d_out.reshape({-1, N}).matmul(dequant_w);
+  std::vector<int64_t> out_shape = d_out.sizes().vec();
+  out_shape.back() = K;
+  return d_x_flat.reshape(out_shape).contiguous(); // d_out[..., :-1] + [K]
+}
+
+//
+// Test function
+//
+
+void test_vulkan_linear_q4gsw_backward_impl(
+    const std::vector<int64_t>& d_out_sizes,
+    const int64_t K,
+    const int64_t group_size) {
+  const int64_t N = d_out_sizes.back();
+  const int64_t num_groups = K / group_size;
+
+  at::Tensor codes =
+      at::randint(0, 16, {N, K}, at::device(at::kCPU).dtype(at::kInt));
+  at::Tensor scales =
+      at::rand({num_groups, N}, at::device(at::kCPU).dtype(at::kFloat)) + 0.5;
+  at::Tensor packed = pack_codes_flat(codes);
+  at::Tensor d_out =
+      at::rand(d_out_sizes, at::device(at::kCPU).dtype(at::kFloat));
+
+  at::Tensor d_x_ref =
+      linear_q4gsw_backward_reference_impl(d_out, codes, scales, group_size);
+
+  using namespace vkcompute;
+
+  GraphConfig config;
+  ComputeGraph graph(config);
+
+  ValueRef r_weights = graph.add_tensorref(
+      packed.sizes().vec(),
+      from_at_scalartype(packed.scalar_type()),
+      packed.const_data_ptr());
+  ValueRef r_scales = graph.add_tensorref(
+      scales.sizes().vec(),
+      from_at_scalartype(scales.scalar_type()),
+      scales.const_data_ptr());
+
+  IOValueRef r_d_out = graph.add_input_tensor(
+      d_out.sizes().vec(),
+      from_at_scalartype(d_out.scalar_type()),
+      utils::kBuffer);
+  const ValueRef r_group_size = graph.add_scalar<int64_t>(group_size);
+  const ValueRef r_d_x = graph.add_tensor(
+      d_x_ref.sizes().vec(),
+      from_at_scalartype(d_x_ref.scalar_type()),
+      utils::kBuffer);
+
+  VK_GET_OP_FN("et_vk.linear_q4gsw_backward.default")
+  (graph, {r_d_out.value, r_weights, r_scales, r_group_size, r_d_x});
+
+  ValueRef staging_out = graph.set_output_tensor(r_d_x);
+
+  graph.prepare();
+  graph.prepack();
+  graph.propagate_resize();
+
+  graph.maybe_cast_and_copy_into_staging(
+      r_d_out.staging,
+      d_out.const_data_ptr(),
+      d_out.numel(),
+      from_at_scalartype(d_out.scalar_type()));
+
+  graph.execute();
+
+  at::Tensor vk_d_x = at::empty_like(d_x_ref);
+  graph.maybe_cast_and_copy_from_staging(
+      staging_out,
+      vk_d_x.mutable_data_ptr(),
+      vk_d_x.numel(),
+      from_at_scalartype(vk_d_x.scalar_type()));
+
+  ASSERT_TRUE(at::allclose(vk_d_x, d_x_ref, 1e-3, 1e-3));
+}
+
+// Tile-aligned single-group shapes.
+TEST(VulkanLinearQ4gswBackwardTest, test_tile_aligned) {
+  test_vulkan_linear_q4gsw_backward_impl(
+      /*d_out_sizes=*/{8, 16}, /*K=*/32, /*group_size=*/32);
+}
+
+// Multiple quantization groups along K.
+TEST(VulkanLinearQ4gswBackwardTest, test_grouped) {
+  test_vulkan_linear_q4gsw_backward_impl(
+      /*d_out_sizes=*/{8, 32}, /*K=*/64, /*group_size=*/32);
+}
+
+// N not a multiple of 8 (odd N4 -> padded W_4X8 stride) plus partial-M tile.
+TEST(VulkanLinearQ4gswBackwardTest, test_odd_n4_partial_m) {
+  test_vulkan_linear_q4gsw_backward_impl(
+      /*d_out_sizes=*/{5, 12}, /*K=*/16, /*group_size=*/16);
+}
+
+// Leading dims > 2D: M is the flattened product of all leading dims.
+TEST(VulkanLinearQ4gswBackwardTest, test_leading_dims_flatten) {
+  test_vulkan_linear_q4gsw_backward_impl(
+      /*d_out_sizes=*/{2, 3, 16}, /*K=*/32, /*group_size=*/32);
+}
diff --git a/backends/vulkan/test/op_tests/targets.bzl b/backends/vulkan/test/op_tests/targets.bzl
index 0a9b0743415..8284c5c9aca 100644
--- a/backends/vulkan/test/op_tests/targets.bzl
+++ b/backends/vulkan/test/op_tests/targets.bzl
@@ -180,6 +180,36 @@ def define_common_targets(is_fbcode = False):
             ":test_utils",
         ]
     )
+    define_test_targets(
+        "adamw_step_test",
+        extra_deps = [
+            ":test_utils",
+        ]
+    )
+    define_test_targets(
+        "linear_dW_test",
+        extra_deps = [
+            ":test_utils",
+        ]
+    )
+    define_test_targets(
+        "fused_ce_test",
+        extra_deps = [
+            ":test_utils",
+        ]
+    )
+    define_test_targets(
+        "quantized_linear_backward_test",
+        extra_deps = [
+            ":test_utils",
+        ]
+    )
+    define_test_targets(
+        "q4gsw_requant_test",
+        extra_deps = [
+            ":test_utils",
+        ]
+    )
     define_test_targets(
         "rms_norm_test",
         extra_deps = [
diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt
index 48337d18da1..78d9c4f30dd 100644
--- a/backends/webgpu/CMakeLists.txt
+++ b/backends/webgpu/CMakeLists.txt
@@ -26,41 +26,18 @@ if(NOT TARGET vulkan_schema)
 endif()
 
 set(WEBGPU_SRCS
-    runtime/WebGPUBackend.cpp
-    runtime/WebGPUGraph.cpp
-    runtime/WebGPUDelegateHeader.cpp
-    runtime/WebGPUDevice.cpp
-    runtime/WebGPUQueryPool.cpp
-    runtime/ops/OperatorRegistry.cpp
-    runtime/ops/add/BinaryOp.cpp
-    runtime/ops/rms_norm/RmsNorm.cpp
-    runtime/ops/update_cache/UpdateCache.cpp
-    runtime/ops/sdpa/Sdpa.cpp
-    runtime/ops/select_as_symint/SelectAsSymint.cpp
-    runtime/ops/quantized_linear/QuantizedLinear.cpp
-    runtime/ops/mul/BinaryOp.cpp
-    runtime/ops/embedding_q4gsw/EmbeddingQ4gsw.cpp
-    runtime/ops/rope/RotaryEmbedding.cpp
-    runtime/ops/prepack/Prepack.cpp
-    runtime/ops/view_copy/ViewCopy.cpp
-    runtime/ops/select/Select.cpp
-    runtime/ops/sigmoid/UnaryOp.cpp
-    runtime/ops/squeeze/Squeeze.cpp
-    runtime/ops/unsqueeze/Unsqueeze.cpp
-    runtime/ops/slice/Slice.cpp
-    runtime/ops/permute/Permute.cpp
-    runtime/ops/cat/Cat.cpp
-    runtime/ops/index/Index.cpp
-    runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp
-    runtime/ops/mm/Mm.cpp
-    runtime/ops/log_softmax/LogSoftmax.cpp
-    runtime/ops/softmax/Softmax.cpp
-    runtime/ops/bmm/Bmm.cpp
-    runtime/ops/div/BinaryOp.cpp
-    runtime/ops/sub/BinaryOp.cpp
-    runtime/ops/linear/Linear.cpp
+    runtime/WebGPUBackend.cpp runtime/WebGPUGraph.cpp
+    runtime/WebGPUDelegateHeader.cpp runtime/WebGPUDevice.cpp
+    runtime/WebGPUQueryPool.cpp runtime/ops/OperatorRegistry.cpp
 )
 
+# Op handlers: glob so adding an op needs no CMakeLists edit. CONFIGURE_DEPENDS
+# re-globs at build time when op sources are added or removed.
+file(GLOB WEBGPU_OP_SRCS CONFIGURE_DEPENDS
+     "${CMAKE_CURRENT_SOURCE_DIR}/runtime/ops/*/*.cpp"
+)
+list(APPEND WEBGPU_SRCS ${WEBGPU_OP_SRCS})
+
 add_library(webgpu_backend ${WEBGPU_SRCS})
 
 # Verify committed *_wgsl.h match their *.wgsl (drift fails the build).
diff --git a/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp
new file mode 100644
index 00000000000..4a78694d40b
--- /dev/null
+++ b/backends/webgpu/runtime/ops/adamw/AdamwStep.cpp
@@ -0,0 +1,181 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct AdamwStepParams {
+  uint32_t numel;
+  uint32_t _pad0;
+  uint32_t _pad1;
+  uint32_t _pad2;
+  float lr;
+  float beta1;
+  float beta2;
+  float eps;
+  float weight_decay;
+  float bias_correction1;
+  float bias_correction2;
+  float _pad3;
+};
+static_assert(sizeof(AdamwStepParams) == 48, "params must be 48 bytes");
+
+// AdamW step over an fp32 latent (elementwise, in place); mirrors torch AdamW.
+void adamw_step_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int param_id = args.at(0);
+  const int m_id = args.at(1);
+  const int v_id = args.at(2);
+  const int grad_id = args.at(3);
+
+  WGPUDevice device = graph.device();
+  const auto& param = graph.get_tensor(param_id);
+  const auto& m = graph.get_tensor(m_id);
+  const auto& v = graph.get_tensor(v_id);
+  const auto& grad = graph.get_tensor(grad_id);
+
+  uint64_t numel = 1;
+  for (int64_t d : param.dims) {
+    numel *= static_cast<uint64_t>(d);
+  }
+  if (param.dims.empty() || numel == 0) {
+    throw std::runtime_error("adamw_step: empty param");
+  }
+  const uint64_t bytes = numel * sizeof(float);
+  if (param.nbytes != bytes || m.nbytes != bytes || v.nbytes != bytes ||
+      grad.nbytes != bytes) {
+    throw std::runtime_error(
+        "adamw_step: param/m/v/grad must be fp32 and same numel");
+  }
+  if (numel > UINT32_MAX) {
+    throw std::runtime_error("adamw_step: numel exceeds u32");
+  }
+
+  auto scalar = [&](int id, const char* name) -> float {
+    if (graph.get_value_type(id) != WebGPUGraph::ValueType::Double) {
+      throw std::runtime_error(
+          std::string("adamw_step: ") + name + " must be a float scalar");
+    }
+    return static_cast<float>(graph.get_double(id));
+  };
+
+  AdamwStepParams params = {};
+  params.numel = static_cast<uint32_t>(numel);
+  params.lr = scalar(args.at(4), "lr");
+  params.beta1 = scalar(args.at(5), "beta1");
+  params.beta2 = scalar(args.at(6), "beta2");
+  params.eps = scalar(args.at(7), "eps");
+  params.weight_decay = scalar(args.at(8), "weight_decay");
+  params.bias_correction1 = scalar(args.at(9), "bias_correction1");
+  params.bias_correction2 = scalar(args.at(10), "bias_correction2");
+  if (params.bias_correction1 == 0.0f || params.bias_correction2 == 0.0f) {
+    throw std::runtime_error("adamw_step: bias corrections must be non-zero");
+  }
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kAdamwStepWorkgroupSizeX);
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, params.numel, wg_size, "adamw_step");
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(params));
+  graph.add_uniform_buffer_bytes(sizeof(params));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kAdamwStepWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[5] = {};
+  for (uint32_t i = 0; i <= 2; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+    entries[i].buffer.type = WGPUBufferBindingType_Storage;
+  }
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[4].binding = 4;
+  entries[4].visibility = WGPUShaderStage_Compute;
+  entries[4].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 5;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[5] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = param.buffer;
+  bg_entries[0].size = param.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = m.buffer;
+  bg_entries[1].size = m.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = v.buffer;
+  bg_entries[2].size = v.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = grad.buffer;
+  bg_entries[3].size = grad.nbytes;
+  bg_entries[4].binding = 4;
+  bg_entries[4].buffer = uniform_buffer;
+  bg_entries[4].size = sizeof(params);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 5;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "adamw_step"});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.adamw_step.default, adamw_step_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl b/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl
new file mode 100644
index 00000000000..f6a5ad58943
--- /dev/null
+++ b/backends/webgpu/runtime/ops/adamw/adamw_step.wgsl
@@ -0,0 +1,41 @@
+
+@group(0) @binding(0) var<storage, read_write> t_param: array<f32>;
+@group(0) @binding(1) var<storage, read_write> t_m: array<f32>;
+@group(0) @binding(2) var<storage, read_write> t_v: array<f32>;
+@group(0) @binding(3) var<storage, read> t_grad: array<f32>;
+
+struct Params {
+  numel: u32,
+  _pad0: u32,
+  _pad1: u32,
+  _pad2: u32,
+  lr: f32,
+  beta1: f32,
+  beta2: f32,
+  eps: f32,
+  weight_decay: f32,
+  bias_correction1: f32,
+  bias_correction2: f32,
+  _pad3: f32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= params.numel) {
+    return;
+  }
+  let g = t_grad[i];
+  var p = t_param[i];
+  p = p - params.lr * params.weight_decay * p;
+  let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g;
+  let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g;
+  t_m[i] = m;
+  t_v[i] = v;
+  let mhat = m / params.bias_correction1;
+  let vhat = v / params.bias_correction2;
+  t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps);
+}
diff --git a/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h b/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h
new file mode 100644
index 00000000000..38cbda98c6b
--- /dev/null
+++ b/backends/webgpu/runtime/ops/adamw/adamw_step_wgsl.h
@@ -0,0 +1,65 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from adamw_step.wgsl - DO NOT EDIT.
+// wgsl-sha256: 0957c04168872db5e2b39cf5f26beefba27f3b5514ec69cece4de5145e97156f
+inline constexpr const char* kAdamwStepWGSL = R"(
+
+@group(0) @binding(0) var<storage, read_write> t_param: array<f32>;
+@group(0) @binding(1) var<storage, read_write> t_m: array<f32>;
+@group(0) @binding(2) var<storage, read_write> t_v: array<f32>;
+@group(0) @binding(3) var<storage, read> t_grad: array<f32>;
+
+struct Params {
+  numel: u32,
+  _pad0: u32,
+  _pad1: u32,
+  _pad2: u32,
+  lr: f32,
+  beta1: f32,
+  beta2: f32,
+  eps: f32,
+  weight_decay: f32,
+  bias_correction1: f32,
+  bias_correction2: f32,
+  _pad3: f32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let i = gid.x;
+  if (i >= params.numel) {
+    return;
+  }
+  let g = t_grad[i];
+  var p = t_param[i];
+  p = p - params.lr * params.weight_decay * p;
+  let m = params.beta1 * t_m[i] + (1.0 - params.beta1) * g;
+  let v = params.beta2 * t_v[i] + (1.0 - params.beta2) * g * g;
+  t_m[i] = m;
+  t_v[i] = v;
+  let mhat = m / params.bias_correction1;
+  let vhat = v / params.bias_correction2;
+  t_param[i] = p - params.lr * mhat / (sqrt(vhat) + params.eps);
+}
+)";
+
+inline constexpr uint32_t kAdamwStepWorkgroupSizeX = 64;
+inline constexpr uint32_t kAdamwStepWorkgroupSizeY = 1;
+inline constexpr uint32_t kAdamwStepWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp
new file mode 100644
index 00000000000..27c87a2bbeb
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/BooleanOp.cpp
@@ -0,0 +1,251 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+#include <string>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Shared uniform for the byte-packed-bool op family. `scalar` is unused (0) for
+// the scalar-less (unary) logical_not variant.
+struct BoolOpParams {
+  uint32_t num_elements;
+  float scalar;
+  uint32_t _pad1;
+  uint32_t _pad2;
+};
+
+float read_scalar(WebGPUGraph& graph, int id, const char* op_name) {
+  if (graph.get_value_type(id) == WebGPUGraph::ValueType::Double) {
+    return static_cast<float>(graph.get_double(id));
+  }
+  if (graph.get_value_type(id) == WebGPUGraph::ValueType::Int) {
+    return static_cast<float>(graph.get_int(id));
+  }
+  throw std::runtime_error(std::string(op_name) + ": scalar is not int/double");
+}
+
+// Dispatch a byte-packed-bool op (scalar compares + logical_not): read an input
+// buffer, write a u32 output packing 4 bool bytes per word, one thread per word
+// (no inter-thread write race). The caller validates dtypes/shapes and supplies
+// `numel` (logical element count = prod(dims)), the scalar (0 when unused), and
+// the per-variant shader from boolean_op.yaml.
+void dispatch_bool_op(
+    WebGPUGraph& graph,
+    int self_id,
+    int out_id,
+    uint32_t numel,
+    float scalar,
+    const char* wgsl,
+    uint32_t wg_size_x,
+    const char* op_name) {
+  WGPUDevice device = graph.device();
+  const auto& self_tensor = graph.get_tensor(self_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+
+  // The output (and logical_not's packed-bool input) is a u32 array; round the
+  // binding up to whole words even when the byte count isn't a multiple of 4.
+  const size_t self_bind_size = (self_tensor.nbytes + 3) & ~size_t(3);
+  const size_t out_bind_size = (out_tensor.nbytes + 3) & ~size_t(3);
+  const uint32_t n_words = (numel + 3u) / 4u;
+
+  uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x);
+  uint32_t workgroup_count =
+      utils::compute_1d_workgroup_count(device, n_words, wg_size, op_name);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  BoolOpParams params = {numel, scalar, 0u, 0u};
+  WGPUBuffer params_buf =
+      utils::make_uniform(device, &params, sizeof(BoolOpParams));
+  graph.add_uniform_buffer_bytes(sizeof(BoolOpParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {wgsl, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[3] = {};
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].buffer.type = WGPUBufferBindingType_Storage;
+  entries[2].buffer.type = WGPUBufferBindingType_Uniform;
+  for (uint32_t i = 0; i < 3; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+  }
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 3;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[3] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = self_tensor.buffer;
+  bg_entries[0].size = self_bind_size;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = out_tensor.buffer;
+  bg_entries[1].size = out_bind_size;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = params_buf;
+  bg_entries[2].size = sizeof(BoolOpParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 3;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  const size_t dispatch_idx =
+      graph.add_dispatch({pipeline, bind_group, workgroup_count});
+
+  WGPUBuffer p_buf = params_buf;
+  auto resize = [self_id, out_id, scalar, wg_size, dispatch_idx, p_buf, op_name](
+                    WebGPUGraph& g) {
+    const auto& d = g.cur_dims(self_id);
+    uint32_t n = 1u;
+    for (auto x : d) {
+      n *= static_cast<uint32_t>(x);
+    }
+    g.set_cur_dims(out_id, d);
+    BoolOpParams p = {n, scalar, 0u, 0u};
+    wgpuQueueWriteBuffer(g.queue(), p_buf, 0, &p, sizeof(p));
+    const uint32_t nw = (n + 3u) / 4u;
+    g.dispatch_at(dispatch_idx).workgroup_count_x =
+        utils::compute_1d_workgroup_count(g.device(), nw, wg_size, op_name);
+  };
+  graph.add_tensor_resize_hook(self_id, resize);
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(params_buf);
+}
+
+// cmp(self[i], scalar) -> byte-packed bool. args: [self, scalar, out].
+void compare_dispatch(
+    WebGPUGraph& graph,
+    const std::vector<int>& args,
+    const char* wgsl,
+    uint32_t wg_size_x,
+    const char* op_name) {
+  const int self_id = args.at(0);
+  const int out_id = args.at(args.size() - 1);
+  const float scalar = read_scalar(graph, args.at(1), op_name);
+
+  const auto& self_tensor = graph.get_tensor(self_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+  if (self_tensor.buffer == nullptr || out_tensor.buffer == nullptr) {
+    throw std::runtime_error(std::string(op_name) + ": null buffer binding");
+  }
+  if (self_tensor.nbytes % sizeof(float) != 0) {
+    throw std::runtime_error(std::string(op_name) + ": self is not fp32");
+  }
+  const uint32_t numel =
+      static_cast<uint32_t>(self_tensor.nbytes / sizeof(float));
+  if (out_tensor.nbytes != static_cast<size_t>(numel)) {
+    throw std::runtime_error(
+        std::string(op_name) + ": out is not a 1-byte (bool) tensor");
+  }
+  dispatch_bool_op(
+      graph, self_id, out_id, numel, scalar, wgsl, wg_size_x, op_name);
+}
+
+void eq_scalar_impl(WebGPUGraph& g, const std::vector<int>& a) {
+  compare_dispatch(g, a, kCompareEqWGSL, kCompareEqWorkgroupSizeX, "eq.Scalar");
+}
+void ne_scalar_impl(WebGPUGraph& g, const std::vector<int>& a) {
+  compare_dispatch(g, a, kCompareNeWGSL, kCompareNeWorkgroupSizeX, "ne.Scalar");
+}
+void le_scalar_impl(WebGPUGraph& g, const std::vector<int>& a) {
+  compare_dispatch(g, a, kCompareLeWGSL, kCompareLeWorkgroupSizeX, "le.Scalar");
+}
+void ge_scalar_impl(WebGPUGraph& g, const std::vector<int>& a) {
+  compare_dispatch(g, a, kCompareGeWGSL, kCompareGeWorkgroupSizeX, "ge.Scalar");
+}
+void lt_scalar_impl(WebGPUGraph& g, const std::vector<int>& a) {
+  compare_dispatch(g, a, kCompareLtWGSL, kCompareLtWorkgroupSizeX, "lt.Scalar");
+}
+void gt_scalar_impl(WebGPUGraph& g, const std::vector<int>& a) {
+  compare_dispatch(g, a, kCompareGtWGSL, kCompareGtWorkgroupSizeX, "gt.Scalar");
+}
+
+// logical_not: byte-packed bool -> byte-packed bool. args: [self, out].
+void logical_not_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int self_id = args.at(0);
+  const int out_id = args.at(args.size() - 1);
+
+  const auto& self_tensor = graph.get_tensor(self_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+  if (self_tensor.buffer == nullptr || out_tensor.buffer == nullptr) {
+    throw std::runtime_error("logical_not: null buffer binding");
+  }
+  if (out_tensor.nbytes != self_tensor.nbytes) {
+    throw std::runtime_error("logical_not: self/out byte-size mismatch");
+  }
+  const uint32_t numel = static_cast<uint32_t>(self_tensor.nbytes);
+  dispatch_bool_op(
+      graph,
+      self_id,
+      out_id,
+      numel,
+      0.0f,
+      kLogicalNotWGSL,
+      kLogicalNotWorkgroupSizeX,
+      "logical_not");
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.eq.Scalar, eq_scalar_impl);
+  WEBGPU_REGISTER_OP(aten.ne.Scalar, ne_scalar_impl);
+  WEBGPU_REGISTER_OP(aten.le.Scalar, le_scalar_impl);
+  WEBGPU_REGISTER_OP(aten.ge.Scalar, ge_scalar_impl);
+  WEBGPU_REGISTER_OP(aten.lt.Scalar, lt_scalar_impl);
+  WEBGPU_REGISTER_OP(aten.gt.Scalar, gt_scalar_impl);
+  WEBGPU_REGISTER_OP(aten.logical_not.default, logical_not_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl b/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl
new file mode 100644
index 00000000000..4f32792d7dc
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/boolean_op.wgsl
@@ -0,0 +1,37 @@
+@group(0) @binding(0) var<storage, read> input: array<${IN_TYPE}>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return ${OP_EXPR};
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
diff --git a/backends/webgpu/runtime/ops/boolean_op/boolean_op.yaml b/backends/webgpu/runtime/ops/boolean_op/boolean_op.yaml
new file mode 100644
index 00000000000..4b14860a7e2
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/boolean_op.yaml
@@ -0,0 +1,20 @@
+boolean_op:
+  parameter_names_with_default_values:
+    IN_TYPE: f32
+    OP_EXPR: input[i] == params.scalar
+  shader_variants:
+    - NAME: compare_eq
+      OP_EXPR: input[i] == params.scalar
+    - NAME: compare_ne
+      OP_EXPR: input[i] != params.scalar
+    - NAME: compare_le
+      OP_EXPR: input[i] <= params.scalar
+    - NAME: compare_ge
+      OP_EXPR: input[i] >= params.scalar
+    - NAME: compare_lt
+      OP_EXPR: input[i] < params.scalar
+    - NAME: compare_gt
+      OP_EXPR: input[i] > params.scalar
+    - NAME: logical_not
+      IN_TYPE: u32
+      OP_EXPR: ((input[i >> 2u] >> ((i & 3u) * 8u)) & 0xFFu) == 0u
diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h
new file mode 100644
index 00000000000..8b7b44e7430
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/compare_eq_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 2a3c203d0255086a6e67a9e7cb08538858e8e6cb6a5542f6acebf04b8ccca6b7
+inline constexpr const char* kCompareEqWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return input[i] == params.scalar;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kCompareEqWorkgroupSizeX = 64;
+inline constexpr uint32_t kCompareEqWorkgroupSizeY = 1;
+inline constexpr uint32_t kCompareEqWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h
new file mode 100644
index 00000000000..bcf7239bd7e
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/compare_ge_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 5fb319a54ed666644f118119639dd8cd33256ea5e5163c4c1392ce6e84b369cc
+inline constexpr const char* kCompareGeWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return input[i] >= params.scalar;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kCompareGeWorkgroupSizeX = 64;
+inline constexpr uint32_t kCompareGeWorkgroupSizeY = 1;
+inline constexpr uint32_t kCompareGeWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h
new file mode 100644
index 00000000000..8f911180c8a
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/compare_gt_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: b4a35be8a34774a47be457130125336b5accf8d24615514f3fe1559159644491
+inline constexpr const char* kCompareGtWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return input[i] > params.scalar;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kCompareGtWorkgroupSizeX = 64;
+inline constexpr uint32_t kCompareGtWorkgroupSizeY = 1;
+inline constexpr uint32_t kCompareGtWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h
new file mode 100644
index 00000000000..87b4ff91c2a
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/compare_le_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 64b8988db9611c1aaff3ba373d32ecb7b276340df337b38a49e12a92727c00cb
+inline constexpr const char* kCompareLeWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return input[i] <= params.scalar;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kCompareLeWorkgroupSizeX = 64;
+inline constexpr uint32_t kCompareLeWorkgroupSizeY = 1;
+inline constexpr uint32_t kCompareLeWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h
new file mode 100644
index 00000000000..930601d229c
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/compare_lt_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: df350a7d0b099d38f9e77a27ffe9e56afa93b2c6a3ff84006227e1c1c0b96521
+inline constexpr const char* kCompareLtWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return input[i] < params.scalar;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kCompareLtWorkgroupSizeX = 64;
+inline constexpr uint32_t kCompareLtWorkgroupSizeY = 1;
+inline constexpr uint32_t kCompareLtWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h
new file mode 100644
index 00000000000..9020fb77671
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/compare_ne_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 4f5492c393e494e6a9a734ea797cbfa85f0f1579451fc46f4291f57e1b5daaa2
+inline constexpr const char* kCompareNeWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return input[i] != params.scalar;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kCompareNeWorkgroupSizeX = 64;
+inline constexpr uint32_t kCompareNeWorkgroupSizeY = 1;
+inline constexpr uint32_t kCompareNeWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h b/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h
new file mode 100644
index 00000000000..65239060454
--- /dev/null
+++ b/backends/webgpu/runtime/ops/boolean_op/logical_not_wgsl.h
@@ -0,0 +1,61 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from boolean_op.wgsl - DO NOT EDIT.
+// wgsl-sha256: 06dc55c61e55c35eb12abf548fafb495097fa8d22daea930af6a32ed724a5b6a
+inline constexpr const char* kLogicalNotWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<u32>;
+@group(0) @binding(1) var<storage, read_write> output: array<u32>;
+
+struct Params {
+  num_elements: u32,
+  scalar: f32,
+  _pad1: u32,
+  _pad2: u32,
+}
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+// Per-variant predicate substituted from boolean_op.yaml (compare / logical_not).
+fn elem_bool(i: u32) -> bool {
+  return ((input[i >> 2u] >> ((i & 3u) * 8u)) & 0xFFu) == 0u;
+}
+
+// One thread per output u32 word packs 4 bool bytes -> no inter-thread race.
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let word_idx = gid.x;
+    let n_words = (params.num_elements + 3u) / 4u;
+    if (word_idx >= n_words) {
+        return;
+    }
+    var packed: u32 = 0u;
+    for (var j: u32 = 0u; j < 4u; j = j + 1u) {
+        let i = word_idx * 4u + j;
+        if (i < params.num_elements) {
+            if (elem_bool(i)) {
+                packed = packed | (1u << (j * 8u));
+            }
+        }
+    }
+    output[word_idx] = packed;
+}
+)";
+
+inline constexpr uint32_t kLogicalNotWorkgroupSizeX = 64;
+inline constexpr uint32_t kLogicalNotWorkgroupSizeY = 1;
+inline constexpr uint32_t kLogicalNotWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/dim_order/DimOrder.cpp b/backends/webgpu/runtime/ops/dim_order/DimOrder.cpp
new file mode 100644
index 00000000000..e57e0375a4f
--- /dev/null
+++ b/backends/webgpu/runtime/ops/dim_order/DimOrder.cpp
@@ -0,0 +1,31 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/view_copy/view_copy.h>
+
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// _clone_dim_order = numel-preserving flat copy (shared DMA helper).
+void clone_dim_order_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_flat_copy(graph, args.at(0), args.at(args.size() - 1));
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(
+      dim_order_ops._clone_dim_order.default, clone_dim_order_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/embedding/Embedding.cpp b/backends/webgpu/runtime/ops/embedding/Embedding.cpp
new file mode 100644
index 00000000000..b05b1af7d9d
--- /dev/null
+++ b/backends/webgpu/runtime/ops/embedding/Embedding.cpp
@@ -0,0 +1,157 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct EmbeddingParams {
+  uint32_t num_elements;
+  uint32_t dim;
+  uint32_t _pad[2];
+};
+
+// aten.embedding: out[row, :] = weight[indices[row], :] (fp32 weight, i32 idx).
+void embedding_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int weight_id = args.at(0);
+  const int indices_id = args.at(1);
+  const int out_id = args.at(args.size() - 1);
+
+  WGPUDevice device = graph.device();
+  const auto& weight = graph.get_tensor(weight_id);
+  const auto& indices = graph.get_tensor(indices_id);
+  const auto& out = graph.get_tensor(out_id);
+
+  if (weight.buffer == nullptr || indices.buffer == nullptr ||
+      out.buffer == nullptr) {
+    throw std::runtime_error("embedding: null buffer binding");
+  }
+  if (weight.dims.size() != 2) {
+    throw std::runtime_error("embedding: weight must be 2D [vocab, dim]");
+  }
+  const uint32_t dim = static_cast<uint32_t>(weight.dims[1]);
+  if (dim == 0) {
+    throw std::runtime_error("embedding: dim == 0");
+  }
+  const size_t out_numel = out.nbytes / sizeof(float);
+  // Index is the int32 downcast of the int64 indices (mirror index op).
+  const size_t index_numel = indices.nbytes / sizeof(int32_t);
+  if (out.nbytes != out_numel * sizeof(float) ||
+      weight.nbytes % sizeof(float) != 0 ||
+      indices.nbytes != index_numel * sizeof(int32_t)) {
+    throw std::runtime_error(
+        "embedding: fp32 weight/out + i32 indices required");
+  }
+  if (out_numel != index_numel * dim) {
+    throw std::runtime_error("embedding: out numel != num_indices * dim");
+  }
+
+  uint32_t num_elements = static_cast<uint32_t>(out_numel);
+  uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kEmbeddingWorkgroupSizeX);
+  utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
+      device, num_elements, wg_size, "embedding");
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  EmbeddingParams params = {};
+  params.num_elements = num_elements;
+  params.dim = dim;
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(EmbeddingParams));
+  graph.add_uniform_buffer_bytes(sizeof(EmbeddingParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kEmbeddingWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[4] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].binding = 1;
+  entries[1].visibility = WGPUShaderStage_Compute;
+  entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[2].binding = 2;
+  entries[2].visibility = WGPUShaderStage_Compute;
+  entries[2].buffer.type = WGPUBufferBindingType_Storage;
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = weight.buffer;
+  bg_entries[0].size = weight.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = indices.buffer;
+  bg_entries[1].size = indices.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = out.buffer;
+  bg_entries[2].size = out.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = uniform_buffer;
+  bg_entries[3].size = sizeof(EmbeddingParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch(
+      {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  wgpuBufferRelease(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.embedding.default, embedding_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/embedding/embedding.wgsl b/backends/webgpu/runtime/ops/embedding/embedding.wgsl
new file mode 100644
index 00000000000..9fece09d293
--- /dev/null
+++ b/backends/webgpu/runtime/ops/embedding/embedding.wgsl
@@ -0,0 +1,25 @@
+@group(0) @binding(0) var<storage, read> weight: array<f32>;
+@group(0) @binding(1) var<storage, read> indices: array<i32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  dim: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    let row = idx / params.dim;
+    let col = idx % params.dim;
+    let row_id = u32(indices[row]);
+    output[idx] = weight[row_id * params.dim + col];
+}
diff --git a/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h b/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h
new file mode 100644
index 00000000000..4a790c74d65
--- /dev/null
+++ b/backends/webgpu/runtime/ops/embedding/embedding_wgsl.h
@@ -0,0 +1,49 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from embedding.wgsl - DO NOT EDIT.
+// wgsl-sha256: e7c94da588a9dd40c0e1be7d3c6ef33310c71390c57f6f330ae23a18866d0ea7
+inline constexpr const char* kEmbeddingWGSL = R"(
+@group(0) @binding(0) var<storage, read> weight: array<f32>;
+@group(0) @binding(1) var<storage, read> indices: array<i32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  dim: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    let row = idx / params.dim;
+    let col = idx % params.dim;
+    let row_id = u32(indices[row]);
+    output[idx] = weight[row_id * params.dim + col];
+}
+)";
+
+inline constexpr uint32_t kEmbeddingWorkgroupSizeX = 256;
+inline constexpr uint32_t kEmbeddingWorkgroupSizeY = 1;
+inline constexpr uint32_t kEmbeddingWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp
new file mode 100644
index 00000000000..507677c7bca
--- /dev/null
+++ b/backends/webgpu/runtime/ops/expand_copy/ExpandCopy.cpp
@@ -0,0 +1,139 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
+#include <executorch/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// out coord -> in coord; size-1 in-dims broadcast via clamp (mirrors mul).
+void expand_copy_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int in_id = args.at(0);
+  const int out_id = args.at(args.size() - 1);
+
+  if (graph.get_value_type(in_id) != WebGPUGraph::ValueType::Tensor ||
+      graph.get_value_type(out_id) != WebGPUGraph::ValueType::Tensor) {
+    throw std::runtime_error("expand_copy: in/out arg is not a tensor");
+  }
+
+  WGPUDevice device = graph.device();
+  const auto& in_tensor = graph.get_tensor(in_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+
+  TensorMeta out_meta;
+  TensorMeta in_meta;
+  fill_tensor_meta(out_tensor, &out_meta);
+  fill_tensor_meta_broadcast(in_tensor, out_meta.ndim, &in_meta);
+  if (out_tensor.nbytes !=
+          static_cast<size_t>(out_meta.numel) * sizeof(float) ||
+      in_tensor.nbytes != static_cast<size_t>(in_meta.numel) * sizeof(float)) {
+    throw std::runtime_error(
+        "expand_copy: non-fp32 operand (nbytes != numel*4)");
+  }
+
+  uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kExpandCopyWorkgroupSizeX);
+  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, out_meta.numel, wg_size, "expand_copy");
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUBuffer out_meta_buf =
+      utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
+  WGPUBuffer in_meta_buf =
+      utils::make_uniform(device, &in_meta, sizeof(TensorMeta));
+  graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kExpandCopyWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[4] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].binding = 1;
+  entries[1].visibility = WGPUShaderStage_Compute;
+  entries[1].buffer.type = WGPUBufferBindingType_Storage;
+  entries[2].binding = 2;
+  entries[2].visibility = WGPUShaderStage_Compute;
+  entries[2].buffer.type = WGPUBufferBindingType_Uniform;
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = in_tensor.buffer;
+  bg_entries[0].size = in_tensor.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = out_tensor.buffer;
+  bg_entries[1].size = out_tensor.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = out_meta_buf;
+  bg_entries[2].size = sizeof(TensorMeta);
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = in_meta_buf;
+  bg_entries[3].size = sizeof(TensorMeta);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  wgpuBufferRelease(out_meta_buf);
+  wgpuBufferRelease(in_meta_buf);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.expand_copy.default, expand_copy_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl
new file mode 100644
index 00000000000..d6c65588978
--- /dev/null
+++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy.wgsl
@@ -0,0 +1,29 @@
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(2) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(3) var<uniform> in_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= out_meta.numel) {
+        return;
+    }
+    var rem = idx;
+    var l: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d];
+    }
+    output[idx] = input[l];
+}
diff --git a/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h
new file mode 100644
index 00000000000..56b3f708767
--- /dev/null
+++ b/backends/webgpu/runtime/ops/expand_copy/expand_copy_wgsl.h
@@ -0,0 +1,53 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from expand_copy.wgsl - DO NOT EDIT.
+// wgsl-sha256: ad996b7fd6eca5c5773715af3a9a117da83ef522042eb0ee918a623096417815
+inline constexpr const char* kExpandCopyWGSL = R"(
+@group(0) @binding(0) var<storage, read> input: array<f32>;
+@group(0) @binding(1) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(2) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(3) var<uniform> in_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= out_meta.numel) {
+        return;
+    }
+    var rem = idx;
+    var l: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        l = l + min(coord, in_meta.sizes[d] - 1u) * in_meta.strides[d];
+    }
+    output[idx] = input[l];
+}
+)";
+
+inline constexpr uint32_t kExpandCopyWorkgroupSizeX = 64;
+inline constexpr uint32_t kExpandCopyWorkgroupSizeY = 1;
+inline constexpr uint32_t kExpandCopyWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/fill/Fill.cpp b/backends/webgpu/runtime/ops/fill/Fill.cpp
new file mode 100644
index 00000000000..2ccd48cc8fa
--- /dev/null
+++ b/backends/webgpu/runtime/ops/fill/Fill.cpp
@@ -0,0 +1,177 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/fill/fill_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Uniform buffer layout matching the WGSL Params struct (16-byte aligned).
+struct FillParams {
+  uint32_t num_elements;
+  float fill_value;
+  uint32_t _pad[2];
+};
+
+float read_scalar(WebGPUGraph& graph, int id, const char* op_name) {
+  if (graph.get_value_type(id) == WebGPUGraph::ValueType::Double) {
+    return static_cast<float>(graph.get_double(id));
+  }
+  if (graph.get_value_type(id) == WebGPUGraph::ValueType::Int) {
+    return static_cast<float>(graph.get_int(id));
+  }
+  throw std::runtime_error(
+      std::string(op_name) + ": fill value is not a scalar");
+}
+
+// Fills the (pre-allocated) output buffer with a constant scalar.
+void add_fill(
+    WebGPUGraph& graph,
+    int out_id,
+    float fill_value,
+    const char* op_name) {
+  WGPUDevice device = graph.device();
+  const auto& out_tensor = graph.get_tensor(out_id);
+  if (out_tensor.buffer == nullptr) {
+    throw std::runtime_error(std::string(op_name) + ": null output buffer");
+  }
+  uint32_t num_elements =
+      static_cast<uint32_t>(out_tensor.nbytes / sizeof(float));
+
+  uint32_t wg_size = utils::clamp_workgroup_size(device, kFillWorkgroupSizeX);
+  utils::WgCount workgroup_count =
+      utils::compute_2d_workgroup_count(device, num_elements, wg_size, op_name);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  FillParams params = {};
+  params.num_elements = num_elements;
+  params.fill_value = fill_value;
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(FillParams));
+  graph.add_uniform_buffer_bytes(sizeof(FillParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kFillWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[2] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  entries[1].binding = 1;
+  entries[1].visibility = WGPUShaderStage_Compute;
+  entries[1].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 2;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[2] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = out_tensor.buffer;
+  bg_entries[0].size = out_tensor.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = uniform_buffer;
+  bg_entries[1].size = sizeof(FillParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 2;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  const size_t dispatch_idx = graph.add_dispatch(
+      {pipeline, bind_group, workgroup_count.x, "", workgroup_count.y});
+
+  // Dynamic shapes: recompute num_elements/dispatch from the live output dims.
+  WGPUBuffer params_buf = uniform_buffer;
+  graph.add_tensor_resize_hook(
+      out_id,
+      [out_id, fill_value, wg_size, dispatch_idx, params_buf](WebGPUGraph& g) {
+        const auto& d = g.cur_dims(out_id);
+        const uint64_t numel = utils::numel_of(d);
+        FillParams p = {};
+        p.num_elements = static_cast<uint32_t>(numel);
+        p.fill_value = fill_value;
+        wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
+        const utils::WgCount wgc = utils::compute_2d_workgroup_count(
+            g.device(), static_cast<uint32_t>(numel), wg_size, "fill(resize)");
+        g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
+        g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
+      });
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+void full_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_fill(
+      graph,
+      args.at(args.size() - 1),
+      read_scalar(graph, args.at(1), "full"),
+      "full");
+}
+
+void full_like_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_fill(
+      graph,
+      args.at(args.size() - 1),
+      read_scalar(graph, args.at(1), "full_like"),
+      "full_like");
+}
+
+void scalar_tensor_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  add_fill(
+      graph,
+      args.at(args.size() - 1),
+      read_scalar(graph, args.at(0), "scalar_tensor"),
+      "scalar_tensor");
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.full.default, full_impl);
+  WEBGPU_REGISTER_OP(aten.full_like.default, full_like_impl);
+  WEBGPU_REGISTER_OP(aten.scalar_tensor.default, scalar_tensor_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/fill/fill.wgsl b/backends/webgpu/runtime/ops/fill/fill.wgsl
new file mode 100644
index 00000000000..684c350e8d6
--- /dev/null
+++ b/backends/webgpu/runtime/ops/fill/fill.wgsl
@@ -0,0 +1,20 @@
+@group(0) @binding(0) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  fill_value: f32,
+}
+@group(0) @binding(1) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    output[idx] = params.fill_value;
+}
diff --git a/backends/webgpu/runtime/ops/fill/fill_wgsl.h b/backends/webgpu/runtime/ops/fill/fill_wgsl.h
new file mode 100644
index 00000000000..a083af10af3
--- /dev/null
+++ b/backends/webgpu/runtime/ops/fill/fill_wgsl.h
@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from fill.wgsl - DO NOT EDIT.
+// wgsl-sha256: 9acfd8079cf404b22856c157b339922911e15bb0e808aa76b5a7fb81a8505cf7
+inline constexpr const char* kFillWGSL = R"(
+@group(0) @binding(0) var<storage, read_write> output: array<f32>;
+
+struct Params {
+  num_elements: u32,
+  fill_value: f32,
+}
+@group(0) @binding(1) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(global_invocation_id) gid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+    let idx = gid.x + gid.y * (num_workgroups.x * wg_size);
+    if (idx >= params.num_elements) {
+        return;
+    }
+    output[idx] = params.fill_value;
+}
+)";
+
+inline constexpr uint32_t kFillWorkgroupSizeX = 256;
+inline constexpr uint32_t kFillWorkgroupSizeY = 1;
+inline constexpr uint32_t kFillWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp b/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp
new file mode 100644
index 00000000000..e51052f978d
--- /dev/null
+++ b/backends/webgpu/runtime/ops/fused_ce/FusedCe.cpp
@@ -0,0 +1,285 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Uniform layout matching the fused_ce.wgsl Params struct (16-byte aligned).
+struct FusedCeParams {
+  uint32_t vocab;
+  uint32_t n_rows;
+  float n_valid;
+  float _pad0;
+};
+static_assert(sizeof(FusedCeParams) == 16, "FusedCeParams must be 16 bytes");
+
+// Mirror reduce.wgsl Params (file-local in Reduce.cpp; re-declared here).
+struct ReduceParams {
+  uint32_t outer;
+  uint32_t r;
+  uint32_t inner;
+  uint32_t is_mean;
+};
+static_assert(sizeof(ReduceParams) == 16, "ReduceParams must be 16 bytes");
+
+WGPUShaderModule make_shader(WGPUDevice device, const char* wgsl) {
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {wgsl, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  return wgpuDeviceCreateShaderModule(device, &shader_desc);
+}
+
+WGPUBuffer create_uniform(
+    WebGPUGraph& graph,
+    WGPUDevice device,
+    const void* data,
+    size_t size) {
+  WGPUBufferDescriptor desc = {};
+  desc.size = size;
+  desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  desc.mappedAtCreation = true;
+  WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &desc);
+  std::memcpy(wgpuBufferGetMappedRange(buffer, 0, size), data, size);
+  wgpuBufferUnmap(buffer);
+  graph.add_uniform_buffer_bytes(size);
+  return buffer;
+}
+
+// out valuelist packs the 2-tuple (loss, dlogits) as one id.
+void fused_ce_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int logits_id = args.at(0);
+  const int labels_id = args.at(1);
+  const int n_valid_id = args.at(2);
+  const std::vector<int>& outs = graph.get_value_list(args.at(3));
+  if (outs.size() != 2) {
+    throw std::runtime_error(
+        "WebGPU fused_ce: expected 2 outputs (loss, dlogits)");
+  }
+  const int loss_id = outs.at(0);
+  const int dlogits_id = outs.at(1);
+
+  WGPUDevice device = graph.device();
+  const auto& logits = graph.get_tensor(logits_id);
+  const auto& labels = graph.get_tensor(labels_id);
+  const auto& dlogits = graph.get_tensor(dlogits_id);
+  const auto& loss = graph.get_tensor(loss_id);
+
+  if (logits.dims.size() != 2) {
+    throw std::runtime_error("WebGPU fused_ce: logits must be 2D [N, V]");
+  }
+  const uint64_t n_rows = static_cast<uint64_t>(logits.dims[0]);
+  const uint64_t vocab = static_cast<uint64_t>(logits.dims[1]);
+  const uint64_t numel = n_rows * vocab;
+
+  if (dlogits.dims != logits.dims) {
+    throw std::runtime_error(
+        "WebGPU fused_ce: dlogits shape must match logits");
+  }
+  if (logits.nbytes != numel * sizeof(float) ||
+      dlogits.nbytes != numel * sizeof(float)) {
+    throw std::runtime_error("WebGPU fused_ce: logits/dlogits fp32-only");
+  }
+  if (labels.nbytes != n_rows * sizeof(int32_t)) {
+    throw std::runtime_error("WebGPU fused_ce: labels must be int32 [N]");
+  }
+  if (loss.nbytes != sizeof(float)) {
+    throw std::runtime_error("WebGPU fused_ce: loss must be a scalar [1]");
+  }
+  if (graph.get_value_type(n_valid_id) != WebGPUGraph::ValueType::Double) {
+    throw std::runtime_error("WebGPU fused_ce: n_valid must be a float scalar");
+  }
+  const double n_valid = graph.get_double(n_valid_id);
+  if (n_valid <= 0.0) {
+    throw std::runtime_error("WebGPU fused_ce: n_valid must be positive");
+  }
+  if (n_rows > utils::queried_max_workgroups(device)) {
+    throw std::runtime_error("WebGPU fused_ce: n_rows exceeds dispatch limit");
+  }
+
+  WGPUBuffer loss_partial = graph.create_scratch_buffer(n_rows * sizeof(float));
+
+  // one workgroup per row
+  FusedCeParams ce_params = {};
+  ce_params.vocab = static_cast<uint32_t>(vocab);
+  ce_params.n_rows = static_cast<uint32_t>(n_rows);
+  ce_params.n_valid = static_cast<float>(n_valid);
+  const uint32_t ce_wg =
+      utils::clamp_workgroup_size(device, kFusedCeWorkgroupSizeX);
+  WGPUBuffer ce_uniform =
+      create_uniform(graph, device, &ce_params, sizeof(ce_params));
+  WGPUShaderModule ce_shader = make_shader(device, kFusedCeWGSL);
+
+  WGPUBindGroupLayoutEntry ce_entries[5] = {};
+  ce_entries[0].binding = 0;
+  ce_entries[0].visibility = WGPUShaderStage_Compute;
+  ce_entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  ce_entries[1].binding = 1;
+  ce_entries[1].visibility = WGPUShaderStage_Compute;
+  ce_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  ce_entries[2].binding = 2;
+  ce_entries[2].visibility = WGPUShaderStage_Compute;
+  ce_entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  ce_entries[3].binding = 3;
+  ce_entries[3].visibility = WGPUShaderStage_Compute;
+  ce_entries[3].buffer.type = WGPUBufferBindingType_Storage;
+  ce_entries[4].binding = 4;
+  ce_entries[4].visibility = WGPUShaderStage_Compute;
+  ce_entries[4].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor ce_bgl_desc = {};
+  ce_bgl_desc.entryCount = 5;
+  ce_bgl_desc.entries = ce_entries;
+  WGPUBindGroupLayout ce_bgl =
+      wgpuDeviceCreateBindGroupLayout(device, &ce_bgl_desc);
+
+  WGPUPipelineLayoutDescriptor ce_pl_desc = {};
+  ce_pl_desc.bindGroupLayoutCount = 1;
+  ce_pl_desc.bindGroupLayouts = &ce_bgl;
+  WGPUPipelineLayout ce_pl =
+      wgpuDeviceCreatePipelineLayout(device, &ce_pl_desc);
+
+  WGPUConstantEntry ce_wg_const = {};
+  ce_wg_const.key = {"wg_size", WGPU_STRLEN};
+  ce_wg_const.value = static_cast<double>(ce_wg);
+
+  WGPUComputePipelineDescriptor ce_pipe_desc = {};
+  ce_pipe_desc.layout = ce_pl;
+  ce_pipe_desc.compute.module = ce_shader;
+  ce_pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  ce_pipe_desc.compute.constantCount = 1;
+  ce_pipe_desc.compute.constants = &ce_wg_const;
+  WGPUComputePipeline ce_pipe =
+      wgpuDeviceCreateComputePipeline(device, &ce_pipe_desc);
+
+  WGPUBindGroupEntry ce_bg[5] = {};
+  ce_bg[0].binding = 0;
+  ce_bg[0].buffer = dlogits.buffer;
+  ce_bg[0].size = dlogits.nbytes;
+  ce_bg[1].binding = 1;
+  ce_bg[1].buffer = logits.buffer;
+  ce_bg[1].size = logits.nbytes;
+  ce_bg[2].binding = 2;
+  ce_bg[2].buffer = labels.buffer;
+  ce_bg[2].size = labels.nbytes;
+  ce_bg[3].binding = 3;
+  ce_bg[3].buffer = loss_partial;
+  ce_bg[3].size = n_rows * sizeof(float);
+  ce_bg[4].binding = 4;
+  ce_bg[4].buffer = ce_uniform;
+  ce_bg[4].size = sizeof(ce_params);
+
+  WGPUBindGroupDescriptor ce_bg_desc = {};
+  ce_bg_desc.layout = ce_bgl;
+  ce_bg_desc.entryCount = 5;
+  ce_bg_desc.entries = ce_bg;
+  WGPUBindGroup ce_bind_group = wgpuDeviceCreateBindGroup(device, &ce_bg_desc);
+
+  graph.add_dispatch(
+      {ce_pipe, ce_bind_group, static_cast<uint32_t>(n_rows), "fused_ce"});
+
+  wgpuShaderModuleRelease(ce_shader);
+  wgpuBindGroupLayoutRelease(ce_bgl);
+  wgpuPipelineLayoutRelease(ce_pl);
+  wgpuBufferRelease(ce_uniform);
+
+  // reduce loss_partial[N] -> loss[1] (reuses reduce.wgsl)
+  ReduceParams r_params = {};
+  r_params.outer = 1u;
+  r_params.r = static_cast<uint32_t>(n_rows);
+  r_params.inner = 1u;
+  r_params.is_mean = 0u; // mean already folded into loss_partial
+  const uint32_t r_wg =
+      utils::clamp_workgroup_size(device, kReduceWorkgroupSizeX);
+  const uint32_t r_wgc =
+      utils::compute_1d_workgroup_count(device, 1u, r_wg, "fused_ce_reduce");
+  WGPUBuffer r_uniform =
+      create_uniform(graph, device, &r_params, sizeof(r_params));
+  WGPUShaderModule r_shader = make_shader(device, kReduceWGSL);
+
+  WGPUBindGroupLayoutEntry r_entries[3] = {};
+  r_entries[0].binding = 0;
+  r_entries[0].visibility = WGPUShaderStage_Compute;
+  r_entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  r_entries[1].binding = 1;
+  r_entries[1].visibility = WGPUShaderStage_Compute;
+  r_entries[1].buffer.type = WGPUBufferBindingType_Storage;
+  r_entries[2].binding = 2;
+  r_entries[2].visibility = WGPUShaderStage_Compute;
+  r_entries[2].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor r_bgl_desc = {};
+  r_bgl_desc.entryCount = 3;
+  r_bgl_desc.entries = r_entries;
+  WGPUBindGroupLayout r_bgl =
+      wgpuDeviceCreateBindGroupLayout(device, &r_bgl_desc);
+
+  WGPUPipelineLayoutDescriptor r_pl_desc = {};
+  r_pl_desc.bindGroupLayoutCount = 1;
+  r_pl_desc.bindGroupLayouts = &r_bgl;
+  WGPUPipelineLayout r_pl = wgpuDeviceCreatePipelineLayout(device, &r_pl_desc);
+
+  WGPUConstantEntry r_wg_const = {};
+  r_wg_const.key = {"wg_size", WGPU_STRLEN};
+  r_wg_const.value = static_cast<double>(r_wg);
+
+  WGPUComputePipelineDescriptor r_pipe_desc = {};
+  r_pipe_desc.layout = r_pl;
+  r_pipe_desc.compute.module = r_shader;
+  r_pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  r_pipe_desc.compute.constantCount = 1;
+  r_pipe_desc.compute.constants = &r_wg_const;
+  WGPUComputePipeline r_pipe =
+      wgpuDeviceCreateComputePipeline(device, &r_pipe_desc);
+
+  WGPUBindGroupEntry r_bg[3] = {};
+  r_bg[0].binding = 0;
+  r_bg[0].buffer = loss_partial;
+  r_bg[0].size = n_rows * sizeof(float);
+  r_bg[1].binding = 1;
+  r_bg[1].buffer = loss.buffer;
+  r_bg[1].size = loss.nbytes;
+  r_bg[2].binding = 2;
+  r_bg[2].buffer = r_uniform;
+  r_bg[2].size = sizeof(r_params);
+
+  WGPUBindGroupDescriptor r_bg_desc = {};
+  r_bg_desc.layout = r_bgl;
+  r_bg_desc.entryCount = 3;
+  r_bg_desc.entries = r_bg;
+  WGPUBindGroup r_bind_group = wgpuDeviceCreateBindGroup(device, &r_bg_desc);
+
+  graph.add_dispatch({r_pipe, r_bind_group, r_wgc, "fused_ce_reduce"});
+
+  wgpuShaderModuleRelease(r_shader);
+  wgpuBindGroupLayoutRelease(r_bgl);
+  wgpuPipelineLayoutRelease(r_pl);
+  wgpuBufferRelease(r_uniform);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.fused_ce.default, fused_ce_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl b/backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl
new file mode 100644
index 00000000000..5910e5e9576
--- /dev/null
+++ b/backends/webgpu/runtime/ops/fused_ce/fused_ce.wgsl
@@ -0,0 +1,74 @@
+
+@group(0) @binding(0) var<storage, read_write> dlogits: array<f32>;
+@group(0) @binding(1) var<storage, read> logits: array<f32>;
+@group(0) @binding(2) var<storage, read> labels: array<i32>;
+@group(0) @binding(3) var<storage, read_write> loss_partial: array<f32>;
+
+struct Params { vocab: u32, n_rows: u32, n_valid: f32, pad: f32 }
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 256u;
+var<workgroup> red_m: array<f32, 256>;
+var<workgroup> red_l: array<f32, 256>;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(workgroup_id) wg: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x;
+  if (row >= params.n_rows) { return; }
+  let tid = lid.x;
+  let V = params.vocab;
+  let base = row * V;
+  let lbl = labels[row];
+
+  if (lbl < 0) {
+    for (var j = tid; j < V; j = j + wg_size) { dlogits[base + j] = 0.0; }
+    if (tid == 0u) { loss_partial[row] = 0.0; }
+    return;
+  }
+
+  // -3.4e38 not 3.4028235e38 (Tint f32 overflow)
+  var m = -3.4e38;
+  var l = 0.0;
+  for (var j = tid; j < V; j = j + wg_size) {
+    let x = logits[base + j];
+    if (x > m) {
+      l = l * exp(m - x) + 1.0;
+      m = x;
+    } else {
+      l = l + exp(x - m);
+    }
+  }
+  red_m[tid] = m;
+  red_l[tid] = l;
+  workgroupBarrier();
+  for (var s = wg_size / 2u; s > 0u; s = s >> 1u) {
+    if (tid < s) {
+      let ma = red_m[tid];
+      let la = red_l[tid];
+      let mb = red_m[tid + s];
+      let lb = red_l[tid + s];
+      let mm = max(ma, mb);
+      red_m[tid] = mm;
+      red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm);
+    }
+    workgroupBarrier();
+  }
+  let row_max = red_m[0];
+  let denom = red_l[0];
+  workgroupBarrier();
+
+  let inv = 1.0 / denom;
+  let scale = 1.0 / params.n_valid;
+  if (tid == 0u) {
+    let lse = row_max + log(denom);
+    loss_partial[row] = (lse - logits[base + u32(lbl)]) * scale;
+  }
+  workgroupBarrier();
+  for (var j = tid; j < V; j = j + wg_size) {
+    var g = exp(logits[base + j] - row_max) * inv * scale;
+    if (j == u32(lbl)) { g = g - scale; }
+    dlogits[base + j] = g;
+  }
+}
diff --git a/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h b/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h
new file mode 100644
index 00000000000..6ed1f709b25
--- /dev/null
+++ b/backends/webgpu/runtime/ops/fused_ce/fused_ce_wgsl.h
@@ -0,0 +1,98 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from fused_ce.wgsl - DO NOT EDIT.
+// wgsl-sha256: 416484a64dc4f31940ca63965e28b5747b34cd6b3a64e141ffaa5835ec28891d
+inline constexpr const char* kFusedCeWGSL = R"(
+
+@group(0) @binding(0) var<storage, read_write> dlogits: array<f32>;
+@group(0) @binding(1) var<storage, read> logits: array<f32>;
+@group(0) @binding(2) var<storage, read> labels: array<i32>;
+@group(0) @binding(3) var<storage, read_write> loss_partial: array<f32>;
+
+struct Params { vocab: u32, n_rows: u32, n_valid: f32, pad: f32 }
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 256u;
+var<workgroup> red_m: array<f32, 256>;
+var<workgroup> red_l: array<f32, 256>;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(
+    @builtin(workgroup_id) wg: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>) {
+  let row = wg.x;
+  if (row >= params.n_rows) { return; }
+  let tid = lid.x;
+  let V = params.vocab;
+  let base = row * V;
+  let lbl = labels[row];
+
+  if (lbl < 0) {
+    for (var j = tid; j < V; j = j + wg_size) { dlogits[base + j] = 0.0; }
+    if (tid == 0u) { loss_partial[row] = 0.0; }
+    return;
+  }
+
+  // -3.4e38 not 3.4028235e38 (Tint f32 overflow)
+  var m = -3.4e38;
+  var l = 0.0;
+  for (var j = tid; j < V; j = j + wg_size) {
+    let x = logits[base + j];
+    if (x > m) {
+      l = l * exp(m - x) + 1.0;
+      m = x;
+    } else {
+      l = l + exp(x - m);
+    }
+  }
+  red_m[tid] = m;
+  red_l[tid] = l;
+  workgroupBarrier();
+  for (var s = wg_size / 2u; s > 0u; s = s >> 1u) {
+    if (tid < s) {
+      let ma = red_m[tid];
+      let la = red_l[tid];
+      let mb = red_m[tid + s];
+      let lb = red_l[tid + s];
+      let mm = max(ma, mb);
+      red_m[tid] = mm;
+      red_l[tid] = la * exp(ma - mm) + lb * exp(mb - mm);
+    }
+    workgroupBarrier();
+  }
+  let row_max = red_m[0];
+  let denom = red_l[0];
+  workgroupBarrier();
+
+  let inv = 1.0 / denom;
+  let scale = 1.0 / params.n_valid;
+  if (tid == 0u) {
+    let lse = row_max + log(denom);
+    loss_partial[row] = (lse - logits[base + u32(lbl)]) * scale;
+  }
+  workgroupBarrier();
+  for (var j = tid; j < V; j = j + wg_size) {
+    var g = exp(logits[base + j] - row_max) * inv * scale;
+    if (j == u32(lbl)) { g = g - scale; }
+    dlogits[base + j] = g;
+  }
+}
+)";
+
+inline constexpr uint32_t kFusedCeWorkgroupSizeX = 256;
+inline constexpr uint32_t kFusedCeWorkgroupSizeY = 1;
+inline constexpr uint32_t kFusedCeWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/gather/Gather.cpp b/backends/webgpu/runtime/ops/gather/Gather.cpp
new file mode 100644
index 00000000000..d7e170129d1
--- /dev/null
+++ b/backends/webgpu/runtime/ops/gather/Gather.cpp
@@ -0,0 +1,182 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
+#include <executorch/backends/webgpu/runtime/ops/gather/gather_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct GatherParams {
+  uint32_t dim;
+  uint32_t _pad[3];
+};
+
+// gather: out[c] = self[c] with c[dim] replaced by index[c].
+void gather_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int self_id = args.at(0);
+  const int dim_id = args.at(1);
+  const int index_id = args.at(2);
+  const int out_id = args.at(args.size() - 1);
+
+  if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::Int) {
+    throw std::runtime_error("gather: dim is not an int");
+  }
+  WGPUDevice device = graph.device();
+  const auto& self_tensor = graph.get_tensor(self_id);
+  const auto& index_tensor = graph.get_tensor(index_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+
+  if (self_tensor.buffer == nullptr || index_tensor.buffer == nullptr ||
+      out_tensor.buffer == nullptr) {
+    throw std::runtime_error("gather: null buffer binding");
+  }
+
+  const int64_t ndim = static_cast<int64_t>(out_tensor.dims.size());
+  int64_t dim = graph.get_int(dim_id);
+  if (dim < 0) {
+    dim += ndim;
+  }
+  if (ndim == 0 || dim < 0 || dim >= ndim) {
+    throw std::runtime_error("gather: dim out of range");
+  }
+
+  TensorMeta out_meta;
+  TensorMeta self_meta;
+  fill_tensor_meta(out_tensor, &out_meta);
+  fill_tensor_meta(self_tensor, &self_meta);
+  if (out_meta.ndim != self_meta.ndim) {
+    throw std::runtime_error("gather: self/out rank mismatch");
+  }
+
+  const size_t out_numel = out_tensor.nbytes / sizeof(float);
+  const size_t index_numel = index_tensor.nbytes / sizeof(int32_t);
+  if (out_tensor.nbytes != out_numel * sizeof(float) ||
+      self_tensor.nbytes % sizeof(float) != 0 ||
+      index_tensor.nbytes != index_numel * sizeof(int32_t)) {
+    throw std::runtime_error("gather: fp32 self/out + i32 index required");
+  }
+  if (out_numel != index_numel) {
+    throw std::runtime_error("gather: out numel != index numel");
+  }
+
+  uint32_t wg_size = utils::clamp_workgroup_size(device, kGatherWorkgroupSizeX);
+  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, out_meta.numel, wg_size, "gather");
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  GatherParams params = {};
+  params.dim = static_cast<uint32_t>(dim);
+
+  WGPUBuffer out_meta_buf =
+      utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
+  WGPUBuffer self_meta_buf =
+      utils::make_uniform(device, &self_meta, sizeof(TensorMeta));
+  WGPUBuffer params_buf =
+      utils::make_uniform(device, &params, sizeof(GatherParams));
+  graph.add_uniform_buffer_bytes(2 * sizeof(TensorMeta) + sizeof(GatherParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kGatherWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[6] = {};
+  entries[0].binding = 0;
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].binding = 1;
+  entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[2].binding = 2;
+  entries[2].buffer.type = WGPUBufferBindingType_Storage;
+  entries[3].binding = 3;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+  entries[4].binding = 4;
+  entries[4].buffer.type = WGPUBufferBindingType_Uniform;
+  entries[5].binding = 5;
+  entries[5].buffer.type = WGPUBufferBindingType_Uniform;
+  for (auto& e : entries) {
+    e.visibility = WGPUShaderStage_Compute;
+  }
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 6;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[6] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = self_tensor.buffer;
+  bg_entries[0].size = self_tensor.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = index_tensor.buffer;
+  bg_entries[1].size = index_tensor.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = out_tensor.buffer;
+  bg_entries[2].size = out_tensor.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = out_meta_buf;
+  bg_entries[3].size = sizeof(TensorMeta);
+  bg_entries[4].binding = 4;
+  bg_entries[4].buffer = self_meta_buf;
+  bg_entries[4].size = sizeof(TensorMeta);
+  bg_entries[5].binding = 5;
+  bg_entries[5].buffer = params_buf;
+  bg_entries[5].size = sizeof(GatherParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 6;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  wgpuBufferRelease(out_meta_buf);
+  wgpuBufferRelease(self_meta_buf);
+  wgpuBufferRelease(params_buf);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.gather.default, gather_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/gather/gather.wgsl b/backends/webgpu/runtime/ops/gather/gather.wgsl
new file mode 100644
index 00000000000..fb62206be64
--- /dev/null
+++ b/backends/webgpu/runtime/ops/gather/gather.wgsl
@@ -0,0 +1,39 @@
+@group(0) @binding(0) var<storage, read> self_: array<f32>;
+@group(0) @binding(1) var<storage, read> indices: array<i32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> self_meta: TensorMeta;
+
+struct GatherParams {
+  dim: u32,
+}
+@group(0) @binding(5) var<uniform> params: GatherParams;
+
+override wg_size: u32 = 256;
+
+@compute @workgroup_size(wg_size)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let o = gid.x;
+    if (o >= out_meta.numel) {
+        return;
+    }
+    var rem = o;
+    var self_idx: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let c = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        var coord = c;
+        if (d == params.dim) {
+            coord = u32(indices[o]);
+        }
+        self_idx = self_idx + coord * self_meta.strides[d];
+    }
+    output[o] = self_[self_idx];
+}
diff --git a/backends/webgpu/runtime/ops/gather/gather_wgsl.h b/backends/webgpu/runtime/ops/gather/gather_wgsl.h
new file mode 100644
index 00000000000..5b6922e21ea
--- /dev/null
+++ b/backends/webgpu/runtime/ops/gather/gather_wgsl.h
@@ -0,0 +1,63 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from gather.wgsl - DO NOT EDIT.
+// wgsl-sha256: cd19a9ed2753e2a97e96fb90f7a72e8f40d08feb6c84fbed7f2c52e71d77a03c
+inline constexpr const char* kGatherWGSL = R"(
+@group(0) @binding(0) var<storage, read> self_: array<f32>;
+@group(0) @binding(1) var<storage, read> indices: array<i32>;
+@group(0) @binding(2) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(3) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(4) var<uniform> self_meta: TensorMeta;
+
+struct GatherParams {
+  dim: u32,
+}
+@group(0) @binding(5) var<uniform> params: GatherParams;
+
+override wg_size: u32 = 256;
+
+@compute @workgroup_size(wg_size)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let o = gid.x;
+    if (o >= out_meta.numel) {
+        return;
+    }
+    var rem = o;
+    var self_idx: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let c = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        var coord = c;
+        if (d == params.dim) {
+            coord = u32(indices[o]);
+        }
+        self_idx = self_idx + coord * self_meta.strides[d];
+    }
+    output[o] = self_[self_idx];
+}
+)";
+
+inline constexpr uint32_t kGatherWorkgroupSizeX = 256;
+inline constexpr uint32_t kGatherWorkgroupSizeY = 1;
+inline constexpr uint32_t kGatherWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp b/backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp
new file mode 100644
index 00000000000..f8c56a08979
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp
@@ -0,0 +1,173 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct LinearDwParams {
+  uint32_t M;
+  uint32_t N;
+  uint32_t K;
+  uint32_t _pad;
+};
+static_assert(sizeof(LinearDwParams) == 16, "params must be 16 bytes");
+
+// STE weight gradient d_W[N,K] = d_out^T @ x.
+void linear_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int dout_id = args.at(0);
+  const int x_id = args.at(1);
+  const int dw_id = args.at(2);
+
+  WGPUDevice device = graph.device();
+  const auto& dout = graph.get_tensor(dout_id);
+  const auto& x = graph.get_tensor(x_id);
+  const auto& dw = graph.get_tensor(dw_id);
+
+  if (dw.dims.size() != 2 || dout.dims.empty() || x.dims.empty()) {
+    throw std::runtime_error("linear_dW: bad tensor ranks");
+  }
+  const uint32_t N = static_cast<uint32_t>(dw.dims[0]);
+  const uint32_t K = static_cast<uint32_t>(dw.dims[1]);
+  if (N == 0 || K == 0) {
+    throw std::runtime_error("linear_dW: N or K == 0");
+  }
+
+  uint64_t dout_numel = 1;
+  for (int64_t d : dout.dims) {
+    dout_numel *= static_cast<uint64_t>(d);
+  }
+  uint64_t x_numel = 1;
+  for (int64_t d : x.dims) {
+    x_numel *= static_cast<uint64_t>(d);
+  }
+  if (static_cast<uint32_t>(dout.dims.back()) != N ||
+      static_cast<uint32_t>(x.dims.back()) != K) {
+    throw std::runtime_error("linear_dW: d_out/x last dim mismatch");
+  }
+  const uint32_t M = static_cast<uint32_t>(dout_numel / N);
+  if (dout_numel % N != 0 || x_numel % K != 0 || x_numel / K != M) {
+    throw std::runtime_error("linear_dW: M mismatch across d_out/x");
+  }
+
+  // fp32-only byte-size guards (mirror the forward's byte checks).
+  if (dw.nbytes != static_cast<uint64_t>(N) * K * sizeof(float) ||
+      dout.nbytes != dout_numel * sizeof(float) ||
+      x.nbytes != x_numel * sizeof(float)) {
+    throw std::runtime_error("linear_dW: fp32-only (byte-size mismatch)");
+  }
+
+  LinearDwParams params = {};
+  params.M = M;
+  params.N = N;
+  params.K = K;
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kLinearDwWorkgroupSizeX);
+  const uint64_t tiles =
+      utils::div_up<uint64_t>(N, 4u) * utils::div_up<uint64_t>(K, 4u);
+  if (tiles > UINT32_MAX) {
+    throw std::runtime_error("linear_dW: tile count exceeds u32");
+  }
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, static_cast<uint32_t>(tiles), wg_size, "linear_dW");
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(params));
+  graph.add_uniform_buffer_bytes(sizeof(params));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kLinearDwWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[4] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  for (uint32_t i = 1; i <= 2; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+    entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  }
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = dw.buffer;
+  bg_entries[0].size = dw.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = dout.buffer;
+  bg_entries[1].size = dout.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = x.buffer;
+  bg_entries[2].size = x.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = uniform_buffer;
+  bg_entries[3].size = sizeof(params);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "linear_dW"});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.linear_dW.default, linear_dW_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearBackward.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearBackward.cpp
new file mode 100644
index 00000000000..9f0ac36dfbd
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearBackward.cpp
@@ -0,0 +1,208 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct Q4gswBackwardParams {
+  uint32_t M;
+  uint32_t N;
+  uint32_t K;
+  uint32_t K_packed;
+  uint32_t group_size;
+  uint32_t padded_N;
+  uint32_t has_bias;
+  uint32_t _pad;
+};
+static_assert(sizeof(Q4gswBackwardParams) == 32, "params must be 32 bytes");
+
+// linear_q4gsw_backward: d_x[M,K] = d_out[M,N] @ dequant(W)[N,K].
+void q4gsw_backward_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int dout_id = args.at(0);
+  const int weight_id = args.at(1);
+  const int scales_id = args.at(2);
+  const int group_size_id = args.at(3);
+  const int dx_id = args.at(4);
+
+  WGPUDevice device = graph.device();
+  const auto& dout = graph.get_tensor(dout_id);
+  const auto& weight = graph.get_tensor(weight_id);
+  const auto& scales = graph.get_tensor(scales_id);
+  const auto& dx = graph.get_tensor(dx_id);
+
+  if (weight.dims.size() != 2 || scales.dims.size() != 2 || dx.dims.empty() ||
+      dout.dims.empty()) {
+    throw std::runtime_error("q4gsw_backward: bad tensor ranks");
+  }
+  const uint32_t N = static_cast<uint32_t>(weight.dims[0]);
+  const uint32_t K_packed = static_cast<uint32_t>(weight.dims[1]);
+  const uint32_t K = static_cast<uint32_t>(dx.dims.back());
+  if (N == 0 || K == 0) {
+    throw std::runtime_error("q4gsw_backward: N or K == 0");
+  }
+  uint64_t dx_numel = 1;
+  for (int64_t d : dx.dims) {
+    dx_numel *= static_cast<uint64_t>(d);
+  }
+  const uint32_t M = static_cast<uint32_t>(dx_numel / K);
+  const uint32_t num_groups = static_cast<uint32_t>(scales.dims[0]);
+  const uint32_t padded_N = static_cast<uint32_t>(scales.dims[1]);
+
+  if (graph.get_value_type(group_size_id) != WebGPUGraph::ValueType::Int) {
+    throw std::runtime_error("q4gsw_backward: group_size must be Int");
+  }
+  const int64_t group_size = graph.get_int(group_size_id);
+  if (group_size <= 0) {
+    throw std::runtime_error("q4gsw_backward: group_size must be positive");
+  }
+  const uint32_t gs = static_cast<uint32_t>(group_size);
+
+  // fp32 + shape guards (mirror the forward's byte checks).
+  if (dx.nbytes != dx_numel * sizeof(float)) {
+    throw std::runtime_error("q4gsw_backward: d_x fp32-only");
+  }
+  if (dout.nbytes != static_cast<uint64_t>(M) * N * sizeof(float)) {
+    throw std::runtime_error("q4gsw_backward: d_out fp32/shape mismatch");
+  }
+  if (scales.nbytes !=
+      static_cast<uint64_t>(num_groups) * padded_N * sizeof(float)) {
+    throw std::runtime_error("q4gsw_backward: scales fp32/shape mismatch");
+  }
+  if (weight.nbytes != static_cast<uint64_t>(N) * K_packed) {
+    throw std::runtime_error("q4gsw_backward: weight byte-size mismatch");
+  }
+  if (K_packed != (K + 1u) / 2u) {
+    throw std::runtime_error("q4gsw_backward: K_packed != ceil(K/2)");
+  }
+  if (num_groups < (K + gs - 1u) / gs || padded_N < N) {
+    throw std::runtime_error("q4gsw_backward: scales too small");
+  }
+
+  Q4gswBackwardParams params = {};
+  params.M = M;
+  params.N = N;
+  params.K = K;
+  params.K_packed = K_packed;
+  params.group_size = gs;
+  params.padded_N = padded_N;
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kQ4gswBackwardWorkgroupSizeX);
+  const uint64_t tiles = static_cast<uint64_t>((M + 3u) / 4u) * ((K + 3u) / 4u);
+  if (tiles > UINT32_MAX) {
+    throw std::runtime_error("q4gsw_backward: tile count exceeds u32");
+  }
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, static_cast<uint32_t>(tiles), wg_size, "q4gsw_backward");
+
+  WGPUBufferDescriptor uniform_desc = {};
+  uniform_desc.size = sizeof(Q4gswBackwardParams);
+  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  uniform_desc.mappedAtCreation = true;
+  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
+  std::memcpy(
+      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(Q4gswBackwardParams)),
+      &params,
+      sizeof(Q4gswBackwardParams));
+  wgpuBufferUnmap(uniform_buffer);
+  graph.add_uniform_buffer_bytes(sizeof(Q4gswBackwardParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kQ4gswBackwardWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[5] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  for (uint32_t i = 1; i <= 3; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+    entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  }
+  entries[4].binding = 4;
+  entries[4].visibility = WGPUShaderStage_Compute;
+  entries[4].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 5;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[5] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = dx.buffer;
+  bg_entries[0].size = dx.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = dout.buffer;
+  bg_entries[1].size = dout.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = weight.buffer;
+  bg_entries[2].size = weight.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = scales.buffer;
+  bg_entries[3].size = scales.nbytes;
+  bg_entries[4].binding = 4;
+  bg_entries[4].buffer = uniform_buffer;
+  bg_entries[4].size = sizeof(Q4gswBackwardParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 5;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_backward"});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  wgpuBufferRelease(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_backward.default, q4gsw_backward_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearRequant.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
new file mode 100644
index 00000000000..3c75d63165a
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
@@ -0,0 +1,192 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+struct Q4gswRequantParams {
+  uint32_t N;
+  uint32_t K;
+  uint32_t K_packed;
+  uint32_t group_size;
+  uint32_t padded_N;
+  uint32_t num_words;
+  uint32_t _pad0;
+  uint32_t _pad1;
+};
+static_assert(sizeof(Q4gswRequantParams) == 32, "params must be 32 bytes");
+
+// STE re-quant + int4 pack at a frozen per-group scale (only the codes move).
+void q4gsw_requant_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  const int latent_id = args.at(0);
+  const int scales_id = args.at(1);
+  const int group_size_id = args.at(2);
+  const int packed_id = args.at(3);
+
+  WGPUDevice device = graph.device();
+  const auto& latent = graph.get_tensor(latent_id);
+  const auto& scales = graph.get_tensor(scales_id);
+  const auto& packed = graph.get_tensor(packed_id);
+
+  if (latent.dims.size() != 2 || scales.dims.size() != 2 ||
+      packed.dims.size() != 2) {
+    throw std::runtime_error("q4gsw_requant: bad tensor ranks");
+  }
+  const uint32_t N = static_cast<uint32_t>(latent.dims[0]);
+  const uint32_t K = static_cast<uint32_t>(latent.dims[1]);
+  const uint32_t K_packed = static_cast<uint32_t>(packed.dims[1]);
+  const uint32_t num_groups = static_cast<uint32_t>(scales.dims[0]);
+  const uint32_t padded_N = static_cast<uint32_t>(scales.dims[1]);
+  if (N == 0 || K == 0) {
+    throw std::runtime_error("q4gsw_requant: N or K == 0");
+  }
+  if (static_cast<uint32_t>(packed.dims[0]) != N) {
+    throw std::runtime_error("q4gsw_requant: packed rows != N");
+  }
+  if (K_packed != (K + 1u) / 2u) {
+    throw std::runtime_error("q4gsw_requant: K_packed != ceil(K/2)");
+  }
+  if ((static_cast<uint64_t>(N) * K_packed) % 4u != 0u) {
+    throw std::runtime_error(
+        "q4gsw_requant: N*K_packed must be a multiple of 4 (u32-packed)");
+  }
+
+  if (graph.get_value_type(group_size_id) != WebGPUGraph::ValueType::Int) {
+    throw std::runtime_error("q4gsw_requant: group_size must be Int");
+  }
+  const int64_t group_size = graph.get_int(group_size_id);
+  if (group_size <= 0) {
+    throw std::runtime_error("q4gsw_requant: group_size must be positive");
+  }
+  const uint32_t gs = static_cast<uint32_t>(group_size);
+  if (num_groups < (K + gs - 1u) / gs || padded_N < N) {
+    throw std::runtime_error("q4gsw_requant: scales dims too small for K/N");
+  }
+
+  if (latent.nbytes != static_cast<uint64_t>(N) * K * sizeof(float)) {
+    throw std::runtime_error("q4gsw_requant: latent fp32-only");
+  }
+  if (scales.nbytes !=
+      static_cast<uint64_t>(num_groups) * padded_N * sizeof(float)) {
+    throw std::runtime_error("q4gsw_requant: scales fp32/shape mismatch");
+  }
+  if (packed.nbytes != static_cast<uint64_t>(N) * K_packed) {
+    throw std::runtime_error("q4gsw_requant: packed byte-size mismatch");
+  }
+
+  const uint32_t num_words =
+      static_cast<uint32_t>((static_cast<uint64_t>(N) * K_packed) / 4u);
+
+  Q4gswRequantParams params = {};
+  params.N = N;
+  params.K = K;
+  params.K_packed = K_packed;
+  params.group_size = gs;
+  params.padded_N = padded_N;
+  params.num_words = num_words;
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kQ4gswRequantWorkgroupSizeX);
+  const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, num_words, wg_size, "q4gsw_requant");
+
+  WGPUBuffer uniform_buffer =
+      utils::make_uniform(device, &params, sizeof(params));
+  graph.add_uniform_buffer_bytes(sizeof(params));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kQ4gswRequantWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[4] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_Storage;
+  for (uint32_t i = 1; i <= 2; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+    entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  }
+  entries[3].binding = 3;
+  entries[3].visibility = WGPUShaderStage_Compute;
+  entries[3].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 4;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[4] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = packed.buffer;
+  bg_entries[0].size = packed.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = latent.buffer;
+  bg_entries[1].size = latent.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = scales.buffer;
+  bg_entries[2].size = scales.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = uniform_buffer;
+  bg_entries[3].size = sizeof(params);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 4;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_requant"});
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(et_vk.q4gsw_requant.default, q4gsw_requant_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/linear_dW.wgsl b/backends/webgpu/runtime/ops/quantized_linear/linear_dW.wgsl
new file mode 100644
index 00000000000..d7a68cd5d40
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/linear_dW.wgsl
@@ -0,0 +1,69 @@
+// STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32).
+
+@group(0) @binding(0) var<storage, read_write> t_dw: array<f32>;    // [N, K]
+@group(0) @binding(1) var<storage, read> t_dout: array<f32>;         // [M, N]
+@group(0) @binding(2) var<storage, read> t_x: array<f32>;            // [M, K]
+
+struct Params {
+  M: u32,
+  N: u32,
+  K: u32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+const TN: u32 = 4u;
+const TK: u32 = 4u;
+const TILE_ELEMS: u32 = TN * TK;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let nnt = (params.N + TN - 1u) / TN;
+  let nkt = (params.K + TK - 1u) / TK;
+  let tiles = nnt * nkt;
+  if (gid.x >= tiles) {
+    return;
+  }
+  let row_tile = gid.x / nkt;
+  let col_tile = gid.x % nkt;
+  let n0 = row_tile * TN;
+  let k0 = col_tile * TK;
+
+  var acc: array<f32, TILE_ELEMS>;
+  for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) {
+    acc[i] = 0.0;
+  }
+
+  var m: u32 = 0u;
+  loop {
+    if (m >= params.M) {
+      break;
+    }
+    // Load the TN d_out values for row m once; reused across all TK k columns.
+    var dout_reg: array<f32, TN>;
+    for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
+      let n_eff = min(n0 + nl, params.N - 1u);
+      dout_reg[nl] = t_dout[m * params.N + n_eff];
+    }
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k_eff = min(k0 + kl, params.K - 1u);
+      let xv = t_x[m * params.K + k_eff];
+      for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
+        acc[nl * TK + kl] = acc[nl * TK + kl] + dout_reg[nl] * xv;
+      }
+    }
+    m = m + 1u;
+  }
+
+  for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
+    let n = n0 + nl;
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k = k0 + kl;
+      if (n < params.N && k < params.K) {
+        t_dw[n * params.K + k] = acc[nl * TK + kl];
+      }
+    }
+  }
+}
diff --git a/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h
new file mode 100644
index 00000000000..9bdf7dc4e44
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/linear_dW_wgsl.h
@@ -0,0 +1,93 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from linear_dW.wgsl - DO NOT EDIT.
+// wgsl-sha256: be50075764ccd87e0d0c86124f13fd753aa618f994c9a5e1b8b631516aa51f84
+inline constexpr const char* kLinearDwWGSL = R"(
+// STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32).
+
+@group(0) @binding(0) var<storage, read_write> t_dw: array<f32>;    // [N, K]
+@group(0) @binding(1) var<storage, read> t_dout: array<f32>;         // [M, N]
+@group(0) @binding(2) var<storage, read> t_x: array<f32>;            // [M, K]
+
+struct Params {
+  M: u32,
+  N: u32,
+  K: u32,
+  _pad: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+const TN: u32 = 4u;
+const TK: u32 = 4u;
+const TILE_ELEMS: u32 = TN * TK;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let nnt = (params.N + TN - 1u) / TN;
+  let nkt = (params.K + TK - 1u) / TK;
+  let tiles = nnt * nkt;
+  if (gid.x >= tiles) {
+    return;
+  }
+  let row_tile = gid.x / nkt;
+  let col_tile = gid.x % nkt;
+  let n0 = row_tile * TN;
+  let k0 = col_tile * TK;
+
+  var acc: array<f32, TILE_ELEMS>;
+  for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) {
+    acc[i] = 0.0;
+  }
+
+  var m: u32 = 0u;
+  loop {
+    if (m >= params.M) {
+      break;
+    }
+    // Load the TN d_out values for row m once; reused across all TK k columns.
+    var dout_reg: array<f32, TN>;
+    for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
+      let n_eff = min(n0 + nl, params.N - 1u);
+      dout_reg[nl] = t_dout[m * params.N + n_eff];
+    }
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k_eff = min(k0 + kl, params.K - 1u);
+      let xv = t_x[m * params.K + k_eff];
+      for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
+        acc[nl * TK + kl] = acc[nl * TK + kl] + dout_reg[nl] * xv;
+      }
+    }
+    m = m + 1u;
+  }
+
+  for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
+    let n = n0 + nl;
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k = k0 + kl;
+      if (n < params.N && k < params.K) {
+        t_dw[n * params.K + k] = acc[nl * TK + kl];
+      }
+    }
+  }
+}
+)";
+
+inline constexpr uint32_t kLinearDwWorkgroupSizeX = 64;
+inline constexpr uint32_t kLinearDwWorkgroupSizeY = 1;
+inline constexpr uint32_t kLinearDwWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward.wgsl
new file mode 100644
index 00000000000..569d9450b5b
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward.wgsl
@@ -0,0 +1,84 @@
+// Mirrors q4gsw_linear.wgsl dequant; contracts over N (forward over K).
+
+@group(0) @binding(0) var<storage, read_write> t_dx: array<f32>;    // [M, K]
+@group(0) @binding(1) var<storage, read> t_dout: array<f32>;         // [M, N]
+@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
+@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
+
+struct Params {
+  M: u32,
+  N: u32,
+  K: u32,
+  K_packed: u32,
+  group_size: u32,
+  padded_N: u32,
+  has_bias: u32,
+  _pad: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+const TM: u32 = 4u;
+const TK: u32 = 4u;
+const TILE_ELEMS: u32 = TM * TK;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let nrt = (params.M + TM - 1u) / TM;
+  let nkt = (params.K + TK - 1u) / TK;
+  let tiles = nrt * nkt;
+  if (gid.x >= tiles) {
+    return;
+  }
+  let row_tile = gid.x / nkt;
+  let col_tile = gid.x % nkt;
+  let m0 = row_tile * TM;
+  let k0 = col_tile * TK;
+
+  var acc: array<f32, TILE_ELEMS>;
+  for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) {
+    acc[i] = 0.0;
+  }
+
+  var n: u32 = 0u;
+  loop {
+    if (n >= params.N) {
+      break;
+    }
+    // Load TM d_out for column n once; reused across TK k columns.
+    var dout_reg: array<f32, TM>;
+    for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
+      let m_eff = min(m0 + ml, params.M - 1u);
+      dout_reg[ml] = t_dout[m_eff * params.N + n];
+    }
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k_eff = min(k0 + kl, params.K - 1u);
+      let byte_idx = n * params.K_packed + (k_eff >> 1u);
+      let word = t_weight[byte_idx >> 2u];
+      let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu;
+      var nib: u32;
+      if ((k_eff & 1u) == 0u) {
+        nib = b & 0x0Fu;
+      } else {
+        nib = (b >> 4u) & 0x0Fu;
+      }
+      let q = f32(i32(nib) - 8);
+      let dq = q * t_scales[(k_eff / params.group_size) * params.padded_N + n];
+      for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
+        acc[ml * TK + kl] = acc[ml * TK + kl] + dout_reg[ml] * dq;
+      }
+    }
+    n = n + 1u;
+  }
+
+  for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
+    let m = m0 + ml;
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k = k0 + kl;
+      if (m < params.M && k < params.K) {
+        t_dx[m * params.K + k] = acc[ml * TK + kl];
+      }
+    }
+  }
+}
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h
new file mode 100644
index 00000000000..8882fff0b09
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_backward_wgsl.h
@@ -0,0 +1,108 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from q4gsw_backward.wgsl - DO NOT EDIT.
+// wgsl-sha256: eb7da436c89fc3cd8c062af5450d11f28fd4fc5905cc38d043fc3584252438b2
+inline constexpr const char* kQ4gswBackwardWGSL = R"(
+// Mirrors q4gsw_linear.wgsl dequant; contracts over N (forward over K).
+
+@group(0) @binding(0) var<storage, read_write> t_dx: array<f32>;    // [M, K]
+@group(0) @binding(1) var<storage, read> t_dout: array<f32>;         // [M, N]
+@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
+@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
+
+struct Params {
+  M: u32,
+  N: u32,
+  K: u32,
+  K_packed: u32,
+  group_size: u32,
+  padded_N: u32,
+  has_bias: u32,
+  _pad: u32,
+}
+@group(0) @binding(4) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+const TM: u32 = 4u;
+const TK: u32 = 4u;
+const TILE_ELEMS: u32 = TM * TK;
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let nrt = (params.M + TM - 1u) / TM;
+  let nkt = (params.K + TK - 1u) / TK;
+  let tiles = nrt * nkt;
+  if (gid.x >= tiles) {
+    return;
+  }
+  let row_tile = gid.x / nkt;
+  let col_tile = gid.x % nkt;
+  let m0 = row_tile * TM;
+  let k0 = col_tile * TK;
+
+  var acc: array<f32, TILE_ELEMS>;
+  for (var i: u32 = 0u; i < TILE_ELEMS; i = i + 1u) {
+    acc[i] = 0.0;
+  }
+
+  var n: u32 = 0u;
+  loop {
+    if (n >= params.N) {
+      break;
+    }
+    // Load TM d_out for column n once; reused across TK k columns.
+    var dout_reg: array<f32, TM>;
+    for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
+      let m_eff = min(m0 + ml, params.M - 1u);
+      dout_reg[ml] = t_dout[m_eff * params.N + n];
+    }
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k_eff = min(k0 + kl, params.K - 1u);
+      let byte_idx = n * params.K_packed + (k_eff >> 1u);
+      let word = t_weight[byte_idx >> 2u];
+      let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu;
+      var nib: u32;
+      if ((k_eff & 1u) == 0u) {
+        nib = b & 0x0Fu;
+      } else {
+        nib = (b >> 4u) & 0x0Fu;
+      }
+      let q = f32(i32(nib) - 8);
+      let dq = q * t_scales[(k_eff / params.group_size) * params.padded_N + n];
+      for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
+        acc[ml * TK + kl] = acc[ml * TK + kl] + dout_reg[ml] * dq;
+      }
+    }
+    n = n + 1u;
+  }
+
+  for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
+    let m = m0 + ml;
+    for (var kl: u32 = 0u; kl < TK; kl = kl + 1u) {
+      let k = k0 + kl;
+      if (m < params.M && k < params.K) {
+        t_dx[m * params.K + k] = acc[ml * TK + kl];
+      }
+    }
+  }
+}
+)";
+
+inline constexpr uint32_t kQ4gswBackwardWorkgroupSizeX = 64;
+inline constexpr uint32_t kQ4gswBackwardWorkgroupSizeY = 1;
+inline constexpr uint32_t kQ4gswBackwardWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant.wgsl
new file mode 100644
index 00000000000..1f1af914630
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant.wgsl
@@ -0,0 +1,54 @@
+// STE re-quant + int4 pack: round(latent/scale) in [-8,7], 2 nibbles/byte.
+
+@group(0) @binding(0) var<storage, read_write> t_packed: array<u32>;
+@group(0) @binding(1) var<storage, read> t_latent: array<f32>;       // [N, K]
+@group(0) @binding(2) var<storage, read> t_scales: array<f32>;
+
+struct Params {
+  N: u32,
+  K: u32,
+  K_packed: u32,
+  group_size: u32,
+  padded_N: u32,
+  num_words: u32,
+  _pad0: u32,
+  _pad1: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+fn quant_nibble(n: u32, k: u32) -> u32 {
+  let s = t_scales[(k / params.group_size) * params.padded_N + n];
+  var q: i32 = 0;
+  if (s != 0.0) {
+    q = i32(round(t_latent[n * params.K + k] / s));
+  }
+  q = clamp(q, -8, 7);
+  return u32(q + 8) & 0xFu;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let word_idx = gid.x;
+  if (word_idx >= params.num_words) {
+    return;
+  }
+  var word: u32 = 0u;
+  for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) {
+    let byte_idx = word_idx * 4u + bi;
+    let n = byte_idx / params.K_packed;
+    let byte_in_row = byte_idx % params.K_packed;
+    let k_lo = byte_in_row * 2u;
+    var b: u32 = 0u;
+    if (n < params.N && k_lo < params.K) {
+      b = quant_nibble(n, k_lo);
+      let k_hi = k_lo + 1u;
+      if (k_hi < params.K) {
+        b = b | (quant_nibble(n, k_hi) << 4u);
+      }
+    }
+    word = word | (b << (bi * 8u));
+  }
+  t_packed[word_idx] = word;
+}
diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h
new file mode 100644
index 00000000000..84f30d7e1b8
--- /dev/null
+++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_requant_wgsl.h
@@ -0,0 +1,78 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from q4gsw_requant.wgsl - DO NOT EDIT.
+// wgsl-sha256: 66cbf8061af4580be1773acd2e9d11d539d1a9e4a3fc8d08bd1d4def42e4d643
+inline constexpr const char* kQ4gswRequantWGSL = R"(
+// STE re-quant + int4 pack: round(latent/scale) in [-8,7], 2 nibbles/byte.
+
+@group(0) @binding(0) var<storage, read_write> t_packed: array<u32>;
+@group(0) @binding(1) var<storage, read> t_latent: array<f32>;       // [N, K]
+@group(0) @binding(2) var<storage, read> t_scales: array<f32>;
+
+struct Params {
+  N: u32,
+  K: u32,
+  K_packed: u32,
+  group_size: u32,
+  padded_N: u32,
+  num_words: u32,
+  _pad0: u32,
+  _pad1: u32,
+}
+@group(0) @binding(3) var<uniform> params: Params;
+
+override wg_size: u32 = 64u;
+
+fn quant_nibble(n: u32, k: u32) -> u32 {
+  let s = t_scales[(k / params.group_size) * params.padded_N + n];
+  var q: i32 = 0;
+  if (s != 0.0) {
+    q = i32(round(t_latent[n * params.K + k] / s));
+  }
+  q = clamp(q, -8, 7);
+  return u32(q + 8) & 0xFu;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+  let word_idx = gid.x;
+  if (word_idx >= params.num_words) {
+    return;
+  }
+  var word: u32 = 0u;
+  for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) {
+    let byte_idx = word_idx * 4u + bi;
+    let n = byte_idx / params.K_packed;
+    let byte_in_row = byte_idx % params.K_packed;
+    let k_lo = byte_in_row * 2u;
+    var b: u32 = 0u;
+    if (n < params.N && k_lo < params.K) {
+      b = quant_nibble(n, k_lo);
+      let k_hi = k_lo + 1u;
+      if (k_hi < params.K) {
+        b = b | (quant_nibble(n, k_hi) << 4u);
+      }
+    }
+    word = word | (b << (bi * 8u));
+  }
+  t_packed[word_idx] = word;
+}
+)";
+
+inline constexpr uint32_t kQ4gswRequantWorkgroupSizeX = 64;
+inline constexpr uint32_t kQ4gswRequantWorkgroupSizeY = 1;
+inline constexpr uint32_t kQ4gswRequantWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/reduce/Reduce.cpp b/backends/webgpu/runtime/ops/reduce/Reduce.cpp
new file mode 100644
index 00000000000..b276354d2a8
--- /dev/null
+++ b/backends/webgpu/runtime/ops/reduce/Reduce.cpp
@@ -0,0 +1,273 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <cstdint>
+#include <cstring>
+#include <stdexcept>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// Uniform layout matching the WGSL Params struct (16-byte aligned).
+struct ReduceParams {
+  uint32_t outer;
+  uint32_t r;
+  uint32_t inner;
+  uint32_t is_mean;
+};
+static_assert(sizeof(ReduceParams) == 16, "ReduceParams must be 16 bytes");
+
+void decompose(
+    const std::vector<int64_t>& dims,
+    int64_t dim,
+    uint32_t& outer,
+    uint32_t& r,
+    uint32_t& inner) {
+  const int64_t ndim = static_cast<int64_t>(dims.size());
+  if (dim < 0) {
+    dim += ndim;
+  }
+  if (ndim == 0 || dim < 0 || dim >= ndim) {
+    throw std::runtime_error("WebGPU reduce: dim out of range");
+  }
+  uint64_t o = 1, in = 1;
+  for (int64_t d = 0; d < dim; ++d) {
+    o *= static_cast<uint64_t>(dims[d]);
+  }
+  for (int64_t d = dim + 1; d < ndim; ++d) {
+    in *= static_cast<uint64_t>(dims[d]);
+  }
+  outer = static_cast<uint32_t>(o);
+  r = static_cast<uint32_t>(dims[dim]);
+  inner = static_cast<uint32_t>(in);
+}
+
+void reduce_impl(
+    WebGPUGraph& graph,
+    const std::vector<int>& args,
+    bool is_mean,
+    const char* op_name) {
+  const int in_id = args.at(0);
+  const int dim_id = args.at(1);
+  const int keepdim_id = args.at(2);
+  const int out_id = args.at(args.size() - 1);
+
+  WGPUDevice device = graph.device();
+  const auto& in = graph.get_tensor(in_id);
+  const auto& out = graph.get_tensor(out_id);
+
+  bool keepdim = false;
+  if (graph.get_value_type(keepdim_id) == WebGPUGraph::ValueType::Int) {
+    keepdim = graph.get_int(keepdim_id) != 0;
+  }
+
+  if (in.dims.empty()) {
+    throw std::runtime_error("WebGPU reduce: scalar input unsupported");
+  }
+  if (graph.get_value_type(dim_id) != WebGPUGraph::ValueType::IntList) {
+    throw std::runtime_error("WebGPU reduce: dim arg is not an IntList");
+  }
+  const std::vector<int64_t>& reduce_dims = graph.get_int_list(dim_id);
+  // Single-dim reduction only for now; multi-dim is a tracked extension.
+  if (reduce_dims.size() != 1) {
+    throw std::runtime_error(
+        "WebGPU reduce: only single-dim reduction is supported");
+  }
+  const int64_t dim = reduce_dims[0];
+
+  uint32_t outer = 0, r = 0, inner = 0;
+  decompose(in.dims, dim, outer, r, inner);
+  if (outer == 0 || r == 0 || inner == 0) {
+    throw std::runtime_error("WebGPU reduce: zero-sized reduction");
+  }
+
+  uint64_t in_numel = 1;
+  for (int64_t d : in.dims) {
+    in_numel *= static_cast<uint64_t>(d);
+  }
+  const uint64_t outputs = static_cast<uint64_t>(outer) * inner;
+  if (in.nbytes != in_numel * sizeof(float) ||
+      out.nbytes != outputs * sizeof(float)) {
+    throw std::runtime_error("WebGPU reduce: fp32-only (byte-size mismatch)");
+  }
+  if (outputs > UINT32_MAX) {
+    throw std::runtime_error(
+        "WebGPU reduce: output count exceeds dispatch limit");
+  }
+
+  const uint32_t wg_size =
+      utils::clamp_workgroup_size(device, kReduceWorkgroupSizeX);
+  // Cooperative reduction: one workgroup per output element (2D-folded grid).
+  const utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(
+      device, static_cast<uint32_t>(outputs), 1u, op_name);
+
+  ReduceParams params = {};
+  params.outer = outer;
+  params.r = r;
+  params.inner = inner;
+  params.is_mean = is_mean ? 1u : 0u;
+
+  WGPUBufferDescriptor uniform_desc = {};
+  uniform_desc.size = sizeof(ReduceParams);
+  uniform_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
+  uniform_desc.mappedAtCreation = true;
+  WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc);
+  void* mapped =
+      wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(ReduceParams));
+  std::memcpy(mapped, &params, sizeof(ReduceParams));
+  wgpuBufferUnmap(uniform_buffer);
+  graph.add_uniform_buffer_bytes(sizeof(ReduceParams));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kReduceWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[3] = {};
+  entries[0].binding = 0;
+  entries[0].visibility = WGPUShaderStage_Compute;
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].binding = 1;
+  entries[1].visibility = WGPUShaderStage_Compute;
+  entries[1].buffer.type = WGPUBufferBindingType_Storage;
+  entries[2].binding = 2;
+  entries[2].visibility = WGPUShaderStage_Compute;
+  entries[2].buffer.type = WGPUBufferBindingType_Uniform;
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 3;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[3] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = in.buffer;
+  bg_entries[0].size = in.nbytes;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = out.buffer;
+  bg_entries[1].size = out.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = uniform_buffer;
+  bg_entries[2].size = sizeof(ReduceParams);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 3;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  const size_t dispatch_idx = graph.add_dispatch(
+      {pipeline, bind_group, workgroup_count.x, op_name, workgroup_count.y});
+
+  // Dynamic shapes: recompute the decomposition for the reduced dim + dispatch.
+  WGPUBuffer params_buf = uniform_buffer;
+  const uint32_t is_mean_u = is_mean ? 1u : 0u;
+  const uint64_t build_outputs = outputs;
+  graph.add_tensor_resize_hook(
+      in_id,
+      [in_id,
+       out_id,
+       dim,
+       keepdim,
+       is_mean_u,
+       build_outputs,
+       dispatch_idx,
+       params_buf](WebGPUGraph& g) {
+        const auto& d = g.cur_dims(in_id);
+        uint32_t o = 0, rr = 0, n = 0;
+        decompose(std::vector<int64_t>(d.begin(), d.end()), dim, o, rr, n);
+        if (o == 0u || rr == 0u || n == 0u) {
+          throw std::runtime_error("WebGPU reduce: live zero-sized reduction");
+        }
+        const uint64_t live_outputs = static_cast<uint64_t>(o) * n;
+        if (live_outputs > build_outputs) {
+          throw std::runtime_error(
+              "WebGPU reduce: live output count exceeds build max");
+        }
+        ReduceParams p = {};
+        p.outer = o;
+        p.r = rr;
+        p.inner = n;
+        p.is_mean = is_mean_u;
+        wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p));
+        const utils::WgCount wgc = utils::compute_2d_workgroup_count(
+            g.device(),
+            static_cast<uint32_t>(live_outputs),
+            1u,
+            "reduce(resize)");
+        g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x;
+        g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y;
+        // Propagate reduced output dims for downstream resize hooks.
+        int64_t nd = static_cast<int64_t>(d.size());
+        int64_t rd = dim < 0 ? dim + nd : dim;
+        std::vector<int64_t> od;
+        for (int64_t i = 0; i < nd; ++i) {
+          if (i == rd) {
+            if (keepdim) {
+              od.push_back(1);
+            }
+          } else {
+            od.push_back(d[i]);
+          }
+        }
+        g.set_cur_dims(out_id, od);
+      });
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  // Graph owns it so the resize hook can rewrite it; freed in the dtor.
+  graph.own_uniform_buffer(uniform_buffer);
+}
+
+void sum_dim_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  reduce_impl(graph, args, /*is_mean=*/false, "sum.dim_IntList");
+}
+
+void mean_dim_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  reduce_impl(graph, args, /*is_mean=*/true, "mean.dim");
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.sum.dim_IntList, sum_dim_impl);
+  WEBGPU_REGISTER_OP(aten.mean.dim, mean_dim_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/reduce/reduce.wgsl b/backends/webgpu/runtime/ops/reduce/reduce.wgsl
new file mode 100644
index 00000000000..81c4f63c753
--- /dev/null
+++ b/backends/webgpu/runtime/ops/reduce/reduce.wgsl
@@ -0,0 +1,56 @@
+struct Params {
+  outer_: u32,
+  r_: u32,
+  inner_: u32,
+  is_mean: u32,
+};
+
+@group(0) @binding(0) var<storage, read> inp: array<f32>;
+@group(0) @binding(1) var<storage, read_write> out: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+
+// Cooperative shared-memory reduction, one workgroup per output element: each
+// thread sums a strided slice of the reduced dim into a shared partial, then
+// thread 0 folds the partials. Same one-workgroup-per-row shared-memory shape as
+// Vulkan's reduce_per_row_buffer.glsl. Fixed 256 upper bound >= any clamped
+// wg_size; only [0, wg_size) is used.
+var<workgroup> partials: array<f32, 256>;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  // One workgroup per output; 2D-fold lifts the 65535 grid cap. `t` is uniform
+  // across the workgroup, so the early return keeps the barrier in uniform flow.
+  let t = wid.x + wid.y * num_workgroups.x;
+  let outs = params.outer_ * params.inner_;
+  if (t >= outs) {
+    return;
+  }
+  let oo = t / params.inner_;
+  let ii = t % params.inner_;
+  let base = oo * params.r_ * params.inner_ + ii;
+
+  var acc: f32 = 0.0;
+  var k: u32 = lid.x;
+  while (k < params.r_) {
+    acc = acc + inp[base + k * params.inner_];
+    k = k + wg_size;
+  }
+  partials[lid.x] = acc;
+  workgroupBarrier();
+
+  if (lid.x == 0u) {
+    var s: f32 = partials[0];
+    for (var i: u32 = 1u; i < wg_size; i = i + 1u) {
+      s = s + partials[i];
+    }
+    if (params.is_mean == 1u) {
+      s = s / f32(params.r_);
+    }
+    out[t] = s;
+  }
+}
diff --git a/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h b/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h
new file mode 100644
index 00000000000..f3b06ca8b33
--- /dev/null
+++ b/backends/webgpu/runtime/ops/reduce/reduce_wgsl.h
@@ -0,0 +1,80 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from reduce.wgsl - DO NOT EDIT.
+// wgsl-sha256: b4a40b67af55986ea5136f06c260cf3625e4db7809975fc616aeace16b479b2d
+inline constexpr const char* kReduceWGSL = R"(
+struct Params {
+  outer_: u32,
+  r_: u32,
+  inner_: u32,
+  is_mean: u32,
+};
+
+@group(0) @binding(0) var<storage, read> inp: array<f32>;
+@group(0) @binding(1) var<storage, read_write> out: array<f32>;
+@group(0) @binding(2) var<uniform> params: Params;
+
+override wg_size: u32 = 256;
+
+// Cooperative shared-memory reduction, one workgroup per output element: each
+// thread sums a strided slice of the reduced dim into a shared partial, then
+// thread 0 folds the partials. Same one-workgroup-per-row shared-memory shape as
+// Vulkan's reduce_per_row_buffer.glsl. Fixed 256 upper bound >= any clamped
+// wg_size; only [0, wg_size) is used.
+var<workgroup> partials: array<f32, 256>;
+
+@compute @workgroup_size(wg_size)
+fn main(
+    @builtin(workgroup_id) wid: vec3<u32>,
+    @builtin(local_invocation_id) lid: vec3<u32>,
+    @builtin(num_workgroups) num_workgroups: vec3<u32>) {
+  // One workgroup per output; 2D-fold lifts the 65535 grid cap. `t` is uniform
+  // across the workgroup, so the early return keeps the barrier in uniform flow.
+  let t = wid.x + wid.y * num_workgroups.x;
+  let outs = params.outer_ * params.inner_;
+  if (t >= outs) {
+    return;
+  }
+  let oo = t / params.inner_;
+  let ii = t % params.inner_;
+  let base = oo * params.r_ * params.inner_ + ii;
+
+  var acc: f32 = 0.0;
+  var k: u32 = lid.x;
+  while (k < params.r_) {
+    acc = acc + inp[base + k * params.inner_];
+    k = k + wg_size;
+  }
+  partials[lid.x] = acc;
+  workgroupBarrier();
+
+  if (lid.x == 0u) {
+    var s: f32 = partials[0];
+    for (var i: u32 = 1u; i < wg_size; i = i + 1u) {
+      s = s + partials[i];
+    }
+    if (params.is_mean == 1u) {
+      s = s / f32(params.r_);
+    }
+    out[t] = s;
+  }
+}
+)";
+
+inline constexpr uint32_t kReduceWorkgroupSizeX = 256;
+inline constexpr uint32_t kReduceWorkgroupSizeY = 1;
+inline constexpr uint32_t kReduceWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/where/Where.cpp b/backends/webgpu/runtime/ops/where/Where.cpp
new file mode 100644
index 00000000000..bde779498dd
--- /dev/null
+++ b/backends/webgpu/runtime/ops/where/Where.cpp
@@ -0,0 +1,241 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
+#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
+#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
+#include <executorch/backends/webgpu/runtime/ops/TensorMeta.h>
+#include <executorch/backends/webgpu/runtime/ops/where/where_wgsl.h>
+
+#include <webgpu/webgpu.h>
+
+#include <algorithm>
+#include <stdexcept>
+#include <vector>
+
+namespace executorch::backends::webgpu {
+
+namespace {
+
+// where.self selects self/other by cond; broadcasts all three.
+void where_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+  // args: [condition, self, other, out]; out is the last value id.
+  const int cond_id = args.at(0);
+  const int a_id = args.at(1);
+  const int b_id = args.at(2);
+  const int out_id = args.at(args.size() - 1);
+
+  WGPUDevice device = graph.device();
+
+  const auto& cond_tensor = graph.get_tensor(cond_id);
+  const auto& a_tensor = graph.get_tensor(a_id);
+  const auto& b_tensor = graph.get_tensor(b_id);
+  const auto& out_tensor = graph.get_tensor(out_id);
+
+  if (cond_tensor.buffer == nullptr || a_tensor.buffer == nullptr ||
+      b_tensor.buffer == nullptr || out_tensor.buffer == nullptr) {
+    throw std::runtime_error("where: null buffer binding");
+  }
+  if (out_tensor.dims.size() > kTensorMetaMaxNdim ||
+      cond_tensor.dims.size() > kTensorMetaMaxNdim ||
+      a_tensor.dims.size() > kTensorMetaMaxNdim ||
+      b_tensor.dims.size() > kTensorMetaMaxNdim) {
+    throw std::runtime_error("where: tensor rank exceeds 4 (MAX_NDIM)");
+  }
+
+  const uint32_t out_ndim = static_cast<uint32_t>(out_tensor.dims.size());
+
+  TensorMeta out_meta;
+  TensorMeta cond_meta;
+  TensorMeta a_meta;
+  TensorMeta b_meta;
+  fill_tensor_meta_broadcast(out_tensor, out_ndim, &out_meta);
+  fill_tensor_meta_broadcast(cond_tensor, out_ndim, &cond_meta);
+  fill_tensor_meta_broadcast(a_tensor, out_ndim, &a_meta);
+  fill_tensor_meta_broadcast(b_tensor, out_ndim, &b_meta);
+
+  // a/b/out are fp32; cond is 1-byte bool (read byte-packed as array<u32>).
+  if (out_tensor.nbytes !=
+          static_cast<size_t>(out_meta.numel) * sizeof(float) ||
+      a_tensor.nbytes != static_cast<size_t>(a_meta.numel) * sizeof(float) ||
+      b_tensor.nbytes != static_cast<size_t>(b_meta.numel) * sizeof(float)) {
+    throw std::runtime_error(
+        "where: non-fp32 self/other (nbytes != numel * 4)");
+  }
+  if (cond_tensor.nbytes != static_cast<size_t>(cond_meta.numel)) {
+    throw std::runtime_error("where: condition is not a 1-byte (bool) tensor");
+  }
+
+  // Buffer is 4-byte-rounded at alloc; bind the padded span for the u32 read.
+  const size_t cond_bind_size = (cond_tensor.nbytes + 3) & ~size_t(3);
+
+  uint32_t wg_size = utils::clamp_workgroup_size(device, kWhereWorkgroupSizeX);
+  uint32_t workgroup_count = utils::compute_1d_workgroup_count(
+      device, out_meta.numel, wg_size, "where");
+
+  WGPUConstantEntry wg_size_constant = {};
+  wg_size_constant.key = {"wg_size", WGPU_STRLEN};
+  wg_size_constant.value = static_cast<double>(wg_size);
+
+  WGPUBuffer out_meta_buf =
+      utils::make_uniform(device, &out_meta, sizeof(TensorMeta));
+  WGPUBuffer cond_meta_buf =
+      utils::make_uniform(device, &cond_meta, sizeof(TensorMeta));
+  WGPUBuffer a_meta_buf =
+      utils::make_uniform(device, &a_meta, sizeof(TensorMeta));
+  WGPUBuffer b_meta_buf =
+      utils::make_uniform(device, &b_meta, sizeof(TensorMeta));
+  graph.add_uniform_buffer_bytes(4 * sizeof(TensorMeta));
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kWhereWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry entries[8] = {};
+  entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[2].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
+  entries[3].buffer.type = WGPUBufferBindingType_Storage;
+  entries[4].buffer.type = WGPUBufferBindingType_Uniform;
+  entries[5].buffer.type = WGPUBufferBindingType_Uniform;
+  entries[6].buffer.type = WGPUBufferBindingType_Uniform;
+  entries[7].buffer.type = WGPUBufferBindingType_Uniform;
+  for (uint32_t i = 0; i < 8; i++) {
+    entries[i].binding = i;
+    entries[i].visibility = WGPUShaderStage_Compute;
+  }
+
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 8;
+  bgl_desc.entries = entries;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pipeline_layout =
+      wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipeline_desc = {};
+  pipeline_desc.layout = pipeline_layout;
+  pipeline_desc.compute.module = shader;
+  pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  pipeline_desc.compute.constantCount = 1;
+  pipeline_desc.compute.constants = &wg_size_constant;
+  WGPUComputePipeline pipeline =
+      wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
+
+  WGPUBindGroupEntry bg_entries[8] = {};
+  bg_entries[0].binding = 0;
+  bg_entries[0].buffer = cond_tensor.buffer;
+  bg_entries[0].size = cond_bind_size;
+  bg_entries[1].binding = 1;
+  bg_entries[1].buffer = a_tensor.buffer;
+  bg_entries[1].size = a_tensor.nbytes;
+  bg_entries[2].binding = 2;
+  bg_entries[2].buffer = b_tensor.buffer;
+  bg_entries[2].size = b_tensor.nbytes;
+  bg_entries[3].binding = 3;
+  bg_entries[3].buffer = out_tensor.buffer;
+  bg_entries[3].size = out_tensor.nbytes;
+  bg_entries[4].binding = 4;
+  bg_entries[4].buffer = out_meta_buf;
+  bg_entries[4].size = sizeof(TensorMeta);
+  bg_entries[5].binding = 5;
+  bg_entries[5].buffer = cond_meta_buf;
+  bg_entries[5].size = sizeof(TensorMeta);
+  bg_entries[6].binding = 6;
+  bg_entries[6].buffer = a_meta_buf;
+  bg_entries[6].size = sizeof(TensorMeta);
+  bg_entries[7].binding = 7;
+  bg_entries[7].buffer = b_meta_buf;
+  bg_entries[7].size = sizeof(TensorMeta);
+
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 8;
+  bg_desc.entries = bg_entries;
+  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  const size_t dispatch_idx =
+      graph.add_dispatch({pipeline, bind_group, workgroup_count});
+
+  // Dynamic shapes: rebuild the 4 broadcast TensorMeta UBOs + dispatch count.
+  WGPUBuffer o_buf = out_meta_buf, c_buf = cond_meta_buf, a_buf = a_meta_buf,
+             bb_buf = b_meta_buf;
+  auto where_resize = [cond_id,
+                       a_id,
+                       b_id,
+                       out_id,
+                       wg_size,
+                       dispatch_idx,
+                       o_buf,
+                       c_buf,
+                       a_buf,
+                       bb_buf](WebGPUGraph& g) {
+    const auto& c = g.cur_dims(cond_id);
+    const auto& a = g.cur_dims(a_id);
+    const auto& b = g.cur_dims(b_id);
+    const size_t r = std::max({c.size(), a.size(), b.size()});
+    auto dim_at = [r](const std::vector<int64_t>& d, size_t i) -> int64_t {
+      return (i + d.size() < r) ? 1 : d[i - (r - d.size())];
+    };
+    std::vector<int64_t> out_d(r, 1);
+    for (size_t i = 0; i < r; i++) {
+      const int64_t cv = dim_at(c, i), av = dim_at(a, i), bv = dim_at(b, i);
+      int64_t m = std::max({cv, av, bv});
+      if ((cv != m && cv != 1) || (av != m && av != 1) ||
+          (bv != m && bv != 1)) {
+        throw std::runtime_error(
+            "where(resize): operands not broadcast-compatible");
+      }
+      out_d[i] = m;
+    }
+    g.set_cur_dims(out_id, out_d);
+    const uint32_t out_ndim = static_cast<uint32_t>(r);
+    WebGPUTensor tc, ta, tb, to;
+    tc.dims = c;
+    ta.dims = a;
+    tb.dims = b;
+    to.dims = out_d;
+    TensorMeta om, cm, am, bm;
+    fill_tensor_meta_broadcast(to, out_ndim, &om);
+    fill_tensor_meta_broadcast(tc, out_ndim, &cm);
+    fill_tensor_meta_broadcast(ta, out_ndim, &am);
+    fill_tensor_meta_broadcast(tb, out_ndim, &bm);
+    wgpuQueueWriteBuffer(g.queue(), o_buf, 0, &om, sizeof(om));
+    wgpuQueueWriteBuffer(g.queue(), c_buf, 0, &cm, sizeof(cm));
+    wgpuQueueWriteBuffer(g.queue(), a_buf, 0, &am, sizeof(am));
+    wgpuQueueWriteBuffer(g.queue(), bb_buf, 0, &bm, sizeof(bm));
+    g.dispatch_at(dispatch_idx).workgroup_count_x =
+        utils::compute_1d_workgroup_count(
+            g.device(), om.numel, wg_size, "where(resize)");
+  };
+  graph.add_tensor_resize_hook(cond_id, where_resize);
+  graph.add_tensor_resize_hook(a_id, where_resize);
+  graph.add_tensor_resize_hook(b_id, where_resize);
+
+  wgpuShaderModuleRelease(shader);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuPipelineLayoutRelease(pipeline_layout);
+  graph.own_uniform_buffer(out_meta_buf);
+  graph.own_uniform_buffer(cond_meta_buf);
+  graph.own_uniform_buffer(a_meta_buf);
+  graph.own_uniform_buffer(b_meta_buf);
+}
+
+} // namespace
+
+WEBGPU_REGISTER_OPERATORS {
+  WEBGPU_REGISTER_OP(aten.where.self, where_impl);
+}
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/runtime/ops/where/where.wgsl b/backends/webgpu/runtime/ops/where/where.wgsl
new file mode 100644
index 00000000000..eb32f314f83
--- /dev/null
+++ b/backends/webgpu/runtime/ops/where/where.wgsl
@@ -0,0 +1,49 @@
+@group(0) @binding(0) var<storage, read> cond: array<u32>;
+@group(0) @binding(1) var<storage, read> input_a: array<f32>;
+@group(0) @binding(2) var<storage, read> input_b: array<f32>;
+@group(0) @binding(3) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(4) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> cond_meta: TensorMeta;
+@group(0) @binding(6) var<uniform> a_meta: TensorMeta;
+@group(0) @binding(7) var<uniform> b_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+
+// 1-byte bool packed 4-per-u32; extract byte i (mirrors q4gsw weight unpack).
+fn cond_is_true(i: u32) -> bool {
+    let word = cond[i >> 2u];
+    return ((word >> ((i & 3u) * 8u)) & 0xFFu) != 0u;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    var rem = idx;
+    var lc: u32 = 0u;
+    var la: u32 = 0u;
+    var lb: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        lc = lc + min(coord, cond_meta.sizes[d] - 1u) * cond_meta.strides[d];
+        la = la + min(coord, a_meta.sizes[d] - 1u) * a_meta.strides[d];
+        lb = lb + min(coord, b_meta.sizes[d] - 1u) * b_meta.strides[d];
+    }
+
+    if (cond_is_true(lc)) {
+        output[idx] = input_a[la];
+    } else {
+        output[idx] = input_b[lb];
+    }
+}
diff --git a/backends/webgpu/runtime/ops/where/where_wgsl.h b/backends/webgpu/runtime/ops/where/where_wgsl.h
new file mode 100644
index 00000000000..a8ef64fb169
--- /dev/null
+++ b/backends/webgpu/runtime/ops/where/where_wgsl.h
@@ -0,0 +1,73 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace executorch::backends::webgpu {
+
+// @generated from where.wgsl - DO NOT EDIT.
+// wgsl-sha256: e02e6b9e7a073d2aac3d2c110ed5c382ab6663ee9656dc24f4b9c9fdd9e98a38
+inline constexpr const char* kWhereWGSL = R"(
+@group(0) @binding(0) var<storage, read> cond: array<u32>;
+@group(0) @binding(1) var<storage, read> input_a: array<f32>;
+@group(0) @binding(2) var<storage, read> input_b: array<f32>;
+@group(0) @binding(3) var<storage, read_write> output: array<f32>;
+
+struct TensorMeta {
+  ndim: u32,
+  numel: u32,
+  sizes: vec4<u32>,
+  strides: vec4<u32>,
+}
+@group(0) @binding(4) var<uniform> out_meta: TensorMeta;
+@group(0) @binding(5) var<uniform> cond_meta: TensorMeta;
+@group(0) @binding(6) var<uniform> a_meta: TensorMeta;
+@group(0) @binding(7) var<uniform> b_meta: TensorMeta;
+
+override wg_size: u32 = 64u;
+
+// 1-byte bool packed 4-per-u32; extract byte i (mirrors q4gsw weight unpack).
+fn cond_is_true(i: u32) -> bool {
+    let word = cond[i >> 2u];
+    return ((word >> ((i & 3u) * 8u)) & 0xFFu) != 0u;
+}
+
+@compute @workgroup_size(wg_size, 1, 1)
+fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
+    let idx = gid.x;
+    if (idx >= out_meta.numel) {
+        return;
+    }
+
+    var rem = idx;
+    var lc: u32 = 0u;
+    var la: u32 = 0u;
+    var lb: u32 = 0u;
+    for (var d: u32 = 0u; d < out_meta.ndim; d = d + 1u) {
+        let coord = rem / out_meta.strides[d];
+        rem = rem % out_meta.strides[d];
+        lc = lc + min(coord, cond_meta.sizes[d] - 1u) * cond_meta.strides[d];
+        la = la + min(coord, a_meta.sizes[d] - 1u) * a_meta.strides[d];
+        lb = lb + min(coord, b_meta.sizes[d] - 1u) * b_meta.strides[d];
+    }
+
+    if (cond_is_true(lc)) {
+        output[idx] = input_a[la];
+    } else {
+        output[idx] = input_b[lb];
+    }
+}
+)";
+
+inline constexpr uint32_t kWhereWorkgroupSizeX = 64;
+inline constexpr uint32_t kWhereWorkgroupSizeY = 1;
+inline constexpr uint32_t kWhereWorkgroupSizeZ = 1;
+
+} // namespace executorch::backends::webgpu
diff --git a/backends/webgpu/test/ops/test_adamw.py b/backends/webgpu/test/ops/test_adamw.py
new file mode 100644
index 00000000000..b15b4d25bce
--- /dev/null
+++ b/backends/webgpu/test/ops/test_adamw.py
@@ -0,0 +1,142 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""AdamW optimizer step (`et_vk.adamw_step`) export + fp64 golden.
+
+`adamw_step(param, m, v, grad, lr, beta1, beta2, eps, weight_decay, bc1, bc2)`
+updates the fp32 latent in place: decoupled weight decay, then the bias-corrected
+Adam moment update. `bc1`/`bc2` (= 1 - beta^t) are host-precomputed so the kernel
+carries no step counter. The op mutates and returns `param`/`m`/`v` (aliased), so
+export wraps it in `auto_functionalized`, which VulkanPartitioner tags by name (the
+same mutating-op path as `update_cache`). Golden is the fp64 reference, computed
+independently so a lossy fp32 op impl cannot fake-pass.
+"""
+
+from __future__ import annotations
+
+import unittest
+from dataclasses import dataclass
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# torch.optim.AdamW defaults; bc1/bc2 are 1 - beta^step (host-precomputed).
+LR = 1e-3
+BETA1 = 0.9
+BETA2 = 0.999
+EPS = 1e-8
+
+
+@dataclass(frozen=True)
+class AdamwConfig:
+    name: str
+    numel: int
+    weight_decay: float = 0.01
+    step: int = 1
+
+
+CONFIGS = [
+    AdamwConfig("small", 64),
+    AdamwConfig("no_wd", 256, weight_decay=0.0),  # wd=0 -> the plain Adam path
+    AdamwConfig("later_step", 1000, step=10),  # bias correction well past t=1
+]
+
+
+def _inputs(cfg: AdamwConfig):
+    """Deterministic fp32 param/m/v/grad + the host bias corrections."""
+    g = torch.Generator().manual_seed(0)
+    param = torch.randn(cfg.numel, generator=g, dtype=torch.float32)
+    m = torch.randn(cfg.numel, generator=g, dtype=torch.float32) * 0.1
+    v = torch.rand(cfg.numel, generator=g, dtype=torch.float32) * 0.01
+    grad = torch.randn(cfg.numel, generator=g, dtype=torch.float32)
+    bc1 = 1.0 - BETA1**cfg.step
+    bc2 = 1.0 - BETA2**cfg.step
+    return param, m, v, grad, bc1, bc2
+
+
+def _fp64_golden(param, m, v, grad, wd, bc1, bc2):
+    """fp64 truth for one AdamW step; mirrors adamw_step.wgsl exactly."""
+    p = param.double()
+    g = grad.double()
+    p = p - LR * wd * p
+    m64 = BETA1 * m.double() + (1.0 - BETA1) * g
+    v64 = BETA2 * v.double() + (1.0 - BETA2) * g * g
+    mhat = m64 / bc1
+    vhat = v64 / bc2
+    p = p - LR * mhat / (torch.sqrt(vhat) + EPS)
+    return p.float(), m64.float(), v64.float()
+
+
+class _AdamwModule(torch.nn.Module):
+    def __init__(self, wd: float, bc1: float, bc2: float) -> None:
+        super().__init__()
+        self.wd = wd
+        self.bc1 = bc1
+        self.bc2 = bc2
+
+    def forward(self, param, m, v, grad):
+        return torch.ops.et_vk.adamw_step(
+            param, m, v, grad, LR, BETA1, BETA2, EPS, self.wd, self.bc1, self.bc2
+        )
+
+
+def _export(cfg: AdamwConfig):
+    param, m, v, grad, bc1, bc2 = _inputs(cfg)
+    ep = torch.export.export(
+        _AdamwModule(cfg.weight_decay, bc1, bc2), (param, m, v, grad)
+    )
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestAdamwStep(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for cfg in CONFIGS:
+            with self.subTest(config=cfg.name):
+                et = _export(cfg)
+                self.assertTrue(
+                    _delegates(et), f"no VulkanBackend delegate in {cfg.name}"
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        for cfg in CONFIGS:
+            with self.subTest(config=cfg.name):
+                param, m, v, grad, bc1, bc2 = _inputs(cfg)
+                g_param, g_m, g_v = _fp64_golden(
+                    param, m, v, grad, cfg.weight_decay, bc1, bc2
+                )
+                # Op mutates in place; clone so the golden saw the originals.
+                out_param, out_m, out_v = torch.ops.et_vk.adamw_step(
+                    param.clone(),
+                    m.clone(),
+                    v.clone(),
+                    grad,
+                    LR,
+                    BETA1,
+                    BETA2,
+                    EPS,
+                    cfg.weight_decay,
+                    bc1,
+                    bc2,
+                )
+                torch.testing.assert_close(out_param, g_param, atol=5e-4, rtol=1e-3)
+                torch.testing.assert_close(out_m, g_m, atol=5e-4, rtol=1e-3)
+                torch.testing.assert_close(out_v, g_v, atol=5e-4, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_compare.py b/backends/webgpu/test/ops/test_compare.py
new file mode 100644
index 00000000000..172c5e75113
--- /dev/null
+++ b/backends/webgpu/test/ops/test_compare.py
@@ -0,0 +1,109 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.{eq,ne,le,ge,lt,gt}.Scalar` export + golden for the WebGPU backend.
+
+Each scalar comparison lowers to a single `aten.<op>.Scalar` node that the
+kernel computes as `cmp(self[i], scalar)` and writes as a byte-packed bool. The
+delegation test locks that every variant partitions to `VulkanBackend`; the
+golden test locks the fp32 module output against the fp64 torch truth. The
+deterministic ramp is exact in fp32 and straddles (and hits) the scalar, so both
+precisions agree bit-for-bit and every op has mixed True/False cases; a
+non-multiple-of-4 numel exercises the kernel tail (4 elems per u32 word).
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+SCALAR = 0.0
+OPS = ("eq", "ne", "le", "ge", "lt", "gt")
+SHAPES = {"tail": (3, 5), "3d": (2, 4, 8)}
+
+_TORCH_OP = {
+    "eq": torch.eq,
+    "ne": torch.ne,
+    "le": torch.le,
+    "ge": torch.ge,
+    "lt": torch.lt,
+    "gt": torch.gt,
+}
+
+
+class CompareModule(torch.nn.Module):
+    def __init__(self, op: str, scalar: float) -> None:
+        super().__init__()
+        self.op = op
+        self.scalar = scalar
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.op == "eq":
+            return x == self.scalar
+        if self.op == "ne":
+            return x != self.scalar
+        if self.op == "le":
+            return x <= self.scalar
+        if self.op == "ge":
+            return x >= self.scalar
+        if self.op == "lt":
+            return x < self.scalar
+        return x > self.scalar
+
+
+def _det_input(shape: tuple[int, ...]) -> torch.Tensor:
+    """Deterministic fp32 spanning [-0.5, 0.5] in 1/16 steps (exact in fp32,
+    hits and straddles 0.0)."""
+    n = 1
+    for d in shape:
+        n *= d
+    flat = torch.arange(n, dtype=torch.float32)
+    return (((flat % 17) - 8) / 16.0).reshape(shape)
+
+
+def _export(m: torch.nn.Module, x: torch.Tensor):
+    ep = torch.export.export(m, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestCompare(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for op in OPS:
+            for name, shape in SHAPES.items():
+                with self.subTest(op=op, shape=name):
+                    x = _det_input(shape)
+                    et = _export(CompareModule(op, SCALAR).eval(), x)
+                    self.assertTrue(
+                        _delegated(et),
+                        f"Expected a VulkanBackend delegate ({op}.Scalar {name})",
+                    )
+
+    def test_module_matches_fp64_golden(self) -> None:
+        for op in OPS:
+            for name, shape in SHAPES.items():
+                with self.subTest(op=op, shape=name):
+                    x = _det_input(shape)
+                    got = CompareModule(op, SCALAR)(x)
+                    golden = _TORCH_OP[op](x.double(), float(SCALAR))
+                    torch.testing.assert_close(got, golden)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_dim_order.py b/backends/webgpu/test/ops/test_dim_order.py
new file mode 100644
index 00000000000..a1d64244c35
--- /dev/null
+++ b/backends/webgpu/test/ops/test_dim_order.py
@@ -0,0 +1,95 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`dim_order_ops._clone_dim_order.default` export delegation for the WebGPU backend.
+
+`torch.clone(x)` lowers to `dim_order_ops._clone_dim_order.default` (the edge
+dialect's dim-order-aware clone). On the buffer-only WebGPU backend it is a
+numel-preserving flat copy handled by the shared `add_flat_copy` DMA helper (no
+WGSL). The partitioner tags it (single-node partitions are allowed), so a
+`VulkanBackend` delegate is formed; `RemoveRedundantOpsTransform` then folds the
+identity clone out of the delegate in preprocess, so it never reaches the native
+runtime (hence no `.golden.bin` / native sweep — like `clone`).
+
+`test_export_delegates` locks the contract that the op is absorbed into the
+delegate and never survives as a top-level portable (CPU-fallback) node;
+`test_golden_matches_eager` locks the fp64 `x.clone()` reference. Configs cover
+1D, 3D, and 4D contiguous inputs.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> input_shape
+CONFIGS = {
+    "flat": (16,),
+    "3d": (2, 3, 4),
+    "4d": (1, 3, 4, 5),
+}
+
+
+class CloneDimOrderModule(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.clone(x)
+
+
+def _det_input(shape):
+    g = torch.Generator().manual_seed(0)
+    return torch.randn(*shape, generator=g, dtype=torch.float32)
+
+
+def _lower(x: torch.Tensor):
+    ep = torch.export.export(CloneDimOrderModule().eval(), (x,))
+    return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()])
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+def _op_absent_from_toplevel(edge, op_substr: str) -> bool:
+    # Delegated/folded ops are absorbed; none may survive as a top-level node.
+    gm = edge.exported_program().graph_module
+    return all(op_substr not in str(getattr(n, "target", "")) for n in gm.graph.nodes)
+
+
+class TestCloneDimOrder(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, shape in CONFIGS.items():
+            with self.subTest(name=name):
+                edge = _lower(_det_input(shape))
+                et = edge.to_executorch()
+                self.assertTrue(
+                    _delegates(et),
+                    f"Expected a VulkanBackend delegate (clone_dim_order {name})",
+                )
+                self.assertTrue(
+                    _op_absent_from_toplevel(edge, "_clone_dim_order"),
+                    f"_clone_dim_order left as a top-level portable op for {name}",
+                )
+
+    def test_golden_matches_eager(self) -> None:
+        for name, shape in CONFIGS.items():
+            with self.subTest(name=name):
+                x = _det_input(shape)
+                ref = x.to(torch.float64).clone()
+                torch.testing.assert_close(
+                    CloneDimOrderModule()(x).to(torch.float64), ref
+                )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_div.py b/backends/webgpu/test/ops/test_div.py
new file mode 100644
index 00000000000..5d2bbfacff4
--- /dev/null
+++ b/backends/webgpu/test/ops/test_div.py
@@ -0,0 +1,94 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.div.Tensor` (broadcast) export + fp64 golden for the WebGPU backend.
+
+Exports single-op divide graphs through VulkanPartitioner and asserts they
+delegate to the Vulkan backend (div is absent from the top-level portable ops),
+then locks the golden math against an fp64 torch reference (`a / b` with PyTorch
+broadcasting). Configs span the same-shape fast path, a last-dim broadcast at
+LLM width, and a mixed-rank left-pad case. Divisors are bounded away from zero
+so the fp32-vs-fp64 comparison stays well-conditioned.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (shape_a, shape_b). Output shape is the broadcast of the two.
+CONFIGS = {
+    "same": ((8, 32), (8, 32)),  # fast path (same-shape elementwise)
+    "bcast_lastdim": ((1, 1, 7, 896), (1, 1, 7, 1)),  # last-dim broadcast, LLM
+    "mixedrank": ((4,), (3, 4)),  # right-aligned left-pad (in.ndim < out.ndim)
+}
+
+
+class DivModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return a / b
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed); divisor bounded away from zero."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32).abs() + 0.5
+    return a, b
+
+
+def _export(a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(DivModule().eval(), (a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+def _top_level_op_names(et) -> set[str]:
+    return {
+        op.name
+        for plan in et.executorch_program.execution_plan
+        for op in plan.operators
+    }
+
+
+class TestDiv(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                et = _export(a, b)
+                self.assertTrue(
+                    _delegated(et), f"Expected a VulkanBackend delegate (div {name})"
+                )
+                self.assertFalse(
+                    any("div" in n for n in _top_level_op_names(et)),
+                    f"div should be delegated, not a top-level portable op (div {name})",
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                got = DivModule()(a, b)
+                golden = (a.double() / b.double()).to(torch.float32)
+                torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_embedding.py b/backends/webgpu/test/ops/test_embedding.py
new file mode 100644
index 00000000000..f90aa1e086e
--- /dev/null
+++ b/backends/webgpu/test/ops/test_embedding.py
@@ -0,0 +1,103 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.embedding.default` (fp32 row-gather) export + golden for the WebGPU backend.
+
+Exports single-op embedding graphs through VulkanPartitioner and checks a torch
+golden. embedding is a pure row-gather -- out[i, :] = weight[idx[i], :] -- on the
+token-embedding path that feeds the transformer (and the fine-tuning training
+window). 1D indices exercise the [S, D] output; 2D indices the batched [B, S, D]
+path. The indices span the full vocab (incl. the first/last rows) so a wrong
+row-stride would miss the golden.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (num_embeddings, embedding_dim, indices_shape).
+CONFIGS = {
+    "rows_1d": (32, 16, (8,)),
+    "batched_2d": (128, 8, (2, 4)),
+}
+
+
+class EmbeddingModule(torch.nn.Module):
+    def __init__(self, weight: torch.Tensor) -> None:
+        super().__init__()
+        self.weight = torch.nn.Parameter(weight, requires_grad=False)
+
+    def forward(self, idx: torch.Tensor) -> torch.Tensor:
+        return torch.nn.functional.embedding(idx, self.weight)
+
+
+def _det_weight(num_embeddings: int, dim: int) -> torch.Tensor:
+    """Deterministic fp32 [num_embeddings, dim] table (distinct per-row values)."""
+    return torch.linspace(-1.0, 1.0, num_embeddings * dim, dtype=torch.float32).reshape(
+        num_embeddings, dim
+    )
+
+
+def _det_indices(num_embeddings: int, shape: tuple[int, ...]) -> torch.Tensor:
+    """Deterministic int64 indices spread across the vocab, forced to hit row 0
+    and the last row so an off-by-one row-stride shows up in the golden."""
+    n = 1
+    for s in shape:
+        n *= s
+    flat = (torch.arange(n, dtype=torch.int64) * 7 + 3) % num_embeddings
+    flat[0] = 0
+    flat[-1] = num_embeddings - 1
+    return flat.reshape(shape)
+
+
+def _export(m: torch.nn.Module, idx: torch.Tensor):
+    ep = torch.export.export(m, (idx,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestEmbedding(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        # aten.embedding must be absorbed into the VulkanBackend delegate.
+        for name, (num_embeddings, dim, shape) in CONFIGS.items():
+            with self.subTest(name=name):
+                weight = _det_weight(num_embeddings, dim)
+                idx = _det_indices(num_embeddings, shape)
+                et = _export(EmbeddingModule(weight).eval(), idx)
+                self.assertTrue(
+                    _delegates(et),
+                    f"Expected a VulkanBackend delegate (embedding {name})",
+                )
+
+    def test_golden_matches_eager(self) -> None:
+        # fp64 gather golden: out[i,:] == weight[idx[i],:], bit-exact.
+        for name, (num_embeddings, dim, shape) in CONFIGS.items():
+            with self.subTest(name=name):
+                weight = _det_weight(num_embeddings, dim)
+                idx = _det_indices(num_embeddings, shape)
+                got = EmbeddingModule(weight)(idx)
+                golden = torch.nn.functional.embedding(idx, weight.double()).to(
+                    torch.float32
+                )
+                torch.testing.assert_close(got, golden)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_expand_copy.py b/backends/webgpu/test/ops/test_expand_copy.py
new file mode 100644
index 00000000000..b3b79cafba7
--- /dev/null
+++ b/backends/webgpu/test/ops/test_expand_copy.py
@@ -0,0 +1,120 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.expand_copy.default` export + golden for the WebGPU backend.
+
+Exports single-op expand-copy graphs through VulkanPartitioner and checks an fp64
+torch golden. expand_copy materializes a broadcasted view (size-1 input dims, and
+rank-increasing leading dims) into the target shape via a pure gather -- no
+arithmetic, so fp64 and fp32 agree exactly. Configs cover a broadcast leading
+dim, a broadcast middle dim, and a rank increase (input rank < output rank) that
+exercises the kernel's right-alignment of input dims into the output rank.
+"""
+
+from __future__ import annotations
+
+import math
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (input shape, expanded shape).
+CONFIGS = {
+    "broadcast_leading": ((1, 4), (3, 4)),
+    "broadcast_middle": ((2, 1, 5), (2, 4, 5)),
+    "rank_increase": ((4,), (3, 4)),
+}
+
+
+class ExpandCopyModule(torch.nn.Module):
+    def __init__(self, shape: tuple[int, ...]) -> None:
+        super().__init__()
+        self.shape = shape
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x.expand(self.shape).clone()
+
+
+def _det_input(shape: tuple[int, ...]) -> torch.Tensor:
+    """Deterministic fp32 ramp; a broadcast dim repeats it so the copy is visible."""
+    return torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape)
+
+
+def _export(m: torch.nn.Module, x: torch.Tensor):
+    ep = torch.export.export(m, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+def _top_level_op_names(et) -> set[str]:
+    return {
+        op.name
+        for plan in et.executorch_program.execution_plan
+        for op in plan.operators
+    }
+
+
+class TestExpandCopy(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (in_shape, out_shape) in CONFIGS.items():
+            with self.subTest(name=name):
+                x = _det_input(in_shape)
+                et = _export(ExpandCopyModule(out_shape).eval(), x)
+                self.assertTrue(
+                    _delegates(et),
+                    f"Expected a VulkanBackend delegate (expand_copy {name})",
+                )
+                # Delegated => expand_copy absent from top-level portable ops.
+                self.assertFalse(
+                    any("expand_copy" in n for n in _top_level_op_names(et)),
+                    f"expand_copy leaked into top-level ops ({name})",
+                )
+
+    def test_golden_matches_torch(self) -> None:
+        for name, (in_shape, out_shape) in CONFIGS.items():
+            with self.subTest(name=name):
+                x = _det_input(in_shape)
+                golden = x.double().expand(out_shape).clone()
+                got = ExpandCopyModule(out_shape)(x)
+                torch.testing.assert_close(got.double(), golden)
+
+
+def export_expand_copy_model(
+    out_shape: tuple[int, ...],
+    in_shape: tuple[int, ...],
+    pte_path: str,
+    golden_path: str,
+    input_path: str,
+) -> None:
+    """Write an expand_copy .pte + torch golden (raw LE fp32) + raw LE fp32 input."""
+    m = ExpandCopyModule(out_shape).eval()
+    x = _det_input(in_shape)
+    golden = m(x).detach().numpy().astype("<f4")
+    et = _export(m, x)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    x.numpy().astype("<f4").tofile(input_path)
+    print(
+        f"Exported {pte_path}; golden {golden_path} ({golden.size} floats); "
+        f"input {input_path} ({x.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_fill.py b/backends/webgpu/test/ops/test_fill.py
new file mode 100644
index 00000000000..b42e30c6663
--- /dev/null
+++ b/backends/webgpu/test/ops/test_fill.py
@@ -0,0 +1,91 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.full` / `aten.full_like` export + fp64 golden for the WebGPU backend.
+
+Both lower to the `fill` handler, which writes a constant scalar into the
+pre-allocated output buffer. `full` and `full_like` are on the training-backward
+critical path (gradient seeds). To keep the op in the graph (a constant-shaped
+`full` would fold away), the fill size is taken from a live input: `full` uses
+`x.shape`, `full_like` uses `x`. Configs span 1D/2D/4D and a non-multiple-of-256
+shape that exercises the dispatch bounds guard.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (kind, shape, fill_value)
+CONFIGS = {
+    "full_1d": ("full", (37,), 0.0),  # non-multiple of 256: bounds-guard path
+    "full_2d": ("full", (4, 8), 3.0),
+    "full_4d": ("full", (2, 3, 4, 5), -1.5),
+    "full_like_2d": ("full_like", (16, 16), 7.25),
+}
+
+
+class FullModule(torch.nn.Module):
+    def __init__(self, val: float) -> None:
+        super().__init__()
+        self.val = val
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.full(x.shape, self.val)
+
+
+class FullLikeModule(torch.nn.Module):
+    def __init__(self, val: float) -> None:
+        super().__init__()
+        self.val = val
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.full_like(x, self.val)
+
+
+def _module(kind: str, val: float) -> torch.nn.Module:
+    return (FullModule(val) if kind == "full" else FullLikeModule(val)).eval()
+
+
+def _export(m: torch.nn.Module, x: torch.Tensor):
+    ep = torch.export.export(m, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestFill(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (kind, shape, val) in CONFIGS.items():
+            with self.subTest(name=name):
+                x = torch.zeros(shape, dtype=torch.float32)
+                et = _export(_module(kind, val), x)
+                self.assertTrue(
+                    _delegates(et), f"Expected a VulkanBackend delegate (fill {name})"
+                )
+
+    def test_golden_matches_eager(self) -> None:
+        for name, (kind, shape, val) in CONFIGS.items():
+            with self.subTest(name=name):
+                x = torch.zeros(shape, dtype=torch.float32)
+                golden = torch.full(shape, val, dtype=torch.float64)
+                torch.testing.assert_close(_module(kind, val)(x).double(), golden)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_fused_ce.py b/backends/webgpu/test/ops/test_fused_ce.py
new file mode 100644
index 00000000000..13d2131c64a
--- /dev/null
+++ b/backends/webgpu/test/ops/test_fused_ce.py
@@ -0,0 +1,130 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Fused cross-entropy training op (`et_vk.fused_ce`) export + fp64 golden.
+
+`fused_ce(logits[M,V], labels[M], n_valid) -> (loss, dlogits[M,V])` computes the
+mean-over-valid CE loss and its gradient in one op (labels < 0 are ignored/pad).
+Golden is the fp64 reference (`logsumexp - picked`, `softmax - onehot`), the same math
+torch's cross_entropy uses; the native test reconstructs the deterministic inputs.
+"""
+
+import os
+import unittest
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from executorch.backends.vulkan import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+
+@dataclass(frozen=True)
+class CeConfig:
+    name: str
+    m: int  # rows (valid + pad positions)
+    v: int  # vocab
+    n_pad: int = 0  # trailing rows with label = -1 (ignored)
+
+
+# Mirrored by the C++ kFusedCeConfigs table. Llama-3.2-1B vocab = 128256.
+CONFIGS = [
+    CeConfig("tiny", 4, 32),
+    CeConfig("masked", 8, 128, n_pad=3),  # some ignored labels
+    CeConfig("llama_vocab", 16, 128256),  # real vocab width
+]
+
+
+def _inputs(cfg: CeConfig):
+    """Deterministic logits [M,V] + labels [M] (last n_pad = -1); reconstructable in C++."""
+    flat = np.arange(cfg.m * cfg.v, dtype=np.int64)
+    logits = torch.from_numpy(
+        (((flat % 23) - 11).astype(np.float32) / np.float32(8.0)).reshape(cfg.m, cfg.v)
+    )
+    labels = torch.from_numpy((np.arange(cfg.m, dtype=np.int64) * 7 + 3) % cfg.v)
+    if cfg.n_pad:
+        labels[cfg.m - cfg.n_pad :] = -1
+    n_valid = float(max(1, cfg.m - cfg.n_pad))
+    return logits, labels, n_valid
+
+
+def _fp64_golden(logits: torch.Tensor, labels: torch.Tensor, n_valid: float):
+    mask = labels >= 0
+    safe = labels.clamp(min=0).long()
+    lg = logits.double()
+    lse = torch.logsumexp(lg, dim=-1)
+    picked = lg.gather(-1, safe[:, None]).squeeze(-1)
+    loss = torch.where(mask, (lse - picked) / n_valid, torch.zeros_like(lse)).sum()
+    softmax = torch.softmax(lg, dim=-1)
+    onehot = torch.nn.functional.one_hot(safe, logits.shape[-1]).double()
+    dlogits = torch.where(
+        mask[:, None], (softmax - onehot) / n_valid, torch.zeros_like(softmax)
+    )
+    return loss.to(torch.float32), dlogits.to(torch.float32)
+
+
+class _CeModule(torch.nn.Module):
+    def forward(self, logits, labels, n_valid):
+        return torch.ops.et_vk.fused_ce(logits, labels, n_valid)
+
+
+def _export(logits, labels, n_valid):
+    ep = torch.export.export(_CeModule(), (logits, labels, n_valid))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+class TestFusedCe(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for cfg in CONFIGS:
+            if cfg.v > 1024:  # width-independent; skip the 128k fixture
+                continue
+            with self.subTest(config=cfg.name):
+                logits, labels, n_valid = _inputs(cfg)
+                et = _export(logits, labels, n_valid)
+                found = any(
+                    d.id == "VulkanBackend"
+                    for plan in et.executorch_program.execution_plan
+                    for d in plan.delegates
+                )
+                self.assertTrue(found, f"no VulkanBackend delegate in {cfg.name}")
+
+    def test_op_matches_fp64_golden(self) -> None:
+        for cfg in CONFIGS:
+            if cfg.v > 1024:
+                continue
+            with self.subTest(config=cfg.name):
+                logits, labels, n_valid = _inputs(cfg)
+                loss, dlogits = torch.ops.et_vk.fused_ce(logits, labels, n_valid)
+                g_loss, g_dlogits = _fp64_golden(logits, labels, n_valid)
+                torch.testing.assert_close(loss, g_loss, atol=5e-4, rtol=1e-3)
+                torch.testing.assert_close(dlogits, g_dlogits, atol=5e-4, rtol=1e-3)
+
+
+def export_fused_ce_model(cfg: CeConfig, pte_path: str, golden_path: str) -> None:
+    logits, labels, n_valid = _inputs(cfg)
+    et = _export(logits, labels, n_valid)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    g_loss, g_dlogits = _fp64_golden(logits, labels, n_valid)
+    # loss scalar then dlogits, both raw LE fp32
+    np.concatenate([g_loss.reshape(1).numpy(), g_dlogits.reshape(-1).numpy()]).astype(
+        "<f4"
+    ).tofile(golden_path)
+    print(f"Exported {pte_path}; golden {golden_path}")
+
+
+def export_all_fused_ce_models(out_dir: str) -> None:
+    for cfg in CONFIGS:
+        pte = os.path.join(out_dir, f"fused_ce_{cfg.name}.pte")
+        golden = os.path.join(out_dir, f"fused_ce_{cfg.name}.golden.bin")
+        export_fused_ce_model(cfg, pte, golden)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_gather.py b/backends/webgpu/test/ops/test_gather.py
new file mode 100644
index 00000000000..3aff9b8626f
--- /dev/null
+++ b/backends/webgpu/test/ops/test_gather.py
@@ -0,0 +1,123 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.gather.default` export + fp64 golden for the WebGPU backend.
+
+Exports single-op gather graphs through VulkanPartitioner and writes a torch-computed
+golden (the native binary has no ATen). gather(self, dim, index) copies self along
+`dim` at the positions named by index (out has index's shape). Configs cover the last
+dim, dim 0, and a rank-3 negative dim (exercises the handler's dim-normalization). The
+native test reconstructs the deterministic inputs bit-for-bit.
+"""
+
+from __future__ import annotations
+
+import os
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (self_shape, dim, index_shape). Ranks stay <= 4 (TensorMeta MAX_NDIM).
+CONFIGS = {
+    "cols": ((4, 8), 1, (4, 3)),
+    "rows": ((5, 6), 0, (3, 6)),
+    "rank3_neg": ((2, 3, 4), -1, (2, 3, 2)),
+}
+
+
+class GatherModule(torch.nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dim = dim
+
+    def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor:
+        return torch.gather(x, self.dim, index)
+
+
+def _det_inputs(self_shape, dim: int, index_shape):
+    """Distinct fp32 source (a wrong pick is visible) + in-range int64 index."""
+    n = 1
+    for s in self_shape:
+        n *= s
+    x = torch.arange(n, dtype=torch.float32).reshape(self_shape)
+    bound = self_shape[dim]
+    m = 1
+    for s in index_shape:
+        m *= s
+    index = (torch.arange(m, dtype=torch.int64) % bound).reshape(index_shape)
+    return x, index
+
+
+def _lower(m: torch.nn.Module, x: torch.Tensor, index: torch.Tensor):
+    ep = torch.export.export(m, (x, index))
+    return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()])
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+def _op_delegated(edge) -> bool:
+    # gather must be absorbed into the delegate, not a top-level CPU node.
+    gm = edge.exported_program().graph_module
+    return all("gather" not in str(getattr(n, "target", "")) for n in gm.graph.nodes)
+
+
+class TestGather(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (self_shape, dim, index_shape) in CONFIGS.items():
+            with self.subTest(name=name):
+                x, index = _det_inputs(self_shape, dim, index_shape)
+                edge = _lower(GatherModule(dim).eval(), x, index)
+                et = edge.to_executorch()
+                self.assertTrue(
+                    _delegated(et),
+                    f"Expected a VulkanBackend delegate (gather {name})",
+                )
+                self.assertTrue(
+                    _op_delegated(edge),
+                    f"gather not delegated (fell back to CPU) for {name}",
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        for name, (self_shape, dim, index_shape) in CONFIGS.items():
+            with self.subTest(name=name):
+                x, index = _det_inputs(self_shape, dim, index_shape)
+                got = GatherModule(dim)(x, index)
+                golden = torch.gather(x.double(), dim, index).to(torch.float32)
+                torch.testing.assert_close(got, golden)
+
+
+def export_gather_model(name: str, pte_path: str, golden_path: str) -> None:
+    """Write one config's gather .pte + fp64 torch golden (raw LE fp32)."""
+    self_shape, dim, index_shape = CONFIGS[name]
+    x, index = _det_inputs(self_shape, dim, index_shape)
+    et = _lower(GatherModule(dim).eval(), x, index).to_executorch()
+    golden = torch.gather(x.double(), dim, index).to(torch.float32)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.numpy().astype("<f4").tofile(golden_path)
+    print(f"Exported {pte_path}; golden {golden_path} ({golden.numel()} floats)")
+
+
+def export_all_gather_models(out_dir: str) -> None:
+    for name in CONFIGS:
+        export_gather_model(
+            name,
+            os.path.join(out_dir, f"gather_{name}.pte"),
+            os.path.join(out_dir, f"gather_{name}.golden.bin"),
+        )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_linear_dW.py b/backends/webgpu/test/ops/test_linear_dW.py
new file mode 100644
index 00000000000..16a6e6a6ec0
--- /dev/null
+++ b/backends/webgpu/test/ops/test_linear_dW.py
@@ -0,0 +1,88 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`et_vk.linear_dW` (STE weight gradient) export + fp64 golden.
+
+The weight gradient of a frozen 4-bit linear: `d_W[N, K] = d_out^T @ x`, the grad
+wrt the dequantized weight (both operands are fp32; no int4 unpack). Reached by a
+direct op call in the on-device training graph. CONFIGS reuse Llama-3.2-1B linear
+shapes plus a non-tile-aligned shape that exercises the kernel's boundary clamp.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (m tokens, k in_features, n out_features).
+CONFIGS = {
+    "q_proj_112": (112, 2048, 2048),
+    "kv_proj_112": (112, 2048, 512),
+    "boundary": (13, 18, 10),  # non-multiple-of-4: exercises the min()-clamp
+}
+
+
+class LinearDwModule(torch.nn.Module):
+    def forward(self, d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+        return torch.ops.et_vk.linear_dW(d_out, x)
+
+
+def _det_inputs(m: int, k: int, n: int):
+    """Deterministic fp32 d_out [m, n] + x [m, k] (fixed seed)."""
+    g = torch.Generator().manual_seed(0)
+    d_out = torch.randn(m, n, generator=g, dtype=torch.float32)
+    x = torch.randn(m, k, generator=g, dtype=torch.float32)
+    return d_out, x
+
+
+def _fp64_golden(d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    """fp64 truth: d_W = d_out^T @ x, [N, M] @ [M, K] = [N, K]."""
+    return (d_out.double().t() @ x.double()).to(torch.float32)
+
+
+def _export(d_out: torch.Tensor, x: torch.Tensor):
+    ep = torch.export.export(LinearDwModule().eval(), (d_out, x))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestLinearDw(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (m, k, n) in CONFIGS.items():
+            with self.subTest(config=name):
+                d_out, x = _det_inputs(m, k, n)
+                et = _export(d_out, x)
+                self.assertTrue(
+                    _delegated(et),
+                    f"Expected a VulkanBackend delegate (linear_dW {name})",
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        # Op (d_out^T @ x) vs fp64 matmul truth: guards the formula.
+        for name, (m, k, n) in CONFIGS.items():
+            with self.subTest(config=name):
+                d_out, x = _det_inputs(m, k, n)
+                got = torch.ops.et_vk.linear_dW(d_out, x)
+                golden = _fp64_golden(d_out, x)
+                self.assertEqual(tuple(got.shape), (n, k))
+                torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_logical_not.py b/backends/webgpu/test/ops/test_logical_not.py
new file mode 100644
index 00000000000..6e9c15eaa08
--- /dev/null
+++ b/backends/webgpu/test/ops/test_logical_not.py
@@ -0,0 +1,116 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.logical_not.default` export + golden for the WebGPU backend.
+
+Exports single-op logical_not graphs through VulkanPartitioner and writes a
+torch-computed golden. logical_not lowers to a byte-packed bool kernel (1 byte /
+elem, one u32 word / thread), so the golden is an EXACT bool match -- the native
+test compares the raw uint8 output byte-for-byte. The bool input is produced by a
+delegated `x >= threshold` compare; the golden uses the De Morgan inverse
+`x < threshold` as an independent reference. A non-multiple-of-4 numel exercises
+the shader's partial-final-word masking (the `i < num_elements` branch).
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> shape. "tail3x7" has numel 21 (not a multiple of 4) => partial word.
+CONFIGS = {
+    "vec1d": (16,),
+    "mat2d": (4, 8),
+    "tail3x7": (3, 7),
+    "cube3d": (2, 4, 8),
+}
+
+
+class LogicalNotModule(torch.nn.Module):
+    def __init__(self, threshold: float) -> None:
+        super().__init__()
+        self.threshold = threshold
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.logical_not(x >= self.threshold)
+
+
+def _det_input(shape: tuple[int, ...]) -> torch.Tensor:
+    """Deterministic fp32 spanning negatives, zero, and positives so the bool
+    input to logical_not carries both True and False values."""
+    g = torch.Generator().manual_seed(0)
+    return torch.randn(*shape, generator=g, dtype=torch.float32)
+
+
+def _lower(m: torch.nn.Module, x: torch.Tensor):
+    ep = torch.export.export(m, (x,))
+    return to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()])
+
+
+def _delegates(edge) -> bool:
+    et = edge.to_executorch()
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+def _top_level_targets(edge) -> list[str]:
+    return [
+        str(n.target)
+        for n in edge.exported_program().graph_module.graph.nodes
+        if n.op == "call_function"
+    ]
+
+
+class TestLogicalNot(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, shape in CONFIGS.items():
+            with self.subTest(config=name):
+                edge = _lower(LogicalNotModule(0.0).eval(), _det_input(shape))
+                targets = _top_level_targets(edge)
+                self.assertFalse(
+                    any("logical_not" in t for t in targets),
+                    f"logical_not must be delegated, not portable ({name}): {targets}",
+                )
+                self.assertTrue(
+                    _delegates(edge),
+                    f"Expected a VulkanBackend delegate (logical_not {name})",
+                )
+
+    def test_golden_matches_eager(self) -> None:
+        for name, shape in CONFIGS.items():
+            with self.subTest(config=name):
+                x = _det_input(shape)
+                got = LogicalNotModule(0.0)(x)
+                golden = x < 0.0  # De Morgan inverse of logical_not(x >= 0)
+                self.assertEqual(got.dtype, torch.bool)
+                torch.testing.assert_close(got, golden)
+
+
+def export_logical_not_model(pte_path: str, golden_path: str, input_path: str) -> None:
+    """Write logical_not .pte + torch golden (raw uint8) + raw LE fp32 input."""
+    m = LogicalNotModule(0.0).eval()
+    x = _det_input(CONFIGS["mat2d"])
+    golden = m(x).detach().numpy().astype("<u1")
+    et = _lower(m, x).to_executorch()
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    x.numpy().astype("<f4").tofile(input_path)
+    print(
+        f"Exported {pte_path}; golden {golden_path} ({golden.size} bytes); "
+        f"input {input_path} ({x.numel()} floats)"
+    )
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_q4gsw_requant.py b/backends/webgpu/test/ops/test_q4gsw_requant.py
new file mode 100644
index 00000000000..a342902ebea
--- /dev/null
+++ b/backends/webgpu/test/ops/test_q4gsw_requant.py
@@ -0,0 +1,123 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`et_vk.q4gsw_requant` (STE re-quant + int4 pack) export + fp32 code golden.
+
+Writes updated fp32 latent weights back to the 4-bit group-symmetric packed codes
+`et_vk.linear_q4gsw` reads (only the codes move; the per-group scale is frozen).
+Reached by a direct op call in the on-device training graph after the optimizer
+step. The golden is computed in fp32 (not fp64) on purpose: the kernel rounds
+`round(latent / scale)` in IEEE-754 fp32, so a bit-exact contract must round in
+fp32 too -- an fp64 reference would flip codes at half-way ties. CONFIGS reuse
+Llama-3.2-1B linear shapes plus an odd-K shape that exercises the final-nibble
+tail guard.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (n out_features, k in_features, group_size).
+CONFIGS = {
+    "kv_proj": (512, 2048, 64),
+    "q_proj_g32": (2048, 2048, 32),
+    "small_odd_k": (6, 129, 64),  # odd K -> trailing low-nibble-only byte
+}
+
+
+class Q4gswRequantModule(torch.nn.Module):
+    def __init__(self, group_size: int) -> None:
+        super().__init__()
+        self.group_size = group_size
+
+    def forward(self, latent: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
+        return torch.ops.et_vk.q4gsw_requant(latent, scales, self.group_size)
+
+
+def _det_inputs(n: int, k: int, gs: int):
+    """Deterministic fp32 latent [N, K] + frozen scales [num_groups, N] (fixed seed).
+
+    Scales are small relative to the latent so `latent / scale` spans well beyond
+    [-8, 7], exercising the clamp on both ends.
+    """
+    num_groups = (k + gs - 1) // gs
+    g = torch.Generator().manual_seed(0)
+    latent = torch.randn(n, k, generator=g, dtype=torch.float32)
+    scales = torch.rand(num_groups, n, generator=g, dtype=torch.float32) * 0.1 + 0.05
+    return latent, scales
+
+
+def _reference_codes(
+    latent: torch.Tensor, scales: torch.Tensor, gs: int
+) -> torch.Tensor:
+    """fp32 truth for the int4 codes: clamp(round(latent / scale), -8, 7), [N, K]."""
+    n, k = latent.shape
+    group_idx = torch.arange(k) // gs  # [K]
+    scale_full = scales.t()[:, group_idx]  # [N, K]: scales[k // gs, n]
+    q = torch.round(latent / scale_full)
+    return torch.clamp(q, -8, 7).to(torch.int64)
+
+
+def _unpack(packed: torch.Tensor, k: int) -> torch.Tensor:
+    """Undo the nibble packing: even k -> low nibble, odd k -> high nibble, code - 8."""
+    n = packed.shape[0]
+    p = packed.to(torch.int64)
+    low = p & 0xF
+    high = (p >> 4) & 0xF
+    codes = torch.zeros(n, k, dtype=torch.int64)
+    n_low = codes[:, 0::2].shape[1]
+    n_high = codes[:, 1::2].shape[1]
+    codes[:, 0::2] = low[:, :n_low] - 8
+    codes[:, 1::2] = high[:, :n_high] - 8
+    return codes
+
+
+def _export(latent: torch.Tensor, scales: torch.Tensor, gs: int):
+    ep = torch.export.export(Q4gswRequantModule(gs).eval(), (latent, scales))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegated(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestQ4gswRequant(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (n, k, gs) in CONFIGS.items():
+            with self.subTest(config=name):
+                latent, scales = _det_inputs(n, k, gs)
+                et = _export(latent, scales, gs)
+                self.assertTrue(
+                    _delegated(et),
+                    f"Expected a VulkanBackend delegate (q4gsw_requant {name})",
+                )
+
+    def test_op_matches_fp32_golden(self) -> None:
+        # Op codes vs fp32 quant truth: guards formula+layout, bit-exact.
+        for name, (n, k, gs) in CONFIGS.items():
+            with self.subTest(config=name):
+                latent, scales = _det_inputs(n, k, gs)
+                packed = torch.ops.et_vk.q4gsw_requant(latent, scales, gs)
+                self.assertEqual(packed.dtype, torch.uint8)
+                self.assertEqual(tuple(packed.shape), (n, (k + 1) // 2))
+                got = _unpack(packed, k)
+                golden = _reference_codes(latent, scales, gs)
+                torch.testing.assert_close(got, golden, atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_quantized_linear_backward.py b/backends/webgpu/test/ops/test_quantized_linear_backward.py
new file mode 100644
index 00000000000..a675be79251
--- /dev/null
+++ b/backends/webgpu/test/ops/test_quantized_linear_backward.py
@@ -0,0 +1,157 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Backward of 4-bit quantized linear (`et_vk.linear_q4gsw_backward`) export + fp64 golden.
+
+Mirrors test_quantized_linear.py. The backward computes `d_x = d_out @ dequant(W)` so
+gradients flow through a frozen 4-bit base into a LoRA/DiReFT adapter. CONFIGS reuse the
+real Llama-3.2-1B linear shapes (the backward's d_out is [M, N] and its output d_x is
+[M, K]). The golden is the fp64 dequant-matmul truth; the native test
+(test_webgpu_native.cpp) reconstructs the identical deterministic ramp bit-for-bit.
+"""
+
+import os
+import unittest
+from dataclasses import dataclass
+
+import numpy as np
+import torch
+
+from executorch.backends.vulkan import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+from torchao.quantization.granularity import PerGroup
+from torchao.quantization.quant_api import IntxWeightOnlyConfig, quantize_
+
+
+@dataclass(frozen=True)
+class BwdConfig:
+    name: str
+    m: int  # rows (tokens)
+    k: int  # in_features (== d_x cols)
+    n: int  # out_features (== d_out cols)
+    group_size: int = 32  # K % group_size == 0, K % 8 == 0, N % 8 == 0
+    heavy: bool = False
+
+
+# Mirrored by the C++ kQ4gswBackwardConfigs table (Llama-3.2-1B shapes).
+CONFIGS = [
+    BwdConfig("q_proj", 1, 2048, 2048),  # also o_proj
+    BwdConfig("kv_proj", 1, 2048, 512),
+    BwdConfig("gate_proj", 1, 2048, 8192),
+    BwdConfig("down_proj", 1, 8192, 2048),
+    BwdConfig("q_proj_112", 112, 2048, 2048),  # S=112 multi-row training window
+]
+
+
+def _make_quantized_model(k: int, n: int, group_size: int) -> torch.nn.Module:
+    torch.manual_seed(0)  # load-bearing: fixes the weights the golden uses
+    m = torch.nn.Linear(k, n, bias=False).eval()
+    quantize_(
+        m,
+        IntxWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(group_size)),
+    )
+    return m
+
+
+def _ramp(m_rows: int, cols: int) -> torch.Tensor:
+    """Deterministic fp32 [rows, cols]; the C++ side reconstructs it bit-for-bit.
+
+    v[flat] = ((flat % 17) - 8) / 16 -- exact in fp32 (small modulus, po2 denominator).
+    """
+    flat = np.arange(m_rows * cols, dtype=np.int64)
+    v = ((flat % 17) - 8).astype(np.float32) / np.float32(16.0)
+    return torch.from_numpy(v).reshape(m_rows, cols)
+
+
+def _packed_qweights(m: torch.nn.Module):
+    """The int4 packed weights + per-group scales `et_vk.linear_q4gsw` consumes.
+
+    Recover them the same way the forward op does: `dequant(W)` is [N, K]; the backward op
+    takes the same (weights, weight_scales, group_size) triple the partitioner extracts.
+    """
+    aqt = m.weight  # AffineQuantizedTensor
+    return aqt
+
+
+def _fp64_golden(m: torch.nn.Module, d_out: torch.Tensor) -> np.ndarray:
+    """fp64 truth: d_x = d_out @ dequant(W), dequant(W) is [N, K] -> [M, N]@[N, K] = [M, K]."""
+    wq = m.weight.dequantize()  # [N, K]
+    d_x = d_out.double() @ wq.double()  # [M, K] in fp64
+    return d_x.to(torch.float32).numpy().astype("<f4")
+
+
+class _BackwardModule(torch.nn.Module):
+    """Wraps the linear so autograd through it exercises the registered backward op."""
+
+    def __init__(self, lin: torch.nn.Module) -> None:
+        super().__init__()
+        self.lin = lin
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.lin(x)
+
+
+def _export_backward(m: torch.nn.Module, x: torch.Tensor):
+    # Training-style export: forward + backward makes the backward reachable.
+    mod = _BackwardModule(m)
+    ep = torch.export.export(mod, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+class TestQuantizedLinearBackward(unittest.TestCase):
+    def test_op_matches_fp64_golden(self) -> None:
+        # Op impl (d_out @ dequant(W)) vs fp64 truth: guards backward formula.
+        for cfg in CONFIGS:
+            if cfg.heavy:
+                continue
+            with self.subTest(config=cfg.name):
+                m = _make_quantized_model(cfg.k, cfg.n, cfg.group_size)
+                d_out = _ramp(cfg.m, cfg.n)
+                got = torch.ops.et_vk.linear_q4gsw_backward(
+                    d_out, _packed_qweights(m), None, cfg.group_size
+                )
+                golden = torch.from_numpy(_fp64_golden(m, d_out))
+                torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+    def test_autograd_backward_matches_golden(self) -> None:
+        # autograd through linear_q4gsw uses the registered backward op.
+        for cfg in CONFIGS:
+            if cfg.heavy:
+                continue
+            with self.subTest(config=cfg.name):
+                m = _make_quantized_model(cfg.k, cfg.n, cfg.group_size)
+                x = _ramp(cfg.m, cfg.k).requires_grad_(True)
+                d_out = _ramp(cfg.m, cfg.n)
+                y = m(x)
+                y.backward(d_out)
+                golden = torch.from_numpy(_fp64_golden(m, d_out))
+                torch.testing.assert_close(x.grad, golden, atol=5e-4, rtol=1e-3)
+
+
+def export_backward_model(cfg: BwdConfig, pte_path: str, golden_path: str) -> None:
+    """Export one config's backward .pte + its fp64 golden (raw LE fp32)."""
+    m = _make_quantized_model(cfg.k, cfg.n, cfg.group_size)
+    x = _ramp(cfg.m, cfg.k)
+    et = _export_backward(m, x)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    _fp64_golden(m, _ramp(cfg.m, cfg.n)).tofile(golden_path)
+    print(f"Exported {pte_path}; golden {golden_path} ({cfg.m * cfg.k} floats)")
+
+
+def export_all_backward_models(out_dir: str, include_heavy: bool = False) -> None:
+    for cfg in CONFIGS:
+        if cfg.heavy and not include_heavy:
+            continue
+        pte = os.path.join(out_dir, f"q4gsw_backward_{cfg.name}.pte")
+        golden = os.path.join(out_dir, f"q4gsw_backward_{cfg.name}.golden.bin")
+        export_backward_model(cfg, pte, golden)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_reduce.py b/backends/webgpu/test/ops/test_reduce.py
new file mode 100644
index 00000000000..39c9dc48506
--- /dev/null
+++ b/backends/webgpu/test/ops/test_reduce.py
@@ -0,0 +1,127 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.sum.dim_IntList` / `aten.mean.dim` single-dim reduction export + fp64 golden.
+
+Exports single-op sum/mean graphs through VulkanPartitioner and checks the kernel
+math against an fp64 torch reference. The handler reduces one dim at a time via an
+outer/r/inner decomposition: `dim=-1` gives inner=1 (unit-stride reduction), a
+middle dim gives inner>1 (the non-unit-stride path); `keepdim` toggles whether the
+reduced dim survives in the output shape.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+
+class ReduceModule(torch.nn.Module):
+    def __init__(self, op: str, dim: int, keepdim: bool) -> None:
+        super().__init__()
+        self.op = op
+        self.dim = dim
+        self.keepdim = keepdim
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.op == "sum":
+            return torch.sum(x, dim=self.dim, keepdim=self.keepdim)
+        return torch.mean(x, dim=self.dim, keepdim=self.keepdim)
+
+
+# (name, shape, dim, keepdim): dim=-1 -> inner=1; middle dim -> inner>1.
+CONFIGS = [
+    ("last_dim_keep", (4, 8), -1, True),
+    ("last_dim_drop", (4, 8), -1, False),
+    ("middle_dim_drop", (2, 3, 4), 1, False),  # inner=4: non-unit-stride reduction
+    ("middle_dim_keep", (2, 3, 4), 1, True),
+]
+
+
+def _det_input(shape) -> torch.Tensor:
+    """Deterministic fp32 [shape]; the C++ side reconstructs it bit-for-bit.
+
+    v[flat] = ((flat % 17) - 8) / 16 -- exact in fp32 (small modulus, po2 denominator).
+    """
+    n = 1
+    for s in shape:
+        n *= s
+    flat = torch.arange(n, dtype=torch.float32)
+    return ((flat % 17) - 8).div(16.0).reshape(shape)
+
+
+def _export(m: torch.nn.Module, x: torch.Tensor):
+    ep = torch.export.export(m, (x,))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+def _fp64_golden(x: torch.Tensor, op: str, dim: int, keepdim: bool) -> torch.Tensor:
+    xd = x.double()
+    if op == "sum":
+        ref = torch.sum(xd, dim=dim, keepdim=keepdim)
+    else:
+        ref = torch.mean(xd, dim=dim, keepdim=keepdim)
+    return ref.to(torch.float32)
+
+
+class TestReduce(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for op in ("sum", "mean"):
+            for name, shape, dim, keepdim in CONFIGS:
+                with self.subTest(op=op, config=name):
+                    x = _det_input(shape)
+                    et = _export(ReduceModule(op, dim, keepdim).eval(), x)
+                    self.assertTrue(
+                        _delegates(et),
+                        f"Expected a VulkanBackend delegate ({op} {name})",
+                    )
+
+    def test_matches_fp64_golden(self) -> None:
+        for op in ("sum", "mean"):
+            for name, shape, dim, keepdim in CONFIGS:
+                with self.subTest(op=op, config=name):
+                    x = _det_input(shape)
+                    got = ReduceModule(op, dim, keepdim)(x)
+                    golden = _fp64_golden(x, op, dim, keepdim)
+                    torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)
+
+
+def export_reduce_model(
+    op: str,
+    shape,
+    dim: int,
+    keepdim: bool,
+    pte_path: str,
+    golden_path: str,
+    input_path: str,
+) -> None:
+    """Write a reduce .pte + torch fp64 golden (raw LE fp32) + raw LE fp32 input."""
+    m = ReduceModule(op, dim, keepdim).eval()
+    x = _det_input(shape)
+    et = _export(m, x)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    _fp64_golden(x, op, dim, keepdim).numpy().astype("<f4").tofile(golden_path)
+    x.numpy().astype("<f4").tofile(input_path)
+    print(f"Exported {pte_path}; golden {golden_path}; input {input_path}")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_sub.py b/backends/webgpu/test/ops/test_sub.py
new file mode 100644
index 00000000000..5966998bd7b
--- /dev/null
+++ b/backends/webgpu/test/ops/test_sub.py
@@ -0,0 +1,114 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.sub.Tensor` export + golden for the WebGPU backend.
+
+Exports single-op subtraction graphs through VulkanPartitioner (the WebGPU runtime
+consumes the Vulkan VK00 delegate directly) and checks an fp64 torch golden
+(`out = in1 - alpha * in2`). The native/etvk numeric oracle compares the GPU kernel
+against this same reference; this suite locks delegation + the reference math.
+Configs span same-shape 2D/3D, trailing- and leading-dim broadcast, and alpha != 1.
+"""
+
+from __future__ import annotations
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+
+# name -> (shape_a, shape_b); b broadcasts into a.
+CONFIGS = {
+    "2d": ((4, 4), (4, 4)),
+    "3d": ((2, 3, 4), (2, 3, 4)),
+    "bcast_last": ((4, 4), (4, 1)),
+    "bcast_row": ((4, 4), (1, 4)),
+}
+
+
+class SubModule(torch.nn.Module):
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return torch.sub(a, b)
+
+
+class SubAlphaModule(torch.nn.Module):
+    def __init__(self, alpha: float) -> None:
+        super().__init__()
+        self.alpha = alpha
+
+    def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+        return torch.sub(a, b, alpha=self.alpha)
+
+
+def _det_inputs(shape_a, shape_b):
+    """Deterministic fp32 inputs (fixed seed) for a config."""
+    g = torch.Generator().manual_seed(0)
+    a = torch.randn(*shape_a, generator=g, dtype=torch.float32)
+    b = torch.randn(*shape_b, generator=g, dtype=torch.float32)
+    return a, b
+
+
+def _export(m: torch.nn.Module, a: torch.Tensor, b: torch.Tensor):
+    ep = torch.export.export(m.eval(), (a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestSub(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        # Delegation => no aten.sub.Tensor left in the top-level portable graph.
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                et = _export(SubModule(), a, b)
+                self.assertTrue(
+                    _delegates(et),
+                    f"Expected a VulkanBackend delegate (sub {name})",
+                )
+
+    def test_golden_matches_fp64(self) -> None:
+        for name, (sa, sb) in CONFIGS.items():
+            with self.subTest(name=name):
+                a, b = _det_inputs(sa, sb)
+                ref = (a.double() - b.double()).to(torch.float32)
+                torch.testing.assert_close(SubModule()(a, b), ref)
+
+    def test_golden_matches_fp64_alpha(self) -> None:
+        # Locks the handler's out = in1 - alpha * in2 path (alpha != 1).
+        alpha = 2.5
+        a, b = _det_inputs((4, 4), (4, 4))
+        ref = (a.double() - alpha * b.double()).to(torch.float32)
+        torch.testing.assert_close(SubAlphaModule(alpha)(a, b), ref)
+
+
+def export_sub_model(pte_path: str, golden_path: str, input_path: str) -> None:
+    """Write sub(a, b) .pte + fp64-computed torch golden + raw LE fp32 inputs (in1, in2)."""
+    a, b = _det_inputs((1024, 1024), (1024, 1024))
+    golden = (a.double() - b.double()).to(torch.float32).numpy().astype("<f4")
+    et = _export(SubModule(), a, b)
+    with open(pte_path, "wb") as f:
+        f.write(et.buffer)
+    golden.tofile(golden_path)
+    with open(input_path, "wb") as f:
+        a.numpy().astype("<f4").tofile(f)
+        b.numpy().astype("<f4").tofile(f)
+    print(f"Exported {pte_path}; golden {golden_path} ({golden.size} floats)")
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/backends/webgpu/test/ops/test_where.py b/backends/webgpu/test/ops/test_where.py
new file mode 100644
index 00000000000..02ccaf2d19a
--- /dev/null
+++ b/backends/webgpu/test/ops/test_where.py
@@ -0,0 +1,86 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""`aten.where.self` export + fp64 golden for the WebGPU backend.
+
+`where(cond, a, b) -> cond ? a : b`, with cond a 1-byte bool and a/b fp32
+(broadcast across all three operands). The kernel reads cond byte-packed as
+`array<u32>` and relinearizes each out coord onto every operand. Configs cover
+the equal-shape path plus broadcasts that exercise the size-1 clamp on cond, a,
+and b. The native binary has no ATen, so the golden is computed with torch here
+and checked in etvk CI.
+"""
+
+import unittest
+
+import torch
+
+from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner
+from executorch.exir import to_edge_transform_and_lower
+
+# name -> (cond_shape, a_shape, b_shape). Output is the broadcast of all three.
+CONFIGS = {
+    "equal": ((4, 8), (4, 8), (4, 8)),
+    "broadcast": ((4, 1), (4, 8), (1, 8)),
+    "cond_row": ((8,), (4, 8), (4, 8)),
+}
+
+
+class WhereModule(torch.nn.Module):
+    def forward(
+        self, cond: torch.Tensor, a: torch.Tensor, b: torch.Tensor
+    ) -> torch.Tensor:
+        return torch.where(cond, a, b)
+
+
+def _det_inputs(cond_shape, a_shape, b_shape):
+    """Deterministic (bool cond, fp32 a, fp32 b) for a config."""
+    g = torch.Generator().manual_seed(0)
+    cond = torch.rand(cond_shape, generator=g) > 0.5
+    a = torch.randn(*a_shape, generator=g, dtype=torch.float32)
+    b = torch.randn(*b_shape, generator=g, dtype=torch.float32)
+    return cond, a, b
+
+
+def _fp64_golden(cond, a, b):
+    return torch.where(cond, a.double(), b.double()).to(torch.float32)
+
+
+def _export(cond, a, b):
+    ep = torch.export.export(WhereModule().eval(), (cond, a, b))
+    return to_edge_transform_and_lower(
+        ep, partitioner=[VulkanPartitioner()]
+    ).to_executorch()
+
+
+def _delegates(et) -> bool:
+    return any(
+        d.id == "VulkanBackend"
+        for plan in et.executorch_program.execution_plan
+        for d in plan.delegates
+    )
+
+
+class TestWhere(unittest.TestCase):
+    def test_export_delegates(self) -> None:
+        for name, (cs, as_, bs) in CONFIGS.items():
+            with self.subTest(config=name):
+                cond, a, b = _det_inputs(cs, as_, bs)
+                et = _export(cond, a, b)
+                self.assertTrue(
+                    _delegates(et), f"Expected a VulkanBackend delegate (where {name})"
+                )
+
+    def test_op_matches_fp64_golden(self) -> None:
+        for name, (cs, as_, bs) in CONFIGS.items():
+            with self.subTest(config=name):
+                cond, a, b = _det_inputs(cs, as_, bs)
+                out = WhereModule()(cond, a, b)
+                torch.testing.assert_close(out, _fp64_golden(cond, a, b))
+
+
+if __name__ == "__main__":
+    unittest.main()