Skip to content

Commit b749d6c

Browse files
JCNTHJCNTH
committed
Update (base update)
[ghstack-poisoned]
1 parent dffd4fc commit b749d6c

6 files changed

Lines changed: 24 additions & 24 deletions

File tree

‎backends/vulkan/custom_ops_lib.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1154,25 +1154,25 @@ def rms_norm_impl(
11541154

11551155

11561156
# STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
1157-
def linear_q4gsw_dw_impl(
1157+
def linear_q4gsw_dW_impl(
11581158
d_out: torch.Tensor,
11591159
x: torch.Tensor,
11601160
) -> torch.Tensor:
11611161
return d_out.reshape(-1, d_out.shape[-1]).t() @ x.reshape(-1, x.shape[-1])
11621162

11631163

1164-
def linear_q4gsw_dw_meta(
1164+
def linear_q4gsw_dW_meta(
11651165
d_out: torch.Tensor,
11661166
x: torch.Tensor,
11671167
) -> torch.Tensor:
11681168
return d_out.new_empty((d_out.shape[-1], x.shape[-1]))
11691169

11701170

1171-
name = "linear_q4gsw_dw"
1171+
name = "linear_q4gsw_dW"
11721172
lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
1173-
lib.impl(name, linear_q4gsw_dw_impl, "CompositeExplicitAutograd")
1174-
lib.impl(name, linear_q4gsw_dw_meta, "Meta")
1175-
linear_q4gsw_dw_op = getattr(getattr(torch.ops, namespace), name)
1173+
lib.impl(name, linear_q4gsw_dW_impl, "CompositeExplicitAutograd")
1174+
lib.impl(name, linear_q4gsw_dW_meta, "Meta")
1175+
linear_q4gsw_dW_op = getattr(getattr(torch.ops, namespace), name)
11761176

11771177

11781178
##################

‎backends/vulkan/op_registry.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1784,8 +1784,8 @@ def register_logical_not():
17841784
)
17851785

17861786

1787-
@update_features(exir_ops.edge.et_vk.linear_q4gsw_dw.default)
1788-
def register_linear_q4gsw_dw():
1787+
@update_features(exir_ops.edge.et_vk.linear_q4gsw_dW.default)
1788+
def register_linear_q4gsw_dW():
17891789
return OpFeatures(
17901790
inputs_storage=utils.CONTIGUOUS_ANY,
17911791
inputs_dtypes=utils.FP_T,

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

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
#include <executorch/backends/webgpu/runtime/WebGPUGraph.h>
1010
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
1111
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
12-
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw_wgsl.h>
12+
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_dW_wgsl.h>
1313

1414
#include <webgpu/webgpu.h>
1515

@@ -30,7 +30,7 @@ struct Q4gswDwParams {
3030
static_assert(sizeof(Q4gswDwParams) == 16, "params must be 16 bytes");
3131

3232
// STE weight gradient d_W[N,K] = d_out^T @ x.
33-
void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
33+
void q4gsw_dW_impl(WebGPUGraph& graph, const std::vector<int>& args) {
3434
const int dout_id = args.at(0);
3535
const int x_id = args.at(1);
3636
const int dw_id = args.at(2);
@@ -41,12 +41,12 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
4141
const auto& dw = graph.get_tensor(dw_id);
4242

4343
if (dw.dims.size() != 2 || dout.dims.empty() || x.dims.empty()) {
44-
throw std::runtime_error("q4gsw_dw: bad tensor ranks");
44+
throw std::runtime_error("q4gsw_dW: bad tensor ranks");
4545
}
4646
const uint32_t N = static_cast<uint32_t>(dw.dims[0]);
4747
const uint32_t K = static_cast<uint32_t>(dw.dims[1]);
4848
if (N == 0 || K == 0) {
49-
throw std::runtime_error("q4gsw_dw: N or K == 0");
49+
throw std::runtime_error("q4gsw_dW: N or K == 0");
5050
}
5151

5252
uint64_t dout_numel = 1;
@@ -59,18 +59,18 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
5959
}
6060
if (static_cast<uint32_t>(dout.dims.back()) != N ||
6161
static_cast<uint32_t>(x.dims.back()) != K) {
62-
throw std::runtime_error("q4gsw_dw: d_out/x last dim mismatch");
62+
throw std::runtime_error("q4gsw_dW: d_out/x last dim mismatch");
6363
}
6464
const uint32_t M = static_cast<uint32_t>(dout_numel / N);
6565
if (dout_numel % N != 0 || x_numel % K != 0 || x_numel / K != M) {
66-
throw std::runtime_error("q4gsw_dw: M mismatch across d_out/x");
66+
throw std::runtime_error("q4gsw_dW: M mismatch across d_out/x");
6767
}
6868

6969
// fp32-only byte-size guards (mirror the forward's byte checks).
7070
if (dw.nbytes != static_cast<uint64_t>(N) * K * sizeof(float) ||
7171
dout.nbytes != dout_numel * sizeof(float) ||
7272
x.nbytes != x_numel * sizeof(float)) {
73-
throw std::runtime_error("q4gsw_dw: fp32-only (byte-size mismatch)");
73+
throw std::runtime_error("q4gsw_dW: fp32-only (byte-size mismatch)");
7474
}
7575

7676
Q4gswDwParams params = {};
@@ -83,10 +83,10 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
8383
const uint64_t tiles =
8484
utils::div_up<uint64_t>(N, 4u) * utils::div_up<uint64_t>(K, 4u);
8585
if (tiles > UINT32_MAX) {
86-
throw std::runtime_error("q4gsw_dw: tile count exceeds u32");
86+
throw std::runtime_error("q4gsw_dW: tile count exceeds u32");
8787
}
8888
const uint32_t workgroup_count = utils::compute_1d_workgroup_count(
89-
device, static_cast<uint32_t>(tiles), wg_size, "q4gsw_dw");
89+
device, static_cast<uint32_t>(tiles), wg_size, "q4gsw_dW");
9090

9191
WGPUBuffer uniform_buffer =
9292
utils::make_uniform(device, &params, sizeof(params));
@@ -156,7 +156,7 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
156156
bg_desc.entries = bg_entries;
157157
WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
158158

159-
graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_dw"});
159+
graph.add_dispatch({pipeline, bind_group, workgroup_count, "q4gsw_dW"});
160160

161161
wgpuShaderModuleRelease(shader);
162162
wgpuBindGroupLayoutRelease(bgl);
@@ -167,7 +167,7 @@ void q4gsw_dw_impl(WebGPUGraph& graph, const std::vector<int>& args) {
167167
} // namespace
168168

169169
WEBGPU_REGISTER_OPERATORS {
170-
WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_dw.default, q4gsw_dw_impl);
170+
WEBGPU_REGISTER_OP(et_vk.linear_q4gsw_dW.default, q4gsw_dW_impl);
171171
}
172172

173173
} // namespace executorch::backends::webgpu

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

File renamed without changes.

backends/webgpu/runtime/ops/quantized_linear/q4gsw_dw_wgsl.h renamed to backends/webgpu/runtime/ops/quantized_linear/q4gsw_dW_wgsl.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,7 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from q4gsw_dw.wgsl - DO NOT EDIT.
15+
// @generated from q4gsw_dW.wgsl - DO NOT EDIT.
1616
// wgsl-sha256: be50075764ccd87e0d0c86124f13fd753aa618f994c9a5e1b8b631516aa51f84
1717
inline constexpr const char* kQ4gswDwWGSL = R"(
1818
// STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32).

‎backends/webgpu/test/ops/test_linear_q4gsw_dw.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7-
"""`et_vk.linear_q4gsw_dw` (STE weight gradient) export + fp64 golden.
7+
"""`et_vk.linear_q4gsw_dW` (STE weight gradient) export + fp64 golden.
88
99
The weight gradient of a frozen 4-bit linear: `d_W[N, K] = d_out^T @ x`, the grad
1010
wrt the dequantized weight (both operands are fp32; no int4 unpack). Reached by a
@@ -31,7 +31,7 @@
3131

3232
class Q4gswDwModule(torch.nn.Module):
3333
def forward(self, d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
34-
return torch.ops.et_vk.linear_q4gsw_dw(d_out, x)
34+
return torch.ops.et_vk.linear_q4gsw_dW(d_out, x)
3535

3636

3737
def _det_inputs(m: int, k: int, n: int):
@@ -70,15 +70,15 @@ def test_export_delegates(self) -> None:
7070
et = _export(d_out, x)
7171
self.assertTrue(
7272
_delegated(et),
73-
f"Expected a VulkanBackend delegate (linear_q4gsw_dw {name})",
73+
f"Expected a VulkanBackend delegate (linear_q4gsw_dW {name})",
7474
)
7575

7676
def test_op_matches_fp64_golden(self) -> None:
7777
# Op (d_out^T @ x) vs fp64 matmul truth: guards the formula.
7878
for name, (m, k, n) in CONFIGS.items():
7979
with self.subTest(config=name):
8080
d_out, x = _det_inputs(m, k, n)
81-
got = torch.ops.et_vk.linear_q4gsw_dw(d_out, x)
81+
got = torch.ops.et_vk.linear_q4gsw_dW(d_out, x)
8282
golden = _fp64_golden(d_out, x)
8383
self.assertEqual(tuple(got.shape), (n, k))
8484
torch.testing.assert_close(got, golden, atol=5e-4, rtol=1e-3)

0 commit comments

Comments
 (0)