webgpu: Fuse FlashAttention decode kernels and extend to any sequence length - #28389
Conversation
The 3-kernel FlashAttention decode path (QKT, SplitVx, VxReduce) previously only handled sequence_length=1. This extends it to work with any sequence length, providing a fallback for devices without Subgroups support. When Subgroups is available and seq_len>1, the subgroup-based prefill path is still preferred. The extended decode path activates when Subgroups is unavailable. Changes: - Workgroup layout now includes new_sequence_length dimension in all 3 kernels - Q offset supports both BSNH and BNSH layouts via q_BNSH template param - Causal masking (is_unidirectional) for self-attention with seq_len>1 - use_seqlen_k support for static KV cache (past_present_share_buffer) - Relaxed CanApplyFlashAttention to allow seq_len>1 without Subgroups Verified: whisper-tiny-int4 correct transcription, phi4-mini correct generation, 16/16 MHA unit tests pass.
…FlashAttention decode path Extend the 3-kernel decode path to process multiple Q rows per workgroup (m_tile=1/2/4) to amortize K/V memory loads for larger sequence lengths. Remove the Subgroups feature requirement from CanApplyFlashAttention so the decode path works on all WebGPU devices. The subgroup-based prefill path is replaced by the extended decode path with m_tile. Fix causal masking to compute past_sequence_length dynamically from total_sequence_length - new_sequence_length, which is correct for both GQA (where past_sequence_length_ is the buffer size) and graph capture (where total_sequence_length_ is on GPU).
…n decode Merge the separate QK^T and SplitVx shaders into a fused QKV shader that computes QK^T, local softmax, and V multiply in one kernel launch, eliminating the intermediate qk tensor and reducing dispatch count from 3 to 2. The VxReduce shader now rescales partial outputs using per-tile online softmax metadata (local_max, local_sum).
Use FlashAttentionProgram (single kernel, subgroup-based) for prefill when sequence is long enough to benefit. Fall back to the split-reduce path (fused QKV + VxReduce) for short sequences, when subgroups are unavailable, or when total_sequence_length is large relative to sequence_length.
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
qjia7
changed the title
There was a problem hiding this comment.
Pull request overview
This PR updates the WebGPU FlashAttention implementation to improve the decode/prefill path by fusing decode shaders (QKᵀ + softmax + V) into a single kernel, adding m_tile to process multiple query rows per workgroup, extending decode beyond seq_len=1, and routing between subgroup-based and non-subgroup paths to broaden device support.
Changes:
- Replace the decode QKT + SplitVx pipeline with a fused QKV shader + VxReduce (online softmax) pipeline, supporting
sequence_length > 1. - Add path routing between subgroup-based prefill and split-reduce decode to enable FlashAttention without subgroup support.
- Extend decode uniforms/shader logic to support causal masking, optional
seqlen_k, andm_tileoptimization.
Reviewed changes
Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/contrib_ops/webgpu/bert/flash_attention.h | Updates program declarations/uniforms: replaces QKT/SplitVx with fused QKV; extends VxReduce options (seqlen_k/head_sink/m_tile). |
| onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc | Implements fused QKV dispatch, split-reduce routing, new metadata/output shapes, and removes subgroup gating from CanApplyFlashAttention. |
| onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template | Updates reduction shader to rescale partial V outputs using online softmax metadata; adds m_tile, head_sink support, and sequence-length indexing. |
| onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template | Deleted (functionality folded into fused QKV shader). |
| onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkv.wgsl.template | New fused shader implementing QKᵀ + bias/mask + local softmax + V multiply, emitting per-tile metadata for reduce. |
| onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkt.wgsl.template | Deleted (replaced by fused QKV shader). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…ce path When graph capture is enabled with static KV cache, total_sequence_length is 0 (GPU-based seqlen_k), causing dispatch size 0 in the split-reduce QKV kernel. Fix by broadening use_indirect_dispatch to remove the sequence_length == 1 restriction, so the indirect buffer computes dispatch sizes on GPU for any sequence length. The indirect buffer now writes a flattened total workgroup count (batch * heads * num_q_tiles * num_kv_tiles) with normalization to stay within WebGPU maxComputeWorkgroupsPerDimension (65535).
…malization Move the indirect dispatch buffer normalization logic (splitting into 2D when exceeding maxComputeWorkgroupsPerDimension) into a shared WGSL function used by both CopyKVCacheProgram and SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram.
# Conflicts: # onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc
Narrow the second clause of the split-reduce routing heuristic from seq_len < 64 to seq_len < 32 to keep more medium-length prefill shapes on the single-kernel FlashAttention path.
Convert raw subscript access on always-registered tensors in flash_attention_decode_qkv and flash_attention_decode_vx_reduce to the shared .getByOffset/.setByOffset accessor helpers, matching the pattern used by matmul_nbits and split_packed_qkv_with_rotary_embedding_and_copykv. The wgsl-template generator emits unconditional `*params.var_X` hoists for every variable accessed via the helpers, so the corresponding C++ bindings now capture the AddInput/AddOutput refs and forward them with WGSL_TEMPLATE_VARIABLE entries. Conditional inputs (seqlens_k, attention_bias, head_sink) keep raw subscript access and stay out of the variable list.
…e num_q_tiles uniform Rename the WGSL helper from write_indirect_dispatch(total) to normalize_dispatch_group_size(x, y, z) so it mirrors ProgramManager::NormalizeDispatchGroupSize on the CPU side. The new form fast-paths the common case where every dim is within maxComputeWorkgroupsPerDimension, avoiding any f32 arithmetic, and keeps the sqrt and cbrt fallbacks for the rare overflow tiers. Drop the duplicated (new_sequence_length, m_tile) uniforms from CopyKVCacheProgram and SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram. Precompute num_q_tiles once in ApplyFlashAttention and pass it as a single uniform. CPU becomes the sole source of truth, the per-shader ceil-div recompute goes away, and uniform layout shrinks by one entry per program.
The previous heuristic also gated on total_sequence_length_ > 1000, but that signal is 0 when graph capture is enabled (seqlen_k lives on the GPU). The carve-out is also not needed: benchmarks show split-reduce is uniformly faster for short Q across all measured KV cache lengths (1.13x-2.07x at total_sequence_length 128 / 500 / 2000 on a representative LLM).
Verdict: Approve — the kernel fusion is sound, the online-softmax math in the new VxReduce is right, and the routing heuristic is now well-justified. A few observations worth flagging, none blocking.What's actually happening here (worth restating for the next reviewer)Three changes ride together in this PR:
The PR also addresses Copilot's one open comment by simplifying the routing rather than papering over the What I checked carefully1. Online-softmax rescaling in
|
Summary
use_seqlen_ksupport for static KV cacheqktensor (B×H×seq×present_seq) and reducing dispatch count from 3 to 2Resolved Issues
Whisper decoding prefill improved from 4.68ms to 1.09ms. Whisper's decoder attention has a small sequence length but large total sequence length (seq_len=4, total_seq_len=1500). The default prefill shader (FlashAttentionProgram) has low parallelism in this case because each workgroup iterates serially over the full KV cache. The split-reduce path tiles the KV dimension across workgroups, achieving much higher GPU occupancy for this workload shape.
Details
Fused QKV kernel: Each workgroup computes QK^T dot products, applies attention bias and causal mask, computes local softmax (per-tile max and sum), normalizes, and multiplies by V — all in one kernel. Per-tile metadata (max, sum) is written for the VxReduce shader to rescale partial outputs using online softmax:
output = Σ(partial_i × local_sum_i × exp(local_max_i - global_max)) / global_sum.Path routing (
use_split_reduce): The split-reduce path is used whensequence_length_ < 32; otherwise the single-kernel FlashAttentionProgram prefill path is used. Microbenchmarks on Phi-4 (32 heads, head_size 128, GQA group 3) show split-reduce is 1.13×-2.07× faster than the fused prefill kernel acrosssequence_length ∈ {16, 30, 31}×total_sequence_length ∈ {128, 500, 2000}. The previous heuristic additionally gated ontotal_sequence_length_ > 1000, but that signal is 0 under graph capture (seqlen_k lives on the GPU) and the carve-out is unnecessary because split-reduce is uniformly faster for short Q.Test plan