Skip to content

[ExecuTorch][WebGPU] Make shader workgroup sizes runtime-configurable #20954

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
JCNTH wants to merge 2 commits into
base: gh/JCNTH/86/base
Choose a base branch
from
Open

[ExecuTorch][WebGPU] Make shader workgroup sizes runtime-configurable #20954

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
65c1fbb
Update
JCNTH Jul 15, 2026
fb95a32
Update
JCNTH Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions backends/webgpu/runtime/WebGPUUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,16 @@ inline uint32_t clamp_workgroup_size(WGPUDevice device, uint32_t desired) {
return desired;
}

// Clamp to device limit, then floor to pow2 (reduction kernels halve stride).
inline uint32_t clamp_workgroup_size_pow2(WGPUDevice device, uint32_t desired) {
uint32_t v = clamp_workgroup_size(device, desired);
uint32_t p = 1u;
while (p <= (v >> 1u)) {
p <<= 1u;
}
return p;
}

struct WgCount {
uint32_t x;
uint32_t y;
Expand Down
12 changes: 9 additions & 3 deletions backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -254,6 +254,11 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
const uint32_t wg_size =
utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX);
const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u);
// GEMV (bicol) is a pow2 tree reduction; compute its size only when used.
const uint32_t gemv_wg_size = use_gemv
? utils::clamp_workgroup_size_pow2(
device, kQ4gswLinearCoop4BicolWorkgroupSizeX)
: 0u;
// steel (256-thread) is the preferred M>1 prefill GEMM; 0 count = ineligible.
const bool use_steel = !use_gemv && steel_supported(device) &&
steel_workgroup_count(device, M, N, K) > 0u;
Expand Down Expand Up @@ -372,16 +377,17 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector<int>& args) {
WGPUPipelineLayout pipeline_layout =
wgpuDeviceCreatePipelineLayout(device, &pl_desc);

// GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed.
const bool fixed_wg = use_steel || use_shmem_gemm;
WGPUConstantEntry wg_size_constant = {};
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
wg_size_constant.value = static_cast<double>(wg_size);
wg_size_constant.value =
static_cast<double>(use_gemv ? gemv_wg_size : wg_size);

WGPUComputePipelineDescriptor pipeline_desc = {};
pipeline_desc.layout = pipeline_layout;
pipeline_desc.compute.module = shader;
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
// Only tiled GEMM overrides wg_size; GEMV/shmem (64) + steel (256) are fixed.
const bool fixed_wg = use_gemv || use_steel || use_shmem_gemm;
pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u;
pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant;
WGPUComputePipeline pipeline =
Expand Down
11 changes: 6 additions & 5 deletions backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,10 +17,11 @@ struct Params {
@group(0) @binding(5) var<uniform> params: Params;

// Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
const WG: u32 = 64u;
var<workgroup> partial: array<vec2<f32>, WG>;
// wg_size must be a power of two (the tree reduction halves the stride).
override wg_size: u32 = 64u;
var<workgroup> partial: array<vec2<f32>, wg_size>;

@compute @workgroup_size(WG, 1, 1)
@compute @workgroup_size(wg_size, 1, 1)
fn main(
@builtin(workgroup_id) wid: vec3<u32>,
@builtin(num_workgroups) ngrp: vec3<u32>,
Expand Down Expand Up @@ -69,12 +70,12 @@ fn main(
acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
}
w = w + WG;
w = w + wg_size;
}

partial[lid.x] = vec2<f32>(acc0, acc1);
workgroupBarrier();
var s: u32 = WG >> 1u;
var s: u32 = wg_size >> 1u;
loop {
if (s == 0u) {
break;
Expand Down
13 changes: 7 additions & 6 deletions backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_coop4_bicol_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from q4gsw_linear_coop4_bicol.wgsl - DO NOT EDIT.
// wgsl-sha256: 44c53ca7879f7a3382a45703a67cbc6fc221ecbcada58f8c413df50abbc6e544
// wgsl-sha256: ef677c2771c3d4a73704f4725bb490174149b971bc62e696595b8d13941c979e
inline constexpr const char* kQ4gswLinearCoop4BicolWGSL = R"(
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@group(0) @binding(1) var<storage, read> t_input: array<f32>;
Expand All @@ -34,10 +34,11 @@ struct Params {
@group(0) @binding(5) var<uniform> params: Params;

// Cooperative-over-K GEMV, 2 output columns/workgroup; input read once, reused.
const WG: u32 = 64u;
var<workgroup> partial: array<vec2<f32>, WG>;
// wg_size must be a power of two (the tree reduction halves the stride).
override wg_size: u32 = 64u;
var<workgroup> partial: array<vec2<f32>, wg_size>;

@compute @workgroup_size(WG, 1, 1)
@compute @workgroup_size(wg_size, 1, 1)
fn main(
@builtin(workgroup_id) wid: vec3<u32>,
@builtin(num_workgroups) ngrp: vec3<u32>,
Expand Down Expand Up @@ -86,12 +87,12 @@ fn main(
acc1 = acc1 + in0 * f32(i32(b1 & 0x0Fu) - 8) * scale1;
acc1 = acc1 + in1 * f32(i32((b1 >> 4u) & 0x0Fu) - 8) * scale1;
}
w = w + WG;
w = w + wg_size;
}

partial[lid.x] = vec2<f32>(acc0, acc1);
workgroupBarrier();
var s: u32 = WG >> 1u;
var s: u32 = wg_size >> 1u;
loop {
if (s == 0u) {
break;
Expand Down
3 changes: 2 additions & 1 deletion backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@ struct Params {
}
@group(0) @binding(5) var<uniform> params: Params;

override wg_size: u32 = 64u;
// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob).
const wg_size: u32 = 64u;

// Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
const WG_M: u32 = 32u; // output rows per workgroup
Expand Down
5 changes: 3 additions & 2 deletions backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_shmem_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from q4gsw_linear_gemm_shmem.wgsl - DO NOT EDIT.
// wgsl-sha256: 0a1219d3b6781315a21066089fca3f92235587e8af8eb734185f35ea4bfc8a52
// wgsl-sha256: f0180cad41701807b34912b57d9715c02588c9447783bbf691f93836168fe981
inline constexpr const char* kQ4gswLinearGemmShmemWGSL = R"(
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@group(0) @binding(1) var<storage, read> t_input: array<f32>;
Expand All @@ -33,7 +33,8 @@ struct Params {
}
@group(0) @binding(5) var<uniform> params: Params;

override wg_size: u32 = 64u;
// Fixed 64: 8x8 grid + 512-elem shared tiles are hard-locked to it (no knob).
const wg_size: u32 = 64u;

// Shmem-staged tiled GEMM (M>1): dequant weight into shmem once per K-tile.
const WG_M: u32 = 32u; // output rows per workgroup
Expand Down
1 change: 1 addition & 0 deletions backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@ struct Params {
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
var<workgroup> As: array<${buffer_scalar_type(DTYPE)}, 1024>; // BM*BK
var<workgroup> Bs: array<${buffer_scalar_type(DTYPE)}, 1024>; // BK*BN
// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand Down
3 changes: 2 additions & 1 deletion backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
// wgsl-sha256: 36b3d3f9dd08a529909c13ec7d66cd0cf392c347ca047a4d38453b3c295f72ce
// wgsl-sha256: fe8b71afc5634e7db5607deb58f30c6517f93e1d66370ea017bcd9071201c6ca
inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"(
enable f16;
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
Expand Down Expand Up @@ -50,6 +50,7 @@ struct Params {
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
var<workgroup> As: array<f16, 1024>; // BM*BK
var<workgroup> Bs: array<f16, 1024>; // BK*BN
// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand Down
3 changes: 2 additions & 1 deletion backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
// wgsl-sha256: 1f916bcd30dbbbcc7eca37e795ecc26e3c72e645ccd2c361fa0ac4e66f1a174a
// wgsl-sha256: 1b3f5e384e09bde50c66d43db4797e01e5610d70af235355dbc4966643bf9103
inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"(
enable f16;
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
Expand Down Expand Up @@ -50,6 +50,7 @@ struct Params {
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
var<workgroup> As: array<f16, 1024>; // BM*BK
var<workgroup> Bs: array<f16, 1024>; // BK*BN
// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand Down
3 changes: 2 additions & 1 deletion backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
// wgsl-sha256: 00cdd2f2fb98a5c7343d16fdf7e59f1b840e180cec3f82bf9b569513c0a45396
// wgsl-sha256: d77542e5975fab44c58644c771390e89598ce82ab25e38d57150b581b4904acd
inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"(
enable f16;
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
Expand Down Expand Up @@ -50,6 +50,7 @@ struct Params {
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
var<workgroup> As: array<f16, 1024>; // BM*BK
var<workgroup> Bs: array<f16, 1024>; // BK*BN
// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand Down
3 changes: 2 additions & 1 deletion backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT.
// wgsl-sha256: dd771b9ab096410f3ad0d9259bef7816e41330a434325ea28baa5abbfb2841d2
// wgsl-sha256: dcb63b04d002899acfed88d273fe658618650d5dcbc108540cef086bb97f5bce
inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"(
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
@group(0) @binding(1) var<storage, read> t_input: array<f32>;
Expand Down Expand Up @@ -49,6 +49,7 @@ struct Params {
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
var<workgroup> As: array<f32, 1024>; // BM*BK
var<workgroup> Bs: array<f32, 1024>; // BK*BN
// 16x16 = 256 threads, bound to the 64x64 tile + 4x4 reg tile (not a knob).
@compute @workgroup_size(16, 16)
fn main(@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand Down
14 changes: 12 additions & 2 deletions backends/webgpu/runtime/ops/rms_norm/RmsNorm.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -178,11 +178,21 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
WGPUPipelineLayout pipeline_layout =
wgpuDeviceCreatePipelineLayout(device, &pl_desc);

// Runtime-overridable workgroup size (mirrors add op); clamp only reduces.
// Pow2 required: the kernel halves the reduction stride (wg_size / 2u).
const uint32_t wg_size =
utils::clamp_workgroup_size_pow2(device, kRmsNormWorkgroupSizeX);
WGPUConstantEntry wg_size_constant = {};
wg_size_constant.key = {"wg_size", WGPU_STRLEN};
wg_size_constant.value = static_cast<double>(wg_size);

// Create compute pipeline
WGPUComputePipelineDescriptor pipeline_desc = {};
pipeline_desc.layout = pipeline_layout;
pipeline_desc.compute.module = shader;
pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN};
pipeline_desc.compute.constantCount = 1;
pipeline_desc.compute.constants = &wg_size_constant;
WGPUComputePipeline pipeline =
wgpuDeviceCreateComputePipeline(device, &pipeline_desc);

Expand Down Expand Up @@ -217,10 +227,10 @@ void rms_norm_impl(WebGPUGraph& graph, const std::vector<int>& args) {
// One workgroup per row (kRmsNormWorkgroupSizeX threads cooperate per row)
static_assert(
kRmsNormWorkgroupSizeX == 64,
"must match @workgroup_size and WG_SIZE in rms_norm.wgsl");
"kRmsNormWorkgroupSizeX must match override wg_size default (64)");
static_assert(
kRmsNormVec4WorkgroupSizeX == 64,
"must match @workgroup_size and WG_SIZE in rms_norm_vec4.wgsl");
"kRmsNormVec4WorkgroupSizeX must match override wg_size default (64)");
const size_t dispatch_idx =
graph.add_dispatch({pipeline, bind_group, num_rows});

Expand Down
19 changes: 10 additions & 9 deletions backends/webgpu/runtime/ops/rms_norm/rms_norm.wgsl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,14 @@ struct Params {
}
@group(0) @binding(3) var<uniform> params: Params;

const WG_SIZE: u32 = 64u;
// wg_size must be a power of two (the tree reduction halves the stride).
override wg_size: u32 = 64u;

var<workgroup> shared_sum: array<${accum_scalar_type(DTYPE)}, WG_SIZE>;
var<workgroup> shared_sum: array<${accum_scalar_type(DTYPE)}, wg_size>;

fn reduce_shared(worker_id: u32) {
workgroupBarrier();
var stride: u32 = WG_SIZE / 2u;
var stride: u32 = wg_size / 2u;
loop {
if (stride == 0u) {
break;
Expand All @@ -32,10 +33,10 @@ fn reduce_shared(worker_id: u32) {
}

$if VEC == 4:
// vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4
// vec4 variant of rms_norm: each lane strides by wg_size over rw4 = row_width/4
// texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq
// divides by it (not rw4). The host selects this only when row_width % 4 == 0.
@compute @workgroup_size(64, 1, 1)
@compute @workgroup_size(wg_size, 1, 1)
fn main(
@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand Down Expand Up @@ -64,7 +65,7 @@ fn main(
local_sq_sum = local_sq_sum + dot(vec4<f32>(v), vec4<f32>(v));
$else:
local_sq_sum = local_sq_sum + dot(v, v);
x4 = x4 + WG_SIZE;
x4 = x4 + wg_size;
}
$else:
var x: u32 = worker_id;
Expand All @@ -77,7 +78,7 @@ fn main(
local_sq_sum = local_sq_sum + f32(v) * f32(v);
$else:
local_sq_sum = local_sq_sum + v * v;
x = x + WG_SIZE;
x = x + wg_size;
}

shared_sum[worker_id] = local_sq_sum;
Expand All @@ -96,7 +97,7 @@ fn main(
t_out[base4 + x4] = vec4<f16>(vec4<f32>(t_in[base4 + x4]) * rstd * vec4<f32>(t_weight[x4]));
$else:
t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4];
x4 = x4 + WG_SIZE;
x4 = x4 + wg_size;
}
$else:
x = worker_id;
Expand All @@ -110,6 +111,6 @@ fn main(
t_out[base + x] = f16(f32(v) * rstd * f32(w));
$else:
t_out[base + x] = v * rstd * w;
x = x + WG_SIZE;
x = x + wg_size;
}
}
17 changes: 9 additions & 8 deletions backends/webgpu/runtime/ops/rms_norm/rms_norm_vec4_wgsl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,7 @@
namespace executorch::backends::webgpu {

// @generated from rms_norm.wgsl - DO NOT EDIT.
// wgsl-sha256: 4c0ba56708bf125a7ec6ea3c51d1288e05ac00a8e2cfa10e38e9a208e230b8df
// wgsl-sha256: 62fdfe03fc67eb44fa17ca5e91e433b3a45a96c76151777856d6f557fd829919
inline constexpr const char* kRmsNormVec4WGSL = R"(
@group(0) @binding(0) var<storage, read_write> t_out: array<vec4<f32>>;
@group(0) @binding(1) var<storage, read> t_in: array<vec4<f32>>;
Expand All @@ -27,13 +27,14 @@ struct Params {
}
@group(0) @binding(3) var<uniform> params: Params;

const WG_SIZE: u32 = 64u;
// wg_size must be a power of two (the tree reduction halves the stride).
override wg_size: u32 = 64u;

var<workgroup> shared_sum: array<f32, WG_SIZE>;
var<workgroup> shared_sum: array<f32, wg_size>;

fn reduce_shared(worker_id: u32) {
workgroupBarrier();
var stride: u32 = WG_SIZE / 2u;
var stride: u32 = wg_size / 2u;
loop {
if (stride == 0u) {
break;
Expand All @@ -46,10 +47,10 @@ fn reduce_shared(worker_id: u32) {
}
}

// vec4 variant of rms_norm: each lane strides by WG_SIZE over rw4 = row_width/4
// vec4 variant of rms_norm: each lane strides by wg_size over rw4 = row_width/4
// texels and accumulates dot(v, v). row_width is the ELEMENT count, so mean_sq
// divides by it (not rw4). The host selects this only when row_width % 4 == 0.
@compute @workgroup_size(64, 1, 1)
@compute @workgroup_size(wg_size, 1, 1)
fn main(
@builtin(workgroup_id) wid: vec3<u32>,
@builtin(local_invocation_id) lid: vec3<u32>) {
Expand All @@ -71,7 +72,7 @@ fn main(
}
let v = t_in[base4 + x4];
local_sq_sum = local_sq_sum + dot(v, v);
x4 = x4 + WG_SIZE;
x4 = x4 + wg_size;
}

shared_sum[worker_id] = local_sq_sum;
Expand All @@ -86,7 +87,7 @@ fn main(
break;
}
t_out[base4 + x4] = t_in[base4 + x4] * rstd * t_weight[x4];
x4 = x4 + WG_SIZE;
x4 = x4 + wg_size;
}
}
)";
Expand Down
Loading
Loading