[ExecuTorch][WebGPU] Add optimized SDPA op

[JCNTH]

Stack from ghstack (oldest at bottom):

• [ExecuTorch][WebGPU] split_with_sizes_copy op tests #20996
• [ExecuTorch][WebGPU] Add split_with_sizes_copy op #20995
• [ExecuTorch][WebGPU] batch_norm op tests #20994
• [ExecuTorch][WebGPU] Add batch_norm op #20993
• [ExecuTorch][WebGPU] upsample_bilinear2d op tests #20992
• [ExecuTorch][WebGPU] Add upsample_bilinear2d op #20991
• [ExecuTorch][WebGPU] leaky_relu op tests #20990
• [ExecuTorch][WebGPU] Add leaky_relu op #20989
• [ExecuTorch][WebGPU] Tests for aten._to_copy.default int<->float convert #20988
• [ExecuTorch][WebGPU] Add aten._to_copy.default with int<->float numeric convert #20987
• [ExecuTorch][WebGPU] Tests for aten.sub.Tensor broadcast #20986
• [ExecuTorch][WebGPU] Add aten.sub.Tensor with TensorMeta broadcast #20985
• [ExecuTorch][WebGPU] apply_rotary_emb_hf op tests #20876
• [ExecuTorch][WebGPU] et_vk.apply_rotary_emb_hf: HuggingFace rotate-half RoPE runtime op (unblocks Qwen3) #20875
• [ExecuTorch][WebGPU] et_vk conv2d im2col-GEMM tests #20874
• [ExecuTorch][WebGPU] et_vk conv2d: route standard (groups==1) conv through an im2col tiled GEMM (1.1-2.4x) #20873
• [ExecuTorch][WebGPU] et_vk.sdpa channel-attention tests #20872
• [ExecuTorch][WebGPU] et_vk.sdpa: shape-route QK to a per-entry kernel for channel attention (15-30x faster) #20871
• [ExecuTorch][WebGPU] Migrate quantized_linear to make_compute_pipeline #20870
• [ExecuTorch][WebGPU] Migrate multi-dispatch ops to make_compute_pipeline #20869
• [ExecuTorch][WebGPU] Migrate 9 ops to make_compute_pipeline (layout:"auto") #20868
• [ExecuTorch][WebGPU] Adopt WebGPU layout:"auto" in make_compute_pipeline #20867
• [ExecuTorch][WebGPU] slice_copy native golden test #20866
• [ExecuTorch][WebGPU] slice_copy Double-arg scalar handling + view_copy/alias glue #20865
• [ExecuTorch][WebGPU] relu op tests #20864
• [ExecuTorch][WebGPU] Add relu op (shared unary handler; sigmoid adopts make_compute_pipeline) #20863
• [ExecuTorch][WebGPU] max_pool2d op tests #20862
• [ExecuTorch][WebGPU] Add max_pool2d op #20861
• [ExecuTorch][WebGPU] upsample_nearest2d op tests #20860
• [ExecuTorch][WebGPU] Add upsample_nearest2d op #20859
• [ExecuTorch][WebGPU] constant_pad_nd op tests #20858
• [ExecuTorch][WebGPU] Add constant_pad_nd op #20857
• [ExecuTorch][WebGPU] addmm op tests #20856
• [ExecuTorch][WebGPU] Add optimized addmm op #20855
• [ExecuTorch][WebGPU] embedding op tests #20854
• [ExecuTorch][WebGPU] Add embedding op #20853
• [ExecuTorch][WebGPU] SDPA op tests #20852
• -> [ExecuTorch][WebGPU] Add optimized SDPA op #20851
• [ExecuTorch][WebGPU] conv2d op tests #20850
• [ExecuTorch][WebGPU] Add optimized conv2d op #20849
• [ExecuTorch][WebGPU] linear_fp32 op tests #20848
• [ExecuTorch][WebGPU] Add optimized linear_fp32 op #20847
• [ExecuTorch][WebGPU] layer_norm op tests #20846
• [ExecuTorch][WebGPU] Add optimized layer_norm op #20845
• [ExecuTorch][WebGPU] gelu op tests #20844
• [ExecuTorch][WebGPU] Add optimized gelu op #20843
• [ExecuTorch][WebGPU] Shared pipeline + dispatch-grid + vec4-align helpers (op-suite foundation) #20842

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 attention softmax(q @ kᵀ * scale + attn_mask) @ v that 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-cache sdpa_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

• Before: no non-causal fused-attention kernel — vision-encoder attention blocks could not lower into the delegate.
• After: et_vk_sdpa_impl chains three compute dispatches over [B, H, S, D] (DSHB) row-major tensors — et_vk_sdpa_qk.wgsl computes the scaled Q·Kᵀ scores plus an optional additive mask, the reused sdpa_softmax.wgsl does the row-wise softmax, and et_vk_sdpa_av.wgsl computes softmax · V into the output.

Implementation

• QK phase (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/k are bound as array<vec4<f32>> over D, and the thread loops over c (key positions) and d4 (D/4) accumulating dot(q4, k4), then multiplies by scale and adds mask[...] when has_mask. Row count B*H*S_q stays well under the 1D dispatch limit for any ViT.
• Softmax phase: reuses the existing 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.
• AV phase (et_vk_sdpa_av.wgsl): one thread per (b, h, s, d4) computing a vec4<f32> of four output elements; v/out are bound as array<vec4<f32>> over D and the thread contracts over c scalar (S_kv is not guaranteed % 4 == 0), accumulating sm[...] * v4.
• Two fp32 scratch buffers (attn, softmax, each B*H*S_q*S_kv) are allocated via graph.create_scratch_buffer; scale defaults to 1/sqrt(D) when the arg is None or takes the Double value; the three uniforms are compact structs (QkParams/AvParams 32 bytes, SoftmaxParams 16 bytes).
• Uses the shared runtime helpers: utils::make_compute_pipeline, utils::make_uniform, utils::check_vec4_aligned (guards D % 4 == 0), utils::clamp_workgroup_size + utils::compute_1d_workgroup_count (QK / AV grids), utils::compute_2d_workgroup_count (softmax grid), and utils::make_optional_binding (a 4-byte dummy satisfies the mask binding when absent; the shader never reads it under has_mask == 0).
• Mirrors Vulkan backends/vulkan/runtime/graph/ops/impl/SDPA.cpp SDPAMode::FUSED (general non-cache SDPA, [B, H, S, D] DSHB layout, optional additive attn_mask, optional scale, unpadded fp32 attention weights).

Constraints — fp32 only (bails on q/out byte mismatch); non-causal (causality is expressed as a baked additive mask input, not a code path); D % 4 == 0 for the vec4 QK/AV kernels (every model in scope uses D=64 or 128); q rank ≥ 3; k.dims == v.dims, q/k/v share H and D and all leading batch dims, and out.dims == q.dims; asymmetric S_q != S_kv is 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.
@exported-using-ghexport

Differential Revision: D110836679

Differential Revision: D110836679

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20851

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit 2ddf136 with merge base c2b273e ():

NEW FAILURE - The following job has failed:

• pull / test-multimodal-linux (gemma3-4b) / linux-job (gh)
  RuntimeError: Command docker exec -t cdf0152634c2ade23433800f254ace833df12ea584e2ec41f003d3ef65e352b6 /exec failed with exit code 139

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.