From 9e9885cb10b340eee11e6655c96578a8d5f9cf0f Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 11:51:37 -0700 Subject: [PATCH 1/3] [ExecuTorch][WebGPU] Fuse attention QKV projections in prefill Pull Request resolved: https://github.com/pytorch/executorch/pull/20880 Fuse the three attention q/k/v `linear_q4gsw` projections into one multi-output GEMM that scatter-writes q/k/v, removing the occupancy-starved small k/v dispatches in prefill on Llama-3.2-1B. Auto-applied graph pass, perplexity-neutral. Problem Attention lowers q/k/v as three separate `linear_q4gsw` GEMMs; the two N=512 k/v projections each dispatch a grid too small to fill the GPU, so they run occupancy-starved. Solution - Before: q/k/v = 3 GEMM dispatches (N=2048/512/512). - After: one [2048->3072] multi-output GEMM that scatter-writes q/k/v to their own buffers, M-gated so it fires only at prefill (M>1) and keeps the three coop4 GEMVs at decode (M=1, where the fused 64x64 tile would waste 63/64 rows). Implementation - New standalone kernel `q4gsw_linear_gemm_qkv_fused.wgsl` (+ generated `_wgsl.h`); build-time pattern detection and an M-gated dispatch/resize hook in `WebGPUGraph::build`. - q/k/v are repointed to fresh scratch buffers so the memory planner's reuse-aliasing cannot overlap the simultaneous fused write with a still-live input. - The fused dispatch owns its one-off compute pipeline (released in the destructor), mirroring `q4gsw_linear_impl`. Constraints Auto-applied when the pattern matches (no flag): the detection gates on shader-f16, a 256-thread workgroup, and the Llama-3.2 GQA projection shape {2048, 512, 512} with no bias, so any graph without that triple gets the byte-identical baseline lowering. Co-authored-with: Claude Code. ghstack-source-id: 405251962 @exported-using-ghexport Differential Revision: [D111489253](https://our.internmc.facebook.com/intern/diff/D111489253/) --- backends/webgpu/runtime/WebGPUGraph.cpp | 497 ++++++++++++++++++ backends/webgpu/runtime/WebGPUGraph.h | 28 + .../q4gsw_linear_gemm_qkv_fused.wgsl | 104 ++++ .../q4gsw_linear_gemm_qkv_fused_wgsl.h | 128 +++++ 4 files changed, 757 insertions(+) create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl create mode 100644 backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index d9ab6467882..94288c8404e 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -16,6 +17,7 @@ #include #include +#include #include #include #include @@ -75,6 +77,20 @@ int normalize_dim(int dim, int rank, const char* op) { return dim; } +// Uniform layout matching the fused-QKV WGSL Params struct (16B-aligned, 32B); +// identical to QuantizedLinear.cpp's Q4gswParams (kept local to this TU). +struct QkvFusedParams { + uint32_t M; + uint32_t N; + uint32_t K; + uint32_t K_packed; + uint32_t group_size; + uint32_t padded_N; + uint32_t has_bias; + uint32_t _pad; +}; +static_assert(sizeof(QkvFusedParams) == 32, "QkvFusedParams must be 32 bytes"); + } // namespace WebGPUGraph::WebGPUGraph() = default; @@ -685,6 +701,175 @@ void WebGPUGraph::build( // Phase 3: Build operator dispatch chain const auto* chain = graph->chain(); + + // QKV-concat fusion detection (auto-applied graph pass, no flag): the maps + // stay empty when no q/k/v triple matches -> the Phase-3 loop below runs + // verbatim. Find each attention q/k/v triple: EXACTLY 3 et_vk.linear_q4gsw + // ops sharing args[0] (the same input activation), in chain order q,k,v with + // the N-pattern {2048,512,512}, on the steel route (K%16==0), + // group_size%16==0, no bias. The fused kernel needs shader-f16 + a 256-thread + // WG, so gate on those (else leave the triple to the normal per-linear + // handlers). qkv_fused_skip holds all 3 op indices; qkv_anchor maps the FIRST + // op index -> its group, so the fused dispatch is emitted IN-PLACE at the + // anchor (correct execution order). + std::vector qkv_groups; + std::unordered_map + qkv_first; // first triple op -> group (repoint buffers) + std::unordered_map + qkv_last; // last triple op -> group (emit fused) + std::unordered_map + qkv_member; // any triple op -> group (record dispatch) + if (chain) { + bool device_ok = false; + { + WGPULimits limits = {}; + const bool have = + wgpuDeviceGetLimits(device_, &limits) == WGPUStatus_Success; + bool f16 = false; + if (auto* ctx = get_default_webgpu_context()) { + f16 = ctx->shader_f16_supported; + } + device_ok = + have && f16 && limits.maxComputeInvocationsPerWorkgroup >= 256u; + } + if (device_ok) { + // Group linear_q4gsw op indices by input id, preserving chain order. + std::unordered_map> by_input; + std::vector input_order; + for (unsigned i = 0; i < chain->size(); i++) { + const auto* oc = chain->Get(i); + if (oc->name()->str() != "et_vk.linear_q4gsw.default") { + continue; + } + const auto* a = oc->args(); + if (!a || a->size() < 6) { + continue; + } + const int inp = static_cast(a->Get(0)); + if (by_input.find(inp) == by_input.end()) { + input_order.push_back(inp); + } + by_input[inp].push_back(i); + } + auto op_arg = [&](unsigned oi, unsigned j) { + return static_cast(chain->Get(oi)->args()->Get(j)); + }; + for (int inp : input_order) { + const auto& ops = by_input[inp]; + if (ops.size() != 3) { + continue; // gate+up is a 2-group; o/down/lm_head are 1 each. + } + // args: [in, weight, scales, group_size, bias, out]. + const int wq = op_arg(ops[0], 1), sqid = op_arg(ops[0], 2), + bq = op_arg(ops[0], 4), oq = op_arg(ops[0], 5); + const int wk = op_arg(ops[1], 1), skid = op_arg(ops[1], 2), + bk = op_arg(ops[1], 4), ok = op_arg(ops[1], 5); + const int wv = op_arg(ops[2], 1), svid = op_arg(ops[2], 2), + bv = op_arg(ops[2], 4), ov = op_arg(ops[2], 5); + const int gsid = op_arg(ops[0], 3); + if (op_arg(ops[1], 3) != gsid || op_arg(ops[2], 3) != gsid) { + continue; // all 3 must share the group_size scalar. + } + if (get_value_type(bq) == ValueType::Tensor || + get_value_type(bk) == ValueType::Tensor || + get_value_type(bv) == ValueType::Tensor) { + continue; // fused kernel path assumes has_bias == 0. + } + const auto& twq = tensors_[wq]; + const auto& twk = tensors_[wk]; + const auto& twv = tensors_[wv]; + if (twq.dims.size() != 2 || twk.dims.size() != 2 || + twv.dims.size() != 2) { + continue; + } + const uint32_t Nq = static_cast(twq.dims[0]); + const uint32_t Nk = static_cast(twk.dims[0]); + const uint32_t Nv = static_cast(twv.dims[0]); + if (Nq != 2048u || Nk != 512u || Nv != 512u) { + continue; // kernel hardcodes N_Q=2048, N_KV=512 (Llama-3.2 GQA). + } + const uint32_t K_packed = static_cast(twq.dims[1]); + if (static_cast(twk.dims[1]) != K_packed || + static_cast(twv.dims[1]) != K_packed) { + continue; + } + const auto& tin = tensors_[inp]; + if (tin.dims.empty()) { + continue; + } + const uint32_t K = static_cast(tin.dims.back()); + if (K == 0 || K % 16u != 0u || K_packed != (K + 1u) / 2u) { + continue; // steel route stages a full BK=16 K-tile with no K-mask. + } + if (get_value_type(gsid) != ValueType::Int) { + continue; + } + const int64_t gsv = get_int(gsid); + if (gsv <= 0 || static_cast(gsv) % 16u != 0u) { + continue; // hoisted scale must be constant across the BK tile. + } + const uint32_t gs = static_cast(gsv); + const auto& tsq = tensors_[sqid]; + const auto& tsk = tensors_[skid]; + const auto& tsv = tensors_[svid]; + if (tsq.dims.size() != 2 || tsk.dims.size() != 2 || + tsv.dims.size() != 2) { + continue; + } + const uint32_t num_groups = static_cast(tsq.dims[0]); + if (static_cast(tsk.dims[0]) != num_groups || + static_cast(tsv.dims[0]) != num_groups) { + continue; + } + const uint32_t pNq = static_cast(tsq.dims[1]); + const uint32_t pNk = static_cast(tsk.dims[1]); + const uint32_t pNv = static_cast(tsv.dims[1]); + if (pNq < Nq || pNk < Nk || pNv < Nv || + num_groups < (K + gs - 1u) / gs) { + continue; + } + // All source + destination buffers must be live (Phase 1/2 allocated). + if (!twq.buffer || !twk.buffer || !twv.buffer || !tsq.buffer || + !tsk.buffer || !tsv.buffer || !tin.buffer || !tensors_[oq].buffer || + !tensors_[ok].buffer || !tensors_[ov].buffer) { + continue; + } + + QkvFusionGroup grp; + grp.input_id = inp; + grp.out_q = oq; + grp.out_k = ok; + grp.out_v = ov; + grp.weight_q = wq; + grp.weight_k = wk; + grp.weight_v = wv; + grp.scales_q = sqid; + grp.scales_k = skid; + grp.scales_v = svid; + grp.Nq = Nq; + grp.Nk = Nk; + grp.Nv = Nv; + grp.K = K; + grp.K_packed = K_packed; + grp.group_size = gs; + grp.num_groups = num_groups; + grp.padded_N_q = pNq; + grp.padded_N_k = pNk; + grp.padded_N_v = pNv; + grp.op_idx[0] = ops[0]; + grp.op_idx[1] = ops[1]; + grp.op_idx[2] = ops[2]; + const size_t gidx = qkv_groups.size(); + qkv_groups.push_back(grp); + qkv_first[ops[0]] = gidx; + qkv_last[ops[2]] = gidx; + qkv_member[ops[0]] = gidx; + qkv_member[ops[1]] = gidx; + qkv_member[ops[2]] = gidx; + } + } + } + if (chain) { for (unsigned i = 0; i < chain->size(); i++) { const auto* op_call = chain->Get(i); @@ -702,7 +887,55 @@ void WebGPUGraph::build( } } + // QKV fusion (M-gated): keep the 3 separate q/k/v linears AND add a fused + // multi-output GEMM; the fused resize hook selects by LIVE M (prefill M>1 + // -> fused runs, the 3 zeroed; decode M==1 -> the 3 coop4 GEMVs run, + // fused zeroed -- the fused 64x64 tile is ~4x slower than coop4 at M=1). + // At the FIRST triple op, repoint the 3 outputs to FRESH distinct + // buffers: the planner reuse-aliases q/k/v (each dies right after RoPE), + // which is fatal for a simultaneous fused write, so BOTH paths use + // non-aliased storage. All maps empty when no triple matches (verbatim + // path). + { + auto fit = qkv_first.find(i); + if (fit != qkv_first.end()) { + const auto& g = qkv_groups[fit->second]; + tensors_[g.out_q].buffer = + create_scratch_buffer(tensors_[g.out_q].nbytes); + tensors_[g.out_k].buffer = + create_scratch_buffer(tensors_[g.out_k].nbytes); + tensors_[g.out_v].buffer = + create_scratch_buffer(tensors_[g.out_v].nbytes); + } + } + webgpu_operator_registry().get_op_fn(op_name)(*this, args); + + { + auto mit = qkv_member.find(i); + if (mit != qkv_member.end()) { + QkvFusionGroup& g = qkv_groups[mit->second]; + const size_t di = num_dispatches() - 1; // this linear's dispatch + if (i == g.op_idx[0]) { + g.sep_dispatch[0] = di; + // Emit the fused dispatch RIGHT AFTER the q-linear (the anchor) so + // at M>1 it writes q/k/v BEFORE any consumer. q/k/v may be + // interleaved with rope in the chain, so emitting it at the LAST + // triple op would let a consumer (rope-q) read still-unwritten + // fresh_q -> garbage. + add_qkv_fused_dispatch(g); + } else if (i == g.op_idx[1]) { + g.sep_dispatch[1] = di; + } else { + g.sep_dispatch[2] = di; + } + } + auto lit = qkv_last.find(i); + if (lit != qkv_last.end()) { + // All 3 sep dispatch indices + the fused index are now known. + add_qkv_fused_hook(qkv_groups[lit->second]); + } + } } } @@ -806,6 +1039,270 @@ WGPUBindGroupLayout WebGPUGraph::get_or_create_bgl( return bgl; } +void WebGPUGraph::add_qkv_fused_dispatch(QkvFusionGroup& g) { + const uint32_t N = g.Nq + g.Nk + g.Nv; // fused output width (3072) + + const auto& in = tensors_[g.input_id]; + const auto& out_q = tensors_[g.out_q]; + const auto& out_k = tensors_[g.out_k]; + const auto& out_v = tensors_[g.out_v]; + const auto& wq = tensors_[g.weight_q]; + const auto& wk = tensors_[g.weight_k]; + const auto& wv = tensors_[g.weight_v]; + const auto& sq = tensors_[g.scales_q]; + const auto& sk = tensors_[g.scales_k]; + const auto& sv = tensors_[g.scales_v]; + + // Buffers were repointed to FRESH distinct slots at the first triple op (see + // the build() op-walk), so out_q/k/v no longer alias. Live M from the shared + // input. + uint64_t in_numel = 1; + for (int64_t d : in.dims) { + in_numel *= static_cast(d); + } + const uint32_t M = static_cast(in_numel / g.K); + + // Fused weight [N, K_packed]: a byte-contiguous row-stack of Wq;Wk;Wv (q4gsw + // packs each output row independently along a shared K_packed, so stacking + // along N is a flat append -- bit-exact). Fused scales [num_groups, N]: a + // strided PER-GROUP-ROW gather (dest row stride N != the per-linear source + // strides padded_N_{q,k,v}), NOT a flat append. Both dtor-freed via scratch. + const uint64_t kp = static_cast(g.K_packed); // packed bytes / row + const uint64_t fs = sizeof(float); + WGPUBuffer fused_weight = create_scratch_buffer(static_cast(N) * kp); + WGPUBuffer fused_scales = create_scratch_buffer( + static_cast(g.num_groups) * N * sizeof(float)); + + // Sources are direct constants materialized in Phase 1 (or prepack outputs + // materialized earlier in Phase 3); all writes are already enqueued on + // queue_, so this build-time copy sees the materialized bytes. + WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device_, nullptr); + wgpuCommandEncoderCopyBufferToBuffer( + enc, wq.buffer, 0, fused_weight, 0, static_cast(g.Nq) * kp); + wgpuCommandEncoderCopyBufferToBuffer( + enc, + wk.buffer, + 0, + fused_weight, + static_cast(g.Nq) * kp, + static_cast(g.Nk) * kp); + wgpuCommandEncoderCopyBufferToBuffer( + enc, + wv.buffer, + 0, + fused_weight, + static_cast(g.Nq + g.Nk) * kp, + static_cast(g.Nv) * kp); + for (uint32_t grp = 0; grp < g.num_groups; grp++) { + const uint64_t dst_row = static_cast(grp) * N * fs; + wgpuCommandEncoderCopyBufferToBuffer( + enc, + sq.buffer, + static_cast(grp) * g.padded_N_q * fs, + fused_scales, + dst_row, + static_cast(g.Nq) * fs); + wgpuCommandEncoderCopyBufferToBuffer( + enc, + sk.buffer, + static_cast(grp) * g.padded_N_k * fs, + fused_scales, + dst_row + static_cast(g.Nq) * fs, + static_cast(g.Nk) * fs); + wgpuCommandEncoderCopyBufferToBuffer( + enc, + sv.buffer, + static_cast(grp) * g.padded_N_v * fs, + fused_scales, + dst_row + static_cast(g.Nq + g.Nk) * fs, + static_cast(g.Nv) * fs); + } + WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, nullptr); + wgpuQueueSubmit(queue_, 1, &cmd); + wgpuCommandBufferRelease(cmd); + wgpuCommandEncoderRelease(enc); + + // Params UBO (owned; rewritten by the resize hook). padded_N == N (fused + // scales row stride); has_bias == 0 (attention q/k/v are bias-less). + QkvFusedParams params = {}; + params.M = M; + params.N = N; + params.K = g.K; + params.K_packed = g.K_packed; + params.group_size = g.group_size; + params.padded_N = N; + params.has_bias = 0; + WGPUBufferDescriptor u_desc = {}; + u_desc.size = sizeof(QkvFusedParams); + u_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + u_desc.mappedAtCreation = true; + WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device_, &u_desc); + std::memcpy( + wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(QkvFusedParams)), + ¶ms, + sizeof(QkvFusedParams)); + wgpuBufferUnmap(uniform_buffer); + add_uniform_buffer_bytes(sizeof(QkvFusedParams)); + + // 4-byte dummy for the fixed bias binding (has_bias == 0). + WGPUBuffer bias_dummy = create_scratch_buffer(4); + + // Bespoke 8-binding layout: 3 rw-storage outputs + 4 ro-storage + 1 uniform. + // One-off shader/bgl/pipeline owned by the dispatch (matches + // q4gsw_linear_impl). + WGPUBindGroupLayoutEntry entries[8] = {}; + for (uint32_t i = 0; i < 3; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_Storage; + } + for (uint32_t i = 3; i < 7; i++) { + entries[i].binding = i; + entries[i].visibility = WGPUShaderStage_Compute; + entries[i].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + } + entries[7].binding = 7; + entries[7].visibility = WGPUShaderStage_Compute; + entries[7].buffer.type = WGPUBufferBindingType_Uniform; + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 8; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kQ4gswLinearGemmQkvFusedWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device_, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device_, &pipeline_desc); + + WGPUBindGroupEntry bg[8] = {}; + bg[0].binding = 0; + bg[0].buffer = out_q.buffer; + bg[0].size = out_q.nbytes; + bg[1].binding = 1; + bg[1].buffer = out_k.buffer; + bg[1].size = out_k.nbytes; + bg[2].binding = 2; + bg[2].buffer = out_v.buffer; + bg[2].size = out_v.nbytes; + bg[3].binding = 3; + bg[3].buffer = in.buffer; + bg[3].size = in.nbytes; + bg[4].binding = 4; + bg[4].buffer = fused_weight; + bg[4].size = static_cast(N) * kp; + bg[5].binding = 5; + bg[5].buffer = fused_scales; + bg[5].size = static_cast(g.num_groups) * N * fs; + bg[6].binding = 6; + bg[6].buffer = bias_dummy; + bg[6].size = 4; + bg[7].binding = 7; + bg[7].buffer = uniform_buffer; + bg[7].size = sizeof(QkvFusedParams); + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 8; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device_, &bg_desc); + + // 1D dispatch over ceil(M/BM) * ceil(N/BN) tiles (BM=BN=64), matching the + // kernel's nbN = ceil(N/64) tile decode (NOT grid-strided). + const uint32_t nbN = (N + 63u) / 64u; + const uint32_t nbM = (M + 63u) / 64u; + const size_t fused_idx = + add_dispatch({pipeline, bind_group, nbN * nbM, "linear_q4gsw_qkv_fused"}); + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + own_uniform_buffer(uniform_buffer); + g.fused_dispatch = fused_idx; // consumed by add_qkv_fused_hook at the last op + g.fused_params = uniform_buffer; +} + +// M-gate coordinator: registered at the LAST triple op (all dispatch indices +// known). Prefill (M>1): run the fused GEMM, zero the 3 separate linears. +// Decode (M==1): zero the fused, leave the 3 coop4 GEMVs (their own hooks set +// the decode wg) -- the fused 64x64 tile wastes 63/64 rows at M=1. Recomputes +// live M + the 3 output cur_dims + fused params. Inert on a static graph; a +// workgroup_count of 0 = no-op. +void WebGPUGraph::add_qkv_fused_hook(const QkvFusionGroup& g) { + const int input_id = g.input_id, out_q_id = g.out_q, out_k_id = g.out_k, + out_v_id = g.out_v; + const uint32_t K = g.K, Kp = g.K_packed, gs = g.group_size, Nq = g.Nq, + Nk = g.Nk, Nv = g.Nv, Nf = g.Nq + g.Nk + g.Nv; + const size_t fused_idx = g.fused_dispatch, sep0 = g.sep_dispatch[0], + sep1 = g.sep_dispatch[1], sep2 = g.sep_dispatch[2]; + WGPUBuffer params_buf = g.fused_params; + add_tensor_resize_hook( + input_id, + [input_id, + out_q_id, + out_k_id, + out_v_id, + K, + Kp, + gs, + Nq, + Nk, + Nv, + Nf, + fused_idx, + sep0, + sep1, + sep2, + params_buf](WebGPUGraph& gr) { + const auto& d = gr.cur_dims(input_id); + uint64_t numel = 1; + for (int64_t v : d) { + numel *= static_cast(v); + } + const uint32_t m = static_cast(numel / K); + std::vector oq = d; + oq.back() = static_cast(Nq); + std::vector ok = d; + ok.back() = static_cast(Nk); + std::vector ov = d; + ov.back() = static_cast(Nv); + gr.set_cur_dims(out_q_id, oq); + gr.set_cur_dims(out_k_id, ok); + gr.set_cur_dims(out_v_id, ov); + QkvFusedParams p = {}; + p.M = m; + p.N = Nf; + p.K = K; + p.K_packed = Kp; + p.group_size = gs; + p.padded_N = Nf; + p.has_bias = 0; + wgpuQueueWriteBuffer(gr.queue(), params_buf, 0, &p, sizeof(p)); + if (m > 1u) { + const uint32_t nbN2 = (Nf + 63u) / 64u; + const uint32_t nbM2 = (m + 63u) / 64u; + gr.dispatch_at(fused_idx).workgroup_count_x = nbN2 * nbM2; + gr.dispatch_at(sep0).workgroup_count_x = 0u; + gr.dispatch_at(sep1).workgroup_count_x = 0u; + gr.dispatch_at(sep2).workgroup_count_x = 0u; + } else { + gr.dispatch_at(fused_idx).workgroup_count_x = 0u; + } + }); +} + void WebGPUGraph::copy_inputs(const std::vector& inputs) { for (size_t i = 0; i < inputs.size() && i < input_ids_.size(); i++) { const InputData& in = inputs[i]; diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 66f0e401de5..d7b81b7b532 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -450,6 +450,34 @@ class WebGPUGraph { std::unordered_map bgl_cache_; size_t uniform_buffer_bytes_ = 0; + + // QKV-concat fusion: one detected attention q/k/v linear + // triple sharing an input activation (value ids + shapes), fused in build() + // into a single multi-output q4gsw GEMM that scatter-writes q/k/v. Only used + // during build(); inert (never populated) when no q/k/v triple matches. + struct QkvFusionGroup { + int input_id = -1; + int out_q = -1, out_k = -1, out_v = -1; + int weight_q = -1, weight_k = -1, weight_v = -1; + int scales_q = -1, scales_k = -1, scales_v = -1; + uint32_t Nq = 0, Nk = 0, Nv = 0; // 2048, 512, 512 + uint32_t K = 0, K_packed = 0, group_size = 0, num_groups = 0; + uint32_t padded_N_q = 0, padded_N_k = 0, padded_N_v = 0; + unsigned op_idx[3] = {0, 0, 0}; // the 3 q/k/v linear op-chain indices + size_t sep_dispatch[3] = { + 0, + 0, + 0}; // their dispatch indices (filled in build()) + size_t fused_dispatch = 0; // the fused GEMM dispatch index + WGPUBuffer fused_params = + nullptr; // the fused params UBO (rewritten by the hook) + }; + // Concat the 3 packed weights (row-stack) + scales (strided gather) into + // fused buffers, then record ONE fused-GEMM dispatch (bespoke 8-binding + // layout) that writes the 3 original q/k/v output buffers, plus a 3-output + // resize hook. + void add_qkv_fused_dispatch(QkvFusionGroup& g); + void add_qkv_fused_hook(const QkvFusionGroup& g); }; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl new file mode 100644 index 00000000000..c89924a67c1 --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused.wgsl @@ -0,0 +1,104 @@ +enable f16; +// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM +// (vec4 activation load) that scatter-writes each output column range to a SEPARATE buffer -- +// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches; +// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned +// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE: +// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the +// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead +// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV). +@group(0) @binding(0) var t_out_q: array; +@group(0) @binding(1) var t_out_k: array; +@group(0) @binding(2) var t_out_v: array; +@group(0) @binding(3) var t_input: array>; +@group(0) @binding(4) var t_weight: array; +@group(0) @binding(5) var t_scales: array; +@group(0) @binding(6) var t_bias: array; +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(7) var params: Params; +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; +const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u; +var As: array; +var Bs: array; +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let nbN = (params.N + BN - 1u) / BN; + let bx = wid.x % nbN; + let by = wid.x / nbN; + let row0 = by * BM; + let col0 = bx * BN; + let tid = lid.y * 16u + lid.x; + var acc: array, 4>; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; } + } + let ar = tid / 4u; + let ac = (tid % 4u) * 4u; + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y); + As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w); + } else { + As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; + As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + } + if (tid < BN) { + let c = tid; + let n = col0 + c; + if (n < params.N) { + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < BK; k = k + 1u) { + var a: array; + var bvec: array; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + } + } + workgroupBarrier(); + k0 = k0 + BK; + } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { + let r = row0 + lid.y * 4u + m; + let c = col0 + lid.x * 4u + n; // global fused column [0, 3072) + if (r < params.M && c < params.N) { + var val = f32(acc[m][n]); + if (params.has_bias != 0u) { val = val + t_bias[c]; } + if (c < N_Q) { t_out_q[r * N_Q + c] = val; } + else if (c < N_QK) { t_out_k[r * N_KV + (c - N_Q)] = val; } + else { t_out_v[r * N_KV + (c - N_QK)] = val; } + } + } + } +} diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h new file mode 100644 index 00000000000..93243698a3b --- /dev/null +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_qkv_fused_wgsl.h @@ -0,0 +1,128 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from q4gsw_linear_gemm_qkv_fused.wgsl - DO NOT EDIT. +// wgsl-sha256: 93e127e8ee4609d846015c8b75a600a29502e19a92bdf3a08e3429635f834085 +inline constexpr const char* kQ4gswLinearGemmQkvFusedWGSL = R"( +enable f16; +// Fused QKV q4gsw GEMM (Llama attention projections): one [M, N=3072] pwdq + f16-accumulate GEMM +// (vec4 activation load) that scatter-writes each output column range to a SEPARATE buffer -- +// c<2048 -> q, [2048,2560) -> k, [2560,3072) -> v. Replaces the 3 separate q/k/v linear dispatches; +// fixes the N=512 K/V occupancy starvation (16 WGs -> 96 WGs at M~128). Boundaries are 64-tile-aligned +// so each 64-col tile maps to exactly one output (uniform branch per workgroup). Per-output ROW STRIDE: +// q=2048, k=v=512. BIT-EXACT to 3 separate pwdqf16acc linears (fusing along N does not change the +// per-column K-accumulation order). Validated on Canary M4 Pro: correct (maxRel ~1e-3), scatter overhead +// 1.02x (free), concat win 1.63x on the QKV block. Boundaries hardcoded for Llama-3.2-1B GQA (32Q/8KV). +@group(0) @binding(0) var t_out_q: array; +@group(0) @binding(1) var t_out_k: array; +@group(0) @binding(2) var t_out_v: array; +@group(0) @binding(3) var t_input: array>; +@group(0) @binding(4) var t_weight: array; +@group(0) @binding(5) var t_scales: array; +@group(0) @binding(6) var t_bias: array; +struct Params { + M: u32, + N: u32, + K: u32, + K_packed: u32, + group_size: u32, + padded_N: u32, + has_bias: u32, + _pad: u32, +} +@group(0) @binding(7) var params: Params; +const BM: u32 = 64u; const BN: u32 = 64u; const BK: u32 = 16u; +const N_Q: u32 = 2048u; const N_QK: u32 = 2560u; const N_KV: u32 = 512u; +var As: array; +var Bs: array; +@compute @workgroup_size(16, 16) +fn main(@builtin(workgroup_id) wid: vec3, + @builtin(local_invocation_id) lid: vec3) { + let nbN = (params.N + BN - 1u) / BN; + let bx = wid.x % nbN; + let by = wid.x / nbN; + let row0 = by * BM; + let col0 = bx * BN; + let tid = lid.y * 16u + lid.x; + var acc: array, 4>; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = 0.0h; } + } + let ar = tid / 4u; + let ac = (tid % 4u) * 4u; + var k0: u32 = 0u; + loop { + if (k0 >= params.K) { break; } + let arow = row0 + ar; + if (arow < params.M) { + let base = arow * params.K + k0 + ac; + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = f16(av.x); As[ar * BK + ac + 1u] = f16(av.y); + As[ar * BK + ac + 2u] = f16(av.z); As[ar * BK + ac + 3u] = f16(av.w); + } else { + As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; + As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; + } + if (tid < BN) { + let c = tid; + let n = col0 + c; + if (n < params.N) { + let scale_row = (k0 / params.group_size) * params.padded_N; + let scale = f16(t_scales[scale_row + n]); + let base_word = n * (params.K_packed >> 2u) + (k0 >> 3u); + let w0 = t_weight[base_word]; + let w1 = t_weight[base_word + 1u]; + for (var br: u32 = 0u; br < BK; br = br + 1u) { + let word = select(w1, w0, br < 8u); + let nib = (word >> ((br & 7u) * 4u)) & 0x0Fu; + Bs[br * BN + c] = f16(i32(nib) - 8) * scale; + } + } else { + for (var br: u32 = 0u; br < BK; br = br + 1u) { Bs[br * BN + c] = 0.0h; } + } + } + workgroupBarrier(); + for (var k: u32 = 0u; k < BK; k = k + 1u) { + var a: array; + var bvec: array; + for (var m: u32 = 0u; m < 4u; m = m + 1u) { a[m] = As[(lid.y * 4u + m) * BK + k]; } + for (var n: u32 = 0u; n < 4u; n = n + 1u) { bvec[n] = Bs[k * BN + lid.x * 4u + n]; } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { acc[m][n] = fma(a[m], bvec[n], acc[m][n]); } + } + } + workgroupBarrier(); + k0 = k0 + BK; + } + for (var m: u32 = 0u; m < 4u; m = m + 1u) { + for (var n: u32 = 0u; n < 4u; n = n + 1u) { + let r = row0 + lid.y * 4u + m; + let c = col0 + lid.x * 4u + n; // global fused column [0, 3072) + if (r < params.M && c < params.N) { + var val = f32(acc[m][n]); + if (params.has_bias != 0u) { val = val + t_bias[c]; } + if (c < N_Q) { t_out_q[r * N_Q + c] = val; } + else if (c < N_QK) { t_out_k[r * N_KV + (c - N_Q)] = val; } + else { t_out_v[r * N_KV + (c - N_QK)] = val; } + } + } + } +} +)"; + +inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeX = 16; +inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeY = 16; +inline constexpr uint32_t kQ4gswLinearGemmQkvFusedWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From d6cf974a4f7b5960c5e6f5fda90266e3b21943b9 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 11:51:38 -0700 Subject: [PATCH 2/3] [ExecuTorch][WebGPU] Fuse SwiGLU sigmoid+muls in prefill Pull Request resolved: https://github.com/pytorch/executorch/pull/20881 Fold the SwiGLU `sigmoid` + two `mul`s into a single elementwise pass computing (gate * sigmoid(gate)) * up, cutting the MLP activation traffic from ~8 units to ~3 on Llama-3.2-1B. Auto-applied graph pass, bit-exact. Problem The SwiGLU MLP lowers `silu(gate)*up` as three memory-bound elementwise dispatches (`sigmoid` -> `mul` -> `mul`) over the [S, 8192] intermediate. Solution - Before: `sigmoid` + 2 `mul`s (~8 tensor-traffic units over [S, 8192]). - After: one `silu_mul_fused` pass computing `(g*sigmoid(g))*up` with the sigmoid in-register (~3 traffic units); the folded sigmoid + 1st-mul dispatches are dropped. Implementation - New standalone kernel `silu_mul_fused.wgsl` (+ generated `_wgsl.h`); build-time triple detection (requiring single-consumer sig/silu intermediates) + an in-place fused dispatch at the 2nd-mul anchor in `WebGPUGraph::build`. - gate is repointed to a private pooled scratch buffer at its producer op, and the fused output to another, so the memory planner's reuse-aliasing (which places up_proj's output onto gate's dead slot) cannot stomp gate before the fused dispatch reads it; both pooled buffers are released after last use for cross-layer reuse. Constraints Auto-applied to fp32 elementwise triples with identical dims and single-consumer sig/silu intermediates (no flag); a graph without a matching triple gets the byte-identical baseline lowering. Stacked on the QKV-concat fusion diff. Co-authored-with: Claude Code. ghstack-source-id: 405251982 @exported-using-ghexport Differential Revision: [D111489277](https://our.internmc.facebook.com/intern/diff/D111489277/) --- backends/webgpu/runtime/WebGPUGraph.cpp | 347 ++++++++++++++++++ backends/webgpu/runtime/WebGPUGraph.h | 9 + .../runtime/ops/mul/silu_mul_fused.wgsl | 24 ++ .../runtime/ops/mul/silu_mul_fused_wgsl.h | 48 +++ 4 files changed, 428 insertions(+) create mode 100644 backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl create mode 100644 backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 94288c8404e..2bb091d391e 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -8,6 +8,7 @@ #include #include +#include #include #include @@ -870,6 +871,178 @@ void WebGPUGraph::build( } } + // SwiGLU fusion detection (auto-applied graph pass, no flag): all sets stay + // empty when no SiLU-gate triple matches + // -> the Phase-3 loop below runs verbatim. Fold each + // SiLU-gate MLP triple sigmoid(g) -> mul(g,sig)=silu -> mul(silu,up)=out + // into ONE elementwise dispatch that computes sigmoid + silu in registers + // (gate + up read once, one output written): 8 traffic units -> 3. Bit-exact + // (same fp op order, and the sigmoid form matches sigmoid.wgsl). swiglu_skip + // holds the sigmoid + the 1st-mul op indices (their dispatches are dropped -- + // sig/silu become dead), and swiglu_anchor maps the 2nd-mul op (where out + + // up are both live) -> its group, so the fused dispatch is emitted IN-PLACE + // there (correct execution order). The sig/silu intermediates must be + // single-consumer (folding them can't strand a second reader). + std::vector> swiglu_groups; // {gate, up, out} + std::unordered_set swiglu_skip; // sigmoid + 1st-mul op indices + std::unordered_map swiglu_anchor; // 2nd-mul op idx -> group + // gate_proj op idx -> group: repoint gate to a PRIVATE pooled buffer there, + // before gate_proj is lowered (root-cause fix, see the detection guard + // below). + std::unordered_map swiglu_gate_acquire; + // out's last-use op idx -> group: release the pooled fused-output buffer + // after its final consumer's dispatch is built (pool recycling for the + // +memory). + std::unordered_map swiglu_out_release; + if (chain) { + // Consumer count: appearances as a NON-output (non-last) arg, ValueList- + // expanded. sig/silu are safe to fold only if each is consumed exactly once + // (by the 1st/2nd mul respectively) and nowhere else. + std::vector consumer_cnt(num_vals, 0); + for (unsigned i = 0; i < chain->size(); i++) { + const auto* a = chain->Get(i)->args(); + if (!a || a->size() == 0) { + continue; + } + for (unsigned j = 0; j + 1 < a->size(); j++) { + const int id = static_cast(a->Get(j)); + if (id < 0 || id >= num_vals) { + continue; + } + consumer_cnt[id]++; + if (value_types_[id] == ValueType::ValueList) { + for (int m : value_lists_[id]) { + if (m >= 0 && m < num_vals) { + consumer_cnt[m]++; + } + } + } + } + } + // Producer (op that writes each value = its last arg) + last-use (last op + // referencing a value in ANY arg, ValueList-expanded). Used to (a) find + // gate's producer op so gate can be repointed to a private buffer before it + // is written, and (b) find out's last consumer so the pooled out buffer is + // released only after it is truly dead. + std::vector producer(num_vals, -1); + std::vector last_use(num_vals, -1); + for (unsigned i = 0; i < chain->size(); i++) { + const auto* a = chain->Get(i)->args(); + if (!a || a->size() == 0) { + continue; + } + for (unsigned j = 0; j < a->size(); j++) { + const int id = static_cast(a->Get(j)); + if (id < 0 || id >= num_vals) { + continue; + } + last_use[id] = static_cast(i); + if (value_types_[id] == ValueType::ValueList) { + for (int m : value_lists_[id]) { + if (m >= 0 && m < num_vals) { + last_use[m] = static_cast(i); + } + } + } + } + const int outv = static_cast(a->Get(a->size() - 1)); + if (outv >= 0 && outv < num_vals) { + producer[outv] = static_cast(i); + } + } + struct SigInfo { + unsigned op; + int g_in; + }; + struct Mul1Info { + unsigned op; + int g_in; + unsigned sig_op; + int sig_out; + }; + std::unordered_map sigmoid_by_out; // sig_out -> {op, g} + std::unordered_map + mul1_by_out; // silu_out -> {op, g, sig...} + for (unsigned i = 0; i < chain->size(); i++) { + const auto* oc = chain->Get(i); + const std::string nm = oc->name()->str(); + const auto* a = oc->args(); + if (!a) { + continue; + } + if (nm == "aten.sigmoid.default" && a->size() >= 2) { + sigmoid_by_out[static_cast(a->Get(1))] = { + i, static_cast(a->Get(0))}; + continue; + } + if (nm != "aten.mul.Tensor" || a->size() < 3) { + continue; + } + const int x = static_cast(a->Get(0)); + const int y = static_cast(a->Get(1)); + const int out = static_cast(a->Get(2)); + // 2nd mul? one operand is a recorded silu (a 1st-mul output). + int silu = -1, up = -1; + if (mul1_by_out.count(x) != 0) { + silu = x; + up = y; + } else if (mul1_by_out.count(y) != 0) { + silu = y; + up = x; + } + if (silu >= 0) { + const Mul1Info& m1 = mul1_by_out[silu]; + const int g = m1.g_in; + const bool tensors_ok = g >= 0 && up >= 0 && out >= 0 && + get_value_type(g) == ValueType::Tensor && + get_value_type(up) == ValueType::Tensor && + get_value_type(out) == ValueType::Tensor; + if (tensors_ok) { + const auto& tg = tensors_[g]; + const auto& tu = tensors_[up]; + const auto& to = tensors_[out]; + // Elementwise, all fp32, identical dims (so the fused output's live + // dims + // == gate's on resize), live buffers, and single-consumer + // intermediates. gate must be consumed by EXACTLY the sigmoid + + // 1st-mul (consumer_cnt + // == 2) so it is dead once the fused dispatch reads it, and have a + // real producer op (so it can be repointed before it is written). + if (tg.buffer && tu.buffer && to.buffer && tg.elem_size == 4 && + tu.elem_size == 4 && to.elem_size == 4 && tg.dims == tu.dims && + tg.dims == to.dims && tg.nbytes == tu.nbytes && + tg.nbytes == to.nbytes && consumer_cnt[m1.sig_out] == 1 && + consumer_cnt[silu] == 1 && consumer_cnt[g] == 2 && + producer[g] >= 0) { + const size_t gidx = swiglu_groups.size(); + swiglu_groups.push_back({g, up, out}); + swiglu_skip.insert(m1.sig_op); + swiglu_skip.insert(m1.op); + swiglu_anchor[i] = gidx; + // The serialized memory planner reuse-aliases up onto gate's slot + // (gate dies at the 1st mul, up_proj is emitted between the muls), + // so up_proj would stomp gate's buffer before the fused dispatch + // (at this 2nd-mul anchor) reads it. Give gate a private buffer at + // its producer op; release it right after the fused read. out is + // likewise pooled + released after last_use[out] (its final + // consumer, e.g. down_proj). + swiglu_gate_acquire[static_cast(producer[g])] = gidx; + swiglu_out_release[static_cast(last_use[out])] = gidx; + } + } + continue; // a 2nd-mul is never also a 1st-mul + } + // 1st mul? inputs are exactly {sigmoid input g, sigmoid output sig}. + auto sx = sigmoid_by_out.find(x); + auto sy = sigmoid_by_out.find(y); + if (sx != sigmoid_by_out.end() && sx->second.g_in == y) { + mul1_by_out[out] = {i, y, sx->second.op, x}; + } else if (sy != sigmoid_by_out.end() && sy->second.g_in == x) { + mul1_by_out[out] = {i, x, sy->second.op, y}; + } + } + } + if (chain) { for (unsigned i = 0; i < chain->size(); i++) { const auto* op_call = chain->Get(i); @@ -887,6 +1060,34 @@ void WebGPUGraph::build( } } + // SwiGLU fusion. At gate_proj repoint gate to a private pooled buffer (so + // up_proj can't stomp its planner-aliased slot before the fused reads + // it); drop the folded sigmoid + 1st-mul; at the 2nd-mul anchor emit ONE + // fused silu*up dispatch then release gate. Sets empty when + // no SwiGLU triple matched (verbatim path). + { + auto ga = swiglu_gate_acquire.find(i); + if (ga != swiglu_gate_acquire.end()) { + // Repoint BEFORE gate_proj is lowered below, so it writes the private + // buffer. gate_proj falls through to normal lowering (not + // skip/anchor). + const int gate_id = swiglu_groups[ga->second][0]; + tensors_[gate_id].buffer = acquire_scratch(tensors_[gate_id].nbytes); + } + if (swiglu_skip.count(i) != 0) { + continue; // sigmoid / 1st-mul: folded into the fused dispatch + } + auto sa = swiglu_anchor.find(i); + if (sa != swiglu_anchor.end()) { + const auto& grp = swiglu_groups[sa->second]; + add_swiglu_fused_dispatch(grp[0], grp[1], grp[2]); + // gate is dead once the fused dispatch has read it + // (consumer_cnt[g]==2, both folded) -> return its buffer to the pool + // for the next layer. + release_scratch(tensors_[grp[0]].buffer); + continue; + } + } // QKV fusion (M-gated): keep the 3 separate q/k/v linears AND add a fused // multi-output GEMM; the fused resize hook selects by LIVE M (prefill M>1 // -> fused runs, the 3 zeroed; decode M==1 -> the 3 coop4 GEMVs run, @@ -936,6 +1137,14 @@ void WebGPUGraph::build( add_qkv_fused_hook(qkv_groups[lit->second]); } } + // SwiGLU: this op is out's last consumer (its dispatch just captured + // out's buffer above) -> return the pooled fused-output buffer for reuse. + { + auto orl = swiglu_out_release.find(i); + if (orl != swiglu_out_release.end()) { + release_scratch(tensors_[swiglu_groups[orl->second][2]].buffer); + } + } } } @@ -1234,6 +1443,144 @@ void WebGPUGraph::add_qkv_fused_dispatch(QkvFusionGroup& g) { g.fused_params = uniform_buffer; } +namespace { +// Uniform layout matching silu_mul_fused.wgsl Params (16B-aligned). +struct SiluMulParams { + uint32_t num_elements; + uint32_t _pad[3]; +}; +} // namespace + +// SwiGLU fusion: emit ONE elementwise dispatch computing +// out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2 muls. +// Elementwise (no M-gate: identical at decode and prefill). +void WebGPUGraph::add_swiglu_fused_dispatch( + int gate_id, + int up_id, + int out_id) { + // Private distinct output buffer (mirrors the QKV aliasing guard): the + // planner reuse-aliases `out` onto a dead slot (e.g. sigmoid's), which + // without this would bind the same buffer as ro `gate` AND rw `output` -> + // Dawn writable-aliasing / all-zeros. Repoint BEFORE the bind group so it + // captures the private buffer; downstream consumers (lowered later) also see + // it. gate is still in_use here (released only after this call), so + // acquire_scratch hands out a DISTINCT slot. Pooled (not dedicated): the + // caller releases it after out's last consumer, so N layers recycle a small + // constant of buffers. tensor_mem_obj_ids_[out] stays + // >= 0, so the dtor never per-tensor-frees it (scratch_pool_ owns it). + tensors_[out_id].buffer = acquire_scratch(tensors_[out_id].nbytes); + + const auto& gate = tensors_[gate_id]; + const auto& up = tensors_[up_id]; + const auto& out = tensors_[out_id]; + const uint32_t num_elements = + static_cast(out.nbytes / sizeof(float)); + + SiluMulParams params = {num_elements, {0u, 0u, 0u}}; + WGPUBufferDescriptor u_desc = {}; + u_desc.size = sizeof(SiluMulParams); + u_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst; + u_desc.mappedAtCreation = true; + WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device_, &u_desc); + std::memcpy( + wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(SiluMulParams)), + ¶ms, + sizeof(SiluMulParams)); + wgpuBufferUnmap(uniform_buffer); + add_uniform_buffer_bytes(sizeof(SiluMulParams)); + + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {kSiluMulFusedWGSL, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device_, &shader_desc); + + // BGL: gate (ro) 0, up (ro) 1, output (rw) 2, params (uniform) 3. + WGPUBindGroupLayoutEntry entries[4] = {}; + entries[0].binding = 0; + entries[0].visibility = WGPUShaderStage_Compute; + entries[0].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[1].binding = 1; + entries[1].visibility = WGPUShaderStage_Compute; + entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage; + entries[2].binding = 2; + entries[2].visibility = WGPUShaderStage_Compute; + entries[2].buffer.type = WGPUBufferBindingType_Storage; + entries[3].binding = 3; + entries[3].visibility = WGPUShaderStage_Compute; + entries[3].buffer.type = WGPUBufferBindingType_Uniform; + WGPUBindGroupLayoutDescriptor bgl_desc = {}; + bgl_desc.entryCount = 4; + bgl_desc.entries = entries; + WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc); + + WGPUPipelineLayoutDescriptor pl_desc = {}; + pl_desc.bindGroupLayoutCount = 1; + pl_desc.bindGroupLayouts = &bgl; + WGPUPipelineLayout pipeline_layout = + wgpuDeviceCreatePipelineLayout(device_, &pl_desc); + + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device_, &pipeline_desc); + + WGPUBindGroupEntry bg[4] = {}; + bg[0].binding = 0; + bg[0].buffer = gate.buffer; + bg[0].size = gate.nbytes; + bg[1].binding = 1; + bg[1].buffer = up.buffer; + bg[1].size = up.nbytes; + bg[2].binding = 2; + bg[2].buffer = out.buffer; + bg[2].size = out.nbytes; + bg[3].binding = 3; + bg[3].buffer = uniform_buffer; + bg[3].size = sizeof(SiluMulParams); + WGPUBindGroupDescriptor bg_desc = {}; + bg_desc.layout = bgl; + bg_desc.entryCount = 4; + bg_desc.entries = bg; + WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device_, &bg_desc); + + const uint32_t wg = kSiluMulFusedWorkgroupSizeX; + const uint32_t workgroup_count = (num_elements + wg - 1) / wg; + if (workgroup_count > 65535) { + throw std::runtime_error( + "silu_mul_fused: workgroup count exceeds 65535 (1D dispatch limit)"); + } + add_dispatch({pipeline, bind_group, workgroup_count, "silu_mul_fused"}); + const size_t dispatch_idx = num_dispatches() - 1; + + wgpuShaderModuleRelease(shader); + wgpuBindGroupLayoutRelease(bgl); + wgpuPipelineLayoutRelease(pipeline_layout); + own_uniform_buffer(uniform_buffer); + + // Dynamic shapes: gate/up/out share dims, so out's live dims == gate's; + // recompute num_elements + dispatch from gate's live shape. Triggers on gate + // -- exactly the input the folded sigmoid's hook keyed on, so it is dirtied + // on every resize. + WGPUBuffer params_buf = uniform_buffer; + add_tensor_resize_hook( + gate_id, [gate_id, out_id, wg, dispatch_idx, params_buf](WebGPUGraph& g) { + const auto& d = g.cur_dims(gate_id); + g.set_cur_dims(out_id, d); + uint64_t numel = 1; + for (int64_t v : d) { + numel *= static_cast(v); + } + SiluMulParams p = {static_cast(numel), {0u, 0u, 0u}}; + wgpuQueueWriteBuffer(g.queue(), params_buf, 0, &p, sizeof(p)); + g.dispatch_at(dispatch_idx).workgroup_count_x = + (static_cast(numel) + wg - 1) / wg; + }); +} + // M-gate coordinator: registered at the LAST triple op (all dispatch indices // known). Prefill (M>1): run the fused GEMM, zero the 3 separate linears. // Decode (M==1): zero the fused, leave the 3 coop4 GEMVs (their own hooks set diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index d7b81b7b532..6f99c6162a4 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -478,6 +478,15 @@ class WebGPUGraph { // resize hook. void add_qkv_fused_dispatch(QkvFusionGroup& g); void add_qkv_fused_hook(const QkvFusionGroup& g); + + // SwiGLU fusion: emit ONE fused elementwise dispatch + // computing out = (gate * sigmoid(gate)) * up, replacing the sigmoid + 2 + // muls. `out` is repointed to a private pooled buffer (aliasing guard); + // `gate` is likewise given a private pooled buffer at its producer op by the + // build() walk (the planner reuse-aliases up onto gate's slot, so up_proj + // would stomp gate before the fused reads it). Only used during build(); the + // detection maps are empty (inert) when no SwiGLU triple matches. + void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id); }; } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl b/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl new file mode 100644 index 00000000000..f8ae8136222 --- /dev/null +++ b/backends/webgpu/runtime/ops/mul/silu_mul_fused.wgsl @@ -0,0 +1,24 @@ +@group(0) @binding(0) var gate: array; +@group(0) @binding(1) var up: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate +// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch. +// sigmoid + silu are computed in registers (never written to memory), so gate + up +// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the +// multiply order match the original ops -> bit-exact. +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let g = gate[idx]; + let sig = 1.0 / (1.0 + exp(-g)); + output[idx] = (g * sig) * up[idx]; +} diff --git a/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h b/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h new file mode 100644 index 00000000000..e847eee4429 --- /dev/null +++ b/backends/webgpu/runtime/ops/mul/silu_mul_fused_wgsl.h @@ -0,0 +1,48 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch::backends::webgpu { + +// @generated from silu_mul_fused.wgsl - DO NOT EDIT. +// wgsl-sha256: 4b8ede66c5dbc9829ff48f745eb9ad48fa5a5200058baa532fbf34f78ec2f560 +inline constexpr const char* kSiluMulFusedWGSL = R"( +@group(0) @binding(0) var gate: array; +@group(0) @binding(1) var up: array; +@group(0) @binding(2) var output: array; + +struct Params { + num_elements: u32, +} +@group(0) @binding(3) var params: Params; + +// Fused SwiGLU activation: output = (g * sigmoid(g)) * up, folding the separate +// sigmoid(gate) -> mul(gate,sig)=silu -> mul(silu,up) triple into one dispatch. +// sigmoid + silu are computed in registers (never written to memory), so gate + up +// are read once and one output is written. The sigmoid form (1/(1+exp(-x))) and the +// multiply order match the original ops -> bit-exact. +@compute @workgroup_size(64) +fn main(@builtin(global_invocation_id) gid: vec3) { + let idx = gid.x; + if (idx >= params.num_elements) { + return; + } + let g = gate[idx]; + let sig = 1.0 / (1.0 + exp(-g)); + output[idx] = (g * sig) * up[idx]; +} +)"; + +inline constexpr uint32_t kSiluMulFusedWorkgroupSizeX = 64; +inline constexpr uint32_t kSiluMulFusedWorkgroupSizeY = 1; +inline constexpr uint32_t kSiluMulFusedWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu From f35b0031ae62cced102f7c5a756386bd728b442e Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 11:51:38 -0700 Subject: [PATCH 3/3] [ExecuTorch][WebGPU] vec4 activation loads in the steel q4gsw prefill GEMM MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Pull Request resolved: https://github.com/pytorch/executorch/pull/20882 Read the activation tile as one `vec4` load instead of 4 scalar loads in the steel q4gsw prefill GEMM — a coalesced 16B fetch that shaves ~1-2% off prefill on Llama-3.2-1B (M4 Pro), bit-exact. Problem The steel GEMM staged the activation tile (As) with 4 scalar `t_input` reads per thread per BK step. Solution - Before: `As[...] = f16(t_input[base + 0/1/2/3])` — 4 scalar loads. - After: `let av = t_input[base >> 2u]; As[...] = f16(av.x/y/z/w)` — one `vec4` load (16B coalesced); `t_input` is bound as `array>`. Implementation Applies across the shared steel template (base / f16 / pwdq / pwdqf16acc variants); generated `_wgsl.h` headers regenerated. Load-only — the vec4 output store was a measured wash on Apple (scalar-ALU) and is deliberately not included. Constraints Bit-exact (identical values — a wider load, not a math change). Requires the activation contiguous with K a multiple of 4, which holds for all Llama q4gsw projections. Co-authored-with: Claude Code. ghstack-source-id: 405251998 @exported-using-ghexport Differential Revision: [D111500432](https://our.internmc.facebook.com/intern/diff/D111500432/) --- .../quantized_linear/q4gsw_linear_gemm_steel.wgsl | 12 +++++++----- ...q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h | 14 ++++++++------ .../q4gsw_linear_gemm_steel_half_pwdq_wgsl.h | 14 ++++++++------ .../q4gsw_linear_gemm_steel_half_wgsl.h | 14 ++++++++------ .../q4gsw_linear_gemm_steel_wgsl.h | 14 ++++++++------ 5 files changed, 39 insertions(+), 29 deletions(-) diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl index f08cb39bd5c..e2dfc610976 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel.wgsl @@ -1,7 +1,7 @@ $if DTYPE == "half": enable f16; @group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; +@group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @group(0) @binding(4) var t_bias: array; @@ -63,10 +63,12 @@ fn main(@builtin(workgroup_id) wid: vec3, let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(t_input[base]); - As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 1u]); - As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 2u]); - As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(t_input[base + 3u]); + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = ${buffer_scalar_type(DTYPE)}(av.x); + As[ar * BK + ac + 1u] = ${buffer_scalar_type(DTYPE)}(av.y); + As[ar * BK + ac + 2u] = ${buffer_scalar_type(DTYPE)}(av.z); + As[ar * BK + ac + 3u] = ${buffer_scalar_type(DTYPE)}(av.w); } else { As[ar * BK + ac + 0u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 1u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 2u] = ${"0.0h" if PWDQ else "0.0"}; As[ar * BK + ac + 3u] = ${"0.0h" if PWDQ else "0.0"}; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h index 6d23ea7ac76..ab4d3c06915 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_f16acc_wgsl.h @@ -13,11 +13,11 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: fe8b71afc5634e7db5607deb58f30c6517f93e1d66370ea017bcd9071201c6ca +// wgsl-sha256: a4346ddef028036f29aa73f0620c586467627626a745159fccef816a32f9475a inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqF16accWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; +@group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @group(0) @binding(4) var t_bias: array; @@ -75,10 +75,12 @@ fn main(@builtin(workgroup_id) wid: vec3, let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - As[ar * BK + ac + 0u] = f16(t_input[base]); - As[ar * BK + ac + 1u] = f16(t_input[base + 1u]); - As[ar * BK + ac + 2u] = f16(t_input[base + 2u]); - As[ar * BK + ac + 3u] = f16(t_input[base + 3u]); + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = f16(av.x); + As[ar * BK + ac + 1u] = f16(av.y); + As[ar * BK + ac + 2u] = f16(av.z); + As[ar * BK + ac + 3u] = f16(av.w); } else { As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h index e684b00a228..251a78b9fc8 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_pwdq_wgsl.h @@ -13,11 +13,11 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: 1b3f5e384e09bde50c66d43db4797e01e5610d70af235355dbc4966643bf9103 +// wgsl-sha256: d772a385edf91547a51fa3f5c4bdb4101da933f4c47828578c2818ff9eadf5fc inline constexpr const char* kQ4gswLinearGemmSteelHalfPwdqWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; +@group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @group(0) @binding(4) var t_bias: array; @@ -75,10 +75,12 @@ fn main(@builtin(workgroup_id) wid: vec3, let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - As[ar * BK + ac + 0u] = f16(t_input[base]); - As[ar * BK + ac + 1u] = f16(t_input[base + 1u]); - As[ar * BK + ac + 2u] = f16(t_input[base + 2u]); - As[ar * BK + ac + 3u] = f16(t_input[base + 3u]); + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = f16(av.x); + As[ar * BK + ac + 1u] = f16(av.y); + As[ar * BK + ac + 2u] = f16(av.z); + As[ar * BK + ac + 3u] = f16(av.w); } else { As[ar * BK + ac + 0u] = 0.0h; As[ar * BK + ac + 1u] = 0.0h; As[ar * BK + ac + 2u] = 0.0h; As[ar * BK + ac + 3u] = 0.0h; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h index d4c78406582..4741e13b74a 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_half_wgsl.h @@ -13,11 +13,11 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: d77542e5975fab44c58644c771390e89598ce82ab25e38d57150b581b4904acd +// wgsl-sha256: 363916e1aacae9635b4573d2ce35d0ee22916ce225537a3717df98a0ed3353da inline constexpr const char* kQ4gswLinearGemmSteelHalfWGSL = R"( enable f16; @group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; +@group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @group(0) @binding(4) var t_bias: array; @@ -78,10 +78,12 @@ fn main(@builtin(workgroup_id) wid: vec3, let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - As[ar * BK + ac + 0u] = f16(t_input[base]); - As[ar * BK + ac + 1u] = f16(t_input[base + 1u]); - As[ar * BK + ac + 2u] = f16(t_input[base + 2u]); - As[ar * BK + ac + 3u] = f16(t_input[base + 3u]); + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = f16(av.x); + As[ar * BK + ac + 1u] = f16(av.y); + As[ar * BK + ac + 2u] = f16(av.z); + As[ar * BK + ac + 3u] = f16(av.w); } else { As[ar * BK + ac + 0u] = 0.0; As[ar * BK + ac + 1u] = 0.0; As[ar * BK + ac + 2u] = 0.0; As[ar * BK + ac + 3u] = 0.0; diff --git a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h index db7a7689a6f..197b4f0ae92 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h +++ b/backends/webgpu/runtime/ops/quantized_linear/q4gsw_linear_gemm_steel_wgsl.h @@ -13,10 +13,10 @@ namespace executorch::backends::webgpu { // @generated from q4gsw_linear_gemm_steel.wgsl - DO NOT EDIT. -// wgsl-sha256: dcb63b04d002899acfed88d273fe658618650d5dcbc108540cef086bb97f5bce +// wgsl-sha256: f8106cb3f18424a9db8e04464e52a7fd522ca78f49e19ba6692f1c8fa83474e6 inline constexpr const char* kQ4gswLinearGemmSteelWGSL = R"( @group(0) @binding(0) var t_out: array; -@group(0) @binding(1) var t_input: array; +@group(0) @binding(1) var t_input: array>; @group(0) @binding(2) var t_weight: array; @group(0) @binding(3) var t_scales: array; @group(0) @binding(4) var t_bias: array; @@ -77,10 +77,12 @@ fn main(@builtin(workgroup_id) wid: vec3, let arow = row0 + ar; if (arow < params.M) { let base = arow * params.K + k0 + ac; - As[ar * BK + ac + 0u] = f32(t_input[base]); - As[ar * BK + ac + 1u] = f32(t_input[base + 1u]); - As[ar * BK + ac + 2u] = f32(t_input[base + 2u]); - As[ar * BK + ac + 3u] = f32(t_input[base + 3u]); + // vec4 coalesced load; base is 4-aligned on the steel route (K%16==0, ac/k0 multiples of 4). + let av = t_input[base >> 2u]; + As[ar * BK + ac + 0u] = f32(av.x); + As[ar * BK + ac + 1u] = f32(av.y); + As[ar * BK + ac + 2u] = f32(av.z); + As[ar * BK + ac + 3u] = f32(av.w); } else { As[ar * BK + ac + 0u] = 0.0; As[ar * BK + ac + 1u] = 0.0; As[ar * BK + ac + 2u] = 0.0; As[ar * BK + ac + 3u] = 0.0;