Skip to content

Commit 5828778

Browse files
pytorchbotJCNTH
pytorchbot
and
JCNTH
authored
[ExecuTorch][WebGPU] Fuse attention QKV projections in prefill (pytorch#21107)
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: pytorch#20880 by @JCNTH ^ 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/JCNTH/52/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/52/head Merge bot PR base: https://github.com/pytorch/executorch/tree/main Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/JCNTH/52/orig @diff-train-skip-merge --------- Co-authored-by: Julian Ng-Thow-Hing <juliannth@meta.com> Co-authored-by: Julian Ng-Thow-Hing <107437036+JCNTH@users.noreply.github.com>
1 parent 6bd0355 commit 5828778

11 files changed

Lines changed: 1224 additions & 29 deletions

‎backends/webgpu/runtime/WebGPUGraph.cpp‎

Lines changed: 844 additions & 0 deletions
Large diffs are not rendered by default.

‎backends/webgpu/runtime/WebGPUGraph.h‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -450,6 +450,43 @@ class WebGPUGraph {
450450
std::unordered_map<std::string, WGPUBindGroupLayout> bgl_cache_;
451451

452452
size_t uniform_buffer_bytes_ = 0;
453+
454+
// QKV-concat fusion: one detected attention q/k/v linear
455+
// triple sharing an input activation (value ids + shapes), fused in build()
456+
// into a single multi-output q4gsw GEMM that scatter-writes q/k/v. Only used
457+
// during build(); inert (never populated) when no q/k/v triple matches.
458+
struct QkvFusionGroup {
459+
int input_id = -1;
460+
int out_q = -1, out_k = -1, out_v = -1;
461+
int weight_q = -1, weight_k = -1, weight_v = -1;
462+
int scales_q = -1, scales_k = -1, scales_v = -1;
463+
uint32_t Nq = 0, Nk = 0, Nv = 0; // 2048, 512, 512
464+
uint32_t K = 0, K_packed = 0, group_size = 0, num_groups = 0;
465+
uint32_t padded_N_q = 0, padded_N_k = 0, padded_N_v = 0;
466+
unsigned op_idx[3] = {0, 0, 0}; // the 3 q/k/v linear op-chain indices
467+
size_t sep_dispatch[3] = {
468+
0,
469+
0,
470+
0}; // their dispatch indices (filled in build())
471+
size_t fused_dispatch = 0; // the fused GEMM dispatch index
472+
WGPUBuffer fused_params =
473+
nullptr; // the fused params UBO (rewritten by the hook)
474+
};
475+
// Concat the 3 packed weights (row-stack) + scales (strided gather) into
476+
// fused buffers, then record ONE fused-GEMM dispatch (bespoke 8-binding
477+
// layout) that writes the 3 original q/k/v output buffers, plus a 3-output
478+
// resize hook.
479+
void add_qkv_fused_dispatch(QkvFusionGroup& g);
480+
void add_qkv_fused_hook(const QkvFusionGroup& g);
481+
482+
// SwiGLU fusion: emit ONE fused elementwise dispatch
483+
// computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2
484+
// muls. `out` is repointed to a private pooled buffer (aliasing guard);
485+
// `gate` is likewise given a private pooled buffer at its producer op by the
486+
// build() walk (the planner reuse-aliases up onto gate's slot, so up_proj
487+
// would stomp gate before the fused reads it). Only used during build(); the
488+
// detection maps are empty (inert) when no SwiGLU triple matches.
489+
void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id);
453490
};
454491

455492
} // namespace executorch::backends::webgpu

‎backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
@group(0) @binding(0) var<storage, read> gate: array<f32>;
2+
@group(0) @binding(1) var<storage, read> up: array<f32>;
3+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
4+
5+
struct Params {
6+
num_elements: u32,
7+
}
8+
@group(0) @binding(3) var<uniform> params: Params;
9+
10+
// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
11+
// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
12+
// sigmoid + silu are computed in registers (never written to memory), so gate + up
13+
// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
14+
// multiply order match the original ops -> bit-exact.
15+
@compute @workgroup_size(64)
16+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
17+
let idx = gid.x;
18+
if (idx >= params.num_elements) {
19+
return;
20+
}
21+
let g = gate[idx];
22+
let sig = 1.0 / (1.0 + exp(-g));
23+
output[idx] = (g * sig) * up[idx];
24+
}

‎backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,48 @@
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 silu_mul_fused.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560
17+
inline constexpr const char* kSiluMulFusedWGSL = R"(
18+
@group(0) @binding(0) var<storage, read> gate: array<f32>;
19+
@group(0) @binding(1) var<storage, read> up: array<f32>;
20+
@group(0) @binding(2) var<storage, read_write> output: array<f32>;
21+
22+
struct Params {
23+
num_elements: u32,
24+
}
25+
@group(0) @binding(3) var<uniform> params: Params;
26+
27+
// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate
28+
// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch.
29+
// sigmoid + silu are computed in registers (never written to memory), so gate + up
30+
// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the
31+
// multiply order match the original ops -> bit-exact.
32+
@compute @workgroup_size(64)
33+
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
34+
let idx = gid.x;
35+
if (idx >= params.num_elements) {
36+
return;
37+
}
38+
let g = gate[idx];
39+
let sig = 1.0 / (1.0 + exp(-g));
40+
output[idx] = (g * sig) * up[idx];
41+
}
42+
)";
43+
44+
inline constexpr uint32_t kSiluMulFusedWorkgroupSizeX = 64;
45+
inline constexpr uint32_t kSiluMulFusedWorkgroupSizeY = 1;
46+
inline constexpr uint32_t kSiluMulFusedWorkgroupSizeZ = 1;
47+
48+
} // namespace executorch::backends::webgpu

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

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
enable f16;
2+
// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM
3+
// (vec4<f32> activation load) that scatter-writes each output column range to a SEPARATE buffer --
4+
// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches;
5+
// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned
6+
// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE:
7+
// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the
8+
// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead
9+
// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV).
10+
@group(0) @binding(0) var<storage, read_write> t_out_q: array<f32>;
11+
@group(0) @binding(1) var<storage, read_write> t_out_k: array<f32>;
12+
@group(0) @binding(2) var<storage, read_write> t_out_v: array<f32>;
13+
@group(0) @binding(3) var<storage, read> t_input: array<vec4<f32>>;
14+
@group(0) @binding(4) var<storage, read> t_weight: array<u32>;
15+
@group(0) @binding(5) var<storage, read> t_scales: array<f32>;
16+
@group(0) @binding(6) var<storage, read> t_bias: array<f32>;
17+
struct Params {
18+
M: u32,
19+
N: u32,
20+
K: u32,
21+
K_packed: u32,
22+
group_size: u32,
23+
padded_N: u32,
24+
has_bias: u32,
25+
_pad: u32,
26+
}
27+
@group(0) @binding(7) var<uniform> params: Params;
28+
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
29+
const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u;
30+
var<workgroup> As: array<f16, 1024>;
31+
var<workgroup> Bs: array<f16, 1024>;
32+
@compute @workgroup_size(16, 16)
33+
fn main(@builtin(workgroup_id) wid: vec3<u32>,
34+
@builtin(local_invocation_id) lid: vec3<u32>) {
35+
let nbN = (params.N + BN - 1u) / BN;
36+
let bx = wid.x % nbN;
37+
let by = wid.x / nbN;
38+
let row0 = by * BM;
39+
let col0 = bx * BN;
40+
let tid = lid.y * 16u + lid.x;
41+
var acc: array<array<f16, 4>, 4>;
42+
for (var m: u32 = 0u; m < 4u; m = m + 1u) {
43+
for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; }
44+
}
45+
let ar = tid / 4u;
46+
let ac = (tid % 4u) * 4u;
47+
var k0: u32 = 0u;
48+
loop {
49+
if (k0 >= params.K) { break; }
50+
let arow = row0 + ar;
51+
if (arow < params.M) {
52+
let base = arow * params.K + k0 + ac;
53+
let av = t_input[base >> 2u];
54+
As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y);
55+
As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w);
56+
} else {
57+
As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h;
58+
As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h;
59+
}
60+
if (tid < BN) {
61+
let c = tid;
62+
let n = col0 + c;
63+
if (n < params.N) {
64+
let scale_row = (k0 / params.group_size) * params.padded_N;
65+
let scale = f16(t_scales[scale_row + n]);
66+
let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u);
67+
let w0 = t_weight[base_word];
68+
let w1 = t_weight[base_word + 1u];
69+
for (var br: u32 = 0u; br < BK; br = br + 1u) {
70+
let word = select(w1, w0, br < 8u);
71+
let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu;
72+
Bs[br * BN + c] = f16(i32(nib) - 8) * scale;
73+
}
74+
} else {
75+
for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; }
76+
}
77+
}
78+
workgroupBarrier();
79+
for (var k: u32 = 0u; k < BK; k = k + 1u) {
80+
var a: array<f16, 4>;
81+
var bvec: array<f16, 4>;
82+
for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; }
83+
for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; }
84+
for (var m: u32 = 0u; m < 4u; m = m + 1u) {
85+
for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); }
86+
}
87+
}
88+
workgroupBarrier();
89+
k0 = k0 + BK;
90+
}
91+
for (var m: u32 = 0u; m < 4u; m = m + 1u) {
92+
for (var n: u32 = 0u; n < 4u; n = n + 1u) {
93+
let r = row0 + lid.y * 4u + m;
94+
let c = col0 + lid.x * 4u + n; // global fused column [0, 3072)
95+
if (r < params.M && c < params.N) {
96+
var val = f32(acc[m][n]);
97+
if (params.has_bias != 0u) { val = val + t_bias[c]; }
98+
if (c < N_Q) { t_out_q[r * N_Q + c] = val; }
99+
else if (c < N_QK) { t_out_k[r * N_KV + (c - N_Q)] = val; }
100+
else { t_out_v[r * N_KV + (c - N_QK)] = val; }
101+
}
102+
}
103+
}
104+
}

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

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
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_qkv_fused.wgsl - DO NOT EDIT.
16+
// wgsl-sha256: 93e127e8ee4609d846015c8b75a600a29502e19a92bdf3a08e3429635f834085
17+
inline constexpr const char* kQ4gswLinearGemmQkvFusedWGSL = R"(
18+
enable f16;
19+
// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM
20+
// (vec4<f32> activation load) that scatter-writes each output column range to a SEPARATE buffer --
21+
// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches;
22+
// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned
23+
// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE:
24+
// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the
25+
// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead
26+
// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV).
27+
@group(0) @binding(0) var<storage, read_write> t_out_q: array<f32>;
28+
@group(0) @binding(1) var<storage, read_write> t_out_k: array<f32>;
29+
@group(0) @binding(2) var<storage, read_write> t_out_v: array<f32>;
30+
@group(0) @binding(3) var<storage, read> t_input: array<vec4<f32>>;
31+
@group(0) @binding(4) var<storage, read> t_weight: array<u32>;
32+
@group(0) @binding(5) var<storage, read> t_scales: array<f32>;
33+
@group(0) @binding(6) var<storage, read> t_bias: array<f32>;
34+
struct Params {
35+
M: u32,
36+
N: u32,
37+
K: u32,
38+
K_packed: u32,
39+
group_size: u32,
40+
padded_N: u32,
41+
has_bias: u32,
42+
_pad: u32,
43+
}
44+
@group(0) @binding(7) var<uniform> params: Params;
45+
const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u;
46+
const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u;
47+
var<workgroup> As: array<f16, 1024>;
48+
var<workgroup> Bs: array<f16, 1024>;
49+
@compute @workgroup_size(16, 16)
50+
fn main(@builtin(workgroup_id) wid: vec3<u32>,
51+
@builtin(local_invocation_id) lid: vec3<u32>) {
52+
let nbN = (params.N + BN - 1u) / BN;
53+
let bx = wid.x % nbN;
54+
let by = wid.x / nbN;
55+
let row0 = by * BM;
56+
let col0 = bx * BN;
57+
let tid = lid.y * 16u + lid.x;
58+
var acc: array<array<f16, 4>, 4>;
59+
for (var m: u32 = 0u; m < 4u; m = m + 1u) {
60+
for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; }
61+
}
62+
let ar = tid / 4u;
63+
let ac = (tid % 4u) * 4u;
64+
var k0: u32 = 0u;
65+
loop {
66+
if (k0 >= params.K) { break; }
67+
let arow = row0 + ar;
68+
if (arow < params.M) {
69+
let base = arow * params.K + k0 + ac;
70+
let av = t_input[base >> 2u];
71+
As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y);
72+
As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w);
73+
} else {
74+
As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h;
75+
As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h;
76+
}
77+
if (tid < BN) {
78+
let c = tid;
79+
let n = col0 + c;
80+
if (n < params.N) {
81+
let scale_row = (k0 / params.group_size) * params.padded_N;
82+
let scale = f16(t_scales[scale_row + n]);
83+
let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u);
84+
let w0 = t_weight[base_word];
85+
let w1 = t_weight[base_word + 1u];
86+
for (var br: u32 = 0u; br < BK; br = br + 1u) {
87+
let word = select(w1, w0, br < 8u);
88+
let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu;
89+
Bs[br * BN + c] = f16(i32(nib) - 8) * scale;
90+
}
91+
} else {
92+
for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; }
93+
}
94+
}
95+
workgroupBarrier();
96+
for (var k: u32 = 0u; k < BK; k = k + 1u) {
97+
var a: array<f16, 4>;
98+
var bvec: array<f16, 4>;
99+
for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; }
100+
for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; }
101+
for (var m: u32 = 0u; m < 4u; m = m + 1u) {
102+
for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); }
103+
}
104+
}
105+
workgroupBarrier();
106+
k0 = k0 + BK;
107+
}
108+
for (var m: u32 = 0u; m < 4u; m = m + 1u) {
109+
for (var n: u32 = 0u; n < 4u; n = n + 1u) {
110+
let r = row0 + lid.y * 4u + m;
111+
let c = col0 + lid.x * 4u + n; // global fused column [0, 3072)
112+
if (r < params.M && c < params.N) {
113+
var val = f32(acc[m][n]);
114+
if (params.has_bias != 0u) { val = val + t_bias[c]; }
115+
if (c < N_Q) { t_out_q[r * N_Q + c] = val; }
116+
else if (c < N_QK) { t_out_k[r * N_KV + (c - N_Q)] = val; }
117+
else { t_out_v[r * N_KV + (c - N_QK)] = val; }
118+
}
119+
}
120+
}
121+
}
122+
)";
123+
124+
inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeX = 16;
125+
inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeY = 16;
126+
inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeZ = 1;
127+
128+
} // namespace executorch::backends::webgpu

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

Lines changed: 7 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
$if DTYPE == "half":
22
enable f16;
33
@group(0) @binding(0) var<storage, read_write> t_out: array<f32>;
4-
@group(0) @binding(1) var<storage, read> t_input: array<f32>;
4+
@group(0) @binding(1) var<storage, read> t_input: array<vec4<f32>>;
55
@group(0) @binding(2) var<storage, read> t_weight: array<u32>;
66
@group(0) @binding(3) var<storage, read> t_scales: array<f32>;
77
@group(0) @binding(4) var<storage, read> t_bias: array<f32>;
@@ -63,10 +63,12 @@ fn main(@builtin(workgroup_id) wid: vec3<u32>,
6363
let arow = row0 + ar;
6464
if (arow < params.M) {
6565
let base = arow * params.K + k0 + ac;
66-
As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(t_input[base]);
67-
As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 1u]);
68-
As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 2u]);
69-
As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 3u]);
66+
// vec4<f32> coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4).
67+
let av = t_input[base >> 2u];
68+
As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(av.x);
69+
As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(av.y);
70+
As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(av.z);
71+
As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(av.w);
7072
} else {
7173
As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"};
7274
As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "0.0"};

0 commit comments

Comments
 (0)