[ExecuTorch][WebGPU] Add gather and embedding ops to the WebGPU backend

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)

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,20 @@ 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,
@@ -1775,6 +1798,31 @@ def register_logical_not():
     )
 
 
+@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 ##
 #######################

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);
+}

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