Update

[ghstack-poisoned]

Author: JCNTH

backends/vulkan/custom_ops_lib.py

@@ -1154,25 +1154,25 @@ def rms_norm_impl(
 
 
 # STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
-def linear_q4gsw_dw_impl(
+def linear_q4gsw_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_q4gsw_dw_meta(
+def linear_q4gsw_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_q4gsw_dw"
+name = "linear_q4gsw_dW"
 lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
-lib.impl(name, linear_q4gsw_dw_impl, "CompositeExplicitAutograd")
-lib.impl(name, linear_q4gsw_dw_meta, "Meta")
-linear_q4gsw_dw_op = getattr(getattr(torch.ops, namespace), name)
+lib.impl(name, linear_q4gsw_dW_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_q4gsw_dW_meta, "Meta")
+linear_q4gsw_dW_op = getattr(getattr(torch.ops, namespace), name)
 
 
 ##################

backends/vulkan/op_registry.py

@@ -1784,8 +1784,8 @@ def register_logical_not():
     )
 
 
-@update_features(exir_ops.edge.et_vk.linear_q4gsw_dw.default)
-def register_linear_q4gsw_dw():
+@update_features(exir_ops.edge.et_vk.linear_q4gsw_dW.default)
+def register_linear_q4gsw_dW():
     return OpFeatures(
         inputs_storage=utils.CONTIGUOUS_ANY,
         inputs_dtypes=utils.FP_T,

backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp

@@ -9,7 +9,7 @@
 #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_dw_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dW_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
@@ -30,7 +30,7 @@ struct Q4gswDwParams {
 static_assert(sizeof(Q4gswDwParams) == 16, "params must be 16 bytes");
 
 // STE weight gradient d_W[N,K] = d_out^T @ x.
-void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
+void q4gsw_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);
@@ -41,12 +41,12 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const auto& dw = graph.get_tensor(dw_id);
 
   if (dw.dims.size() != 2 || dout.dims.empty() || x.dims.empty()) {
-    throw std::runtime_error("q4gsw_dw: bad tensor ranks");
+    throw std::runtime_error("q4gsw_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("q4gsw_dw: N or K == 0");
+    throw std::runtime_error("q4gsw_dW: N or K == 0");
   }
 
   uint64_t dout_numel = 1;
@@ -59,18 +59,18 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   }
   if (static_cast<uint32_t>(dout.dims.back()) != N ||
       static_cast<uint32_t>(x.dims.back()) != K) {
-    throw std::runtime_error("q4gsw_dw: d_out/x last dim mismatch");
+    throw std::runtime_error("q4gsw_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("q4gsw_dw: M mismatch across d_out/x");
+    throw std::runtime_error("q4gsw_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("q4gsw_dw: fp32-only (byte-size mismatch)");
+    throw std::runtime_error("q4gsw_dW: fp32-only (byte-size mismatch)");
   }
 
   Q4gswDwParams params = {};
@@ -83,10 +83,10 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   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("q4gsw_dw: tile count exceeds u32");
+    throw std::runtime_error("q4gsw_dW: tile count exceeds u32");
   }
   const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
-      device, static_cast<uint32_t>(tiles), wg_size, "q4gsw_dw");
+      device, static_cast<uint32_t>(tiles), wg_size, "q4gsw_dW");
 
   WGPUBuffer uniform_buffer =
       utils::make_uniform(device, &params, sizeof(params));
@@ -156,7 +156,7 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   bg_desc.entries = bg_entries;
   WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
 
-  graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_dw"});
+  graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_dW"});
 
   wgpuShaderModuleRelease(shader);
   wgpuBindGroupLayoutRelease(bgl);
@@ -167,7 +167,7 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
 } // namespace
 
 WEBGPU_REGISTER_OPERATORS {
-  WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_dw.default, q4gsw_dw_impl);
+  WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_dW.default, q4gsw_dW_impl);
 }
 
 } // namespace executorch::backends::webgpu

…ntime/ops/quantized_linear/q4gsw_dw.wgsl …ntime/ops/quantized_linear/q4gsw_dW.wgslbackends/webgpu/runtime/ops/quantized_linear/q4gsw_dw.wgsl renamed to backends/webgpu/runtime/ops/quantized_linear/q4gsw_dW.wgsl backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw.wgsl renamed to backends/webgpu/runtime/ops/quantized_linear/q4gsw_dW.wgsl

@@ -12,7 +12,7 @@
 
 namespace executorch::backends::webgpu {
 
-// @generated from q4gsw_dw.wgsl - DO NOT EDIT.
+// @generated from q4gsw_dW.wgsl - DO NOT EDIT.
 // wgsl-sha256: be50075764ccd87e0d0c86124f13fd753aa618f994c9a5e1b8b631516aa51f84
 inline constexpr const char* kQ4gswDwWGSL = R"(
 // STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32).