Skip to content

Commit 55debb5

Browse files
pytorchbotJCNTH
pytorchbot
and
JCNTH
authored
[ExecuTorch][WebGPU] q4gsw prefill: shared-memory tiled GEMM, shape-routed (+150–303% large-K/N) (#20672)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #20605 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/79/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JulianCloudNTH/79/head Merge bot PR base: https://github.com/pytorch/executorch/tree/main Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JulianCloudNTH/79/orig @diff-train-skip-merge --------- Co-authored-by: Julian Ng-Thow-Hing <juliannth@meta.com>
1 parent 5909e23 commit 55debb5

5 files changed

Lines changed: 321 additions & 6 deletions

File tree

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

Lines changed: 36 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
#include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
1111
#include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
1212
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h>
13+
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h>
1314
#include <executorch/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_wgsl.h>
1415

1516
#include <webgpu/webgpu.h>
@@ -39,6 +40,16 @@ static_assert(sizeof(Q4gswParams) == 32, "Q4gswParams must be 32 bytes");
3940
constexpr int64_t kQ4gswTileM = 4;
4041
constexpr int64_t kQ4gswTileN = 4;
4142

43+
// Shmem-GEMM tile dims; MUST match WG_M/WG_N in q4gsw_linear_gemm_shmem.wgsl.
44+
constexpr int64_t kQ4gswShmemTileM = 32;
45+
constexpr int64_t kQ4gswShmemTileN = 32;
46+
// Prefill route: shmem GEMM wins large K/N; the square 2048^2 (q/o proj,
47+
// N=2048) also wins on shmem (+~50% Canary, 10-warm/50-run; the earlier -27%
48+
// did not reproduce), so route it via a lower N threshold while k/v (N=512)
49+
// stays on register-tiled.
50+
constexpr uint32_t kQ4gswShmemMinDim = 4096u;
51+
constexpr uint32_t kQ4gswShmemNMinDim = 2048u;
52+
4253
// et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out].
4354
void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
4455
const int in_id = args.at(0);
@@ -117,12 +128,15 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
117128
"WebGPU linear_q4gsw: scales dims too small for K/N");
118129
}
119130

120-
// M==1 decode -> bicol GEMV (needs K%8==0 && gs%8==0); else tiled GEMM.
131+
// M==1 -> bicol GEMV; M>1 -> shmem GEMM (large K/N) else tiled GEMM.
121132
const uint32_t wg_size =
122133
utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX);
123134
const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u);
124-
const char* shader_src =
125-
use_gemv ? kQ4gswLinearCoop4BicolWGSL : kQ4gswLinearWGSL;
135+
const bool use_shmem_gemm =
136+
!use_gemv && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim);
137+
const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL
138+
: use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL
139+
: kQ4gswLinearWGSL;
126140
uint32_t workgroup_count;
127141
if (use_gemv) {
128142
// bicol: fixed 64 lanes, 2 output columns/workgroup, grid-strided over
@@ -136,6 +150,21 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
136150
if (workgroup_count == 0u) {
137151
throw std::runtime_error("WebGPU linear_q4gsw: zero GEMV dispatch");
138152
}
153+
} else if (use_shmem_gemm) {
154+
// shmem GEMM: one workgroup per tile, no grid-stride -> throw over limit.
155+
const int64_t total_wgs = utils::div_up<int64_t>(M, kQ4gswShmemTileM) *
156+
utils::div_up<int64_t>(N, kQ4gswShmemTileN);
157+
WGPULimits limits = {};
158+
const uint32_t max_wgs =
159+
wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success &&
160+
limits.maxComputeWorkgroupsPerDimension > 0
161+
? limits.maxComputeWorkgroupsPerDimension
162+
: 65535u;
163+
if (total_wgs > static_cast<int64_t>(max_wgs)) {
164+
throw std::runtime_error(
165+
"WebGPU linear_q4gsw: shmem GEMM tile count exceeds the 1D dispatch limit");
166+
}
167+
workgroup_count = static_cast<uint32_t>(total_wgs);
139168
} else {
140169
const int64_t total_tiles = utils::div_up<int64_t>(M, kQ4gswTileM) *
141170
utils::div_up<int64_t>(N, kQ4gswTileN);
@@ -225,9 +254,10 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
225254
pipeline_desc.layout = pipeline_layout;
226255
pipeline_desc.compute.module = shader;
227256
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
228-
// bicol GEMV uses fixed @workgroup_size(64); only the GEMM has an override.
229-
pipeline_desc.compute.constantCount = use_gemv ? 0u : 1u;
230-
pipeline_desc.compute.constants = use_gemv ? nullptr : &wg_size_constant;
257+
// Only the tiled GEMM has a wg_size override; GEMV + shmem are fixed 64.
258+
const bool fixed_wg = use_gemv || use_shmem_gemm;
259+
pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u;
260+
pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant;
231261
WGPUComputePipeline pipeline =
232262
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
233263

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

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,127 @@
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+
override wg_size: u32 = 64u;
20+
21+
// Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
22+
const WG_M: u32 = 32u; // output rows per workgroup
23+
const WG_N: u32 = 32u; // output cols per workgroup
24+
const TM: u32 = 4u; // rows per thread
25+
const TN: u32 = 4u; // cols per thread
26+
const TK: u32 = 16u; // K-tile depth
27+
const THREADS: u32 = 64u; // 8x8 thread grid (must equal wg_size)
28+
29+
var<workgroup> in_sh: array<f32, 512>; // WG_M * TK = 32 * 16
30+
var<workgroup> w_sh: array<f32, 512>; // TK * WG_N = 16 * 32
31+
32+
@compute @workgroup_size(wg_size, 1, 1)
33+
fn main(@builtin(workgroup_id) wid: vec3<u32>,
34+
@builtin(local_invocation_index) lid: u32) {
35+
let nct = (params.N + WG_N - 1u) / WG_N; // output-col tiles
36+
let tile_row = wid.x / nct;
37+
let tile_col = wid.x % nct;
38+
let m_base = tile_row * WG_M; // first output row of this WG's tile
39+
let n_base = tile_col * WG_N; // first output col of this WG's tile
40+
41+
// This thread's 4x4 sub-tile origin within the 32x32 WG tile (8x8 thread grid).
42+
let tr = lid / 8u; // thread row in [0,8)
43+
let tc = lid % 8u; // thread col in [0,8)
44+
let m_thread = m_base + tr * TM;
45+
let n_thread = n_base + tc * TN;
46+
47+
var acc: array<f32, 16>; // TM * TN
48+
for (var i: u32 = 0u; i < TM * TN; i = i + 1u) {
49+
acc[i] = 0.0;
50+
}
51+
52+
var k0: u32 = 0u;
53+
loop {
54+
if (k0 >= params.K) { break; }
55+
56+
// (a) Stage the 32xTK input block into shmem; 64 threads x 8 elems; OOB->0.
57+
for (var t: u32 = lid; t < WG_M * TK; t = t + THREADS) {
58+
let r = t / TK; // row within the WG tile [0,32)
59+
let c = t % TK; // col within the K-tile [0,16)
60+
let m = m_base + r;
61+
let k = k0 + c;
62+
var v: f32 = 0.0;
63+
if (m < params.M && k < params.K) {
64+
v = t_input[m * params.K + k];
65+
}
66+
in_sh[t] = v;
67+
}
68+
69+
// (b) Dequant the TKx32 weight block into shmem once; dq = (nibble-8)*scale.
70+
for (var t: u32 = lid; t < TK * WG_N; t = t + THREADS) {
71+
let kk = t / WG_N; // k within the K-tile [0,16)
72+
let nn = t % WG_N; // col within the WG tile [0,32)
73+
let n = n_base + nn;
74+
let k = k0 + kk;
75+
var dq: f32 = 0.0;
76+
if (n < params.N && k < params.K) {
77+
let byte_idx = n * params.K_packed + (k >> 1u);
78+
let word = t_weight[byte_idx >> 2u];
79+
let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu;
80+
var nib: u32;
81+
if ((k & 1u) == 0u) {
82+
nib = b & 0x0Fu; // even k -> low nibble
83+
} else {
84+
nib = (b >> 4u) & 0x0Fu; // odd k -> high nibble
85+
}
86+
let q = f32(i32(nib) - 8); // +8-shifted on pack; recover [-8,7]
87+
dq = q * t_scales[(k / params.group_size) * params.padded_N + n];
88+
}
89+
w_sh[t] = dq;
90+
}
91+
92+
workgroupBarrier();
93+
94+
// Accumulate this thread's 4x4 tile from shared memory.
95+
let k_lim = min(TK, params.K - k0);
96+
for (var kk: u32 = 0u; kk < k_lim; kk = kk + 1u) {
97+
var a: array<f32, 4>;
98+
for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
99+
a[ml] = in_sh[(tr * TM + ml) * TK + kk];
100+
}
101+
for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
102+
let wv = w_sh[kk * WG_N + (tc * TN + nl)];
103+
for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
104+
acc[ml * TN + nl] = acc[ml * TN + nl] + a[ml] * wv;
105+
}
106+
}
107+
}
108+
109+
workgroupBarrier();
110+
k0 = k0 + TK;
111+
}
112+
113+
// Write the 4x4 register tile out (bounds-guarded; bias added per column).
114+
for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
115+
let m = m_thread + ml;
116+
for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
117+
let n = n_thread + nl;
118+
if (m < params.M && n < params.N) {
119+
var v = acc[ml * TN + nl];
120+
if (params.has_bias != 0u) {
121+
v = v + t_bias[n];
122+
}
123+
t_out[m * params.N + n] = v;
124+
}
125+
}
126+
}
127+
}

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

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
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_gemm_shmem.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 0a1219d3b6781315a21066089fca3f92235587e8af8eb734185f35ea4bfc8a52
17+
inline constexpr const char* kQ4gswLinearGemmShmemWGSL = 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+
override wg_size: u32 = 64u;
37+
38+
// Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
39+
const WG_M: u32 = 32u; // output rows per workgroup
40+
const WG_N: u32 = 32u; // output cols per workgroup
41+
const TM: u32 = 4u; // rows per thread
42+
const TN: u32 = 4u; // cols per thread
43+
const TK: u32 = 16u; // K-tile depth
44+
const THREADS: u32 = 64u; // 8x8 thread grid (must equal wg_size)
45+
46+
var<workgroup> in_sh: array<f32, 512>; // WG_M * TK = 32 * 16
47+
var<workgroup> w_sh: array<f32, 512>; // TK * WG_N = 16 * 32
48+
49+
@compute @workgroup_size(wg_size, 1, 1)
50+
fn main(@builtin(workgroup_id) wid: vec3<u32>,
51+
@builtin(local_invocation_index) lid: u32) {
52+
let nct = (params.N + WG_N - 1u) / WG_N; // output-col tiles
53+
let tile_row = wid.x / nct;
54+
let tile_col = wid.x % nct;
55+
let m_base = tile_row * WG_M; // first output row of this WG's tile
56+
let n_base = tile_col * WG_N; // first output col of this WG's tile
57+
58+
// This thread's 4x4 sub-tile origin within the 32x32 WG tile (8x8 thread grid).
59+
let tr = lid / 8u; // thread row in [0,8)
60+
let tc = lid % 8u; // thread col in [0,8)
61+
let m_thread = m_base + tr * TM;
62+
let n_thread = n_base + tc * TN;
63+
64+
var acc: array<f32, 16>; // TM * TN
65+
for (var i: u32 = 0u; i < TM * TN; i = i + 1u) {
66+
acc[i] = 0.0;
67+
}
68+
69+
var k0: u32 = 0u;
70+
loop {
71+
if (k0 >= params.K) { break; }
72+
73+
// (a) Stage the 32xTK input block into shmem; 64 threads x 8 elems; OOB->0.
74+
for (var t: u32 = lid; t < WG_M * TK; t = t + THREADS) {
75+
let r = t / TK; // row within the WG tile [0,32)
76+
let c = t % TK; // col within the K-tile [0,16)
77+
let m = m_base + r;
78+
let k = k0 + c;
79+
var v: f32 = 0.0;
80+
if (m < params.M && k < params.K) {
81+
v = t_input[m * params.K + k];
82+
}
83+
in_sh[t] = v;
84+
}
85+
86+
// (b) Dequant the TKx32 weight block into shmem once; dq = (nibble-8)*scale.
87+
for (var t: u32 = lid; t < TK * WG_N; t = t + THREADS) {
88+
let kk = t / WG_N; // k within the K-tile [0,16)
89+
let nn = t % WG_N; // col within the WG tile [0,32)
90+
let n = n_base + nn;
91+
let k = k0 + kk;
92+
var dq: f32 = 0.0;
93+
if (n < params.N && k < params.K) {
94+
let byte_idx = n * params.K_packed + (k >> 1u);
95+
let word = t_weight[byte_idx >> 2u];
96+
let b = (word >> ((byte_idx & 3u) * 8u)) & 0xFFu;
97+
var nib: u32;
98+
if ((k & 1u) == 0u) {
99+
nib = b & 0x0Fu; // even k -> low nibble
100+
} else {
101+
nib = (b >> 4u) & 0x0Fu; // odd k -> high nibble
102+
}
103+
let q = f32(i32(nib) - 8); // +8-shifted on pack; recover [-8,7]
104+
dq = q * t_scales[(k / params.group_size) * params.padded_N + n];
105+
}
106+
w_sh[t] = dq;
107+
}
108+
109+
workgroupBarrier();
110+
111+
// Accumulate this thread's 4x4 tile from shared memory.
112+
let k_lim = min(TK, params.K - k0);
113+
for (var kk: u32 = 0u; kk < k_lim; kk = kk + 1u) {
114+
var a: array<f32, 4>;
115+
for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
116+
a[ml] = in_sh[(tr * TM + ml) * TK + kk];
117+
}
118+
for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
119+
let wv = w_sh[kk * WG_N + (tc * TN + nl)];
120+
for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
121+
acc[ml * TN + nl] = acc[ml * TN + nl] + a[ml] * wv;
122+
}
123+
}
124+
}
125+
126+
workgroupBarrier();
127+
k0 = k0 + TK;
128+
}
129+
130+
// Write the 4x4 register tile out (bounds-guarded; bias added per column).
131+
for (var ml: u32 = 0u; ml < TM; ml = ml + 1u) {
132+
let m = m_thread + ml;
133+
for (var nl: u32 = 0u; nl < TN; nl = nl + 1u) {
134+
let n = n_thread + nl;
135+
if (m < params.M && n < params.N) {
136+
var v = acc[ml * TN + nl];
137+
if (params.has_bias != 0u) {
138+
v = v + t_bias[n];
139+
}
140+
t_out[m * params.N + n] = v;
141+
}
142+
}
143+
}
144+
}
145+
)";
146+
147+
inline constexpr uint32_t kQ4gswLinearGemmShmemWorkgroupSizeX = 64;
148+
inline constexpr uint32_t kQ4gswLinearGemmShmemWorkgroupSizeY = 1;
149+
inline constexpr uint32_t kQ4gswLinearGemmShmemWorkgroupSizeZ = 1;
150+
151+
} // namespace executorch::backends::webgpu

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -59,6 +59,10 @@ class Q4gswConfig:
5959
# requires N % 8 == 0 (torchao pads N for the scale layout), so odd-N / N=1 are
6060
# not exportable -- bicol's has1 odd-N guard is defensive (mirrors coop4's
6161
# general-N robustness) and unreachable through this op.
62+
# Prefill shapes routing to the shmem GEMM (K>=4096 or N>=2048); M=128.
63+
Q4gswConfig("gate_proj_pf", 128, 2048, 8192), # gate/up prefill (shmem via N)
64+
Q4gswConfig("down_proj_pf", 128, 8192, 2048), # down prefill (shmem via K)
65+
Q4gswConfig("shmem_edge", 130, 4096, 2056), # partial 32-tile bounds
6266
]
6367

6468

‎backends/webgpu/test/test_webgpu_native.cpp‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -293,6 +293,9 @@ static const Q4gswConfig kQ4gswConfigs[] = {
293293
// scale over 64-256 K-groups). q4gsw requires N % 8 == 0, so odd-N is not
294294
// exportable; bicol's has1 odd-N guard is defensive (mirrors coop4
295295
// general-N robustness).
296+
{"gate_proj_pf", 128, 2048, 8192, 1e-4f, 1e-3f, true, false}, // shmem via N
297+
{"down_proj_pf", 128, 8192, 2048, 1e-3f, 1e-2f, true, false}, // shmem via K
298+
{"shmem_edge", 130, 4096, 2056, 1e-4f, 1e-3f, true, false}, // partial tiles
296299
};
297300

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

0 commit comments

Comments
 (0)