[ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster)

backends/webgpu/runtime/ops/et_vk_sdpa/EtVkSdpa.cpp

@@ -10,6 +10,7 @@
 #include <executorch/backends/webgpu/runtime/WebGPUUtils.h>
 #include <executorch/backends/webgpu/runtime/ops/OperatorRegistry.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_av_wgsl.h>
+#include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_wgsl.h>
 #include <executorch/backends/webgpu/runtime/ops/sdpa/sdpa_softmax_wgsl.h>
 
@@ -149,13 +150,33 @@ void et_vk_sdpa_impl(WebGPUGraph& graph, const std::vector<int>& args) {
   const size_t aw_bytes = static_cast<size_t>(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<uint32_t>(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<uint32_t>(aw_numel),
+        qk_wg_size,
+        "et_vk_sdpa_qk_entry");
+  } else {
+    qk_wg_count = utils::compute_1d_workgroup_count(
+        device, static_cast<uint32_t>(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<int>& 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<int>& 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<int>& 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) {

backends/webgpu/runtime/ops/et_vk_sdpa/et_vk_sdpa_qk_entry.wgsl

@@ -0,0 +1,54 @@
+@group(0) @binding(0) var<storage, read_write> attn: array<f32>;
+@group(0) @binding(1) var<storage, read> q: array<f32>;
+@group(0) @binding(2) var<storage, read> k: array<f32>;
+@group(0) @binding(3) var<storage, read> mask: array<f32>;
+
+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<uniform> 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<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
+  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
+}

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 <cstdint>
+
+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<storage, read_write> attn: array<f32>;
+@group(0) @binding(1) var<storage, read> q: array<f32>;
+@group(0) @binding(2) var<storage, read> k: array<f32>;
+@group(0) @binding(3) var<storage, read> mask: array<f32>;
+
+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<uniform> 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<u32>,
+    @builtin(num_workgroups) nwg: vec3<u32>) {
+  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