[ExecuTorch][WebGPU] Add optimized SDPA op#20851
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20851
Note: Links to docs will display an error until the docs builds have been completed. ✅ You can merge normally! (1 Unrelated Failure)As of commit 0313e86 with merge base aceeb40 ( BROKEN TRUNK - The following job failed but were 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):
Adds a non-causal fused scaled-dot-product-attention kernel to the WebGPU backend, enabling the attention blocks of vision encoders (Florence-2 DaViT / SigLIP, SAM2) and non-causal cross-attention (BART) to run end-to-end on GPU.
Problem — the delegate had no kernel for
et_vk.sdpa.default, the plain non-causal attentionsoftmax(q @ kᵀ * scale + attn_mask) @ vthat the et_vk source transform plugs into every vision-encoder attention block. Without it the attention layers broke the graph. This is distinct from the causal KV-cachesdpa_with_kv_cache: no cache, an optional additive mask, and it must handle asymmetric sequence lengths (S_q != S_kv, e.g. a Hiera pooled query or cross-attention).Solution
et_vk_sdpa_implchains three compute dispatches over[B, H, S, D](DSHB) row-major tensors —et_vk_sdpa_qk.wgslcomputes the scaledQ·Kᵀscores plus an optional additive mask, the reusedsdpa_softmax.wgsldoes the row-wise softmax, andet_vk_sdpa_av.wgslcomputessoftmax · Vinto the output.Implementation
et_vk_sdpa_qk.wgsl): one GPU thread per(b, h, s)row of the[B, H, S_q, S_kv]attention-weight buffer;q/kare bound asarray<vec4<f32>>overD, and the thread loops overc(key positions) andd4(D/4) accumulatingdot(q4, k4), then multiplies byscaleand addsmask[...]whenhas_mask. Row countB*H*S_qstays well under the 1D dispatch limit for any ViT.sdpa_softmax.wgsl(one workgroup per row, hardcoded@workgroup_size(64,1,1)) dispatched on a near-square 2D workgroup grid past the 65535 ceiling; the shader recovers the flat row index from@builtin(num_workgroups), so no override constant is needed.et_vk_sdpa_av.wgsl): one thread per(b, h, s, d4)computing avec4<f32>of four output elements;v/outare bound asarray<vec4<f32>>overDand the thread contracts overcscalar (S_kvis not guaranteed% 4 == 0), accumulatingsm[...] * v4.attn,softmax, eachB*H*S_q*S_kv) are allocated viagraph.create_scratch_buffer;scaledefaults to1/sqrt(D)when the arg isNoneor takes theDoublevalue; the three uniforms are compact structs (QkParams/AvParams32 bytes,SoftmaxParams16 bytes).utils::make_compute_pipeline,utils::make_uniform,utils::check_vec4_aligned(guardsD % 4 == 0),utils::clamp_workgroup_size+utils::compute_1d_workgroup_count(QK / AV grids),utils::compute_2d_workgroup_count(softmax grid), andutils::make_optional_binding(a 4-byte dummy satisfies the mask binding when absent; the shader never reads it underhas_mask == 0).backends/vulkan/runtime/graph/ops/impl/SDPA.cppSDPAMode::FUSED(general non-cache SDPA,[B, H, S, D]DSHB layout, optional additiveattn_mask, optionalscale, unpadded fp32 attention weights).Constraints — fp32 only (bails on
q/outbyte mismatch); non-causal (causality is expressed as a baked additive mask input, not a code path);D % 4 == 0for the vec4 QK/AV kernels (every model in scope usesD=64or128);qrank ≥ 3;k.dims == v.dims,q/k/vshareHandDand all leading batch dims, andout.dims == q.dims; asymmetricS_q != S_kvis supported (reduces bit-identically to self-attention when equal); a supplied mask must be[B, H, S_q, S_kv]fp32.Co-authored-with: Claude Code.
Differential Revision: D110836679