Skip to content

Commit 5b001f4

Browse files
pytorchbotJCNTH
pytorchbot
and
JCNTH
authored
[ExecuTorch][WebGPU] q4gsw decode: 2-column "bicol" GEMV (conservative -15..-35% decode latency) (#20665)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #20597 by @JulianCloudNTH ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/JulianCloudNTH/77/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JulianCloudNTH/77/head Merge bot PR base: https://github.com/pytorch/executorch/tree/main Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JulianCloudNTH/77/orig Differential Revision: [D110017692](https://our.internmc.facebook.com/intern/diff/D110017692/) @diff-train-skip-merge --------- Co-authored-by: Julian Ng-Thow-Hing <juliannth@meta.com>
1 parent 7b6d3bf commit 5b001f4

5 files changed

Lines changed: 256 additions & 10 deletions

File tree

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

Lines changed: 11 additions & 10 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_linear_coop4_wgsl.h>
12+
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h>
1313
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
1414

1515
#include <webgpu/webgpu.h>
@@ -117,21 +117,22 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
117117
"WebGPU linear_q4gsw: scales dims too small for K/N");
118118
}
119119

120-
// M==1 decode -> coop4 GEMV (needs K%8==0 && gs%8==0); else tiled GEMM.
120+
// M==1 decode -> bicol GEMV (needs K%8==0 && gs%8==0); else tiled GEMM.
121121
const uint32_t wg_size =
122122
utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX);
123123
const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u);
124-
const char* shader_src = use_gemv ? kQ4gswLinearCoop4WGSL : kQ4gswLinearWGSL;
124+
const char* shader_src =
125+
use_gemv ? kQ4gswLinearCoop4BicolWGSL : kQ4gswLinearWGSL;
125126
uint32_t workgroup_count;
126127
if (use_gemv) {
127-
// coop4: fixed 64 lanes, 1 workgroup per output, grid-strided over M*N.
128-
const uint64_t outputs =
129-
static_cast<uint64_t>(M) * static_cast<uint64_t>(N);
130-
if (outputs == 0u || outputs > UINT32_MAX) {
131-
throw std::runtime_error("WebGPU linear_q4gsw: M*N out of range");
128+
// bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over
129+
// ceil(N/2) column-pairs (M == 1 on this decode path).
130+
const uint64_t pairs = (static_cast<uint64_t>(N) + 1u) / 2u;
131+
if (pairs == 0u || pairs > UINT32_MAX) {
132+
throw std::runtime_error("WebGPU linear_q4gsw: N/2 out of range");
132133
}
133134
workgroup_count =
134-
utils::clamp_workgroup_count(device, static_cast<uint32_t>(outputs));
135+
utils::clamp_workgroup_count(device, static_cast<uint32_t>(pairs));
135136
if (workgroup_count == 0u) {
136137
throw std::runtime_error("WebGPU linear_q4gsw: zero GEMV dispatch");
137138
}
@@ -224,7 +225,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
224225
pipeline_desc.layout = pipeline_layout;
225226
pipeline_desc.compute.module = shader;
226227
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
227-
// coop4 GEMV uses fixed @workgroup_size(64); only the GEMM has an override.
228+
// bicol GEMV uses fixed @workgroup_size(64); only the GEMM has an override.
228229
pipeline_desc.compute.constantCount = use_gemv ? 0u : 1u;
229230
pipeline_desc.compute.constants = use_gemv ? nullptr : &wg_size_constant;
230231
WGPUComputePipeline pipeline =

‎backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl‎

Lines changed: 105 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,105 @@
1+
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
2+
@group(0) @binding(1) var<storage, read> t_input: array<f32>;
3+
@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
4+
@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
5+
@group(0) @binding(4) var<storage, read> t_bias: array<f32>;
6+
7+
struct Params {
8+
M: u32,
9+
N: u32,
10+
K: u32,
11+
K_packed: u32,
12+
group_size: u32,
13+
padded_N: u32,
14+
has_bias: u32,
15+
_pad: u32,
16+
}
17+
@group(0) @binding(5) var<uniform> params: Params;
18+
19+
// Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
20+
const WG: u32 = 64u;
21+
var<workgroup> partial: array<vec2<f32>, WG>;
22+
23+
@compute @workgroup_size(WG, 1, 1)
24+
fn main(
25+
@builtin(workgroup_id) wid: vec3<u32>,
26+
@builtin(num_workgroups) ngrp: vec3<u32>,
27+
@builtin(local_invocation_id) lid: vec3<u32>) {
28+
let num_pairs = (params.N + 1u) >> 1u; // ceil(N/2) pairs (M == 1 decode)
29+
let stride = ngrp.x;
30+
let num_words = params.K >> 3u; // K / 8 words per row
31+
let row_words = params.K_packed >> 2u; // u32s per weight row (= K/8)
32+
var p = wid.x;
33+
loop {
34+
if (p >= num_pairs) {
35+
break;
36+
}
37+
let col0 = p << 1u;
38+
let col1 = col0 + 1u;
39+
let has1 = col1 < params.N;
40+
let wbase0 = col0 * row_words;
41+
let wbase1 = col1 * row_words; // load guarded by has1
42+
43+
var acc0: f32 = 0.0;
44+
var acc1: f32 = 0.0;
45+
var w: u32 = lid.x;
46+
loop {
47+
if (w >= num_words) {
48+
break;
49+
}
50+
let k0 = w << 3u;
51+
let ib = k0; // in_base = 0 (m == 0 for M == 1 decode)
52+
let sidx = (k0 / params.group_size) * params.padded_N;
53+
let word0 = t_weight[wbase0 + w];
54+
let scale0 = t_scales[sidx + col0];
55+
var word1: u32 = 0u;
56+
var scale1: f32 = 0.0;
57+
if (has1) {
58+
word1 = t_weight[wbase1 + w];
59+
scale1 = t_scales[sidx + col1];
60+
}
61+
for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) {
62+
let kk = bi << 1u;
63+
let in0 = t_input[ib + kk]; // shared across both columns
64+
let in1 = t_input[ib + kk + 1u];
65+
let b0 = (word0 >> (bi * 8u)) & 0xFFu;
66+
acc0 = acc0 + in0 * f32(i32(b0 & 0x0Fu) - 8) * scale0;
67+
acc0 = acc0 + in1 * f32(i32((b0 >> 4u) & 0x0Fu) - 8) * scale0;
68+
let b1 = (word1 >> (bi * 8u)) & 0xFFu;
69+
acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
70+
acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
71+
}
72+
w = w + WG;
73+
}
74+
75+
partial[lid.x] = vec2<f32>(acc0, acc1);
76+
workgroupBarrier();
77+
var s: u32 = WG >> 1u;
78+
loop {
79+
if (s == 0u) {
80+
break;
81+
}
82+
if (lid.x < s) {
83+
partial[lid.x] = partial[lid.x] + partial[lid.x + s];
84+
}
85+
workgroupBarrier();
86+
s = s >> 1u;
87+
}
88+
if (lid.x == 0u) {
89+
var o0 = partial[0].x;
90+
var o1 = partial[0].y;
91+
if (params.has_bias != 0u) {
92+
o0 = o0 + t_bias[col0];
93+
if (has1) {
94+
o1 = o1 + t_bias[col1];
95+
}
96+
}
97+
t_out[col0] = o0;
98+
if (has1) {
99+
t_out[col1] = o1;
100+
}
101+
}
102+
workgroupBarrier();
103+
p = p + stride;
104+
}
105+
}

‎backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,129 @@
1+
/*
2+
* Copyright (c) Meta Platforms, Inc. and affiliates.
3+
* All rights reserved.
4+
*
5+
* This source code is licensed under the BSD-style license found in the
6+
* LICENSE file in the root directory of this source tree.
7+
*/
8+
9+
#pragma once
10+
11+
#include <cstdint>
12+
13+
namespace executorch::backends::webgpu {
14+
15+
// @generated from q4gsw_linear_coop4_bicol.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 44c53ca7879f7a3382a45703a67cbc6fc221ecbcada58f8c413df50abbc6e544
17+
inline constexpr const char* kQ4gswLinearCoop4BicolWGSL = R"(
18+
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
19+
@group(0) @binding(1) var<storage, read> t_input: array<f32>;
20+
@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
21+
@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
22+
@group(0) @binding(4) var<storage, read> t_bias: array<f32>;
23+
24+
struct Params {
25+
M: u32,
26+
N: u32,
27+
K: u32,
28+
K_packed: u32,
29+
group_size: u32,
30+
padded_N: u32,
31+
has_bias: u32,
32+
_pad: u32,
33+
}
34+
@group(0) @binding(5) var<uniform> params: Params;
35+
36+
// Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
37+
const WG: u32 = 64u;
38+
var<workgroup> partial: array<vec2<f32>, WG>;
39+
40+
@compute @workgroup_size(WG, 1, 1)
41+
fn main(
42+
@builtin(workgroup_id) wid: vec3<u32>,
43+
@builtin(num_workgroups) ngrp: vec3<u32>,
44+
@builtin(local_invocation_id) lid: vec3<u32>) {
45+
let num_pairs = (params.N + 1u) >> 1u; // ceil(N/2) pairs (M == 1 decode)
46+
let stride = ngrp.x;
47+
let num_words = params.K >> 3u; // K / 8 words per row
48+
let row_words = params.K_packed >> 2u; // u32s per weight row (= K/8)
49+
var p = wid.x;
50+
loop {
51+
if (p >= num_pairs) {
52+
break;
53+
}
54+
let col0 = p << 1u;
55+
let col1 = col0 + 1u;
56+
let has1 = col1 < params.N;
57+
let wbase0 = col0 * row_words;
58+
let wbase1 = col1 * row_words; // load guarded by has1
59+
60+
var acc0: f32 = 0.0;
61+
var acc1: f32 = 0.0;
62+
var w: u32 = lid.x;
63+
loop {
64+
if (w >= num_words) {
65+
break;
66+
}
67+
let k0 = w << 3u;
68+
let ib = k0; // in_base = 0 (m == 0 for M == 1 decode)
69+
let sidx = (k0 / params.group_size) * params.padded_N;
70+
let word0 = t_weight[wbase0 + w];
71+
let scale0 = t_scales[sidx + col0];
72+
var word1: u32 = 0u;
73+
var scale1: f32 = 0.0;
74+
if (has1) {
75+
word1 = t_weight[wbase1 + w];
76+
scale1 = t_scales[sidx + col1];
77+
}
78+
for (var bi: u32 = 0u; bi < 4u; bi = bi + 1u) {
79+
let kk = bi << 1u;
80+
let in0 = t_input[ib + kk]; // shared across both columns
81+
let in1 = t_input[ib + kk + 1u];
82+
let b0 = (word0 >> (bi * 8u)) & 0xFFu;
83+
acc0 = acc0 + in0 * f32(i32(b0 & 0x0Fu) - 8) * scale0;
84+
acc0 = acc0 + in1 * f32(i32((b0 >> 4u) & 0x0Fu) - 8) * scale0;
85+
let b1 = (word1 >> (bi * 8u)) & 0xFFu;
86+
acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
87+
acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
88+
}
89+
w = w + WG;
90+
}
91+
92+
partial[lid.x] = vec2<f32>(acc0, acc1);
93+
workgroupBarrier();
94+
var s: u32 = WG >> 1u;
95+
loop {
96+
if (s == 0u) {
97+
break;
98+
}
99+
if (lid.x < s) {
100+
partial[lid.x] = partial[lid.x] + partial[lid.x + s];
101+
}
102+
workgroupBarrier();
103+
s = s >> 1u;
104+
}
105+
if (lid.x == 0u) {
106+
var o0 = partial[0].x;
107+
var o1 = partial[0].y;
108+
if (params.has_bias != 0u) {
109+
o0 = o0 + t_bias[col0];
110+
if (has1) {
111+
o1 = o1 + t_bias[col1];
112+
}
113+
}
114+
t_out[col0] = o0;
115+
if (has1) {
116+
t_out[col1] = o1;
117+
}
118+
}
119+
workgroupBarrier();
120+
p = p + stride;
121+
}
122+
}
123+
)";
124+
125+
inline constexpr uint32_t kQ4gswLinearCoop4BicolWorkgroupSizeX = 64;
126+
inline constexpr uint32_t kQ4gswLinearCoop4BicolWorkgroupSizeY = 1;
127+
inline constexpr uint32_t kQ4gswLinearCoop4BicolWorkgroupSizeZ = 1;
128+
129+
} // namespace executorch::backends::webgpu

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,12 @@ class Q4gswConfig:
5353
Q4gswConfig("kv_proj_4k", 4096, 2048, 512),
5454
Q4gswConfig("q_proj_8k", 8192, 2048, 2048, heavy=True), # 67MB golden
5555
Q4gswConfig("kv_proj_8k", 8192, 2048, 512, heavy=True),
56+
# The M==1 decode configs above (q/kv/gate/down_proj) exercise the bicol 2-col
57+
# decode GEMV: the handler routes M==1 -> bicol, so each reads its own per-
58+
# column scale (col0/col1) across many K-groups (down_proj: 256 groups). q4gsw
59+
# requires N % 8 == 0 (torchao pads N for the scale layout), so odd-N / N=1 are
60+
# not exportable -- bicol's has1 odd-N guard is defensive (mirrors coop4's
61+
# general-N robustness) and unreachable through this op.
5662
]
5763

5864

‎backends/webgpu/test/test_webgpu_native.cpp‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -288,6 +288,11 @@ static const Q4gswConfig kQ4gswConfigs[] = {
288288
{"kv_proj_4k", 4096, 2048, 512, 1e-4f, 1e-3f, true, false},
289289
{"q_proj_8k", 8192, 2048, 2048, 1e-4f, 1e-3f, false, true},
290290
{"kv_proj_8k", 8192, 2048, 512, 1e-4f, 1e-3f, false, true},
291+
// The M==1 configs above (q/kv/gate/down_proj) exercise the bicol 2-col
292+
// decode GEMV (handler routes M==1 -> bicol; each reads its own per-column
293+
// scale over 64-256 K-groups). q4gsw requires N % 8 == 0, so odd-N is not
294+
// exportable; bicol's has1 odd-N guard is defensive (mirrors coop4
295+
// general-N robustness).
291296
};
292297

293298
// /16 ramp over the flat index; mirrors test_quantized_linear.py _ramp_input.

0 commit comments

Comments
 (0)