[ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster)#20871
[ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster)#20871JCNTH wants to merge 1 commit into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20871
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (2 Unrelated Failures)As of commit 269817d with merge base aceeb40 ( FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
meta-cla
Bot
added
the
CLA Signed
2 participants
Stack from ghstack (oldest at bottom):
Problem: the fused
et_vk.sdpaQK kernel runs one thread per (b,h,s) row with vec4 loads — ideal for standard attention, but on channel attention (DaViT/Florence, whereS_q = head_dim ~= 32)num_rows = B*H*S_qis tiny, so only a handful of workgroups run serially over a hugeS_kv*D, starving the GPU (the (2,1,1)@103ms dispatch).Solution: add a per-entry QK kernel (one thread per (b,h,s,c) attention entry, 2D-folded) and host-route to it when
num_rowsis below an occupancy floor (4096); standard attention keeps the per-row + vec4 path unchanged.Before:
et_vk_sdpa_qk(per-row, vec4) — the only QK kernel; channel-attn shapes are occupancy-starved.After: router picks
et_vk_sdpa_qk_entry(per-entry, scalar, 2D-folded) for smallnum_rows, else the unchanged per-row kernel.Implementation:
et_vk_sdpa_qk_entry.wgsl(+ generated header) — same bindings andParamsas the per-row kernel, so it is a drop-in underlayout:"auto"; writes a layout-identicalattn[B,H,S_q,S_kv](attn[idx]), so softmax/AV are unchanged and either branch is numerically correct — the floor is a pure perf knob.EtVkSdpa.cppselects the shader and a 2D dispatch (compute_2d_workgroup_count, mirroring the softmax grid) when routed, else the existing 1D per-row dispatch; the grid + dispatch-limit check is computed up front (throw before any buffer alloc -> no leak).LinearFp32.cppK%4vec4 selection,Sdpa.cppvariant selection).Constraints: per-entry drops vec4, so it only wins when the per-row path is occupancy-starved (small
num_rows); the 4096 floor is Canary-tuned.Co-authored-with: Claude Code.
Differential Revision: D110994975