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_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_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_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_dW_impl, "CompositeExplicitAutograd")
+lib.impl(name, linear_dW_meta, "Meta")
+linear_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_dW.default)
+def register_linear_dW():
     return OpFeatures(
         inputs_storage=utils.CONTIGUOUS_ANY,
         inputs_dtypes=utils.FP_T,

backends/webgpu/CMakeLists.txt

@@ -68,7 +68,7 @@ set(WEBGPU_SRCS
     runtime/ops/dim_order/DimOrder.cpp
     runtime/ops/linear/Linear.cpp
     runtime/ops/embedding/Embedding.cpp
-    runtime/ops/quantized_linear/QuantizedLinearDw.cpp
+    runtime/ops/quantized_linear/LinearDw.cpp
     runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
 )
 

…s/quantized_linear/QuantizedLinearDw.cpp …untime/ops/quantized_linear/LinearDw.cppbackends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp renamed to backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp renamed to backends/webgpu/runtime/ops/quantized_linear/LinearDw.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/linear_dW_wgsl.h>
 
 #include <webgpu/webgpu.h>
 
@@ -21,16 +21,16 @@ namespace executorch::backends::webgpu {
 
 namespace {
 
-struct Q4gswDwParams {
+struct LinearDwParams {
   uint32_t M;
   uint32_t N;
   uint32_t K;
   uint32_t _pad;
 };
-static_assert(sizeof(Q4gswDwParams) == 16, "params must be 16 bytes");
+static_assert(sizeof(LinearDwParams) == 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 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);
@@ -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("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("q4gsw_dW: N or K == 0");
+    throw std::runtime_error("linear_dW: N or K == 0");
   }
 
   uint64_t dout_numel = 1;
@@ -59,42 +59,42 @@ 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("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("q4gsw_dW: M mismatch across d_out/x");
+    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("q4gsw_dW: fp32-only (byte-size mismatch)");
+    throw std::runtime_error("linear_dW: fp32-only (byte-size mismatch)");
   }
 
-  Q4gswDwParams params = {};
+  LinearDwParams params = {};
   params.M = M;
   params.N = N;
   params.K = K;
 
   const uint32_t wg_size =
-      utils::clamp_workgroup_size(device, kQ4gswDwWorkgroupSizeX);
+      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("q4gsw_dW: tile count exceeds u32");
+    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, "q4gsw_dW");
+      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 = {kQ4gswDwWGSL, WGPU_STRLEN};
+  wgsl_desc.code = {kLinearDwWGSL, WGPU_STRLEN};
   WGPUShaderModuleDescriptor shader_desc = {};
   shader_desc.nextInChain = &wgsl_desc.chain;
   WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
@@ -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, "linear_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_dW.default, linear_dW_impl);
 }
 
 } // namespace executorch::backends::webgpu

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

@@ -12,9 +12,9 @@
 
 namespace executorch::backends::webgpu {
 
-// @generated from q4gsw_dW.wgsl - DO NOT EDIT.
+// @generated from linear_dW.wgsl - DO NOT EDIT.
 // wgsl-sha256: be50075764ccd87e0d0c86124f13fd753aa618f994c9a5e1b8b631516aa51f84
-inline constexpr const char* kQ4gswDwWGSL = R"(
+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]
@@ -86,8 +86,8 @@ fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
 }
 )";
 
-inline constexpr uint32_t kQ4gswDwWorkgroupSizeX = 64;
-inline constexpr uint32_t kQ4gswDwWorkgroupSizeY = 1;
-inline constexpr uint32_t kQ4gswDwWorkgroupSizeZ = 1;
+inline constexpr uint32_t kLinearDwWorkgroupSizeX = 64;
+inline constexpr uint32_t kLinearDwWorkgroupSizeY = 1;
+inline constexpr uint32_t kLinearDwWorkgroupSizeZ = 1;
 
 } // namespace executorch::backends::webgpu