Skip to content

Commit f2c8a09

Browse files
qjia7qjia7
authored
webgpu: Fuse FlashAttention decode kernels and extend to any sequence length (#28389)
## Summary - Extend the FlashAttention decode path to work with any sequence length (not just seq_len=1), with causal masking and `use_seqlen_k` support for static KV cache - Add m_tile optimization to process multiple Q rows per workgroup (m_tile=1/2/4), amortizing K/V loads - Fuse the separate QKT and SplitVx shaders into a single QKV kernel using online softmax, eliminating the intermediate `qk` tensor (`B×H×seq×present_seq`) and reducing dispatch count from 3 to 2 - Route between prefill (FlashAttentionProgram) and split-reduce (fused QKV + VxReduce) paths based on sequence length ## Resolved 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 when `sequence_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 across `sequence_length ∈ {16, 30, 31}` × `total_sequence_length ∈ {128, 500, 2000}`. The previous heuristic additionally gated on `total_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 - [x] 30/30 MHA unit tests pass - [x] phi4-graph-prune produces correct output - [x] whisper-tiny-int4 produces correct transcription - [x] clang-format clean
1 parent 2cf6c6c commit f2c8a09

7 files changed

Lines changed: 426 additions & 426 deletions

‎onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc‎

Lines changed: 135 additions & 132 deletions
Large diffs are not rendered by default.

‎onnxruntime/contrib_ops/webgpu/bert/flash_attention.h‎

Lines changed: 26 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,9 @@ class SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram final : public Program<S
3535
{"half_rotary_dim", ProgramUniformVariableDataType::Uint32},
3636
{"present_sequence_length", ProgramUniformVariableDataType::Uint32},
3737
{"tile_size", ProgramUniformVariableDataType::Uint32},
38-
{"dispatch_size", ProgramUniformVariableDataType::Uint32});
38+
{"dispatch_size", ProgramUniformVariableDataType::Uint32},
39+
{"batch_size", ProgramUniformVariableDataType::Uint32},
40+
{"num_q_tiles", ProgramUniformVariableDataType::Uint32});
3941

4042
private:
4143
const bool interleaved_;
@@ -56,7 +58,9 @@ class CopyKVCacheProgram final : public Program<CopyKVCacheProgram> {
5658
{"total_sequence_length", ProgramUniformVariableDataType::Uint32},
5759
{"kv_sequence_length", ProgramUniformVariableDataType::Uint32},
5860
{"tile_size", ProgramUniformVariableDataType::Uint32},
59-
{"num_heads", ProgramUniformVariableDataType::Uint32});
61+
{"num_heads", ProgramUniformVariableDataType::Uint32},
62+
{"batch_size", ProgramUniformVariableDataType::Uint32},
63+
{"num_q_tiles", ProgramUniformVariableDataType::Uint32});
6064

6165
private:
6266
bool has_past_;
@@ -138,11 +142,14 @@ class FlashAttentionProgram final : public Program<FlashAttentionProgram> {
138142
int max_k_step_;
139143
};
140144

141-
class FlashAttentionDecodeQKTProgram final : public Program<FlashAttentionDecodeQKTProgram> {
145+
class FlashAttentionDecodeQKVProgram final : public Program<FlashAttentionDecodeQKVProgram> {
142146
public:
143-
FlashAttentionDecodeQKTProgram(const std::string& kernel_name,
144-
bool has_attention_bias, uint32_t tile_size, bool use_indirect_dispatch)
145-
: Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size), use_indirect_dispatch_(use_indirect_dispatch) {
147+
FlashAttentionDecodeQKVProgram(const std::string& kernel_name,
148+
bool has_attention_bias, uint32_t tile_size, int head_size_vec,
149+
bool use_indirect_dispatch, bool q_BNSH = false,
150+
bool is_unidirectional = false,
151+
uint32_t m_tile = 1)
152+
: Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size), head_size_vec_(head_size_vec), use_indirect_dispatch_(use_indirect_dispatch), q_BNSH_(q_BNSH), is_unidirectional_(is_unidirectional), m_tile_(m_tile) {
146153
}
147154

148155
Status GenerateShaderCode(ShaderHelper& sh) const override;
@@ -156,41 +163,23 @@ class FlashAttentionDecodeQKTProgram final : public Program<FlashAttentionDecode
156163
{"num_heads", ProgramUniformVariableDataType::Uint32},
157164
{"batch_size", ProgramUniformVariableDataType::Uint32},
158165
{"attn_bias_dim0", ProgramUniformVariableDataType::Uint32},
159-
{"attn_bias_dim1", ProgramUniformVariableDataType::Uint32});
166+
{"attn_bias_dim1", ProgramUniformVariableDataType::Uint32},
167+
{"new_sequence_length", ProgramUniformVariableDataType::Uint32});
160168

161169
private:
162170
bool has_attention_bias_;
163-
uint32_t tile_size_;
164-
bool use_indirect_dispatch_;
165-
};
166-
167-
class FlashAttentionDecodeSplitVxProgram final : public Program<FlashAttentionDecodeSplitVxProgram> {
168-
public:
169-
FlashAttentionDecodeSplitVxProgram(const std::string& kernel_name, uint32_t tile_size, int head_size_vec, bool use_indirect_dispatch, bool has_head_sink = false)
170-
: Program{kernel_name}, tile_size_(tile_size), head_size_vec_(head_size_vec), use_indirect_dispatch_(use_indirect_dispatch), has_head_sink_(has_head_sink) {
171-
}
172-
173-
Status GenerateShaderCode(ShaderHelper& sh) const override;
174-
175-
WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"total_sequence_length", ProgramUniformVariableDataType::Uint32},
176-
{"head_size_vec", ProgramUniformVariableDataType::Uint32},
177-
{"present_sequence_length", ProgramUniformVariableDataType::Uint32},
178-
{"n_reps", ProgramUniformVariableDataType::Uint32},
179-
{"num_present_sequence_length_tile", ProgramUniformVariableDataType::Uint32},
180-
{"batch_heads", ProgramUniformVariableDataType::Uint32},
181-
{"num_heads", ProgramUniformVariableDataType::Uint32});
182-
183-
private:
184171
uint32_t tile_size_;
185172
int head_size_vec_;
186173
bool use_indirect_dispatch_;
187-
bool has_head_sink_;
174+
bool q_BNSH_;
175+
bool is_unidirectional_;
176+
uint32_t m_tile_;
188177
};
189178

190179
class FlashAttentionDecodeVxReduceProgram final : public Program<FlashAttentionDecodeVxReduceProgram> {
191180
public:
192-
FlashAttentionDecodeVxReduceProgram(const std::string& kernel_name, uint32_t tile_size, uint32_t seq_tile_size, bool use_indirect_dispatch)
193-
: Program{kernel_name}, tile_size_(tile_size), seq_tile_size_(seq_tile_size), use_indirect_dispatch_(use_indirect_dispatch) {
181+
FlashAttentionDecodeVxReduceProgram(const std::string& kernel_name, uint32_t tile_size, uint32_t seq_tile_size, bool use_indirect_dispatch, bool has_head_sink = false, uint32_t m_tile = 1)
182+
: Program{kernel_name}, tile_size_(tile_size), seq_tile_size_(seq_tile_size), use_indirect_dispatch_(use_indirect_dispatch), has_head_sink_(has_head_sink), m_tile_(m_tile) {
194183
}
195184

196185
Status GenerateShaderCode(ShaderHelper& sh) const override;
@@ -199,12 +188,16 @@ class FlashAttentionDecodeVxReduceProgram final : public Program<FlashAttentionD
199188
{"num_total_seq_length_tile", ProgramUniformVariableDataType::Uint32},
200189
{"num_present_sequence_length_tile", ProgramUniformVariableDataType::Uint32},
201190
{"num_head_size_tile", ProgramUniformVariableDataType::Uint32},
202-
{"batch_heads", ProgramUniformVariableDataType::Uint32});
191+
{"batch_heads", ProgramUniformVariableDataType::Uint32},
192+
{"new_sequence_length", ProgramUniformVariableDataType::Uint32},
193+
{"num_heads", ProgramUniformVariableDataType::Uint32});
203194

204195
private:
205196
uint32_t tile_size_;
206197
uint32_t seq_tile_size_;
207198
bool use_indirect_dispatch_;
199+
bool has_head_sink_;
200+
uint32_t m_tile_;
208201
};
209202

210203
Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias,
@@ -225,7 +218,7 @@ Status RunSplitPackedQKVWithRotaryEmbeddingAndCopyKV(onnxruntime::webgpu::Comput
225218
Tensor* present_key,
226219
Tensor* present_value,
227220
Tensor* indirect_buffer,
228-
uint32_t tile_size);
221+
uint32_t tile_size, uint32_t num_q_tiles);
229222
} // namespace webgpu
230223
} // namespace contrib
231224
} // namespace onnxruntime

‎onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkt.wgsl.template‎

Lines changed: 0 additions & 118 deletions
This file was deleted.

0 commit comments

Comments
 (0)