Skip to content

Commit 6d91872

Browse files
JCNTHJCNTH
committed
Update (base update)
[ghstack-poisoned]
1 parent 851bcef commit 6d91872

7 files changed

Lines changed: 37 additions & 37 deletions

File tree

‎backends/vulkan/custom_ops_lib.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1256,25 +1256,25 @@ def adamw_step_meta(
12561256

12571257

12581258
# STE weight gradient d_out^T @ x through the frozen 4-bit linear_q4gsw base.
1259-
def linear_q4gsw_dW_impl(
1259+
def linear_dW_impl(
12601260
d_out: torch.Tensor,
12611261
x: torch.Tensor,
12621262
) -> torch.Tensor:
12631263
return d_out.reshape(-1, d_out.shape[-1]).t() @ x.reshape(-1, x.shape[-1])
12641264

12651265

1266-
def linear_q4gsw_dW_meta(
1266+
def linear_dW_meta(
12671267
d_out: torch.Tensor,
12681268
x: torch.Tensor,
12691269
) -> torch.Tensor:
12701270
return d_out.new_empty((d_out.shape[-1], x.shape[-1]))
12711271

12721272

1273-
name = "linear_q4gsw_dW"
1273+
name = "linear_dW"
12741274
lib.define(f"{name}(Tensor d_out, Tensor x) -> Tensor")
1275-
lib.impl(name, linear_q4gsw_dW_impl, "CompositeExplicitAutograd")
1276-
lib.impl(name, linear_q4gsw_dW_meta, "Meta")
1277-
linear_q4gsw_dW_op = getattr(getattr(torch.ops, namespace), name)
1275+
lib.impl(name, linear_dW_impl, "CompositeExplicitAutograd")
1276+
lib.impl(name, linear_dW_meta, "Meta")
1277+
linear_dW_op = getattr(getattr(torch.ops, namespace), name)
12781278

12791279

12801280
##################

‎backends/vulkan/op_registry.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1806,8 +1806,8 @@ def register_adamw_step():
18061806
)
18071807

18081808

1809-
@update_features(exir_ops.edge.et_vk.linear_q4gsw_dW.default)
1810-
def register_linear_q4gsw_dW():
1809+
@update_features(exir_ops.edge.et_vk.linear_dW.default)
1810+
def register_linear_dW():
18111811
return OpFeatures(
18121812
inputs_storage=utils.CONTIGUOUS_ANY,
18131813
inputs_dtypes=utils.FP_T,

‎backends/webgpu/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ set(WEBGPU_SRCS
7070
runtime/ops/linear/Linear.cpp
7171
runtime/ops/embedding/Embedding.cpp
7272
runtime/ops/adamw/AdamwStep.cpp
73-
runtime/ops/quantized_linear/QuantizedLinearDw.cpp
73+
runtime/ops/quantized_linear/LinearDw.cpp
7474
runtime/ops/quantized_linear/QuantizedLinearRequant.cpp
7575
)
7676

backends/webgpu/runtime/ops/quantized_linear/QuantizedLinearDw.cpp renamed to backends/webgpu/runtime/ops/quantized_linear/LinearDw.cpp

Lines changed: 16 additions & 16 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/linear_dW_wgsl.h>
1313

1414
#include <webgpu/webgpu.h>
1515

@@ -21,16 +21,16 @@ namespace executorch::backends::webgpu {
2121

2222
namespace {
2323

24-
struct Q4gswDwParams {
24+
struct LinearDwParams {
2525
uint32_t M;
2626
uint32_t N;
2727
uint32_t K;
2828
uint32_t _pad;
2929
};
30-
static_assert(sizeof(Q4gswDwParams) == 16, "params must be 16 bytes");
30+
static_assert(sizeof(LinearDwParams) == 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 linear_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("linear_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("linear_dW: N or K == 0");
5050
}
5151

5252
uint64_t dout_numel = 1;
@@ -59,42 +59,42 @@ 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("linear_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("linear_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("linear_dW: fp32-only (byte-size mismatch)");
7474
}
7575

76-
Q4gswDwParams params = {};
76+
LinearDwParams params = {};
7777
params.M = M;
7878
params.N = N;
7979
params.K = K;
8080

8181
const uint32_t wg_size =
82-
utils::clamp_workgroup_size(device, kQ4gswDwWorkgroupSizeX);
82+
utils::clamp_workgroup_size(device, kLinearDwWorkgroupSizeX);
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("linear_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, "linear_dW");
9090

9191
WGPUBuffer uniform_buffer =
9292
utils::make_uniform(device, &params, sizeof(params));
9393
graph.add_uniform_buffer_bytes(sizeof(params));
9494

9595
WGPUShaderSourceWGSL wgsl_desc = {};
9696
wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
97-
wgsl_desc.code = {kQ4gswDwWGSL, WGPU_STRLEN};
97+
wgsl_desc.code = {kLinearDwWGSL, WGPU_STRLEN};
9898
WGPUShaderModuleDescriptor shader_desc = {};
9999
shader_desc.nextInChain = &wgsl_desc.chain;
100100
WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
@@ -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, "linear_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_dW.default, linear_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/linear_dW.wgsl

File renamed without changes.

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,9 +12,9 @@
1212

1313
namespace executorch::backends::webgpu {
1414

15-
// @generated from q4gsw_dW.wgsl - DO NOT EDIT.
15+
// @generated from linear_dW.wgsl - DO NOT EDIT.
1616
// wgsl-sha256: be50075764ccd87e0d0c86124f13fd753aa618f994c9a5e1b8b631516aa51f84
17-
inline constexpr const char* kQ4gswDwWGSL = R"(
17+
inline constexpr const char* kLinearDwWGSL = R"(
1818
// STE weight-gradient d_W[N,K] = sum_m d_out[m,N]*x[m,K] (operands f32).
1919
2020
@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>) {
8686
}
8787
)";
8888

89-
inline constexpr uint32_t kQ4gswDwWorkgroupSizeX = 64;
90-
inline constexpr uint32_t kQ4gswDwWorkgroupSizeY = 1;
91-
inline constexpr uint32_t kQ4gswDwWorkgroupSizeZ = 1;
89+
inline constexpr uint32_t kLinearDwWorkgroupSizeX = 64;
90+
inline constexpr uint32_t kLinearDwWorkgroupSizeY = 1;
91+
inline constexpr uint32_t kLinearDwWorkgroupSizeZ = 1;
9292

9393
} // namespace executorch::backends::webgpu

backends/webgpu/test/ops/test_linear_q4gsw_dw.py renamed to backends/webgpu/test/ops/test_linear_dW.py

Lines changed: 7 additions & 7 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_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
@@ -29,9 +29,9 @@
2929
}
3030

3131

32-
class Q4gswDwModule(torch.nn.Module):
32+
class LinearDwModule(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_dW(d_out, x)
3535

3636

3737
def _det_inputs(m: int, k: int, n: int):
@@ -48,7 +48,7 @@ def _fp64_golden(d_out: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
4848

4949

5050
def _export(d_out: torch.Tensor, x: torch.Tensor):
51-
ep = torch.export.export(Q4gswDwModule().eval(), (d_out, x))
51+
ep = torch.export.export(LinearDwModule().eval(), (d_out, x))
5252
return to_edge_transform_and_lower(
5353
ep, partitioner=[VulkanPartitioner()]
5454
).to_executorch()
@@ -62,23 +62,23 @@ def _delegated(et) -> bool:
6262
)
6363

6464

65-
class TestLinearQ4gswDw(unittest.TestCase):
65+
class TestLinearDw(unittest.TestCase):
6666
def test_export_delegates(self) -> None:
6767
for name, (m, k, n) in CONFIGS.items():
6868
with self.subTest(config=name):
6969
d_out, x = _det_inputs(m, k, n)
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_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_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)