diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 717abe8274b2b..f0a87986267cb 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" @@ -10,6 +12,7 @@ #include "contrib_ops/webgpu/webgpu_contrib_kernels.h" #include "core/providers/webgpu/webgpu_supported_types.h" +#include "core/providers/webgpu/webgpu_utils.h" using namespace onnxruntime::webgpu; using namespace ::onnxruntime::common; @@ -538,6 +541,149 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont return context.RunProgram(program); } +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); + 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.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_unidirectional, is_unidirectional_), + WGSL_TEMPLATE_PARAMETER(kv_step_param, kv_step_), + WGSL_TEMPLATE_PARAMETER(q_BNSH, q_BNSH_), + 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 CanApplyFlashAttentionPrefill(bool is_intel, + 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 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 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)) { + workgroup_size = 128u; + kv_step = 64u; + } else if (is_intel_xe2) { + 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. + // 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) { + 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_), workgroup_size); + const uint32_t dispatch_size = + 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}; + 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(dispatch_size) + .SetWorkgroupSize(workgroup_size) + .CacheHint(has_attention_bias, parameters.head_size_, parameters.num_heads_, parameters.is_unidirectional_, + 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_)}, + {static_cast(present_sequence_length)}, + {static_cast(parameters.batch_size_)}, + {static_cast(parameters.n_reps)}, + {alpha}, + {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 +1193,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(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, + 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..1f30813ce63f8 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -205,6 +205,88 @@ Status ApplyFlashAttentionPagedPrefill(onnxruntime::webgpu::ComputeContext& cont uint32_t max_num_blocks_per_seq, const Tensor* cumulative_seqlens_q = nullptr); +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) + : Program{kernel_name}, + has_attention_bias_(has_attention_bias), + 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), + shuffle_chunk_(shuffle_chunk), + 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_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 CanApplyFlashAttentionPrefill(bool is_intel, + bool 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); + class FlashAttentionDecodeQKVProgram final : public Program { public: FlashAttentionDecodeQKVProgram(const std::string& kernel_name, diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template new file mode 100644 index 0000000000000..c6dccbcd76796 --- /dev/null +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_prefill.wgsl.template @@ -0,0 +1,308 @@ + +#param has_attention_bias +#param has_head_sink +#param has_local_window +#param is_unidirectional +#param kv_step_param +#param q_BNSH +#param qk_f32 +#param qkv_head_size +#param qkv_num_heads +#param shuffle_chunk_param +#param use_seqlen_k +#param use_seqlens_q + +// One query row per thread, with full-head K staging and late V staging in +// slices of up to 16 vec4 values. Shuffle bands fit the adapter's guaranteed +// minimum subgroup size; below 8, use direct shared-memory reads. + +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 + +// Accumulate scores and softmax statistics in f32 to avoid fp16 overflow. +#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 +// 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; +} + +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; +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>; + +// Cooperatively load k_step key columns starting at k_start into shared memory. +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; + } +} + +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; + } +} + +// 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) -> qk_precision { + // Stored as [batch_size, num_heads, new_seq_length, total_sequence_length]. + if (k_idx_global >= total_seq) { + 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); + 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 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) -> qk_precision { + return qk_precision(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; + } + + 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 = qk_precision(head_sink[head_idx]); + var previous_max : qk_precision = sink_value; + var previous_denom : qk_precision = 1; +#else + var previous_max : qk_precision = min_value; + var previous_denom : qk_precision = 0; +#endif + let total_sequence_length = get_total_sequence_length(batch_idx); + +#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 + +#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); + 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 + // 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; +#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; + + for (var k_start = aligned_window_start; k_start < loop_bound; k_start += kv_step) { + loadk(k_start, batch_head_idx, local_idx, kv_step, total_sequence_length); + workgroupBarrier(); + + var qk : array, num_qk_vec>; + + // 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]; +#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. + 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 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]) - exp_base)); + 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 * 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); + } + previous_max = new_max; + previous_denom = d; + // Factor to rescale the already-accumulated output for the new denominator. + let o_ratio = q_element_t(dleft / d); + // 4. Accumulate softmax(qk) * V into the output tile, rescaling the prior + // partial sum by o_ratio. + 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; + } + workgroupBarrier(); + } + } + + 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 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); }