From 7ffb492bf42d2dab6ad333cc51fc6ae63f094ff4 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Tue, 21 Jul 2026 15:28:04 +0800 Subject: [PATCH 1/7] [WebGPU]: Add Intel subgroup prefill flash-attention kernel Introduce FlashAttentionPrefillSimpleProgram and its template (flash_attention_prefill_simple.wgsl.template). This is a simplified version of the FlashAttentionProgram kernel. (TODO: Merge back into the original FlashAttentionProgram) Key implementation details: - Latency Hiding: Only K is loaded before the QK^T + online-softmax stage. V is fetched just before the PV matmul, allowing global-memory latency to overlap with softmax compute. - Avoids Bank Conflicts: K and V tiles are stored in transposed layout ([head_size_vec][kv_step]), ensuring the subgroup column is the inner/ fastest dimension. Adjacent lanes hit adjacent shared-memory addresses. - Static Shuffle Codepath: Unconditionally uses a subgroup-shuffle path based on the host-configured minimum subgroup size. This avoids dynamic runtime SubgroupSize detection inside the WGSL kernel. - Configurations: The host determines kv_step (8 or 16) based on the adapter's guaranteed minimum subgroup size (needs minimum of 8). Workgroup size is 128 for causal and 256 for non-causal prefill. - Unsupported Features: Does not support quantized KV cache (int8/ TurboQuant) or the shared-memory fallback path (to be supported). Integration & Selection: - Adds CanApplyFlashAttentionPrefillSimple and ApplyFlashAttentionPrefillSimple, mirroring the MatMulNBits pattern. - ApplyFlashAttention selects this kernel only if 'is_intel' is true, subgroup_min_size >= 8, and the KV cache is unquantized. Otherwise, it falls back to the existing FlashAttentionProgram. **Test on Phi-4-mini, EdgeLLMOnDeviceModel\2026.5.8.1, Panther Lake:** | Metric | Baseline Prefill TPS | Optimized Prefill TPS | Improve | |---|---|---|---| | Prefill-128 | 1289.89 | 1264.72 | -1.95% | | Prefill-256 | 1647.14 | 1631.21 | -0.97% | | Prefill-512 | 1915.64 | 2012.29 | 5.05% | | Prefill-1024 | 1834.62 | 1935.33 | 5.49% | | Prefill-2048 | 1745.07 | 1947.23 | 11.58% | | Prefill-4096 | 1508.13 | 1812.86 | 20.21% | --- .../webgpu/bert/flash_attention.cc | 137 ++++++++ .../contrib_ops/webgpu/bert/flash_attention.h | 84 +++++ ...ash_attention_prefill_simple.wgsl.template | 296 ++++++++++++++++++ 3 files changed, 517 insertions(+) create mode 100644 onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill_simple.wgsl.template diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 717abe8274b2b..f1431992c8548 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -11,6 +11,8 @@ #include "core/providers/webgpu/webgpu_supported_types.h" +#include + using namespace onnxruntime::webgpu; using namespace ::onnxruntime::common; using namespace ONNX_NAMESPACE; @@ -538,6 +540,130 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont return context.RunProgram(program); } +Status FlashAttentionPrefillSimpleProgram::GenerateShaderCode(ShaderHelper& shader) const { + shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + shader.AddInput("present_key", ShaderUsage::UseUniform); + shader.AddInput("present_value", ShaderUsage::UseUniform); + if (has_attention_bias_) { + shader.AddInput("attention_bias", ShaderUsage::UseUniform); + } + if (use_seqlen_k_) { + shader.AddInput("seqlens_k", ShaderUsage::None); + } + if (use_seqlens_q_) { + shader.AddInput("seqlens_q", ShaderUsage::None); + } + if (has_head_sink_) { + shader.AddInput("head_sink", ShaderUsage::UseUniform); + } + shader.AddOutput("output", ShaderUsage::UseUniform); + + return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_prefill_simple.wgsl.template", + WGSL_TEMPLATE_PARAMETER(has_attention_bias, has_attention_bias_), + WGSL_TEMPLATE_PARAMETER(has_head_sink, has_head_sink_), + WGSL_TEMPLATE_PARAMETER(has_local_window, has_local_window_), + WGSL_TEMPLATE_PARAMETER(is_fp16, is_fp16_), + WGSL_TEMPLATE_PARAMETER(is_unidirectional, is_unidirectional_), + WGSL_TEMPLATE_PARAMETER(kv_step_param, kv_step_), + WGSL_TEMPLATE_PARAMETER(q_BNSH, q_BNSH_), + WGSL_TEMPLATE_PARAMETER(qkv_head_size, qkv_head_size_), + WGSL_TEMPLATE_PARAMETER(qkv_num_heads, qkv_num_heads_), + WGSL_TEMPLATE_PARAMETER(use_seqlen_k, use_seqlen_k_), + WGSL_TEMPLATE_PARAMETER(use_seqlens_q, use_seqlens_q_)); +} + +bool CanApplyFlashAttentionPrefillSimple(bool is_intel, + uint32_t subgroup_min_size, + bool kv_cache_quantization_enabled) { + return is_intel && subgroup_min_size >= 8u && !kv_cache_quantization_enabled; +} + +Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& context, + const Tensor* Q, + const Tensor* present_key, + const Tensor* present_value, + const Tensor* attention_bias, + const Tensor* seqlen_k, + Tensor* attn_output, + const WebgpuAttentionParameters& parameters, + bool has_attention_bias, + bool is_fp16, + bool use_seqlen_k, + bool has_local_window, + int local_window_size, + uint32_t present_sequence_length, + uint32_t subgroup_min_size, + float alpha, + uint32_t attn_bias_dim0, + uint32_t attn_bias_dim1, + uint32_t attn_bias_dim3, + bool use_seqlens_q, + const Tensor* seqlens_q, + bool q_BNSH, + bool has_head_sink, + const Tensor* head_sink) { + // The simple prefill kernel tiles the query sequence by its workgroup size + // (one query row per thread), so it can pick a different workgroup size than + // the full kernel. Larger tiles amortize the shared K/V load across more + // query rows (256 is the WebGPU maxComputeInvocationsPerWorkgroup). Causal + // attention does ~half the score work per row, so 128 fits it better, while + // non-causal uses 256. num_seq_tile / dispatch are recomputed from this size + // so the shader's tiling stays consistent. + const uint32_t prefill_simple_wg_size = parameters.is_unidirectional_ ? 128u : 256u; + // The shuffle loop broadcasts kv_step columns across the subgroup, so kv_step + // must not exceed the guaranteed minimum subgroup size: 16 when that minimum + // is >= 16, otherwise 8. + const uint32_t prefill_simple_kv_step = subgroup_min_size >= 16u ? 16u : 8u; + const uint32_t prefill_simple_num_seq_tile = + (parameters.sequence_length_ + prefill_simple_wg_size - 1) / prefill_simple_wg_size; + const uint32_t prefill_simple_dispatch_size = + parameters.batch_size_ * parameters.num_heads_ * prefill_simple_num_seq_tile; + + FlashAttentionPrefillSimpleProgram program{"FlashAttentionPrefillSimple", + has_attention_bias, + is_fp16, + parameters.is_unidirectional_, + parameters.head_size_, + parameters.num_heads_, + use_seqlen_k, + prefill_simple_kv_step, + has_local_window, + use_seqlens_q, + q_BNSH, + has_head_sink}; + program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {present_key, ProgramTensorMetadataDependency::TypeAndRank, 4}, + {present_value, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + if (has_attention_bias) { + program.AddInputs({{attention_bias, ProgramTensorMetadataDependency::TypeAndRank}}); + } + if (use_seqlen_k) { + program.AddInputs({{seqlen_k, ProgramTensorMetadataDependency::None}}); + } + if (use_seqlens_q) { + program.AddInputs({{seqlens_q, ProgramTensorMetadataDependency::None}}); + } + if (has_head_sink) { + program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}}); + } + program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); + program.SetDispatchGroupSize(prefill_simple_dispatch_size) + .SetWorkgroupSize(prefill_simple_wg_size) + .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, use_seqlen_k, prefill_simple_wg_size, prefill_simple_kv_step, has_local_window, use_seqlens_q, q_BNSH, has_head_sink) + .AddUniformVariables({{static_cast(parameters.sequence_length_)}, + {static_cast(parameters.total_sequence_length_)}, + {static_cast(present_sequence_length)}, + {static_cast(parameters.batch_size_)}, + {static_cast(parameters.n_reps)}, + {alpha}, + {prefill_simple_num_seq_tile}, + {attn_bias_dim0}, + {attn_bias_dim1}, + {attn_bias_dim3}, + {static_cast(has_local_window ? local_window_size : 0)}}); + return context.RunProgram(program); +} + Status FlashAttentionDecodeQKVProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& q = shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); const auto& present_key = shader.AddInput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); @@ -1047,6 +1173,17 @@ Status ApplyDensePrefill(const Tensor* Q, Tensor* present_key, Tensor* present_v attn_bias_dim3 = static_cast(bias_shape[3]); } + const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; + const uint32_t subgroup_min_size = has_subgroups ? context.AdapterInfo().subgroupMinSize : 0u; + // TODO: Merge back into `FlashAttentionProgram`. + if (CanApplyFlashAttentionPrefillSimple(is_intel, subgroup_min_size, kv_cache_quantization_enabled)) { + return ApplyFlashAttentionPrefillSimple( + context, Q, present_key, present_value, attention_bias, seqlen_k, attn_output, + parameters, has_attention_bias, is_fp16, use_seqlen_k, has_local_window, local_window_size, + present_sequence_length, subgroup_min_size, alpha, attn_bias_dim0, attn_bias_dim1, attn_bias_dim3, + use_seqlens_q, seqlens_q, q_BNSH, has_head_sink, head_sink); + } + program.SetDispatchGroupSize(parameters.batch_size_ * parameters.num_heads_ * num_seq_tile) .SetWorkgroupSize(prefill_tile_size) .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 7c0eaf624718e..f93657b62f006 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -205,6 +205,90 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont uint32_t max_num_blocks_per_seq, const Tensor* cumulative_seqlens_q = nullptr); +class FlashAttentionPrefillSimpleProgram final : public Program { + public: + FlashAttentionPrefillSimpleProgram(const std::string& kernel_name, + bool has_attention_bias, + bool is_fp16, + bool is_unidirectional, + int qkv_head_size, + int qkv_num_heads, + bool use_seqlen_k, + uint32_t kv_step, + bool has_local_window, + bool use_seqlens_q, + bool q_BNSH, + bool has_head_sink) + : Program{kernel_name}, + has_attention_bias_(has_attention_bias), + is_fp16_(is_fp16), + is_unidirectional_(is_unidirectional), + qkv_head_size_(qkv_head_size), + qkv_num_heads_(qkv_num_heads), + use_seqlen_k_(use_seqlen_k), + kv_step_(kv_step), + has_local_window_(has_local_window), + use_seqlens_q_(use_seqlens_q), + q_BNSH_(q_BNSH), + has_head_sink_(has_head_sink) {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"new_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"batch_size", ProgramUniformVariableDataType::Uint32}, + {"n_reps", ProgramUniformVariableDataType::Uint32}, + {"alpha", ProgramUniformVariableDataType::Float32}, + {"num_seq_tile", ProgramUniformVariableDataType::Uint32}, + {"attn_bias_dim0", ProgramUniformVariableDataType::Uint32}, + {"attn_bias_dim1", ProgramUniformVariableDataType::Uint32}, + {"attn_bias_dim3", ProgramUniformVariableDataType::Uint32}, + {"local_window_size", ProgramUniformVariableDataType::Uint32}); + + private: + bool has_attention_bias_; + bool is_fp16_; + bool is_unidirectional_; + int qkv_head_size_; + int qkv_num_heads_; + bool use_seqlen_k_; + uint32_t kv_step_; + bool has_local_window_; + bool use_seqlens_q_; + bool q_BNSH_; + bool has_head_sink_; +}; + +bool CanApplyFlashAttentionPrefillSimple(bool is_intel, + uint32_t subgroup_min_size, + bool kv_cache_quantization_enabled); + +Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& context, + const Tensor* Q, + const Tensor* present_key, + const Tensor* present_value, + const Tensor* attention_bias, + const Tensor* seqlen_k, + Tensor* attn_output, + const WebgpuAttentionParameters& parameters, + bool has_attention_bias, + bool is_fp16, + bool use_seqlen_k, + bool has_local_window, + int local_window_size, + uint32_t present_sequence_length, + uint32_t subgroup_min_size, + float alpha, + uint32_t attn_bias_dim0, + uint32_t attn_bias_dim1, + uint32_t attn_bias_dim3, + bool use_seqlens_q, + const Tensor* seqlens_q, + bool q_BNSH, + bool has_head_sink, + const Tensor* head_sink); + class FlashAttentionDecodeQKVProgram final : public Program { public: FlashAttentionDecodeQKVProgram(const std::string& kernel_name, diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill_simple.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill_simple.wgsl.template new file mode 100644 index 0000000000000..105c0289c71b5 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill_simple.wgsl.template @@ -0,0 +1,296 @@ + +#param has_attention_bias +#param has_head_sink +#param has_local_window +#param is_fp16 +#param is_unidirectional +#param kv_step_param +#param q_BNSH +#param qkv_head_size +#param qkv_num_heads +#param use_seqlen_k +#param use_seqlens_q + +// Bare, single-path flash-attention prefill. One workgroup per (batch, head, +// query tile); workgroup_size_x threads (host-tunable, 128 or 256), each owning +// one query row. The K/V sequence streams in tiles of kv_step columns broadcast +// across the subgroup via subgroupShuffle, and an online (streaming) softmax +// keeps the full score matrix from ever materializing. +// +// The V tile is loaded only just before the PV matmul (deferred V load) so the +// QK^T + softmax compute overlaps V's global-memory fetch. + +const head_size : u32 = qkv_head_size; +const num_heads : u32 = qkv_num_heads; + +#if use_seqlen_k +// With seqlens_k, total_sequence_length is read per batch from the GPU buffer. +fn get_total_sequence_length(batch_idx: u32) -> u32 { + return u32(seqlens_k[batch_idx]) + 1u; +} +#else +// Without seqlens_k, total_sequence_length comes from uniforms (max across batches). +fn get_total_sequence_length(batch_idx: u32) -> u32 { + return uniforms.total_sequence_length; +} +#endif + +#if is_fp16 +const min_value = q_element_t(-65504.0); +#else +const min_value = q_element_t(-3.4028234663852886e+38f); +#endif + +// Local (sliding) attention window: key k_idx is attended to only when +// local_window_start <= k_idx < causal_end. +fn is_key_visible(k_idx: u32, local_window_start: u32, causal_end: u32) -> bool { + return k_idx >= local_window_start && k_idx < causal_end; +} + +// Max K/V columns handled per iteration (one per subgroup lane); the host picks +// 8 or 16 to fit the adapter's guaranteed minimum subgroup size. The head +// dimension is expressed in vec4 units. +const kv_step : u32 = kv_step_param; +const vec_factor : u32 = 4u; +const head_size_vec : u32 = head_size / vec_factor; +// Score vec4 lanes per tile: kv_step columns packed 4 per vec4. +const num_qk_vec : u32 = kv_step / vec_factor; + +// K/V tiles shared and shuffled across the subgroup so each lane contributes the +// column it loaded. Stored as [head_size_vec][kv_step] so the subgroup column +// (kv_step) is the inner/fastest dimension: adjacent subgroup lanes hit adjacent +// addresses, avoiding shared-memory bank conflicts. +var k_tile : array, head_size_vec>; +var v_tile : array, head_size_vec>; + +// Per-thread query row and output accumulator live as local vars in $MAIN +// (o_tile is zero-initialized). + +// Cooperatively load k_step key columns starting at k_start into shared memory. +// Columns past total_seq are zeroed so out-of-range lanes contribute nothing. +fn loadk(k_start : u32, batch_head_idx : u32, local_idx : u32, k_step : u32, total_seq : u32) { + let offset = batch_head_idx / uniforms.n_reps * uniforms.present_sequence_length * head_size_vec + + k_start * head_size_vec; + for (var idx : u32 = local_idx; idx < head_size_vec * k_step; idx += workgroup_size_x) { + let slot = u32(idx / head_size_vec); + let val = select(q_value_t(0), present_key[offset + idx], k_start + slot < total_seq); + k_tile[idx % head_size_vec][slot] = val; + } +} + +// Same as loadk, for the value columns. +fn loadv(v_start : u32, batch_head_idx : u32, local_idx : u32, v_step : u32, total_seq : u32) { + let offset = batch_head_idx / uniforms.n_reps * uniforms.present_sequence_length * head_size_vec + + v_start * head_size_vec; + for (var idx : u32 = local_idx; idx < head_size_vec * v_step; idx += workgroup_size_x) { + let slot = u32(idx / head_size_vec); + let val = select(q_value_t(0), present_value[offset + idx], v_start + slot < total_seq); + v_tile[idx % head_size_vec][slot] = val; + } +} + +// Additive attention bias for score (q_idx, k_idx). +#if has_attention_bias +fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, head_idx : u32, total_seq : u32) -> q_element_t { + // Stored as [batch_size, num_heads, new_seq_length, total_sequence_length]. + if (k_idx_global >= total_seq) { + return q_element_t(0); + } + // Broadcasting: if a dimension size is 1, use index 0. + let bias_batch_idx = select(batch_idx, 0u, batch_idx >= uniforms.attn_bias_dim0); + let bias_head_idx = select(head_idx, 0u, head_idx >= uniforms.attn_bias_dim1); + // Stride along the last dim of attention_bias matches its actual shape, which may + // differ from the per-step total_sequence_length (e.g. graph-capture sets that + // uniform to 0, or batches use a global max). The host passes attn_bias_dim3. + let stride_total_seq = uniforms.attn_bias_dim3; + let offset_base = bias_batch_idx * uniforms.attn_bias_dim1 * uniforms.new_sequence_length * stride_total_seq + + bias_head_idx * uniforms.new_sequence_length * stride_total_seq + q_idx_global * stride_total_seq; + return q_element_t(attention_bias[min(offset_base + k_idx_global, offset_base + stride_total_seq - 1u)]); +} +#else +fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, head_idx : u32, total_seq : u32) -> q_element_t { + return q_element_t(0); +} +#endif + +$MAIN { + // Decode (batch, head) from the flattened workgroup index. + let batch_head_idx = u32(workgroup_idx / uniforms.num_seq_tile); + let head_idx = batch_head_idx % num_heads; + let batch_idx = batch_head_idx / num_heads; + + if (batch_idx >= uniforms.batch_size) { + return; + } + + // Per-thread query row and output accumulator (zero-initialized). + var q_tile : array; + var o_tile : array; + + // This thread's query row within the tile; may run past the sequence end. + let seq_tile_idx = workgroup_idx % uniforms.num_seq_tile; + let q_idx_global = seq_tile_idx * workgroup_size_x + local_idx; + let valid_q = q_idx_global < uniforms.new_sequence_length; + if (valid_q) { + // Load this thread's query row into q_tile, pre-scaled by alpha. + let alpha = q_element_t(uniforms.alpha); +#if q_BNSH + let q_offset = batch_idx * num_heads * uniforms.new_sequence_length * head_size_vec + + head_idx * uniforms.new_sequence_length * head_size_vec + + q_idx_global * head_size_vec; +#else + let q_offset = batch_idx * uniforms.new_sequence_length * head_size_vec * num_heads + + q_idx_global * head_size_vec * num_heads + + head_idx * head_size_vec; +#endif + for (var idx = 0u; idx < head_size_vec; idx++) { + q_tile[idx] = q[idx + q_offset] * alpha; + } + } + + // Running statistics for the online (streaming) softmax. +#if has_head_sink + let sink_value = q_element_t(head_sink[head_idx]); + var previous_max : q_element_t = sink_value; + var previous_denom : q_element_t = 1; +#else + var previous_max : q_element_t = min_value; + var previous_denom : q_element_t = 0; +#endif + let total_sequence_length = get_total_sequence_length(batch_idx); + +#if is_unidirectional + // Causal: query i attends to keys [0, i]. +#if use_seqlens_q + // LEFT-aligned per-batch Q length: past_len_b = total_kv_b - q_len_b (>= 0). + let q_len_b = u32(seqlens_q[batch_idx]); + let past_sequence_length = total_sequence_length - q_len_b; +#else + // Uniform q_len path. Right-padded batches with prompt shorter than + // new_sequence_length would underflow u32, so clamp to 0. Variable-q_len + // callers with past > 0 must set use_seqlens_q instead. + let past_sequence_length = select(total_sequence_length - uniforms.new_sequence_length, + 0u, + total_sequence_length <= uniforms.new_sequence_length); +#endif + let max_causal_len_for_workgroup = past_sequence_length + + (seq_tile_idx + 1) * workgroup_size_x; + let loop_bound = min(total_sequence_length, max_causal_len_for_workgroup); + let seq_causal_length = past_sequence_length + q_idx_global + 1; +#else + // Non-causal: every query attends to the whole sequence. + let loop_bound = total_sequence_length; + let seq_causal_length = total_sequence_length; +#endif + +#if has_local_window +#if is_unidirectional + let first_query_causal_length = past_sequence_length + seq_tile_idx * workgroup_size_x + 1u; +#else + let first_query_causal_length = total_sequence_length; +#endif + let first_query_window_start = first_query_causal_length - + min(first_query_causal_length, uniforms.local_window_size); + let local_window_start = seq_causal_length - min(seq_causal_length, uniforms.local_window_size); +#else + let first_query_window_start = 0u; + let local_window_start = 0u; +#endif + // Tiles entirely before the sliding window for every query in this workgroup + // are never visible, so skip straight to the first tile the window can reach. + let aligned_window_start = (first_query_window_start / kv_step) * kv_step; + + // sg_size >= kv_step is guaranteed by the host gate, so each iteration advances + // by a full kv_step and no step cap is needed. sg_size may still exceed kv_step + // (e.g. 32): the QK/PV loads index k_tile/v_tile with the raw sg_id, so lanes + // >= kv_step read past kv_step (relies on WGSL bounds clamping), but their + // k_local is never read by the 0..kv_step-1 shuffles. + + // Stream the K/V sequence one tile at a time. + for (var k_start = aligned_window_start; k_start < loop_bound; k_start += kv_step) { + // Load only K here; V is deferred until just before the PV matmul so the + // QK^T + softmax compute overlaps V's global-memory fetch. The prior + // iteration's trailing barrier (after loadv) still guards k_tile from being + // overwritten while its readers are in flight, so no leading barrier is needed. + loadk(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); + workgroupBarrier(); + + // Scores for this tile, held as num_qk_vec vec4 lanes (kv_step columns); + // qk[j / 4u][j % 4u] is this query's score against column j. The inner loops + // iterate over the compile-time const kv_step so the compiler unrolls them + // and each subgroupShuffle lane index folds to a constant. Requires + // sg_size >= kv_step, which the host guarantees before selecting this shader + // (subgroup_min_size_ge_8 gate; kv_step is 8 or 16 to match). + var qk : array, num_qk_vec>; + + // 1. QK^T. Each lane owns one K column of this tile; subgroup shuffles let + // every thread read all columns and accumulate the dot products. + for (var i = 0u; i < head_size_vec; i++) { + let q_own = q_tile[i]; + let k_local = k_tile[i][sg_id]; + for (var j = 0u; j < kv_step; j++) { + qk[j / 4u][j % 4u] += dot(q_own, subgroupShuffle(k_local, j)); + } + } + + // 2. Add attention bias, then mask columns outside the causal/window boundary. + for (var j = 0u; j < kv_step; j++) { + qk[j / 4u][j % 4u] += loadAttentionBias(batch_idx, q_idx_global, k_start + j, head_idx, total_sequence_length); + qk[j / 4u][j % 4u] = select(min_value, qk[j / 4u][j % 4u], is_key_visible(k_start + j, local_window_start, seq_causal_length)); + } + + // 3. Online softmax. Refresh the running max over this tile's scores. + var local_max_vec = qk[0]; + for (var m = 1u; m < num_qk_vec; m++) { + local_max_vec = max(local_max_vec, qk[m]); + } + let local_max = max(max(local_max_vec.x, local_max_vec.y), max(local_max_vec.z, local_max_vec.w)); + let new_max = max(previous_max, local_max); + + // Exponentiate against the new max (in f32 for stability) and accumulate + // this tile's denominator contribution. + var sum_vec = q_value_t(0); + for (var m = 0u; m < num_qk_vec; m++) { + qk[m] = q_value_t(exp(vec4(qk[m]) - f32(new_max))); + sum_vec += qk[m]; + } + let sum = sum_vec.x + sum_vec.y + sum_vec.z + sum_vec.w; + + // Rescale the carried-over denominator to the new max, then fold in the + // tile. Guard against a zero denominator. + let dleft = previous_denom * q_element_t(exp(f32(previous_max) - f32(new_max))); + var d = dleft + sum; + d = select(d, q_element_t(0.0000001), d == 0); + for (var m = 0u; m < num_qk_vec; m++) { + qk[m] = qk[m] / d; + } + previous_max = new_max; + previous_denom = d; + // Factor to rescale the already-accumulated output for the new denominator. + let o_ratio = dleft / d; + + // Deferred V load: fetch this tile's V now, right before it is consumed. + loadv(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); + workgroupBarrier(); + + // 4. Accumulate softmax(qk) * V into the output tile, rescaling the prior + // partial sum by o_ratio. + for (var i = 0u; i < head_size_vec; i++) { + let v_local = v_tile[i][sg_id]; + var acc = q_value_t(0); + for (var j = 0u; j < kv_step; j++) { + acc += subgroupShuffle(v_local, j) * qk[j / 4u][j % 4u]; + } + o_tile[i] = o_tile[i] * o_ratio + acc; + } + } + + if (valid_q) { + // Write this thread's normalized output row. + let o_offset = batch_idx * uniforms.new_sequence_length * num_heads * head_size_vec + + q_idx_global * num_heads * head_size_vec + head_idx * head_size_vec; + for (var idx = 0u; idx < head_size_vec; idx++) { + output[o_offset + idx] = o_tile[idx]; + } + } +} // MAIN From 811d7d0e66523c9fb67a8bedf097c1801ee43f12 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Tue, 22 Sep 2026 13:14:30 +0800 Subject: [PATCH 2/7] =?UTF-8?q?=EF=BB=BFFixes=20GroupQueryAttentionTest.We?= =?UTF-8?q?bGPU=5FFP16=5FHighMagnitude=5FFlashPrefill.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../webgpu/bert/flash_attention.cc | 6 +- .../contrib_ops/webgpu/bert/flash_attention.h | 4 -- ...ash_attention_prefill_simple.wgsl.template | 56 ++++++++++--------- 3 files changed, 33 insertions(+), 33 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index f1431992c8548..2129db650eb0e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -562,10 +562,10 @@ Status FlashAttentionPrefillSimpleProgram::GenerateShaderCode(ShaderHelper& shad WGSL_TEMPLATE_PARAMETER(has_attention_bias, has_attention_bias_), WGSL_TEMPLATE_PARAMETER(has_head_sink, has_head_sink_), WGSL_TEMPLATE_PARAMETER(has_local_window, has_local_window_), - WGSL_TEMPLATE_PARAMETER(is_fp16, is_fp16_), WGSL_TEMPLATE_PARAMETER(is_unidirectional, is_unidirectional_), WGSL_TEMPLATE_PARAMETER(kv_step_param, kv_step_), WGSL_TEMPLATE_PARAMETER(q_BNSH, q_BNSH_), + WGSL_TEMPLATE_PARAMETER(qk_is_fp16, false), WGSL_TEMPLATE_PARAMETER(qkv_head_size, qkv_head_size_), WGSL_TEMPLATE_PARAMETER(qkv_num_heads, qkv_num_heads_), WGSL_TEMPLATE_PARAMETER(use_seqlen_k, use_seqlen_k_), @@ -587,7 +587,6 @@ Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& con Tensor* attn_output, const WebgpuAttentionParameters& parameters, bool has_attention_bias, - bool is_fp16, bool use_seqlen_k, bool has_local_window, int local_window_size, @@ -621,7 +620,6 @@ Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& con FlashAttentionPrefillSimpleProgram program{"FlashAttentionPrefillSimple", has_attention_bias, - is_fp16, parameters.is_unidirectional_, parameters.head_size_, parameters.num_heads_, @@ -1179,7 +1177,7 @@ Status ApplyDensePrefill(const Tensor* Q, Tensor* present_key, Tensor* present_v if (CanApplyFlashAttentionPrefillSimple(is_intel, subgroup_min_size, kv_cache_quantization_enabled)) { return ApplyFlashAttentionPrefillSimple( context, Q, present_key, present_value, attention_bias, seqlen_k, attn_output, - parameters, has_attention_bias, is_fp16, use_seqlen_k, has_local_window, local_window_size, + parameters, has_attention_bias, use_seqlen_k, has_local_window, local_window_size, present_sequence_length, subgroup_min_size, alpha, attn_bias_dim0, attn_bias_dim1, attn_bias_dim3, use_seqlens_q, seqlens_q, q_BNSH, has_head_sink, head_sink); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index f93657b62f006..9877c3036145d 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -209,7 +209,6 @@ class FlashAttentionPrefillSimpleProgram final : public Program u32 { } #endif -#if is_fp16 -const min_value = q_element_t(-65504.0); +// Default to f32 scores and online-softmax statistics to avoid fp16 overflow +// with high-magnitude inputs, independently of the Q/K/V storage type. +#if qk_is_fp16 +alias qk_precision = f16; +const min_value = qk_precision(-65504.0); #else -const min_value = q_element_t(-3.4028234663852886e+38f); +alias qk_precision = f32; +const min_value = qk_precision(-3.4028234663852886e+38f); #endif // Local (sliding) attention window: key k_idx is attended to only when @@ -91,10 +95,10 @@ fn loadv(v_start : u32, batch_head_idx : u32, local_idx : u32, v_step : u32, tot // Additive attention bias for score (q_idx, k_idx). #if has_attention_bias -fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, head_idx : u32, total_seq : u32) -> q_element_t { +fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, head_idx : u32, total_seq : u32) -> qk_precision { // Stored as [batch_size, num_heads, new_seq_length, total_sequence_length]. if (k_idx_global >= total_seq) { - return q_element_t(0); + return qk_precision(0); } // Broadcasting: if a dimension size is 1, use index 0. let bias_batch_idx = select(batch_idx, 0u, batch_idx >= uniforms.attn_bias_dim0); @@ -105,11 +109,11 @@ fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, he let stride_total_seq = uniforms.attn_bias_dim3; let offset_base = bias_batch_idx * uniforms.attn_bias_dim1 * uniforms.new_sequence_length * stride_total_seq + bias_head_idx * uniforms.new_sequence_length * stride_total_seq + q_idx_global * stride_total_seq; - return q_element_t(attention_bias[min(offset_base + k_idx_global, offset_base + stride_total_seq - 1u)]); + return qk_precision(attention_bias[min(offset_base + k_idx_global, offset_base + stride_total_seq - 1u)]); } #else -fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, head_idx : u32, total_seq : u32) -> q_element_t { - return q_element_t(0); +fn loadAttentionBias(batch_idx : u32, q_idx_global : u32, k_idx_global : u32, head_idx : u32, total_seq : u32) -> qk_precision { + return qk_precision(0); } #endif @@ -150,12 +154,12 @@ $MAIN { // Running statistics for the online (streaming) softmax. #if has_head_sink - let sink_value = q_element_t(head_sink[head_idx]); - var previous_max : q_element_t = sink_value; - var previous_denom : q_element_t = 1; + let sink_value = qk_precision(head_sink[head_idx]); + var previous_max : qk_precision = sink_value; + var previous_denom : qk_precision = 1; #else - var previous_max : q_element_t = min_value; - var previous_denom : q_element_t = 0; + var previous_max : qk_precision = min_value; + var previous_denom : qk_precision = 0; #endif let total_sequence_length = get_total_sequence_length(batch_idx); @@ -221,15 +225,16 @@ $MAIN { // and each subgroupShuffle lane index folds to a constant. Requires // sg_size >= kv_step, which the host guarantees before selecting this shader // (subgroup_min_size_ge_8 gate; kv_step is 8 or 16 to match). - var qk : array, num_qk_vec>; + var qk : array, num_qk_vec>; // 1. QK^T. Each lane owns one K column of this tile; subgroup shuffles let - // every thread read all columns and accumulate the dot products. + // every thread read all columns and accumulate the dot products. Native + // precision vec4 dots are converted to qk_precision before accumulation. for (var i = 0u; i < head_size_vec; i++) { let q_own = q_tile[i]; let k_local = k_tile[i][sg_id]; for (var j = 0u; j < kv_step; j++) { - qk[j / 4u][j % 4u] += dot(q_own, subgroupShuffle(k_local, j)); + qk[j / 4u][j % 4u] += qk_precision(dot(q_own, subgroupShuffle(k_local, j))); } } @@ -247,27 +252,28 @@ $MAIN { let local_max = max(max(local_max_vec.x, local_max_vec.y), max(local_max_vec.z, local_max_vec.w)); let new_max = max(previous_max, local_max); - // Exponentiate against the new max (in f32 for stability) and accumulate + // Exponentiate against the new max in qk_precision and accumulate // this tile's denominator contribution. - var sum_vec = q_value_t(0); + var sum_vec = vec4(0); for (var m = 0u; m < num_qk_vec; m++) { - qk[m] = q_value_t(exp(vec4(qk[m]) - f32(new_max))); + qk[m] = exp(qk[m] - new_max); sum_vec += qk[m]; } let sum = sum_vec.x + sum_vec.y + sum_vec.z + sum_vec.w; // Rescale the carried-over denominator to the new max, then fold in the // tile. Guard against a zero denominator. - let dleft = previous_denom * q_element_t(exp(f32(previous_max) - f32(new_max))); + let dleft = previous_denom * exp(previous_max - new_max); var d = dleft + sum; - d = select(d, q_element_t(0.0000001), d == 0); + d = select(d, qk_precision(0.0000001), d == 0); + var qk_normalized : array, num_qk_vec>; for (var m = 0u; m < num_qk_vec; m++) { - qk[m] = qk[m] / d; + qk_normalized[m] = q_value_t(qk[m] / d); } previous_max = new_max; previous_denom = d; // Factor to rescale the already-accumulated output for the new denominator. - let o_ratio = dleft / d; + let o_ratio = q_element_t(dleft / d); // Deferred V load: fetch this tile's V now, right before it is consumed. loadv(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); @@ -279,7 +285,7 @@ $MAIN { let v_local = v_tile[i][sg_id]; var acc = q_value_t(0); for (var j = 0u; j < kv_step; j++) { - acc += subgroupShuffle(v_local, j) * qk[j / 4u][j % 4u]; + acc += subgroupShuffle(v_local, j) * qk_normalized[j / 4u][j % 4u]; } o_tile[i] = o_tile[i] * o_ratio + acc; } From 92de4dcd43b07be74d9d78dd450c0366411e4d49 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Fri, 16 Oct 2026 19:17:27 +0800 Subject: [PATCH 3/7] webgpu: optimize simple flash attention prefill with tiled K/V staging --- .../webgpu/bert/flash_attention.cc | 56 ++++--- .../contrib_ops/webgpu/bert/flash_attention.h | 4 +- ...ash_attention_prefill_simple.wgsl.template | 154 +++++++++--------- 3 files changed, 115 insertions(+), 99 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 2129db650eb0e..e66a85eb6face 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -565,17 +565,17 @@ Status FlashAttentionPrefillSimpleProgram::GenerateShaderCode(ShaderHelper& shad WGSL_TEMPLATE_PARAMETER(is_unidirectional, is_unidirectional_), WGSL_TEMPLATE_PARAMETER(kv_step_param, kv_step_), WGSL_TEMPLATE_PARAMETER(q_BNSH, q_BNSH_), - WGSL_TEMPLATE_PARAMETER(qk_is_fp16, false), + WGSL_TEMPLATE_PARAMETER(qk_f32, true), WGSL_TEMPLATE_PARAMETER(qkv_head_size, qkv_head_size_), WGSL_TEMPLATE_PARAMETER(qkv_num_heads, qkv_num_heads_), + WGSL_TEMPLATE_PARAMETER(shuffle_chunk_param, shuffle_chunk_), WGSL_TEMPLATE_PARAMETER(use_seqlen_k, use_seqlen_k_), WGSL_TEMPLATE_PARAMETER(use_seqlens_q, use_seqlens_q_)); } bool CanApplyFlashAttentionPrefillSimple(bool is_intel, - uint32_t subgroup_min_size, bool kv_cache_quantization_enabled) { - return is_intel && subgroup_min_size >= 8u && !kv_cache_quantization_enabled; + return is_intel && !kv_cache_quantization_enabled; } Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& context, @@ -601,22 +601,25 @@ Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& con bool q_BNSH, bool has_head_sink, const Tensor* head_sink) { - // The simple prefill kernel tiles the query sequence by its workgroup size - // (one query row per thread), so it can pick a different workgroup size than - // the full kernel. Larger tiles amortize the shared K/V load across more - // query rows (256 is the WebGPU maxComputeInvocationsPerWorkgroup). Causal - // attention does ~half the score work per row, so 128 fits it better, while - // non-causal uses 256. num_seq_tile / dispatch are recomputed from this size - // so the shader's tiling stays consistent. - const uint32_t prefill_simple_wg_size = parameters.is_unidirectional_ ? 128u : 256u; - // The shuffle loop broadcasts kv_step columns across the subgroup, so kv_step - // must not exceed the guaranteed minimum subgroup size: 16 when that minimum - // is >= 16, otherwise 8. - const uint32_t prefill_simple_kv_step = subgroup_min_size >= 16u ? 16u : 8u; - const uint32_t prefill_simple_num_seq_tile = - (parameters.sequence_length_ + prefill_simple_wg_size - 1) / prefill_simple_wg_size; - const uint32_t prefill_simple_dispatch_size = - parameters.batch_size_ * parameters.num_heads_ * prefill_simple_num_seq_tile; + constexpr uint32_t wg_size = 128u; + const uint32_t head_size_vec = static_cast(parameters.head_size_) / 4u; + const uint32_t v_head_vec_step = std::min(16u, head_size_vec); + // Full-head K staging and sliced V staging share the workgroup storage budget. + const uint64_t bytes_per_column = + (uint64_t{head_size_vec} + v_head_vec_step) * 4u * Q->DataType()->Size(); + const uint64_t max_workgroup_storage = context.DeviceLimits().maxComputeWorkgroupStorageSize; + uint32_t kv_step = 64u; + while (kv_step > 8u && bytes_per_column * kv_step > max_workgroup_storage) { + kv_step /= 2u; + } + ORT_RETURN_IF(bytes_per_column * kv_step > max_workgroup_storage, + "Simple prefill K/V tiles exceed workgroup storage limits."); + // Values below 8 select direct shared-memory reads, including without subgroups. + const uint32_t shuffle_chunk = std::min(kv_step, subgroup_min_size); + const uint32_t num_seq_tile = + (parameters.sequence_length_ + wg_size - 1) / wg_size; + const uint32_t dispatch_size = + parameters.batch_size_ * parameters.num_heads_ * num_seq_tile; FlashAttentionPrefillSimpleProgram program{"FlashAttentionPrefillSimple", has_attention_bias, @@ -624,7 +627,8 @@ Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& con parameters.head_size_, parameters.num_heads_, use_seqlen_k, - prefill_simple_kv_step, + kv_step, + shuffle_chunk, has_local_window, use_seqlens_q, q_BNSH, @@ -645,16 +649,18 @@ Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& con program.AddInputs({{head_sink, ProgramTensorMetadataDependency::Type}}); } program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); - program.SetDispatchGroupSize(prefill_simple_dispatch_size) - .SetWorkgroupSize(prefill_simple_wg_size) - .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, use_seqlen_k, prefill_simple_wg_size, prefill_simple_kv_step, has_local_window, use_seqlens_q, q_BNSH, has_head_sink) + program.SetDispatchGroupSize(dispatch_size) + .SetWorkgroupSize(wg_size) + .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, + use_seqlen_k, wg_size, kv_step, shuffle_chunk, + has_local_window, use_seqlens_q, q_BNSH, has_head_sink) .AddUniformVariables({{static_cast(parameters.sequence_length_)}, {static_cast(parameters.total_sequence_length_)}, {static_cast(present_sequence_length)}, {static_cast(parameters.batch_size_)}, {static_cast(parameters.n_reps)}, {alpha}, - {prefill_simple_num_seq_tile}, + {num_seq_tile}, {attn_bias_dim0}, {attn_bias_dim1}, {attn_bias_dim3}, @@ -1174,7 +1180,7 @@ Status ApplyDensePrefill(const Tensor* Q, Tensor* present_key, Tensor* present_v const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; const uint32_t subgroup_min_size = has_subgroups ? context.AdapterInfo().subgroupMinSize : 0u; // TODO: Merge back into `FlashAttentionProgram`. - if (CanApplyFlashAttentionPrefillSimple(is_intel, subgroup_min_size, kv_cache_quantization_enabled)) { + if (CanApplyFlashAttentionPrefillSimple(is_intel, kv_cache_quantization_enabled)) { return ApplyFlashAttentionPrefillSimple( context, Q, present_key, present_value, attention_bias, seqlen_k, attn_output, parameters, has_attention_bias, use_seqlen_k, has_local_window, local_window_size, diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 9877c3036145d..daf51abd1d75c 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -214,6 +214,7 @@ class FlashAttentionPrefillSimpleProgram final : public Program u32 { } #endif -// Default to f32 scores and online-softmax statistics to avoid fp16 overflow -// with high-magnitude inputs, independently of the Q/K/V storage type. -#if qk_is_fp16 -alias qk_precision = f16; -const min_value = qk_precision(-65504.0); -#else +// Accumulate scores and softmax statistics in f32 to avoid fp16 overflow. +// Individual dot products retain the Q/K storage precision. +#if qk_f32 alias qk_precision = f32; const min_value = qk_precision(-3.4028234663852886e+38f); +#else +alias qk_precision = q_element_t; +const min_value = q_element_t(-65504.0); #endif // Local (sliding) attention window: key k_idx is attended to only when @@ -51,24 +47,19 @@ fn is_key_visible(k_idx: u32, local_window_start: u32, causal_end: u32) -> bool return k_idx >= local_window_start && k_idx < causal_end; } -// Max K/V columns handled per iteration (one per subgroup lane); the host picks -// 8 or 16 to fit the adapter's guaranteed minimum subgroup size. The head -// dimension is expressed in vec4 units. const kv_step : u32 = kv_step_param; +#if shuffle_chunk_param >= 8 +const shuffle_chunk : u32 = shuffle_chunk_param; +#endif const vec_factor : u32 = 4u; const head_size_vec : u32 = head_size / vec_factor; -// Score vec4 lanes per tile: kv_step columns packed 4 per vec4. const num_qk_vec : u32 = kv_step / vec_factor; +const k_head_vec_step : u32 = head_size_vec; +const v_head_vec_step : u32 = min(16u, head_size_vec); -// K/V tiles shared and shuffled across the subgroup so each lane contributes the -// column it loaded. Stored as [head_size_vec][kv_step] so the subgroup column -// (kv_step) is the inner/fastest dimension: adjacent subgroup lanes hit adjacent -// addresses, avoiding shared-memory bank conflicts. -var k_tile : array, head_size_vec>; -var v_tile : array, head_size_vec>; - -// Per-thread query row and output accumulator live as local vars in $MAIN -// (o_tile is zero-initialized). +// Adjacent subgroup lanes access adjacent workgroup-memory addresses. +var k_tile : array, k_head_vec_step>; +var v_tile : array, v_head_vec_step>; // Cooperatively load k_step key columns starting at k_start into shared memory. // Columns past total_seq are zeroed so out-of-range lanes contribute nothing. @@ -82,14 +73,16 @@ fn loadk(k_start : u32, batch_head_idx : u32, local_idx : u32, k_step : u32, tot } } -// Same as loadk, for the value columns. -fn loadv(v_start : u32, batch_head_idx : u32, local_idx : u32, v_step : u32, total_seq : u32) { - let offset = batch_head_idx / uniforms.n_reps * uniforms.present_sequence_length * head_size_vec + - v_start * head_size_vec; - for (var idx : u32 = local_idx; idx < head_size_vec * v_step; idx += workgroup_size_x) { - let slot = u32(idx / head_size_vec); - let val = select(q_value_t(0), present_value[offset + idx], v_start + slot < total_seq); - v_tile[idx % head_size_vec][slot] = val; +fn loadv_head(k_start: u32, head_start: u32, batch_head_idx: u32, + local_idx: u32, total_seq: u32) { + let base = batch_head_idx / uniforms.n_reps * uniforms.present_sequence_length * head_size_vec; + for (var idx = local_idx; idx < v_head_vec_step * kv_step; idx += workgroup_size_x) { + let slot = idx / v_head_vec_step; + let h = idx % v_head_vec_step; + let val = select(q_value_t(0), + present_value[base + (k_start + slot) * head_size_vec + head_start + h], + k_start + slot < total_seq && head_start + h < head_size_vec); + v_tile[h][slot] = val; } } @@ -204,38 +197,41 @@ $MAIN { // are never visible, so skip straight to the first tile the window can reach. let aligned_window_start = (first_query_window_start / kv_step) * kv_step; - // sg_size >= kv_step is guaranteed by the host gate, so each iteration advances - // by a full kv_step and no step cap is needed. sg_size may still exceed kv_step - // (e.g. 32): the QK/PV loads index k_tile/v_tile with the raw sg_id, so lanes - // >= kv_step read past kv_step (relies on WGSL bounds clamping), but their - // k_local is never read by the 0..kv_step-1 shuffles. - - // Stream the K/V sequence one tile at a time. for (var k_start = aligned_window_start; k_start < loop_bound; k_start += kv_step) { - // Load only K here; V is deferred until just before the PV matmul so the - // QK^T + softmax compute overlaps V's global-memory fetch. The prior - // iteration's trailing barrier (after loadv) still guards k_tile from being - // overwritten while its readers are in flight, so no leading barrier is needed. + // The previous iteration's PV barriers protect K reads before this overwrite. loadk(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); workgroupBarrier(); - // Scores for this tile, held as num_qk_vec vec4 lanes (kv_step columns); - // qk[j / 4u][j % 4u] is this query's score against column j. The inner loops - // iterate over the compile-time const kv_step so the compiler unrolls them - // and each subgroupShuffle lane index folds to a constant. Requires - // sg_size >= kv_step, which the host guarantees before selecting this shader - // (subgroup_min_size_ge_8 gate; kv_step is 8 or 16 to match). var qk : array, num_qk_vec>; - // 1. QK^T. Each lane owns one K column of this tile; subgroup shuffles let - // every thread read all columns and accumulate the dot products. Native - // precision vec4 dots are converted to qk_precision before accumulation. + // 1. QK^T. Native-precision dots are widened before score accumulation. for (var i = 0u; i < head_size_vec; i++) { let q_own = q_tile[i]; - let k_local = k_tile[i][sg_id]; - for (var j = 0u; j < kv_step; j++) { - qk[j / 4u][j % 4u] += qk_precision(dot(q_own, subgroupShuffle(k_local, j))); +#if shuffle_chunk_param >= 8 + for (var band = 0u; band < kv_step / shuffle_chunk; band++) { + let k_local = k_tile[i][band * shuffle_chunk + (sg_id % shuffle_chunk)]; + for (var m = 0u; m < shuffle_chunk / 4u; m++) { + let j = m * 4u; + qk[band * (shuffle_chunk / 4u) + m] += vec4( + qk_precision(dot(q_own, subgroupShuffle(k_local, j))), + qk_precision(dot(q_own, subgroupShuffle(k_local, j + 1u))), + qk_precision(dot(q_own, subgroupShuffle(k_local, j + 2u))), + qk_precision(dot(q_own, subgroupShuffle(k_local, j + 3u))) + ); + } + } +#else + // Read four K columns directly from shared memory per update. + for (var m = 0u; m < num_qk_vec; m++) { + let j = m * 4u; + qk[m] += vec4( + qk_precision(dot(q_own, k_tile[i][j])), + qk_precision(dot(q_own, k_tile[i][j + 1u])), + qk_precision(dot(q_own, k_tile[i][j + 2u])), + qk_precision(dot(q_own, k_tile[i][j + 3u])) + ); } +#endif } // 2. Add attention bias, then mask columns outside the causal/window boundary. @@ -252,20 +248,22 @@ $MAIN { let local_max = max(max(local_max_vec.x, local_max_vec.y), max(local_max_vec.z, local_max_vec.w)); let new_max = max(previous_max, local_max); - // Exponentiate against the new max in qk_precision and accumulate + // Exponentiate against the new max (in f32 for stability) and accumulate // this tile's denominator contribution. var sum_vec = vec4(0); for (var m = 0u; m < num_qk_vec; m++) { - qk[m] = exp(qk[m] - new_max); + qk[m] = vec4(exp(vec4(qk[m]) - f32(new_max))); sum_vec += qk[m]; } let sum = sum_vec.x + sum_vec.y + sum_vec.z + sum_vec.w; // Rescale the carried-over denominator to the new max, then fold in the // tile. Guard against a zero denominator. - let dleft = previous_denom * exp(previous_max - new_max); + let dleft = previous_denom * qk_precision(exp(f32(previous_max) - f32(new_max))); var d = dleft + sum; d = select(d, qk_precision(0.0000001), d == 0); + // Only the normalized probabilities are cast to q_element_t, right before + // the PV matmul consumes them below. var qk_normalized : array, num_qk_vec>; for (var m = 0u; m < num_qk_vec; m++) { qk_normalized[m] = q_value_t(qk[m] / d); @@ -274,20 +272,30 @@ $MAIN { previous_denom = d; // Factor to rescale the already-accumulated output for the new denominator. let o_ratio = q_element_t(dleft / d); - - // Deferred V load: fetch this tile's V now, right before it is consumed. - loadv(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); - workgroupBarrier(); - // 4. Accumulate softmax(qk) * V into the output tile, rescaling the prior // partial sum by o_ratio. - for (var i = 0u; i < head_size_vec; i++) { - let v_local = v_tile[i][sg_id]; - var acc = q_value_t(0); - for (var j = 0u; j < kv_step; j++) { - acc += subgroupShuffle(v_local, j) * qk_normalized[j / 4u][j % 4u]; + for (var head_start = 0u; head_start < head_size_vec; head_start += v_head_vec_step) { + loadv_head(k_start, head_start, batch_head_idx, local_idx, total_sequence_length); + workgroupBarrier(); + for (var i = 0u; i < v_head_vec_step; i++) { + var acc = q_value_t(0); +#if shuffle_chunk_param >= 8 + for (var band = 0u; band < kv_step / shuffle_chunk; band++) { + let v_local = v_tile[i][band * shuffle_chunk + (sg_id % shuffle_chunk)]; + for (var j = 0u; j < shuffle_chunk; j++) { + let col = band * shuffle_chunk + j; + acc += subgroupShuffle(v_local, j) * qk_normalized[col / 4u][col % 4u]; + } + } +#else + for (var j = 0u; j < kv_step; j++) { + acc += v_tile[i][j] * qk_normalized[j / 4u][j % 4u]; + } +#endif + o_tile[head_start + i] = o_tile[head_start + i] * o_ratio + acc; } - o_tile[i] = o_tile[i] * o_ratio + acc; + // Protect V reads before the next head slice or K/V tile. + workgroupBarrier(); } } From d462e023383b52b8a3ffac13f8757f1b2dd7bd04 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Mon, 19 Oct 2026 12:54:39 +0800 Subject: [PATCH 4/7] Update --- .../webgpu/bert/flash_attention.cc | 53 +++++++++++-------- .../contrib_ops/webgpu/bert/flash_attention.h | 8 +-- ... => flash_attention_prefill.wgsl.template} | 10 +--- 3 files changed, 36 insertions(+), 35 deletions(-) rename onnxruntime/contrib_ops/webgpu/bert/{flash_attention_prefill_simple.wgsl.template => flash_attention_prefill.wgsl.template} (95%) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index e66a85eb6face..1790c14828663 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -10,6 +10,7 @@ #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/webgpu_utils.h" #include @@ -540,7 +541,7 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont return context.RunProgram(program); } -Status FlashAttentionPrefillSimpleProgram::GenerateShaderCode(ShaderHelper& shader) const { +Status FlashAttentionPrefillProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); shader.AddInput("present_key", ShaderUsage::UseUniform); shader.AddInput("present_value", ShaderUsage::UseUniform); @@ -558,7 +559,7 @@ Status FlashAttentionPrefillSimpleProgram::GenerateShaderCode(ShaderHelper& shad } shader.AddOutput("output", ShaderUsage::UseUniform); - return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_prefill_simple.wgsl.template", + return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_prefill.wgsl.template", WGSL_TEMPLATE_PARAMETER(has_attention_bias, has_attention_bias_), WGSL_TEMPLATE_PARAMETER(has_head_sink, has_head_sink_), WGSL_TEMPLATE_PARAMETER(has_local_window, has_local_window_), @@ -573,12 +574,12 @@ Status FlashAttentionPrefillSimpleProgram::GenerateShaderCode(ShaderHelper& shad WGSL_TEMPLATE_PARAMETER(use_seqlens_q, use_seqlens_q_)); } -bool CanApplyFlashAttentionPrefillSimple(bool is_intel, +bool CanApplyFlashAttentionPrefill(bool is_intel, bool kv_cache_quantization_enabled) { return is_intel && !kv_cache_quantization_enabled; } -Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& context, +Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, const Tensor* Q, const Tensor* present_key, const Tensor* present_value, @@ -601,27 +602,33 @@ Status ApplyFlashAttentionPrefillSimple(onnxruntime::webgpu::ComputeContext& con bool q_BNSH, bool has_head_sink, const Tensor* head_sink) { - constexpr uint32_t wg_size = 128u; - const uint32_t head_size_vec = static_cast(parameters.head_size_) / 4u; - const uint32_t v_head_vec_step = std::min(16u, head_size_vec); - // Full-head K staging and sliced V staging share the workgroup storage budget. - const uint64_t bytes_per_column = - (uint64_t{head_size_vec} + v_head_vec_step) * 4u * Q->DataType()->Size(); - const uint64_t max_workgroup_storage = context.DeviceLimits().maxComputeWorkgroupStorageSize; - uint32_t kv_step = 64u; - while (kv_step > 8u && bytes_per_column * kv_step > max_workgroup_storage) { - kv_step /= 2u; - } - ORT_RETURN_IF(bytes_per_column * kv_step > max_workgroup_storage, - "Simple prefill K/V tiles exceed workgroup storage limits."); + const auto& adapter_info = context.AdapterInfo(); + const bool is_intel_xe2 = adapter_info.vendor == std::string_view{"intel"} && + (adapter_info.architecture == std::string_view{"xe-2lpg"} || + adapter_info.architecture == std::string_view{"xe-2hpg"}); + const bool is_intel_xe3 = adapter_info.vendor == std::string_view{"intel"} && + (adapter_info.architecture == std::string_view{"xe-3lpg"} || + adapter_info.architecture == std::string_view{"xe-3lpg-xs"}); + // Default: 64 threads, 16 KV rows per step. + uint32_t wg_size = 64u; + uint32_t kv_step = 16u; + // Xe2: query length <= 128 uses 128 threads and 64 KV rows per step; + // longer queries use 256 threads and 32 KV rows per step. + // Xe3: 128 threads, 64 KV rows per step for all query lengths. + if (is_intel_xe3 || (is_intel_xe2 && parameters.sequence_length_ <= 128)) { + wg_size = 128u; + kv_step = 64u; + } else if (is_intel_xe2) { + wg_size = 256u; + kv_step = 32u; + } // Values below 8 select direct shared-memory reads, including without subgroups. const uint32_t shuffle_chunk = std::min(kv_step, subgroup_min_size); - const uint32_t num_seq_tile = - (parameters.sequence_length_ + wg_size - 1) / wg_size; + const uint32_t num_seq_tile = CeilDiv(static_cast(parameters.sequence_length_), wg_size); const uint32_t dispatch_size = parameters.batch_size_ * parameters.num_heads_ * num_seq_tile; - FlashAttentionPrefillSimpleProgram program{"FlashAttentionPrefillSimple", + FlashAttentionPrefillProgram program{"FlashAttentionPrefill", has_attention_bias, parameters.is_unidirectional_, parameters.head_size_, @@ -1179,9 +1186,9 @@ Status ApplyDensePrefill(const Tensor* Q, Tensor* present_key, Tensor* present_v const bool is_intel = context.AdapterInfo().vendor == std::string_view{"intel"}; const uint32_t subgroup_min_size = has_subgroups ? context.AdapterInfo().subgroupMinSize : 0u; - // TODO: Merge back into `FlashAttentionProgram`. - if (CanApplyFlashAttentionPrefillSimple(is_intel, kv_cache_quantization_enabled)) { - return ApplyFlashAttentionPrefillSimple( + // TODO(daijh): Merge back into `FlashAttentionProgram`. + if (CanApplyFlashAttentionPrefill(is_intel, kv_cache_quantization_enabled)) { + return ApplyFlashAttentionPrefill( context, Q, present_key, present_value, attention_bias, seqlen_k, attn_output, parameters, has_attention_bias, use_seqlen_k, has_local_window, local_window_size, present_sequence_length, subgroup_min_size, alpha, attn_bias_dim0, attn_bias_dim1, attn_bias_dim3, diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index daf51abd1d75c..0558a0bf67bcc 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -205,9 +205,9 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont uint32_t max_num_blocks_per_seq, const Tensor* cumulative_seqlens_q = nullptr); -class FlashAttentionPrefillSimpleProgram final : public Program { +class FlashAttentionPrefillProgram final : public Program { public: - FlashAttentionPrefillSimpleProgram(const std::string& kernel_name, + FlashAttentionPrefillProgram(const std::string& kernel_name, bool has_attention_bias, bool is_unidirectional, int qkv_head_size, @@ -260,10 +260,10 @@ class FlashAttentionPrefillSimpleProgram final : public Program u32 { #endif // Accumulate scores and softmax statistics in f32 to avoid fp16 overflow. -// Individual dot products retain the Q/K storage precision. #if qk_f32 alias qk_precision = f32; const min_value = qk_precision(-3.4028234663852886e+38f); @@ -57,12 +56,10 @@ const num_qk_vec : u32 = kv_step / vec_factor; const k_head_vec_step : u32 = head_size_vec; const v_head_vec_step : u32 = min(16u, head_size_vec); -// Adjacent subgroup lanes access adjacent workgroup-memory addresses. var k_tile : array, k_head_vec_step>; var v_tile : array, v_head_vec_step>; // Cooperatively load k_step key columns starting at k_start into shared memory. -// Columns past total_seq are zeroed so out-of-range lanes contribute nothing. fn loadk(k_start : u32, batch_head_idx : u32, local_idx : u32, k_step : u32, total_seq : u32) { let offset = batch_head_idx / uniforms.n_reps * uniforms.present_sequence_length * head_size_vec + k_start * head_size_vec; @@ -120,7 +117,6 @@ $MAIN { return; } - // Per-thread query row and output accumulator (zero-initialized). var q_tile : array; var o_tile : array; @@ -198,7 +194,6 @@ $MAIN { let aligned_window_start = (first_query_window_start / kv_step) * kv_step; for (var k_start = aligned_window_start; k_start < loop_bound; k_start += kv_step) { - // The previous iteration's PV barriers protect K reads before this overwrite. loadk(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); workgroupBarrier(); @@ -248,8 +243,8 @@ $MAIN { let local_max = max(max(local_max_vec.x, local_max_vec.y), max(local_max_vec.z, local_max_vec.w)); let new_max = max(previous_max, local_max); - // Exponentiate against the new max (in f32 for stability) and accumulate - // this tile's denominator contribution. + // Exponentiate against the new max and accumulate this tile's denominator + // contribution. var sum_vec = vec4(0); for (var m = 0u; m < num_qk_vec; m++) { qk[m] = vec4(exp(vec4(qk[m]) - f32(new_max))); @@ -294,7 +289,6 @@ $MAIN { #endif o_tile[head_start + i] = o_tile[head_start + i] * o_ratio + acc; } - // Protect V reads before the next head slice or K/V tile. workgroupBarrier(); } } From 8276331241babbe7459ff3123e9322ff6124edc8 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Mon, 19 Oct 2026 15:02:53 +0800 Subject: [PATCH 5/7] Up --- .../webgpu/bert/flash_attention.cc | 89 ++++++++++--------- .../contrib_ops/webgpu/bert/flash_attention.h | 68 +++++++------- .../flash_attention_prefill.wgsl.template | 28 +++--- 3 files changed, 99 insertions(+), 86 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 1790c14828663..9e6e810d31b37 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -1,6 +1,8 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include + #include "contrib_ops/cpu/bert/multihead_attention_helper.h" #include "contrib_ops/webgpu/bert/kv_cache_block_quant_int8.h" #include "contrib_ops/webgpu/bert/flash_attention.h" @@ -12,8 +14,6 @@ #include "core/providers/webgpu/webgpu_supported_types.h" #include "core/providers/webgpu/webgpu_utils.h" -#include - using namespace onnxruntime::webgpu; using namespace ::onnxruntime::common; using namespace ONNX_NAMESPACE; @@ -575,40 +575,40 @@ Status FlashAttentionPrefillProgram::GenerateShaderCode(ShaderHelper& shader) co } bool CanApplyFlashAttentionPrefill(bool is_intel, - bool kv_cache_quantization_enabled) { + bool kv_cache_quantization_enabled) { return is_intel && !kv_cache_quantization_enabled; } Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, - const Tensor* Q, - const Tensor* present_key, - const Tensor* present_value, - const Tensor* attention_bias, - const Tensor* seqlen_k, - Tensor* attn_output, - const WebgpuAttentionParameters& parameters, - bool has_attention_bias, - bool use_seqlen_k, - bool has_local_window, - int local_window_size, - uint32_t present_sequence_length, - uint32_t subgroup_min_size, - float alpha, - uint32_t attn_bias_dim0, - uint32_t attn_bias_dim1, - uint32_t attn_bias_dim3, - bool use_seqlens_q, - const Tensor* seqlens_q, - bool q_BNSH, - bool has_head_sink, - const Tensor* head_sink) { + const Tensor* Q, + const Tensor* present_key, + const Tensor* present_value, + const Tensor* attention_bias, + const Tensor* seqlen_k, + Tensor* attn_output, + const WebgpuAttentionParameters& parameters, + bool has_attention_bias, + bool use_seqlen_k, + bool has_local_window, + int local_window_size, + uint32_t present_sequence_length, + uint32_t subgroup_min_size, + float alpha, + uint32_t attn_bias_dim0, + uint32_t attn_bias_dim1, + uint32_t attn_bias_dim3, + bool use_seqlens_q, + const Tensor* seqlens_q, + bool q_BNSH, + bool has_head_sink, + const Tensor* head_sink) { const auto& adapter_info = context.AdapterInfo(); const bool is_intel_xe2 = adapter_info.vendor == std::string_view{"intel"} && - (adapter_info.architecture == std::string_view{"xe-2lpg"} || - adapter_info.architecture == std::string_view{"xe-2hpg"}); + (adapter_info.architecture == std::string_view{"xe-2lpg"} || + adapter_info.architecture == std::string_view{"xe-2hpg"}); const bool is_intel_xe3 = adapter_info.vendor == std::string_view{"intel"} && - (adapter_info.architecture == std::string_view{"xe-3lpg"} || - adapter_info.architecture == std::string_view{"xe-3lpg-xs"}); + (adapter_info.architecture == std::string_view{"xe-3lpg"} || + adapter_info.architecture == std::string_view{"xe-3lpg-xs"}); // Default: 64 threads, 16 KV rows per step. uint32_t wg_size = 64u; uint32_t kv_step = 16u; @@ -622,6 +622,15 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, wg_size = 256u; kv_step = 32u; } + // The K tile (full head) and V slice (up to 16 vec4) are staged in workgroup memory for kv_step rows. + // If that exceeds the device limit, fall back to the default config (64 threads, 16 KV rows per step). + const uint64_t head_size_vec = static_cast(parameters.head_size_) / 4; + const uint64_t vec4_bytes = (Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ? 2u : 4u) * 4u; + if ((head_size_vec + std::min(head_size_vec, 16)) * kv_step * vec4_bytes > + context.DeviceLimits().maxComputeWorkgroupStorageSize) { + wg_size = 64u; + kv_step = 16u; + } // Values below 8 select direct shared-memory reads, including without subgroups. const uint32_t shuffle_chunk = std::min(kv_step, subgroup_min_size); const uint32_t num_seq_tile = CeilDiv(static_cast(parameters.sequence_length_), wg_size); @@ -629,17 +638,17 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, parameters.batch_size_ * parameters.num_heads_ * num_seq_tile; FlashAttentionPrefillProgram program{"FlashAttentionPrefill", - has_attention_bias, - parameters.is_unidirectional_, - parameters.head_size_, - parameters.num_heads_, - use_seqlen_k, - kv_step, - shuffle_chunk, - has_local_window, - use_seqlens_q, - q_BNSH, - has_head_sink}; + has_attention_bias, + parameters.is_unidirectional_, + parameters.head_size_, + parameters.num_heads_, + use_seqlen_k, + kv_step, + shuffle_chunk, + has_local_window, + use_seqlens_q, + q_BNSH, + has_head_sink}; program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, 4}, {present_key, ProgramTensorMetadataDependency::TypeAndRank, 4}, {present_value, ProgramTensorMetadataDependency::TypeAndRank, 4}}); diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 0558a0bf67bcc..1f30813ce63f8 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -208,17 +208,17 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont class FlashAttentionPrefillProgram final : public Program { public: FlashAttentionPrefillProgram(const std::string& kernel_name, - bool has_attention_bias, - bool is_unidirectional, - int qkv_head_size, - int qkv_num_heads, - bool use_seqlen_k, - uint32_t kv_step, - uint32_t shuffle_chunk, - bool has_local_window, - bool use_seqlens_q, - bool q_BNSH, - bool has_head_sink) + bool has_attention_bias, + bool is_unidirectional, + int qkv_head_size, + int qkv_num_heads, + bool use_seqlen_k, + uint32_t kv_step, + uint32_t shuffle_chunk, + bool has_local_window, + bool use_seqlens_q, + bool q_BNSH, + bool has_head_sink) : Program{kernel_name}, has_attention_bias_(has_attention_bias), is_unidirectional_(is_unidirectional), @@ -261,31 +261,31 @@ class FlashAttentionPrefillProgram final : public Program { public: diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template index 1acc5d0b706b5..c6dccbcd76796 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template @@ -55,6 +55,8 @@ const head_size_vec : u32 = head_size / vec_factor; const num_qk_vec : u32 = kv_step / vec_factor; const k_head_vec_step : u32 = head_size_vec; const v_head_vec_step : u32 = min(16u, head_size_vec); +// Padded to whole V slices so the last, partial slice can run a full v_head_vec_step iterations. +const o_tile_vec : u32 = ((head_size_vec + v_head_vec_step - 1u) / v_head_vec_step) * v_head_vec_step; var k_tile : array, k_head_vec_step>; var v_tile : array, v_head_vec_step>; @@ -118,7 +120,7 @@ $MAIN { } var q_tile : array; - var o_tile : array; + var o_tile : array; // This thread's query row within the tile; may run past the sequence end. let seq_tile_idx = workgroup_idx % uniforms.num_seq_tile; @@ -152,8 +154,6 @@ $MAIN { #endif let total_sequence_length = get_total_sequence_length(batch_idx); -#if is_unidirectional - // Causal: query i attends to keys [0, i]. #if use_seqlens_q // LEFT-aligned per-batch Q length: past_len_b = total_kv_b - q_len_b (>= 0). let q_len_b = u32(seqlens_q[batch_idx]); @@ -166,6 +166,9 @@ $MAIN { 0u, total_sequence_length <= uniforms.new_sequence_length); #endif + +#if is_unidirectional + // Causal: query i attends to keys [0, past + i]. let max_causal_len_for_workgroup = past_sequence_length + (seq_tile_idx + 1) * workgroup_size_x; let loop_bound = min(total_sequence_length, max_causal_len_for_workgroup); @@ -177,14 +180,11 @@ $MAIN { #endif #if has_local_window -#if is_unidirectional - let first_query_causal_length = past_sequence_length + seq_tile_idx * workgroup_size_x + 1u; -#else - let first_query_causal_length = total_sequence_length; -#endif - let first_query_window_start = first_query_causal_length - - min(first_query_causal_length, uniforms.local_window_size); - let local_window_start = seq_causal_length - min(seq_causal_length, uniforms.local_window_size); + // The window's left bound is query-relative, with or without the causal right bound. + let first_query_end = past_sequence_length + seq_tile_idx * workgroup_size_x + 1u; + let query_end = past_sequence_length + q_idx_global + 1u; + let first_query_window_start = first_query_end - min(first_query_end, uniforms.local_window_size); + let local_window_start = query_end - min(query_end, uniforms.local_window_size); #else let first_query_window_start = 0u; let local_window_start = 0u; @@ -245,9 +245,13 @@ $MAIN { // Exponentiate against the new max and accumulate this tile's denominator // contribution. + // Masked scores are already min_value. If nothing in this row is visible yet, + // new_max is min_value too and exp(min_value - new_max) would be 1, so + // exponentiate against 0 instead to keep masked entries at 0. + let exp_base = select(f32(new_max), 0.0, new_max <= min_value); var sum_vec = vec4(0); for (var m = 0u; m < num_qk_vec; m++) { - qk[m] = vec4(exp(vec4(qk[m]) - f32(new_max))); + qk[m] = vec4(exp(vec4(qk[m]) - exp_base)); sum_vec += qk[m]; } let sum = sum_vec.x + sum_vec.y + sum_vec.z + sum_vec.w; From 80b831064947b5b7fecf396751f583f8cd2bf446 Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Mon, 19 Oct 2026 15:45:41 +0800 Subject: [PATCH 6/7] Add test cases for head size 80/96 --- .../contrib_ops/group_query_attention_op_test.cc | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 8eff8420fe527..23863edd45d33 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -7263,7 +7263,7 @@ TEST(GroupQueryAttentionTest, WebGPU_BlockQuantInt8_CrossValidate_FlashPrefill) /*max_absolute_error=*/0.03f); } -static void RunFp16HighMagnitudeAttention(int sequence_length, uint32_t bit_width) { +static void RunFp16HighMagnitudeAttention(int sequence_length, uint32_t bit_width, int head_size = 128) { auto ep = bit_width == 0 ? DefaultWebGpuExecutionProvider() : WebGpuEPWithKVCacheQuantization(bit_width); if (!ep) { GTEST_SKIP() << "WebGPU EP not available"; @@ -7272,9 +7272,8 @@ static void RunFp16HighMagnitudeAttention(int sequence_length, uint32_t bit_widt constexpr int batch_size = 1; constexpr int num_heads = 2; constexpr int kv_num_heads = 1; - constexpr int head_size = 128; - constexpr int hidden_size = num_heads * head_size; - constexpr int kv_hidden_size = kv_num_heads * head_size; + const int hidden_size = num_heads * head_size; + const int kv_hidden_size = kv_num_heads * head_size; std::mt19937 rng(1234); std::bernoulli_distribution sign_dist(0.5); @@ -7320,6 +7319,15 @@ TEST(GroupQueryAttentionTest, WebGPU_FP16_HighMagnitude_FlashPrefill) { RunFp16HighMagnitudeAttention(/*sequence_length=*/40, /*bit_width=*/0); } +// head_size / 4 is neither a power of two nor a multiple of 16, which exercises partial V slices in flash prefill. +TEST(GroupQueryAttentionTest, WebGPU_FP16_HighMagnitude_FlashPrefill_HeadSize80) { + RunFp16HighMagnitudeAttention(/*sequence_length=*/40, /*bit_width=*/0, /*head_size=*/80); +} + +TEST(GroupQueryAttentionTest, WebGPU_FP16_HighMagnitude_FlashPrefill_HeadSize96) { + RunFp16HighMagnitudeAttention(/*sequence_length=*/40, /*bit_width=*/0, /*head_size=*/96); +} + TEST(GroupQueryAttentionTest, WebGPU_FP16_HighMagnitude_SplitReduce) { RunFp16HighMagnitudeAttention(/*sequence_length=*/4, /*bit_width=*/0); } From 65c8761dc4c222022b9c2c2746409084ea1808eb Mon Sep 17 00:00:00 2001 From: Jianhui Dai Date: Tue, 27 Oct 2026 13:44:09 +0800 Subject: [PATCH 7/7] Up --- .../contrib_ops/webgpu/bert/flash_attention.cc | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 9e6e810d31b37..f0a87986267cb 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -610,16 +610,16 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, (adapter_info.architecture == std::string_view{"xe-3lpg"} || adapter_info.architecture == std::string_view{"xe-3lpg-xs"}); // Default: 64 threads, 16 KV rows per step. - uint32_t wg_size = 64u; + uint32_t workgroup_size = 64u; uint32_t kv_step = 16u; // Xe2: query length <= 128 uses 128 threads and 64 KV rows per step; // longer queries use 256 threads and 32 KV rows per step. // Xe3: 128 threads, 64 KV rows per step for all query lengths. if (is_intel_xe3 || (is_intel_xe2 && parameters.sequence_length_ <= 128)) { - wg_size = 128u; + workgroup_size = 128u; kv_step = 64u; } else if (is_intel_xe2) { - wg_size = 256u; + workgroup_size = 256u; kv_step = 32u; } // The K tile (full head) and V slice (up to 16 vec4) are staged in workgroup memory for kv_step rows. @@ -628,12 +628,12 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, const uint64_t vec4_bytes = (Q->GetElementType() == ONNX_NAMESPACE::TensorProto_DataType_FLOAT16 ? 2u : 4u) * 4u; if ((head_size_vec + std::min(head_size_vec, 16)) * kv_step * vec4_bytes > context.DeviceLimits().maxComputeWorkgroupStorageSize) { - wg_size = 64u; + workgroup_size = 64u; kv_step = 16u; } // Values below 8 select direct shared-memory reads, including without subgroups. const uint32_t shuffle_chunk = std::min(kv_step, subgroup_min_size); - const uint32_t num_seq_tile = CeilDiv(static_cast(parameters.sequence_length_), wg_size); + const uint32_t num_seq_tile = CeilDiv(static_cast(parameters.sequence_length_), workgroup_size); const uint32_t dispatch_size = parameters.batch_size_ * parameters.num_heads_ * num_seq_tile; @@ -666,9 +666,9 @@ Status ApplyFlashAttentionPrefill(onnxruntime::webgpu::ComputeContext& context, } program.AddOutputs({{attn_output, ProgramTensorMetadataDependency::TypeAndRank, 4}}); program.SetDispatchGroupSize(dispatch_size) - .SetWorkgroupSize(wg_size) + .SetWorkgroupSize(workgroup_size) .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, - use_seqlen_k, wg_size, kv_step, shuffle_chunk, + use_seqlen_k, workgroup_size, kv_step, shuffle_chunk, has_local_window, use_seqlens_q, q_BNSH, has_head_sink) .AddUniformVariables({{static_cast(parameters.sequence_length_)}, {static_cast(parameters.total_sequence_length_)},