diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp index dbe6b03f30c..92a53cb7574 100644 --- a/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include @@ -149,13 +150,33 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { const size_t aw_bytes = static_cast(aw_numel) * sizeof(float); // Up-front dispatch-limit checks (throw BEFORE any buffer alloc → no leak). - // QK = one thread per (b,h,s) row; AV = one per (b,h,s,d4) vec4 (4 output - // elements/thread); softmax = one workgroup per row. compute_1d_workgroup_ - // count throws past the limit. - const uint32_t qk_wg_size = - utils::clamp_workgroup_size(device, kEtVkSdpaQkWorkgroupSizeX); - const uint32_t qk_wg_count = utils::compute_1d_workgroup_count( - device, static_cast(num_rows), qk_wg_size, "et_vk_sdpa_qk"); + // QK: per-row (one thread per (b,h,s) row, vec4 loads) is fastest for + // standard attention, but starves the GPU on channel attention (S_q = + // head_dim, so num_rows is tiny → few workgroups serial over a huge S_kv*D). + // Route to the per-entry kernel (one thread per (b,h,s,c) attn entry, + // 2D-folded) below an occupancy floor; both write a layout-identical + // attn[B,H,S_q,S_kv] so softmax/AV are unchanged, and either branch is + // numerically correct, so the floor is a perf knob only (Canary M4 Pro: + // per-entry ~15-30x faster at num_rows <= 256, per-row wins at num_rows >= + // 8192). AV = one per (b,h,s,d4) vec4; softmax = one workgroup per row. + constexpr uint32_t kQkEntryOccupancyFloor = 4096u; + const bool qk_per_entry = num_rows < kQkEntryOccupancyFloor; + const uint32_t qk_wg_size = utils::clamp_workgroup_size( + device, + qk_per_entry ? kEtVkSdpaQkEntryWorkgroupSizeX + : kEtVkSdpaQkWorkgroupSizeX); + uint32_t qk_wg_count = 0; + utils::WgCount qk_entry_grid = {}; + if (qk_per_entry) { + qk_entry_grid = utils::compute_2d_workgroup_count( + device, + static_cast(aw_numel), + qk_wg_size, + "et_vk_sdpa_qk_entry"); + } else { + qk_wg_count = utils::compute_1d_workgroup_count( + device, static_cast(num_rows), qk_wg_size, "et_vk_sdpa_qk"); + } const uint32_t av_wg_size = utils::clamp_workgroup_size(device, kEtVkSdpaAvWorkgroupSizeX); const uint64_t out_numel4 = @@ -179,7 +200,7 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer attn_buf = graph.create_scratch_buffer(aw_bytes); WGPUBuffer softmax_buf = graph.create_scratch_buffer(aw_bytes); - // ---- Dispatch 1: QK (one thread per (b,h,s) row) ---- + // ---- Dispatch 1: QK (per-row for standard attn, per-entry for channel) ---- { QkParams p = {}; p.B = B; @@ -205,7 +226,7 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( device, - kEtVkSdpaQkWGSL, + qk_per_entry ? kEtVkSdpaQkEntryWGSL : kEtVkSdpaQkWGSL, { {0, WGPUBufferBindingType_Storage, attn_buf, aw_bytes}, {1, WGPUBufferBindingType_ReadOnlyStorage, q.buffer, q.nbytes}, @@ -222,7 +243,12 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector& args) { &wg_const, 1); - graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count}); + if (qk_per_entry) { + graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, qk_entry_grid.x, qk_entry_grid.y); + } else { + graph.add_dispatch({bundle.pipeline, bundle.bind_group, qk_wg_count}); + } wgpuBufferRelease(uniform_buffer); if (mask.owned_dummy != nullptr) { diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl new file mode 100644 index 00000000000..605ac52c89e --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl @@ -0,0 +1,54 @@ +@group(0) @binding(0) var attn: array; +@group(0) @binding(1) var q: array; +@group(0) @binding(2) var k: array; +@group(0) @binding(3) var mask: array; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + has_mask: u32, + _pad0: u32, + scale: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D], +// k [B, H, S_kv, D]. ONE thread per ENTRY (b,h,s,c) of attn_weights +// [B, H, S_q, S_kv] = one D-length dot. Previously this was one thread per ROW +// (looping c) — fine for window/self attention but catastrophic for DaViT +// CHANNEL attention where S_q = head_dim (~32), so B*H*S_q was 128/256/512/1024 +// → only 2/4/8/16 workgroups serial over the huge spatial D (the (2,1,1)@103ms +// dispatch). Parallelizing over all B*H*S_q*S_kv entries (2D-folded past the +// 65535 ceiling, mirroring the softmax phase) gives S_kv× more threads. +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) nwg: vec3) { + let aw_numel = params.B * params.H * params.S_q * params.S_kv; + let idx = gid.x + gid.y * (nwg.x * wg_size); // 2D-folded linear entry id + if (idx >= aw_numel) { + return; + } + let c = idx % params.S_kv; + let row = idx / params.S_kv; // (b,h,s) flattened + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let qbase = ((b * params.H + h) * params.S_q + s) * params.D; + let kbase = ((b * params.H + h) * params.S_kv + c) * params.D; + var acc: f32 = 0.0; + for (var d: u32 = 0u; d < params.D; d = d + 1u) { + acc = acc + q[qbase + d] * k[kbase + d]; + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[idx]; + } + attn[idx] = acc; // attn is [B,H,S_q,S_kv] row-major -> index == idx +} diff --git a/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h new file mode 100644 index 00000000000..6918c6e2621 --- /dev/null +++ b/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h @@ -0,0 +1,78 @@ +/* + * 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 et_vk_sdpa_qk_entry.wgsl - DO NOT EDIT. +// wgsl-sha256: 95ff9f19757d23c41d82054b779c02fb8d83531eef3e8cda5f3a78dde42a4e58 +inline constexpr const char* kEtVkSdpaQkEntryWGSL = R"( +@group(0) @binding(0) var attn: array; +@group(0) @binding(1) var q: array; +@group(0) @binding(2) var k: array; +@group(0) @binding(3) var mask: array; + +struct Params { + B: u32, + H: u32, + S_q: u32, + S_kv: u32, + D: u32, + has_mask: u32, + _pad0: u32, + scale: f32, +} +@group(0) @binding(4) var params: Params; + +override wg_size: u32 = 64; + +// Non-causal fused SDPA, QK phase. DSHB layout, row-major: q [B, H, S_q, D], +// k [B, H, S_kv, D]. ONE thread per ENTRY (b,h,s,c) of attn_weights +// [B, H, S_q, S_kv] = one D-length dot. Previously this was one thread per ROW +// (looping c) — fine for window/self attention but catastrophic for DaViT +// CHANNEL attention where S_q = head_dim (~32), so B*H*S_q was 128/256/512/1024 +// → only 2/4/8/16 workgroups serial over the huge spatial D (the (2,1,1)@103ms +// dispatch). Parallelizing over all B*H*S_q*S_kv entries (2D-folded past the +// 65535 ceiling, mirroring the softmax phase) gives S_kv× more threads. +@compute @workgroup_size(wg_size) +fn main( + @builtin(global_invocation_id) gid: vec3, + @builtin(num_workgroups) nwg: vec3) { + let aw_numel = params.B * params.H * params.S_q * params.S_kv; + let idx = gid.x + gid.y * (nwg.x * wg_size); // 2D-folded linear entry id + if (idx >= aw_numel) { + return; + } + let c = idx % params.S_kv; + let row = idx / params.S_kv; // (b,h,s) flattened + let s = row % params.S_q; + let h = (row / params.S_q) % params.H; + let b = row / (params.S_q * params.H); + + let qbase = ((b * params.H + h) * params.S_q + s) * params.D; + let kbase = ((b * params.H + h) * params.S_kv + c) * params.D; + var acc: f32 = 0.0; + for (var d: u32 = 0u; d < params.D; d = d + 1u) { + acc = acc + q[qbase + d] * k[kbase + d]; + } + acc = acc * params.scale; + if (params.has_mask != 0u) { + acc = acc + mask[idx]; + } + attn[idx] = acc; // attn is [B,H,S_q,S_kv] row-major -> index == idx +} +)"; + +inline constexpr uint32_t kEtVkSdpaQkEntryWorkgroupSizeX = 64; +inline constexpr uint32_t kEtVkSdpaQkEntryWorkgroupSizeY = 1; +inline constexpr uint32_t kEtVkSdpaQkEntryWorkgroupSizeZ = 1; + +} // namespace executorch::backends::webgpu