From f21efe67e39144d6090783d4622f8820a8fbc27b Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 19 Aug 2025 16:11:24 +0800 Subject: [PATCH 01/16] Add dispatchWorkgroupsIndirect support Add seqlen_k to dynamically compute total_seq_length Add Indirect buffer usage fuse PrepareIndirectDispatch shader into CopyKVCache code reuse Update the conditions --- .../webgpu/bert/flash_attention.cc | 145 ++++++++++++++---- .../contrib_ops/webgpu/bert/flash_attention.h | 21 +-- .../flash_attention_decode_qkt.wgsl.template | 9 +- ...sh_attention_decode_split_vx.wgsl.template | 13 +- ...h_attention_decode_vx_reduce.wgsl.template | 6 +- .../webgpu/bert/group_query_attention.cc | 2 +- .../core/providers/webgpu/allocator.cc | 2 +- .../core/providers/webgpu/compute_context.h | 5 +- onnxruntime/core/providers/webgpu/program.cc | 6 + onnxruntime/core/providers/webgpu/program.h | 7 + .../core/providers/webgpu/webgpu_context.cc | 38 ++++- .../core/providers/webgpu/webgpu_context.h | 4 +- 12 files changed, 196 insertions(+), 62 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index ab611a8e5a7c0..2e805ea3da55b 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -32,12 +32,31 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& present_value = shader.AddOutput("present_value", ShaderUsage::UseUniform); const auto& copy_kv_shape = shader.AddIndices("copy_kv_shape"); + // If prepare_indirect_dispatch is enabled, add seqlen_k input and indirect_buffer output + if (prepare_indirect_dispatch_) { + shader.AddInput("seqlen_k"); + shader.AddOutput("indirect_buffer", ShaderUsage::UseUniform); + } + shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.copy_size") << " let output_indices = " << copy_kv_shape.OffsetToIndices("global_idx") << ";\n" << " let head_size_id = output_indices[3];\n" " let sequence_id = output_indices[2];\n" " let num_head_id = output_indices[1];\n" " let batch = output_indices[0];\n"; + + // Add indirect dispatch logic for thread 0 + if (prepare_indirect_dispatch_) { + shader.MainFunctionBody() << " // Prepare indirect dispatch buffer for thread 0\n" + << " if (global_idx == 0u) {\n" + << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n" + << " let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n" + << " indirect_buffer[0] = num_total_seq_length_tile;\n" + << " indirect_buffer[1] = uniforms.num_heads;\n" + << " indirect_buffer[2] = 1u;\n" + << " }\n\n"; + } + if (has_past_) { shader.MainFunctionBody() << "let past_sequence_length = uniforms.past_sequence_length;\n"; if (past_present_share_buffer_) { @@ -70,10 +89,12 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& parameters, const Tensor* K, const Tensor* past_key, Tensor* present_key, - const Tensor* V, const Tensor* past_value, Tensor* present_value) { + const Tensor* V, const Tensor* past_value, Tensor* present_value, + const Tensor* seqlen_k, Tensor* indirect_buffer) { // CopyKVCache takes past key/value and current key/value and copies them to present key and value. // This makes it so that FlashAttention only needs to look at present key and value, and saves // number of input buffers in the shader, which we run out of (<=8) without this optimization. + // If indirect_buffer is provided, also prepare indirect dispatch buffer for flash attention. const int components = parameters.head_size_ % 4 == 0 ? 4 : (parameters.head_size_ % 2 == 0 ? 2 : 1); bool has_past = (parameters.total_sequence_length_ - parameters.kv_sequence_length_) > 0; // parameters.total_sequence_length_ is past_sequence_length + kv_sequence_length. @@ -83,7 +104,13 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt int copy_sequence_length = has_past && parameters.past_present_share_buffer_ ? parameters.kv_sequence_length_ : parameters.total_sequence_length_; TensorShape copy_kv_shape{parameters.batch_size_, num_heads, copy_sequence_length, parameters.head_size_ / components}; int64_t copy_size = copy_kv_shape.Size(); - CopyKVCacheProgram program{"CopyKVCache", has_past, parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH, parameters.past_present_share_buffer_}; + + // Determine if we need to prepare indirect dispatch + bool prepare_indirect_dispatch = (indirect_buffer != nullptr); + constexpr uint32_t tile_size = 64; + + CopyKVCacheProgram program{"CopyKVCache", has_past, parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH, parameters.past_present_share_buffer_, + prepare_indirect_dispatch, tile_size, static_cast(parameters.num_heads_)}; if (parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH) { program.AddInputs({{K, ProgramTensorMetadataDependency::TypeAndRank, components}, {V, ProgramTensorMetadataDependency::TypeAndRank, components}}); @@ -94,24 +121,45 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt program.AddInputs({{K, ProgramTensorMetadataDependency::TypeAndRank, reshaped_KV_shape, components}, {V, ProgramTensorMetadataDependency::TypeAndRank, reshaped_KV_shape, components}}); } + + // Add seqlen_k input if preparing indirect dispatch + if (prepare_indirect_dispatch) { + program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); + } + if (has_past && !parameters.past_present_share_buffer_) { program.AddInputs({{past_key, ProgramTensorMetadataDependency::TypeAndRank, components}, {past_value, ProgramTensorMetadataDependency::TypeAndRank, components}}); } program.AddOutputs({{present_key, ProgramTensorMetadataDependency::Rank, components}, - {present_value, ProgramTensorMetadataDependency::Rank, components}}) - .AddIndices(std::move(copy_kv_shape)); + {present_value, ProgramTensorMetadataDependency::Rank, components}}); + + // Add indirect_buffer output if preparing indirect dispatch + if (prepare_indirect_dispatch) { + program.AddOutput({indirect_buffer, ProgramTensorMetadataDependency::None}); + } + + program.AddIndices(std::move(copy_kv_shape)); program.SetDispatchGroupSize(static_cast((copy_size + 63) / 64)) .SetWorkgroupSize(64) - .CacheHint(has_past, parameters.qkv_format_, parameters.past_present_share_buffer_) + .CacheHint(has_past, parameters.qkv_format_, parameters.past_present_share_buffer_, prepare_indirect_dispatch) .AddUniformVariables({{static_cast(copy_size)}, // Note that when parameters.past_present_share_buffer_ is true, parameters.past_sequence_length_ will become to // max_sequence_length. To get a valid past_sequence_length, we use total_sequence_length - kv_sequence_length. - {static_cast(parameters.total_sequence_length_ - parameters.kv_sequence_length_)}}); + {static_cast(parameters.total_sequence_length_ - parameters.kv_sequence_length_)}, + {tile_size}, + {static_cast(parameters.num_heads_)}}); return context.RunProgram(program); } +// Overloaded version for backward compatibility +Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& parameters, + const Tensor* K, const Tensor* past_key, Tensor* present_key, + const Tensor* V, const Tensor* past_value, Tensor* present_value) { + return CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, nullptr, nullptr); +} + Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const { // Expectations are // qkv have same number of heads and hidden dimension (head size). @@ -147,6 +195,7 @@ Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const { Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); shader.AddInput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); + shader.AddInput("seqlens_k"); if (has_attention_bias_) { shader.AddInput("attention_bias", ShaderUsage::UseUniform); } @@ -163,8 +212,8 @@ Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) } Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& context, const Tensor* Q, - const Tensor* attention_bias, Tensor* output, Tensor* present_key, Tensor* metadata, - const WebgpuAttentionParameters& parameters, uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, uint32_t tile_size) { + const Tensor* attention_bias, Tensor* output, Tensor* present_key, Tensor* metadata, const Tensor* seqlen_k, + const WebgpuAttentionParameters& parameters, const Tensor* indirect_buffer, uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, uint32_t tile_size, bool use_indirect_dispatch) { const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size_)) : parameters.scale_; @@ -173,7 +222,8 @@ Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& conte FlashAttentionDecodeQKTProgram program{"FlashAttentionDecodeQKT", has_attention_bias, tile_size}; program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, components}, - {present_key, ProgramTensorMetadataDependency::TypeAndRank, components}}); + {present_key, ProgramTensorMetadataDependency::TypeAndRank, components}, + {seqlen_k, ProgramTensorMetadataDependency::None}}); if (has_attention_bias) { program.AddInput({attention_bias, ProgramTensorMetadataDependency::TypeAndRank}); } @@ -181,16 +231,18 @@ Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& conte {metadata, ProgramTensorMetadataDependency::Rank, 2}}); const uint32_t vectorized_head_size = parameters.head_size_ / components; - program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile) - .SetWorkgroupSize(64) + if (use_indirect_dispatch) { + program.SetIndirectDispatchTensor(indirect_buffer); + } else { + program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile); + } + program.SetWorkgroupSize(64) .CacheHint(tile_size, has_attention_bias) .AddUniformVariables({{static_cast(vectorized_head_size)}, - {static_cast(parameters.total_sequence_length_)}, {static_cast(alpha)}, // present_sequence_length is used to index into the KV cache, for static kv cache it is the max sequence length. {static_cast(parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_)}, {static_cast(parameters.n_reps)}, - {num_total_seq_length_tile}, {num_present_sequence_length_tile}, {static_cast(parameters.num_heads_)}}); @@ -201,6 +253,7 @@ Status FlashAttentionDecodeSplitVxProgram::GenerateShaderCode(ShaderHelper& shad shader.AddInput("metadata", ShaderUsage::UseUniform); shader.AddInput("qk", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); shader.AddInput("present_value", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); + shader.AddInput("seqlens_k"); shader.AddOutput("out_split_vx", ShaderUsage::UseUniform); const uint32_t tile_size_k_vec = 8u; @@ -217,25 +270,31 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte const Tensor* qk, Tensor* out_split_vx, Tensor* present_value, + const Tensor* seqlen_k, const WebgpuAttentionParameters& parameters, + const Tensor* indirect_buffer, uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, - uint32_t tile_size) { + uint32_t tile_size, + bool use_indirect_dispatch) { const int components = 4; int head_size_vec = parameters.v_head_size_ / components; FlashAttentionDecodeSplitVxProgram program{"FlashAttentionDecodeSplitVx", tile_size, head_size_vec}; program.AddInputs({{metadata, ProgramTensorMetadataDependency::TypeAndRank, 2}, {qk, ProgramTensorMetadataDependency::TypeAndRank}, - {present_value, ProgramTensorMetadataDependency::TypeAndRank, components}}); + {present_value, ProgramTensorMetadataDependency::TypeAndRank, components}, + {seqlen_k, ProgramTensorMetadataDependency::None}}); program.AddOutputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}}); // [B, N, split_k, head_size] - program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile) - .CacheHint(tile_size, head_size_vec) + if (use_indirect_dispatch) { + program.SetIndirectDispatchTensor(indirect_buffer); + } else { + program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile); + } + program.CacheHint(tile_size, head_size_vec) .SetWorkgroupSize(64) - .AddUniformVariables({{static_cast(parameters.total_sequence_length_)}, - {static_cast(head_size_vec)}, + .AddUniformVariables({{static_cast(head_size_vec)}, {static_cast(parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_)}, {static_cast(parameters.n_reps)}, - num_total_seq_length_tile, num_present_sequence_length_tile, {static_cast(parameters.num_heads_)}}); @@ -244,6 +303,7 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte Status FlashAttentionDecodeVxReduceProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("input", ShaderUsage::UseUniform); + shader.AddInput("seqlens_k"); shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_decode_vx_reduce.wgsl.template", @@ -253,21 +313,21 @@ Status FlashAttentionDecodeVxReduceProgram::GenerateShaderCode(ShaderHelper& sha Status ComputeFlashAttentionDecodeVxReduce(onnxruntime::webgpu::ComputeContext& context, const Tensor* out_split_vx, Tensor* output, + const Tensor* seqlen_k, const WebgpuAttentionParameters& parameters, - uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile) { const int components = 4; constexpr int tile_size = 8; int tile_head_size = tile_size * components; FlashAttentionDecodeVxReduceProgram program{"FlashAttentionDecodeVxReduce", tile_size}; - program.AddInputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}}); + program.AddInputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}, + {seqlen_k, ProgramTensorMetadataDependency::None}}); program.AddOutputs({{output, ProgramTensorMetadataDependency::TypeAndRank, components}}); const uint32_t num_head_size_tile = static_cast((parameters.v_head_size_ + tile_head_size - 1) / tile_head_size); program.SetDispatchGroupSize(parameters.num_heads_ * num_head_size_tile) .CacheHint(tile_size) .SetWorkgroupSize(tile_size * tile_size) .AddUniformVariables({{static_cast(parameters.v_head_size_ / components)}, - num_total_seq_length_tile, num_present_sequence_length_tile, {num_head_size_tile}, {static_cast(parameters.num_heads_)}}); @@ -277,10 +337,13 @@ Status ComputeFlashAttentionDecodeVxReduce(onnxruntime::webgpu::ComputeContext& Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value, - const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context) { - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); + const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k) { const int present_sequence_length = parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_; + if (parameters.sequence_length_ > 1) { + // For encode path, use the original CopyKVCache without indirect dispatch preparation + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); + const uint32_t tile_size = 64; bool has_attention_bias = attention_bias != nullptr; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; @@ -318,7 +381,7 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co return context.RunProgram(program); } - // Use present_sequence_length instead of total_sequence_length to make sure the |qk| buffer is static when static qv cache is enabled. + // For decode path (sequence_length == 1) const TensorShapeVector qk_dims({parameters.batch_size_, parameters.num_heads_, parameters.sequence_length_, present_sequence_length}); const TensorShape qk_shape(qk_dims); @@ -326,19 +389,41 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co constexpr uint32_t tile_size = 64; const uint32_t num_total_seq_length_tile = (parameters.total_sequence_length_ + tile_size - 1) / tile_size; const uint32_t num_present_sequence_length_tile = (present_sequence_length + tile_size - 1) / tile_size; + + // Determine if we should use indirect dispatch + const bool use_indirect_dispatch = parameters.past_present_share_buffer_ && + seqlen_k != nullptr && + context.IsGraphCaptureEnabled(); + + // Create indirect dispatch buffer if using indirect dispatch + Tensor* indirect_buffer_ptr = nullptr; + Tensor indirect_buffer; + if (use_indirect_dispatch) { + const TensorShape indirect_buffer_shape{3}; // 3 uint32 values for dispatch dimensions + indirect_buffer = context.CreateGPUTensor(DataTypeImpl::GetType(), indirect_buffer_shape); + indirect_buffer_ptr = &indirect_buffer; + // Use the fused CopyKVCache that also prepares the indirect dispatch buffer + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, indirect_buffer_ptr)); + } else { + // Use the original CopyKVCache without indirect dispatch preparation + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); + } + // The metadata is used to store the max and sum of each tile. const TensorShapeVector metadata_dims({parameters.batch_size_, parameters.num_heads_, num_present_sequence_length_tile, 2}); const TensorShape metadata_shape(metadata_dims); Tensor metadata = context.CreateGPUTensor(DataTypeImpl::GetType(), metadata_shape); - ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeQKT(context, Q, attention_bias, &qk, present_key, &metadata, - parameters, num_total_seq_length_tile, num_present_sequence_length_tile, tile_size)); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeQKT(context, Q, attention_bias, &qk, present_key, &metadata, seqlen_k, + parameters, indirect_buffer_ptr, num_total_seq_length_tile, + num_present_sequence_length_tile, tile_size, use_indirect_dispatch)); const TensorShapeVector out_split_vx_dims({parameters.batch_size_, parameters.num_heads_, num_present_sequence_length_tile, parameters.head_size_}); const TensorShape out_split_vx_shape(out_split_vx_dims); Tensor out_split_vx = context.CreateGPUTensor(Q->DataType(), out_split_vx_shape); - ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeSplitVxScore(context, &metadata, &qk, &out_split_vx, present_value, parameters, num_total_seq_length_tile, num_present_sequence_length_tile, tile_size)); - ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, output, parameters, num_total_seq_length_tile, num_present_sequence_length_tile)); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeSplitVxScore(context, &metadata, &qk, &out_split_vx, present_value, + seqlen_k, parameters, indirect_buffer_ptr, num_total_seq_length_tile, num_present_sequence_length_tile, tile_size, use_indirect_dispatch)); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, output, seqlen_k, parameters, num_present_sequence_length_tile)); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index c75494df253c1..0c6d831c619a7 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -17,19 +17,25 @@ using namespace onnxruntime::webgpu; class CopyKVCacheProgram final : public Program { public: - CopyKVCacheProgram(const std::string& kernel_name, bool has_past, bool kv_BNSH, bool past_present_share_buffer) - : Program{kernel_name}, has_past_(has_past), kv_BNSH_(kv_BNSH), past_present_share_buffer_(past_present_share_buffer) { + CopyKVCacheProgram(const std::string& kernel_name, bool has_past, bool kv_BNSH, bool past_present_share_buffer, + bool prepare_indirect_dispatch = false, uint32_t tile_size = 0, uint32_t num_heads = 0) + : Program{kernel_name}, has_past_(has_past), kv_BNSH_(kv_BNSH), past_present_share_buffer_(past_present_share_buffer), prepare_indirect_dispatch_(prepare_indirect_dispatch), tile_size_(tile_size), num_heads_(num_heads) { } Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"copy_size", ProgramUniformVariableDataType::Uint32}, - {"past_sequence_length", ProgramUniformVariableDataType::Uint32}); + {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"tile_size", ProgramUniformVariableDataType::Uint32}, + {"num_heads", ProgramUniformVariableDataType::Uint32}); private: bool has_past_; bool kv_BNSH_; bool past_present_share_buffer_; + bool prepare_indirect_dispatch_; + uint32_t tile_size_; + uint32_t num_heads_; }; class FlashAttentionProgram final : public Program { @@ -82,11 +88,9 @@ class FlashAttentionDecodeQKTProgram final : public Program tile_qk: array; $MAIN { let local_row = u32(local_idx / tile_size_k_vec); let local_col = local_idx % tile_size_k_vec; - let total_seq_offset = (workgroup_idx % uniforms.num_total_seq_length_tile) * tile_size; - let head_idx = u32(workgroup_idx / uniforms.num_total_seq_length_tile); + let total_sequence_length = u32(seqlens_k[0]) + 1u; + let num_total_seq_length_tile = (total_sequence_length + tile_size - 1) / tile_size; + let total_seq_offset = (workgroup_idx % num_total_seq_length_tile) * tile_size; + let head_idx = u32(workgroup_idx / num_total_seq_length_tile); let q_offset = head_idx * uniforms.head_size_vec; - var total_sequence_length = uniforms.total_sequence_length; let present_offset = u32(head_idx / uniforms.n_reps) * uniforms.present_sequence_length * uniforms.head_size_vec; for (var k: u32 = 0u; k < uniforms.head_size_vec; k += tile_size_k_vec) { if (local_idx < tile_size_k_vec && k + local_idx < uniforms.head_size_vec) { @@ -95,7 +96,7 @@ $MAIN { for (var i = 0u; i < tile_size && (total_seq_offset + i) < total_sequence_length; i++) { l_sum += exp(f32(tile_qk[i]) - l_max); } - let meta_offset = head_idx * uniforms.num_present_sequence_length_tile + workgroup_idx % uniforms.num_total_seq_length_tile; + let meta_offset = head_idx * uniforms.num_present_sequence_length_tile + workgroup_idx % num_total_seq_length_tile; metadata[meta_offset] = metadata_value_t(l_max, l_sum); } } diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template index c7593af311ce2..dd01343164499 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template @@ -40,9 +40,10 @@ var qkv_values: array, $MAIN { let local_row = u32(local_idx / tile_size_k_vec); let local_col = local_idx % tile_size_k_vec; - let total_seq_offset = (workgroup_idx % uniforms.num_total_seq_length_tile) * tile_size; - let head_idx = u32(workgroup_idx / uniforms.num_total_seq_length_tile); - var total_sequence_length = uniforms.total_sequence_length; + let total_sequence_length = u32(seqlens_k[0]) + 1u; + let num_total_seq_length_tile = (total_sequence_length + tile_size - 1) / tile_size; + let total_seq_offset = (workgroup_idx % num_total_seq_length_tile) * tile_size; + let head_idx = u32(workgroup_idx / num_total_seq_length_tile); let present_offset = u32(head_idx / uniforms.n_reps) * head_size_vec * uniforms.present_sequence_length; // Calculate the global max and sum in qk. @@ -50,12 +51,12 @@ $MAIN { { var g_max = f32(-3.402823e+38f); var g_sum = f32(0); - for (var i = 0u; i < uniforms.num_total_seq_length_tile; i++) + for (var i = 0u; i < num_total_seq_length_tile; i++) { let meta_offset = head_idx * uniforms.num_present_sequence_length_tile + i; g_max = max(g_max, metadata[meta_offset].x); } - for (var i = 0u; i < uniforms.num_total_seq_length_tile; i++) + for (var i = 0u; i < num_total_seq_length_tile; i++) { let meta_offset = head_idx * uniforms.num_present_sequence_length_tile + i; let m_value = metadata[meta_offset]; @@ -95,7 +96,7 @@ $MAIN { } for (var i = local_idx; i < head_size_vec; i += workgroup_size_x) { - let out_offset = head_idx * uniforms.num_present_sequence_length_tile * head_size_vec + (workgroup_idx % uniforms.num_total_seq_length_tile) * head_size_vec + i; + let out_offset = head_idx * uniforms.num_present_sequence_length_tile * head_size_vec + (workgroup_idx % num_total_seq_length_tile) * head_size_vec + i; out_split_vx[out_offset] = tile_output[i]; } } diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template index a4381baa638ce..f010aa9781519 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template @@ -23,10 +23,12 @@ $MAIN { var value = output_value_t(0); let local_row = u32(local_idx / tile_size); let local_col = local_idx % tile_size; + let total_sequence_length = u32(seqlens_k[0]) + 1u; + let num_total_seq_length_tile = (total_sequence_length + 63u) / 64u; if (head_size_offset + local_col < uniforms.head_size_vec) { - for (var r = 0u; r < uniforms.num_total_seq_length_tile; r += tile_size) { - if (r + local_row < uniforms.num_total_seq_length_tile) { + for (var r = 0u; r < num_total_seq_length_tile; r += tile_size) { + if (r + local_row < num_total_seq_length_tile) { value += input[in_offset + (r + local_row) * uniforms.head_size_vec + head_size_offset + local_col]; } } diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index 8b7b257dd2852..cb845061404f3 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -206,7 +206,7 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& !use_sliding_window && CanApplyFlashAttention(attention_bias, present_key, present_value, parameters, context)) { return ApplyFlashAttention(query, key, value, attention_bias, output, past_key, present_key, past_value, - present_value, parameters, context); + present_value, parameters, context, seqlen_k); } Tensor qSplit; diff --git a/onnxruntime/core/providers/webgpu/allocator.cc b/onnxruntime/core/providers/webgpu/allocator.cc index d0e8aba92aa4f..b3eb4b5061423 100644 --- a/onnxruntime/core/providers/webgpu/allocator.cc +++ b/onnxruntime/core/providers/webgpu/allocator.cc @@ -27,7 +27,7 @@ void* GpuBufferAllocator::Alloc(size_t size) { stats_.num_allocs++; wgpu::BufferUsage usage = mapped_at_creation_ ? wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapWrite - : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst; + : wgpu::BufferUsage::Storage | wgpu::BufferUsage::CopySrc | wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::Indirect; return buffer_manager_.Create(size, usage); } diff --git a/onnxruntime/core/providers/webgpu/compute_context.h b/onnxruntime/core/providers/webgpu/compute_context.h index fe95917e4e906..6bf7df74ea861 100644 --- a/onnxruntime/core/providers/webgpu/compute_context.h +++ b/onnxruntime/core/providers/webgpu/compute_context.h @@ -8,6 +8,7 @@ #include #include "core/framework/execution_provider.h" +#include "core/providers/webgpu/webgpu_execution_provider.h" #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/webgpu_context.h" @@ -16,7 +17,6 @@ namespace onnxruntime { class Tensor; -class WebGpuExecutionProvider; namespace webgpu { @@ -42,6 +42,9 @@ class ComputeContext { inline bool HasFeature(wgpu::FeatureName feature) const { return webgpu_context_.DeviceHasFeature(feature); } + inline bool IsGraphCaptureEnabled() const { + return ep_.IsGraphCaptureEnabled(); + } #if !defined(__wasm__) inline const wgpu::AdapterPropertiesSubgroupMatrixConfigs& SubgroupMatrixConfigs() const { return webgpu_context_.SubgroupMatrixConfigs(); diff --git a/onnxruntime/core/providers/webgpu/program.cc b/onnxruntime/core/providers/webgpu/program.cc index bd64416401fad..4a515cb988f1f 100644 --- a/onnxruntime/core/providers/webgpu/program.cc +++ b/onnxruntime/core/providers/webgpu/program.cc @@ -319,6 +319,7 @@ ProgramBase::ProgramBase(std::string_view name, ProgramMetadata&& metadata) dispatch_group_size_x_{0}, dispatch_group_size_y_{0}, dispatch_group_size_z_{0}, + indirect_dispatch_tensor_{nullptr}, workgroup_size_x_{0}, workgroup_size_y_{0}, workgroup_size_z_{0} { @@ -359,6 +360,11 @@ ProgramBase& ProgramBase::SetDispatchGroupSize(uint32_t x, uint32_t y, uint32_t return *this; } +ProgramBase& ProgramBase::SetIndirectDispatchTensor(const Tensor* indirect_dispatch_tensor) { + indirect_dispatch_tensor_ = indirect_dispatch_tensor; + return *this; +} + ProgramBase& ProgramBase::SetWorkgroupSize(uint32_t x) { return SetWorkgroupSize(x, 1, 1); } diff --git a/onnxruntime/core/providers/webgpu/program.h b/onnxruntime/core/providers/webgpu/program.h index 02d3bd28f96fc..77fee2631bfc0 100644 --- a/onnxruntime/core/providers/webgpu/program.h +++ b/onnxruntime/core/providers/webgpu/program.h @@ -305,6 +305,9 @@ class ProgramBase { // set the size of dispatch groups. ProgramBase& SetDispatchGroupSize(uint32_t x, uint32_t y, uint32_t z); + // set indirect dispatch tensor for indirect dispatch + ProgramBase& SetIndirectDispatchTensor(const Tensor* indirect_dispatch_tensor); + // set the size of a workgroup grid. Y and Z are 1 if not specified. ProgramBase& SetWorkgroupSize(uint32_t x); // set the size of a workgroup grid. Z is 1 if not specified. @@ -348,6 +351,8 @@ class ProgramBase { inline uint32_t DispatchGroupSizeX() const { return dispatch_group_size_x_; } inline uint32_t DispatchGroupSizeY() const { return dispatch_group_size_y_; } inline uint32_t DispatchGroupSizeZ() const { return dispatch_group_size_z_; } + inline const Tensor* IndirectDispatchTensor() const { return indirect_dispatch_tensor_; } + inline bool UseIndirectDispatch() const { return indirect_dispatch_tensor_ != nullptr; } inline uint32_t WorkgroupSizeX() const { return workgroup_size_x_; } inline uint32_t WorkgroupSizeY() const { return workgroup_size_y_; } inline uint32_t WorkgroupSizeZ() const { return workgroup_size_z_; } @@ -374,6 +379,8 @@ class ProgramBase { uint32_t dispatch_group_size_y_; uint32_t dispatch_group_size_z_; + const Tensor* indirect_dispatch_tensor_; + uint32_t workgroup_size_x_; uint32_t workgroup_size_y_; uint32_t workgroup_size_z_; diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 34bdfe7fa068a..7f5e6c2d05891 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -279,7 +279,11 @@ Status WebGpuContext::Run(ComputeContext& context, const ProgramBase& program) { uint32_t x = program.DispatchGroupSizeX(); uint32_t y = program.DispatchGroupSizeY(); uint32_t z = program.DispatchGroupSizeZ(); - ORT_RETURN_IF_ERROR(program_mgr_->NormalizeDispatchGroupSize(x, y, z)); + + // Skip normalization for indirect dispatch since dimensions are determined by the indirect buffer + if (!program.UseIndirectDispatch()) { + ORT_RETURN_IF_ERROR(program_mgr_->NormalizeDispatchGroupSize(x, y, z)); + } bool is_1d_dispatch = (y == 1 && z == 1); @@ -466,7 +470,7 @@ Status WebGpuContext::Run(ComputeContext& context, const ProgramBase& program) { bind_buffers.push_back(uniform_buffer); } - LaunchComputePipeline(compute_pass_encoder, bind_buffers, *program_artifact, x, y, z); + LaunchComputePipeline(compute_pass_encoder, bind_buffers, *program_artifact, x, y, z, program.IndirectDispatchTensor()); if (uniform_buffer) { buffer_mgr.Release(uniform_buffer); } @@ -746,7 +750,8 @@ void WebGpuContext::OnRunEnd() { void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& compute_pass_encoder, const std::vector& bind_buffers, const ProgramArtifact& program_artifact, - uint32_t x, uint32_t y, uint32_t z) { + uint32_t x, uint32_t y, uint32_t z, + const Tensor* indirect_dispatch_tensor) { uint32_t entry_index = 0; std::vector bind_group_entries; for (WGPUBuffer buffer : bind_buffers) { @@ -762,14 +767,27 @@ void WebGpuContext::LaunchComputePipeline(const wgpu::ComputePassEncoder& comput auto bind_group = wgpuDeviceCreateBindGroup(Device().Get(), &bind_group_desc); if (graph_capture_state_ == GraphCaptureState::Capturing) { + WGPUBuffer indirect_buffer = nullptr; + if (indirect_dispatch_tensor != nullptr) { + indirect_buffer = reinterpret_cast(const_cast(indirect_dispatch_tensor->DataRaw())); + } external_captured_commands_->push_back({program_artifact.compute_pipeline, bind_group, bind_group_layout, - {x, y, z}}); + {x, y, z}, + indirect_buffer}); } else { compute_pass_encoder.SetPipeline(program_artifact.compute_pipeline); wgpuComputePassEncoderSetBindGroup(compute_pass_encoder.Get(), 0, bind_group, 0, nullptr); - compute_pass_encoder.DispatchWorkgroups(x, y, z); + + if (indirect_dispatch_tensor != nullptr) { + // Use indirect dispatch + WGPUBuffer indirect_buffer = reinterpret_cast(const_cast(indirect_dispatch_tensor->DataRaw())); + compute_pass_encoder.DispatchWorkgroupsIndirect(indirect_buffer, 0); + } else { + // Use direct dispatch + compute_pass_encoder.DispatchWorkgroups(x, y, z); + } wgpuBindGroupRelease(bind_group); wgpuBindGroupLayoutRelease(bind_group_layout); @@ -805,7 +823,15 @@ void WebGpuContext::Replay(const std::vector& captu WriteTimestamp(num_pending_dispatches_ * 2); compute_pass_encoder.SetPipeline(command.compute_pipeline); wgpuComputePassEncoderSetBindGroup(compute_pass_encoder.Get(), 0, command.bind_group, 0, nullptr); - compute_pass_encoder.DispatchWorkgroups(command.dispatch_group[0], command.dispatch_group[1], command.dispatch_group[2]); + + if (command.indirect_buffer != nullptr) { + // Use indirect dispatch + compute_pass_encoder.DispatchWorkgroupsIndirect(command.indirect_buffer, 0); + } else { + // Use direct dispatch + compute_pass_encoder.DispatchWorkgroups(command.dispatch_group[0], command.dispatch_group[1], command.dispatch_group[2]); + } + WriteTimestamp(num_pending_dispatches_ * 2 + 1); ++num_pending_dispatches_; if (num_pending_dispatches_ >= max_num_pending_dispatches_ || diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 5ee085c59d2af..353a5e94fcb9c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -31,6 +31,7 @@ struct CapturedCommandInfo { WGPUBindGroup bind_group; WGPUBindGroupLayout bind_group_layout; std::array dispatch_group; + WGPUBuffer indirect_buffer; // WGPUBuffer for indirect dispatch, nullptr if not using indirect dispatch }; struct WebGpuContextConfig { @@ -183,7 +184,8 @@ class WebGpuContext final { void LaunchComputePipeline(const wgpu::ComputePassEncoder& compute_pass_encoder, const std::vector& bind_buffers, const ProgramArtifact& program_artifact, - uint32_t x, uint32_t y, uint32_t z); + uint32_t x, uint32_t y, uint32_t z, + const Tensor* indirect_dispatch_tensor = nullptr); std::vector GetEnabledAdapterToggles() const; std::vector GetEnabledDeviceToggles() const; From 83ef24ca18e28990a8998b84a1482a04366519f4 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 26 Aug 2025 17:23:15 +0800 Subject: [PATCH 02/16] unify the present_sequence_length --- .../webgpu/bert/flash_attention.cc | 27 ++++++++++++------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 2e805ea3da55b..3458b5f178406 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -213,7 +213,7 @@ Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& context, const Tensor* Q, const Tensor* attention_bias, Tensor* output, Tensor* present_key, Tensor* metadata, const Tensor* seqlen_k, - const WebgpuAttentionParameters& parameters, const Tensor* indirect_buffer, uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, uint32_t tile_size, bool use_indirect_dispatch) { + const WebgpuAttentionParameters& parameters, const Tensor* indirect_buffer, uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, uint32_t tile_size, bool use_indirect_dispatch, uint32_t present_sequence_length) { const float alpha = parameters.scale_ == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size_)) : parameters.scale_; @@ -241,7 +241,7 @@ Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& conte .AddUniformVariables({{static_cast(vectorized_head_size)}, {static_cast(alpha)}, // present_sequence_length is used to index into the KV cache, for static kv cache it is the max sequence length. - {static_cast(parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_)}, + {static_cast(present_sequence_length)}, {static_cast(parameters.n_reps)}, {num_present_sequence_length_tile}, {static_cast(parameters.num_heads_)}}); @@ -276,7 +276,8 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte uint32_t num_total_seq_length_tile, uint32_t num_present_sequence_length_tile, uint32_t tile_size, - bool use_indirect_dispatch) { + bool use_indirect_dispatch, + uint32_t present_sequence_length) { const int components = 4; int head_size_vec = parameters.v_head_size_ / components; FlashAttentionDecodeSplitVxProgram program{"FlashAttentionDecodeSplitVx", tile_size, head_size_vec}; @@ -293,7 +294,7 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte program.CacheHint(tile_size, head_size_vec) .SetWorkgroupSize(64) .AddUniformVariables({{static_cast(head_size_vec)}, - {static_cast(parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_)}, + {static_cast(present_sequence_length)}, {static_cast(parameters.n_reps)}, num_present_sequence_length_tile, {static_cast(parameters.num_heads_)}}); @@ -338,7 +339,9 @@ Status ComputeFlashAttentionDecodeVxReduce(onnxruntime::webgpu::ComputeContext& Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, Tensor* output, const Tensor* past_key, Tensor* present_key, const Tensor* past_value, Tensor* present_value, const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k) { - const int present_sequence_length = parameters.is_gqa_ ? parameters.seqlen_present_kv_cache_ : parameters.total_sequence_length_; + // Extract present_sequence_length directly from present_key tensor shape: + // (batch_size, num_heads, total_sequence_length/max_sequence_length, head_size) + const int present_sequence_length = static_cast(present_key->Shape()[2]); if (parameters.sequence_length_ > 1) { // For encode path, use the original CopyKVCache without indirect dispatch preparation @@ -416,14 +419,20 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co Tensor metadata = context.CreateGPUTensor(DataTypeImpl::GetType(), metadata_shape); ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeQKT(context, Q, attention_bias, &qk, present_key, &metadata, seqlen_k, parameters, indirect_buffer_ptr, num_total_seq_length_tile, - num_present_sequence_length_tile, tile_size, use_indirect_dispatch)); + num_present_sequence_length_tile, tile_size, use_indirect_dispatch, + present_sequence_length)); - const TensorShapeVector out_split_vx_dims({parameters.batch_size_, parameters.num_heads_, num_present_sequence_length_tile, parameters.head_size_}); + const TensorShapeVector out_split_vx_dims({parameters.batch_size_, parameters.num_heads_, + num_present_sequence_length_tile, parameters.head_size_}); const TensorShape out_split_vx_shape(out_split_vx_dims); Tensor out_split_vx = context.CreateGPUTensor(Q->DataType(), out_split_vx_shape); ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeSplitVxScore(context, &metadata, &qk, &out_split_vx, present_value, - seqlen_k, parameters, indirect_buffer_ptr, num_total_seq_length_tile, num_present_sequence_length_tile, tile_size, use_indirect_dispatch)); - ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, output, seqlen_k, parameters, num_present_sequence_length_tile)); + seqlen_k, parameters, indirect_buffer_ptr, + num_total_seq_length_tile, + num_present_sequence_length_tile, tile_size, + use_indirect_dispatch, present_sequence_length)); + ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, output, seqlen_k, parameters, + num_present_sequence_length_tile)); return Status::OK(); } From 12a4559714febf38d2d0e80652fad8fc443b70b1 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 26 Aug 2025 20:23:24 +0800 Subject: [PATCH 03/16] update CopyKVCache --- .../webgpu/bert/flash_attention.cc | 28 ++++++------------- .../contrib_ops/webgpu/bert/flash_attention.h | 2 +- 2 files changed, 9 insertions(+), 21 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 3458b5f178406..f6c7b7843b156 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -31,10 +31,9 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& present_key = shader.AddOutput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); const auto& present_value = shader.AddOutput("present_value", ShaderUsage::UseUniform); const auto& copy_kv_shape = shader.AddIndices("copy_kv_shape"); - + shader.AddInput("seqlen_k"); // If prepare_indirect_dispatch is enabled, add seqlen_k input and indirect_buffer output if (prepare_indirect_dispatch_) { - shader.AddInput("seqlen_k"); shader.AddOutput("indirect_buffer", ShaderUsage::UseUniform); } @@ -43,13 +42,13 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { << " let head_size_id = output_indices[3];\n" " let sequence_id = output_indices[2];\n" " let num_head_id = output_indices[1];\n" - " let batch = output_indices[0];\n"; + " let batch = output_indices[0];\n" + " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n"; // Add indirect dispatch logic for thread 0 if (prepare_indirect_dispatch_) { shader.MainFunctionBody() << " // Prepare indirect dispatch buffer for thread 0\n" << " if (global_idx == 0u) {\n" - << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n" << " let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n" << " indirect_buffer[0] = num_total_seq_length_tile;\n" << " indirect_buffer[1] = uniforms.num_heads;\n" @@ -58,7 +57,7 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { } if (has_past_) { - shader.MainFunctionBody() << "let past_sequence_length = uniforms.past_sequence_length;\n"; + shader.MainFunctionBody() << "let past_sequence_length = total_seq_length - uniforms.kv_sequence_length;\n"; if (past_present_share_buffer_) { shader.MainFunctionBody() << " let present_offset = " << present_key.IndicesToOffset("present_key_indices_t(batch, num_head_id, past_sequence_length + sequence_id, head_size_id)") << ";\n" << " let offset = " << key.IndicesToOffset(kv_BNSH_ ? "key_indices_t(batch, num_head_id, sequence_id, head_size_id)" : "key_indices_t(batch, sequence_id, num_head_id, head_size_id)") << ";\n" @@ -122,8 +121,7 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt {V, ProgramTensorMetadataDependency::TypeAndRank, reshaped_KV_shape, components}}); } - // Add seqlen_k input if preparing indirect dispatch - if (prepare_indirect_dispatch) { + if (seqlen_k != nullptr) { program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); } @@ -144,22 +142,13 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt .SetWorkgroupSize(64) .CacheHint(has_past, parameters.qkv_format_, parameters.past_present_share_buffer_, prepare_indirect_dispatch) .AddUniformVariables({{static_cast(copy_size)}, - // Note that when parameters.past_present_share_buffer_ is true, parameters.past_sequence_length_ will become to - // max_sequence_length. To get a valid past_sequence_length, we use total_sequence_length - kv_sequence_length. - {static_cast(parameters.total_sequence_length_ - parameters.kv_sequence_length_)}, + {static_cast(parameters.kv_sequence_length_)}, {tile_size}, {static_cast(parameters.num_heads_)}}); return context.RunProgram(program); } -// Overloaded version for backward compatibility -Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& parameters, - const Tensor* K, const Tensor* past_key, Tensor* present_key, - const Tensor* V, const Tensor* past_value, Tensor* present_value) { - return CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, nullptr, nullptr); -} - Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const { // Expectations are // qkv have same number of heads and hidden dimension (head size). @@ -240,7 +229,6 @@ Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& conte .CacheHint(tile_size, has_attention_bias) .AddUniformVariables({{static_cast(vectorized_head_size)}, {static_cast(alpha)}, - // present_sequence_length is used to index into the KV cache, for static kv cache it is the max sequence length. {static_cast(present_sequence_length)}, {static_cast(parameters.n_reps)}, {num_present_sequence_length_tile}, @@ -345,7 +333,7 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co if (parameters.sequence_length_ > 1) { // For encode path, use the original CopyKVCache without indirect dispatch preparation - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, nullptr)); const uint32_t tile_size = 64; bool has_attention_bias = attention_bias != nullptr; @@ -409,7 +397,7 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, indirect_buffer_ptr)); } else { // Use the original CopyKVCache without indirect dispatch preparation - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value)); + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, nullptr)); } // The metadata is used to store the max and sum of each tile. diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 0c6d831c619a7..a8fc06590ad4f 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -25,7 +25,7 @@ class CopyKVCacheProgram final : public Program { Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"copy_size", ProgramUniformVariableDataType::Uint32}, - {"past_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"kv_sequence_length", ProgramUniformVariableDataType::Uint32}, {"tile_size", ProgramUniformVariableDataType::Uint32}, {"num_heads", ProgramUniformVariableDataType::Uint32}); From bc4b41ed52dbd915a9b7ae71a39cee3d9f411fda Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 25 Aug 2025 15:53:25 +0800 Subject: [PATCH 04/16] Add int64 support for ReduceSum/Sub/Cast --- .../webgpu/math/binary_elementwise_ops.cc | 50 ++++++++++++++++--- .../webgpu/reduction/reduction_ops.cc | 42 ++++++++-------- .../core/providers/webgpu/shader_variable.cc | 2 +- .../core/providers/webgpu/tensor/cast.cc | 39 ++++++++++++--- .../core/providers/webgpu/tensor/cast.h | 3 +- .../providers/webgpu/webgpu_supported_types.h | 13 +++++ 6 files changed, 111 insertions(+), 38 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc index 82645e30082e6..e52c0da282f2e 100644 --- a/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc @@ -21,6 +21,17 @@ Status BinaryElementwiseProgram::GenerateShaderCode(ShaderHelper& shader) const shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size"); + + if (c.NumComponents() == 1) { + shader.MainFunctionBody() << "let outputIndices = " << c.OffsetToIndices("global_idx") << ";\n" + << "let offset_a = " << a.BroadcastedIndicesToOffset("outputIndices", c) << ";\n" + << "let offset_b = " << b.BroadcastedIndicesToOffset("outputIndices", c) << ";\n" + << "let a = " << a.GetByOffset("offset_a") << ";\n" + << "let b = " << b.GetByOffset("offset_b") << ";\n" + << c.SetByOffset("global_idx", expression_); + return Status::OK(); + } + // check whether can use element-wise mode. // If either A or B is scalar, or A and B have the same shape, element-wise mode can be used. // In element-wise mode, no indices calculation is needed. @@ -133,10 +144,38 @@ Status BinaryElementwise::ComputeInternal(ComputeContext& context) const { return Status::OK(); } + std::string additional_impl; + if (get_additional_impl_) { + additional_impl = get_additional_impl_(lhs_tensor->GetElementType(), rhs_tensor->GetElementType()); + } + bool is_broadcast = lhs_shape != rhs_shape; bool is_lhs_scalar = lhs_shape.IsScalar(); bool is_rhs_scalar = rhs_shape.IsScalar(); + + if (output_tensor->DataType() == DataTypeImpl::GetType()) { + BinaryElementwiseProgram program{kernel_name_, + expression_, + additional_impl, + is_broadcast, + is_lhs_scalar, + is_rhs_scalar, + false, + false, + false}; + program + .SetDispatchGroupSize((size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({ + {static_cast(size)}, + }) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::TypeAndRank}); + program + .AddInputs({{lhs_tensor, ProgramTensorMetadataDependency::TypeAndRank}, + {rhs_tensor, ProgramTensorMetadataDependency::TypeAndRank}}); + return context.RunProgram(program); + } + // Check if either input is boolean // For boolean inputs, we need to handle them differently in the shader. This is because `bool` is not a valid type in // storage buffer. We have to use a `u32` to represent 4 boolean values. @@ -175,11 +214,6 @@ Status BinaryElementwise::ComputeInternal(ComputeContext& context) const { uint32_t vec_size = onnxruntime::narrow((size + 3) / 4); - std::string additional_impl; - if (get_additional_impl_) { - additional_impl = get_additional_impl_(lhs_tensor->GetElementType(), rhs_tensor->GetElementType()); - } - BinaryElementwiseProgram program{kernel_name_, expression_, additional_impl, @@ -312,9 +346,9 @@ WEBGPU_BINARY_VERSIONED_KERNEL(Mul, 13, 13, Mul, WebGpuSupportedNumberTypes()) WEBGPU_BINARY_KERNEL(Mul, 14, Mul, WebGpuSupportedNumberTypes()) WEBGPU_BINARY_IMPL(Sub, "a - b") -WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 7, 12, Sub, WebGpuSupportedNumberTypes()) -WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 13, 13, Sub, WebGpuSupportedNumberTypes()) -WEBGPU_BINARY_KERNEL(Sub, 14, Sub, WebGpuSupportedNumberTypes()) +WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 7, 12, Sub, WebGpuSupportedNumberTypesExtended()) +WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 13, 13, Sub, WebGpuSupportedNumberTypesExtended()) +WEBGPU_BINARY_KERNEL(Sub, 14, Sub, WebGpuSupportedNumberTypesExtended()) std::string GetPowImpl(int lhs_element_type, int /* rhs_element_type */) { SS(s, 1024); diff --git a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc index 4b23f6fc67669..bf51a40087e2b 100644 --- a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc @@ -12,31 +12,31 @@ namespace onnxruntime { namespace webgpu { -#define REGISTER_REDUCE_VERSIONED_KERNEL(ReduceOp, begin, end) \ - ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ - ReduceOp, \ - kOnnxDomain, \ - begin, end, \ - kWebGpuExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypes()), \ +#define REGISTER_REDUCE_VERSIONED_KERNEL(ReduceOp, begin, end) \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + ReduceOp, \ + kOnnxDomain, \ + begin, end, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypesExtended()), \ ReduceOp); -#define REGISTER_REDUCE_VERSIONED_KERNEL_WITH_AXIS_IN_INPUT(ReduceOp, begin, end) \ - ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ - ReduceOp, \ - kOnnxDomain, \ - begin, end, \ - kWebGpuExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypes()).InputMemoryType(OrtMemTypeCPUInput, 1), \ +#define REGISTER_REDUCE_VERSIONED_KERNEL_WITH_AXIS_IN_INPUT(ReduceOp, begin, end) \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + ReduceOp, \ + kOnnxDomain, \ + begin, end, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypesExtended()).InputMemoryType(OrtMemTypeCPUInput, 1), \ ReduceOp); -#define REGISTER_REDUCE_KERNEL(ReduceOp, version) \ - ONNX_OPERATOR_KERNEL_EX( \ - ReduceOp, \ - kOnnxDomain, \ - version, \ - kWebGpuExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypes()).InputMemoryType(OrtMemTypeCPUInput, 1), \ +#define REGISTER_REDUCE_KERNEL(ReduceOp, version) \ + ONNX_OPERATOR_KERNEL_EX( \ + ReduceOp, \ + kOnnxDomain, \ + version, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypesExtended()).InputMemoryType(OrtMemTypeCPUInput, 1), \ ReduceOp); REGISTER_REDUCE_VERSIONED_KERNEL(ReduceMean, 1, 10); diff --git a/onnxruntime/core/providers/webgpu/shader_variable.cc b/onnxruntime/core/providers/webgpu/shader_variable.cc index c197e227e2a8c..ac8b3d07e6e80 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.cc +++ b/onnxruntime/core/providers/webgpu/shader_variable.cc @@ -308,7 +308,7 @@ std::string ShaderVariableHelper::SetByOffsetImpl(std::string_view offset, std:: ORT_THROW("Invalid type"); break; case onnxruntime::webgpu::ProgramVariableDataType::Int64: - ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, " << value << " < 0));"; + ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, i32(" << value << ") < 0));"; break; case onnxruntime::webgpu::ProgramVariableDataType::Uint64: ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), 0u);"; diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 313a96ba25509..6f6bb47a79292 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -20,7 +20,8 @@ const std::vector& CastOpTypeConstraints() { DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}; + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}; return types; } } // namespace @@ -87,12 +88,17 @@ Status Cast::ComputeInternal(ComputeContext& context) const { if (size == 0) { return Status::OK(); } + bool is_from_int64 = input_tensor->DataType() == DataTypeImpl::GetType(); + const int in_components = is_from_int64 ? 1 : 4; + const int out_components = to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64 ? 1 : 4; uint32_t vec_size = onnxruntime::narrow((size + 3) / 4); + uint32_t in_vec_size = onnxruntime::narrow(in_components == 1 ? size : vec_size); + uint32_t out_vec_size = onnxruntime::narrow(out_components == 1 ? size : vec_size); - CastProgram program{to_}; + CastProgram program{to_, is_from_int64}; program - .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) - .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {vec_size}, 4}) + .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {in_vec_size}, in_components}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {out_vec_size}, out_components}) .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({ {static_cast(vec_size)}, @@ -121,12 +127,31 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { case ONNX_NAMESPACE::TensorProto_DataType_BOOL: expression = "vec4(a)"; break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + expression = "int32(a)"; + break; default: ORT_NOT_IMPLEMENTED("Cast to type ", to_, " is not supported."); } - sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size") - << " let a = " << input.GetByOffset("global_idx") << ";\n " - << output.SetByOffset("global_idx", expression); + + sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size"); + if (is_from_int64_) { + sh.MainFunctionBody() << " let a0 = " << input.GetByOffset("global_idx * 4") << ";\n" + << " let a1 = " << input.GetByOffset("global_idx * 4 + 1") << ";\n" + << " let a2 = " << input.GetByOffset("global_idx * 4 + 2") << ";\n" + << " let a3 = " << input.GetByOffset("global_idx * 4 + 3") << ";\n" + << " let a = vec4(a0, a1, a2, a3);\n"; + } else { + sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n"; + } + if (to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64) { + sh.MainFunctionBody() << output.SetByOffset("global_idx * 4", "a.x") << "\n" + << output.SetByOffset("global_idx * 4 + 1", "a.y") << "\n" + << output.SetByOffset("global_idx * 4 + 2", "a.z") << "\n" + << output.SetByOffset("global_idx * 4 + 3", "a.w") << "\n"; + } else { + sh.MainFunctionBody() << output.SetByOffset("global_idx", expression); + } return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 925cd200f0aba..11e2a5b32fc7b 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -10,7 +10,7 @@ namespace webgpu { class CastProgram final : public Program { public: - CastProgram(int32_t to) : Program{"Cast"}, to_{to} {} + CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -18,6 +18,7 @@ class CastProgram final : public Program { private: int32_t to_; + bool is_from_int64_; }; class Cast final : public WebGpuKernel { diff --git a/onnxruntime/core/providers/webgpu/webgpu_supported_types.h b/onnxruntime/core/providers/webgpu/webgpu_supported_types.h index ff66cd535399e..c31c2bafee53b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_supported_types.h +++ b/onnxruntime/core/providers/webgpu/webgpu_supported_types.h @@ -15,6 +15,14 @@ using SupportedNumberTypes = int32_t, uint32_t>; +using SupportedNumberTypesExtended = + TypeList< + float, + MLFloat16, + int32_t, + uint32_t, + int64_t>; + using SupportedFloats = TypeList< float, @@ -25,6 +33,11 @@ inline const std::vector& WebGpuSupportedNumberTypes() { return supportedDataTypes; } +inline const std::vector& WebGpuSupportedNumberTypesExtended() { + static const std::vector supportedDataTypes = BuildKernelDefConstraintsFromTypeList(); + return supportedDataTypes; +} + inline const std::vector& WebGpuSupportedFloatTypes() { static const std::vector supportedDataTypes = BuildKernelDefConstraintsFromTypeList(); return supportedDataTypes; From c8231b66fc1f850fce5e48b2422e4ca8e18ed84b Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 29 Aug 2025 20:26:17 +0800 Subject: [PATCH 05/16] fix --- onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc | 10 +++++----- onnxruntime/core/session/inference_session.cc | 2 +- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index f6c7b7843b156..acfa9ba8051b9 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -31,10 +31,10 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& present_key = shader.AddOutput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); const auto& present_value = shader.AddOutput("present_value", ShaderUsage::UseUniform); const auto& copy_kv_shape = shader.AddIndices("copy_kv_shape"); - shader.AddInput("seqlen_k"); + shader.AddInput("seqlen_k", ShaderUsage::None); // If prepare_indirect_dispatch is enabled, add seqlen_k input and indirect_buffer output if (prepare_indirect_dispatch_) { - shader.AddOutput("indirect_buffer", ShaderUsage::UseUniform); + shader.AddOutput("indirect_buffer", ShaderUsage::None); } shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.copy_size") @@ -184,7 +184,7 @@ Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const { Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); shader.AddInput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - shader.AddInput("seqlens_k"); + shader.AddInput("seqlens_k", ShaderUsage::None); if (has_attention_bias_) { shader.AddInput("attention_bias", ShaderUsage::UseUniform); } @@ -241,7 +241,7 @@ Status FlashAttentionDecodeSplitVxProgram::GenerateShaderCode(ShaderHelper& shad shader.AddInput("metadata", ShaderUsage::UseUniform); shader.AddInput("qk", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); shader.AddInput("present_value", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); - shader.AddInput("seqlens_k"); + shader.AddInput("seqlens_k", ShaderUsage::None); shader.AddOutput("out_split_vx", ShaderUsage::UseUniform); const uint32_t tile_size_k_vec = 8u; @@ -292,7 +292,7 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte Status FlashAttentionDecodeVxReduceProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("input", ShaderUsage::UseUniform); - shader.AddInput("seqlens_k"); + shader.AddInput("seqlens_k", ShaderUsage::None); shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_decode_vx_reduce.wgsl.template", diff --git a/onnxruntime/core/session/inference_session.cc b/onnxruntime/core/session/inference_session.cc index 17b7f9af372bc..095dc858c4f51 100644 --- a/onnxruntime/core/session/inference_session.cc +++ b/onnxruntime/core/session/inference_session.cc @@ -145,7 +145,7 @@ static bool HasControlflowNodes(const Graph& graph) { static bool HasMemcpyNodes(const Graph& graph) { for (const auto& node : graph.Nodes()) { - if (node.OpType() == "MemcpyFromHost" || node.OpType() == "MemcpyToHost") { + if (node.OpType() == "MemcpyFromHost") { return true; } } From 1197a17c37d7e43ebe652f0256df0685b60bd5ee Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 2 Sep 2025 16:44:49 +0800 Subject: [PATCH 06/16] keep the dispatch path unchanged --- .../webgpu/bert/flash_attention.cc | 91 ++++++++++++------- .../contrib_ops/webgpu/bert/flash_attention.h | 27 +++--- .../flash_attention_decode_qkt.wgsl.template | 5 + ...sh_attention_decode_split_vx.wgsl.template | 5 + ...h_attention_decode_vx_reduce.wgsl.template | 5 + 5 files changed, 87 insertions(+), 46 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index acfa9ba8051b9..417a4044b795d 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -31,9 +31,9 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& present_key = shader.AddOutput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseIndicesTypeAlias); const auto& present_value = shader.AddOutput("present_value", ShaderUsage::UseUniform); const auto& copy_kv_shape = shader.AddIndices("copy_kv_shape"); - shader.AddInput("seqlen_k", ShaderUsage::None); // If prepare_indirect_dispatch is enabled, add seqlen_k input and indirect_buffer output if (prepare_indirect_dispatch_) { + shader.AddInput("seqlen_k", ShaderUsage::None); shader.AddOutput("indirect_buffer", ShaderUsage::None); } @@ -42,8 +42,12 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { << " let head_size_id = output_indices[3];\n" " let sequence_id = output_indices[2];\n" " let num_head_id = output_indices[1];\n" - " let batch = output_indices[0];\n" - " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n"; + " let batch = output_indices[0];\n"; + if (prepare_indirect_dispatch_) { + shader.MainFunctionBody() << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n"; + } else { + shader.MainFunctionBody() << " let total_seq_length = uniforms.total_sequence_length;\n"; + } // Add indirect dispatch logic for thread 0 if (prepare_indirect_dispatch_) { @@ -89,7 +93,7 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAttentionParameters& parameters, const Tensor* K, const Tensor* past_key, Tensor* present_key, const Tensor* V, const Tensor* past_value, Tensor* present_value, - const Tensor* seqlen_k, Tensor* indirect_buffer) { + uint32_t tile_size, const Tensor* seqlen_k, Tensor* indirect_buffer) { // CopyKVCache takes past key/value and current key/value and copies them to present key and value. // This makes it so that FlashAttention only needs to look at present key and value, and saves // number of input buffers in the shader, which we run out of (<=8) without this optimization. @@ -106,10 +110,9 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt // Determine if we need to prepare indirect dispatch bool prepare_indirect_dispatch = (indirect_buffer != nullptr); - constexpr uint32_t tile_size = 64; CopyKVCacheProgram program{"CopyKVCache", has_past, parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH, parameters.past_present_share_buffer_, - prepare_indirect_dispatch, tile_size, static_cast(parameters.num_heads_)}; + prepare_indirect_dispatch}; if (parameters.qkv_format_ == Q_K_V_BSNH_BNSH_BNSH) { program.AddInputs({{K, ProgramTensorMetadataDependency::TypeAndRank, components}, {V, ProgramTensorMetadataDependency::TypeAndRank, components}}); @@ -121,7 +124,7 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt {V, ProgramTensorMetadataDependency::TypeAndRank, reshaped_KV_shape, components}}); } - if (seqlen_k != nullptr) { + if (prepare_indirect_dispatch) { program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); } @@ -132,7 +135,6 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt program.AddOutputs({{present_key, ProgramTensorMetadataDependency::Rank, components}, {present_value, ProgramTensorMetadataDependency::Rank, components}}); - // Add indirect_buffer output if preparing indirect dispatch if (prepare_indirect_dispatch) { program.AddOutput({indirect_buffer, ProgramTensorMetadataDependency::None}); } @@ -142,6 +144,7 @@ Status CopyKVCache(onnxruntime::webgpu::ComputeContext& context, const WebgpuAtt .SetWorkgroupSize(64) .CacheHint(has_past, parameters.qkv_format_, parameters.past_present_share_buffer_, prepare_indirect_dispatch) .AddUniformVariables({{static_cast(copy_size)}, + {static_cast(parameters.total_sequence_length_)}, {static_cast(parameters.kv_sequence_length_)}, {tile_size}, {static_cast(parameters.num_heads_)}}); @@ -184,7 +187,9 @@ Status FlashAttentionProgram::GenerateShaderCode(ShaderHelper& shader) const { Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("q", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); shader.AddInput("present_key", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); - shader.AddInput("seqlens_k", ShaderUsage::None); + if (use_indirect_dispatch_) { + shader.AddInput("seqlens_k", ShaderUsage::None); + } if (has_attention_bias_) { shader.AddInput("attention_bias", ShaderUsage::UseUniform); } @@ -197,7 +202,8 @@ Status FlashAttentionDecodeQKTProgram::GenerateShaderCode(ShaderHelper& shader) WGSL_TEMPLATE_PARAMETER(has_attention_bias, has_attention_bias_), WGSL_TEMPLATE_PARAMETER(sub_tile_count, sub_tile_count), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec)); + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec), + WGSL_TEMPLATE_PARAMETER(use_indirect_dispatch, use_indirect_dispatch_)); } Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& context, const Tensor* Q, @@ -209,10 +215,12 @@ Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& conte const bool has_attention_bias = attention_bias != nullptr; const int components = 4; - FlashAttentionDecodeQKTProgram program{"FlashAttentionDecodeQKT", has_attention_bias, tile_size}; + FlashAttentionDecodeQKTProgram program{"FlashAttentionDecodeQKT", has_attention_bias, tile_size, use_indirect_dispatch}; program.AddInputs({{Q, ProgramTensorMetadataDependency::TypeAndRank, components}, - {present_key, ProgramTensorMetadataDependency::TypeAndRank, components}, - {seqlen_k, ProgramTensorMetadataDependency::None}}); + {present_key, ProgramTensorMetadataDependency::TypeAndRank, components}}); + if (use_indirect_dispatch) { + program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); + } if (has_attention_bias) { program.AddInput({attention_bias, ProgramTensorMetadataDependency::TypeAndRank}); } @@ -226,8 +234,9 @@ Status ComputeFlashAttentionDecodeQKT(onnxruntime::webgpu::ComputeContext& conte program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile); } program.SetWorkgroupSize(64) - .CacheHint(tile_size, has_attention_bias) + .CacheHint(tile_size, has_attention_bias, use_indirect_dispatch) .AddUniformVariables({{static_cast(vectorized_head_size)}, + {static_cast(parameters.total_sequence_length_)}, {static_cast(alpha)}, {static_cast(present_sequence_length)}, {static_cast(parameters.n_reps)}, @@ -241,7 +250,9 @@ Status FlashAttentionDecodeSplitVxProgram::GenerateShaderCode(ShaderHelper& shad shader.AddInput("metadata", ShaderUsage::UseUniform); shader.AddInput("qk", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); shader.AddInput("present_value", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias | ShaderUsage::UseElementTypeAlias); - shader.AddInput("seqlens_k", ShaderUsage::None); + if (use_indirect_dispatch_) { + shader.AddInput("seqlens_k", ShaderUsage::None); + } shader.AddOutput("out_split_vx", ShaderUsage::UseUniform); const uint32_t tile_size_k_vec = 8u; @@ -250,7 +261,8 @@ Status FlashAttentionDecodeSplitVxProgram::GenerateShaderCode(ShaderHelper& shad WGSL_TEMPLATE_PARAMETER(head_size_vec, head_size_vec_), WGSL_TEMPLATE_PARAMETER(sub_tile_count, WorkgroupSizeX() / tile_size_k_vec), WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), - WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec)); + WGSL_TEMPLATE_PARAMETER(tile_size_k_vec, tile_size_k_vec), + WGSL_TEMPLATE_PARAMETER(use_indirect_dispatch, use_indirect_dispatch_)); } Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeContext& context, @@ -268,20 +280,21 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte uint32_t present_sequence_length) { const int components = 4; int head_size_vec = parameters.v_head_size_ / components; - FlashAttentionDecodeSplitVxProgram program{"FlashAttentionDecodeSplitVx", tile_size, head_size_vec}; + FlashAttentionDecodeSplitVxProgram program{"FlashAttentionDecodeSplitVx", tile_size, head_size_vec, use_indirect_dispatch}; program.AddInputs({{metadata, ProgramTensorMetadataDependency::TypeAndRank, 2}, {qk, ProgramTensorMetadataDependency::TypeAndRank}, - {present_value, ProgramTensorMetadataDependency::TypeAndRank, components}, - {seqlen_k, ProgramTensorMetadataDependency::None}}); + {present_value, ProgramTensorMetadataDependency::TypeAndRank, components}}); program.AddOutputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}}); // [B, N, split_k, head_size] if (use_indirect_dispatch) { - program.SetIndirectDispatchTensor(indirect_buffer); + program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}) + .SetIndirectDispatchTensor(indirect_buffer); } else { program.SetDispatchGroupSize(parameters.num_heads_ * num_total_seq_length_tile); } - program.CacheHint(tile_size, head_size_vec) + program.CacheHint(tile_size, head_size_vec, use_indirect_dispatch) .SetWorkgroupSize(64) - .AddUniformVariables({{static_cast(head_size_vec)}, + .AddUniformVariables({{static_cast(parameters.total_sequence_length_)}, + {static_cast(head_size_vec)}, {static_cast(present_sequence_length)}, {static_cast(parameters.n_reps)}, num_present_sequence_length_tile, @@ -292,11 +305,14 @@ Status ComputeFlashAttentionDecodeSplitVxScore(onnxruntime::webgpu::ComputeConte Status FlashAttentionDecodeVxReduceProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("input", ShaderUsage::UseUniform); - shader.AddInput("seqlens_k", ShaderUsage::None); + if (use_indirect_dispatch_) { + shader.AddInput("seqlens_k", ShaderUsage::None); + } shader.AddOutput("output", ShaderUsage::UseUniform | ShaderUsage::UseValueTypeAlias); return WGSL_TEMPLATE_APPLY(shader, "bert/flash_attention_decode_vx_reduce.wgsl.template", - WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_)); + WGSL_TEMPLATE_PARAMETER(tile_size, tile_size_), + WGSL_TEMPLATE_PARAMETER(use_indirect_dispatch, use_indirect_dispatch_)); } Status ComputeFlashAttentionDecodeVxReduce(onnxruntime::webgpu::ComputeContext& context, @@ -304,19 +320,24 @@ Status ComputeFlashAttentionDecodeVxReduce(onnxruntime::webgpu::ComputeContext& Tensor* output, const Tensor* seqlen_k, const WebgpuAttentionParameters& parameters, - uint32_t num_present_sequence_length_tile) { + uint32_t num_total_seq_length_tile, + uint32_t num_present_sequence_length_tile, + bool use_indirect_dispatch) { const int components = 4; constexpr int tile_size = 8; int tile_head_size = tile_size * components; - FlashAttentionDecodeVxReduceProgram program{"FlashAttentionDecodeVxReduce", tile_size}; - program.AddInputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}, - {seqlen_k, ProgramTensorMetadataDependency::None}}); + FlashAttentionDecodeVxReduceProgram program{"FlashAttentionDecodeVxReduce", tile_size, use_indirect_dispatch}; + program.AddInputs({{out_split_vx, ProgramTensorMetadataDependency::TypeAndRank, components}}); + if (use_indirect_dispatch) { + program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); + } program.AddOutputs({{output, ProgramTensorMetadataDependency::TypeAndRank, components}}); const uint32_t num_head_size_tile = static_cast((parameters.v_head_size_ + tile_head_size - 1) / tile_head_size); program.SetDispatchGroupSize(parameters.num_heads_ * num_head_size_tile) - .CacheHint(tile_size) + .CacheHint(tile_size, use_indirect_dispatch) .SetWorkgroupSize(tile_size * tile_size) .AddUniformVariables({{static_cast(parameters.v_head_size_ / components)}, + num_total_seq_length_tile, num_present_sequence_length_tile, {num_head_size_tile}, {static_cast(parameters.num_heads_)}}); @@ -332,10 +353,9 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co const int present_sequence_length = static_cast(present_key->Shape()[2]); if (parameters.sequence_length_ > 1) { - // For encode path, use the original CopyKVCache without indirect dispatch preparation - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, nullptr)); - const uint32_t tile_size = 64; + // For encode path, use the original CopyKVCache without indirect dispatch preparation + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, tile_size, seqlen_k, nullptr)); bool has_attention_bias = attention_bias != nullptr; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; bool is_nvidia = context.AdapterInfo().vendor == std::string_view{"nvidia"}; @@ -394,10 +414,10 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co indirect_buffer = context.CreateGPUTensor(DataTypeImpl::GetType(), indirect_buffer_shape); indirect_buffer_ptr = &indirect_buffer; // Use the fused CopyKVCache that also prepares the indirect dispatch buffer - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, indirect_buffer_ptr)); + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, tile_size, seqlen_k, indirect_buffer_ptr)); } else { // Use the original CopyKVCache without indirect dispatch preparation - ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, seqlen_k, nullptr)); + ORT_RETURN_IF_ERROR(CopyKVCache(context, parameters, K, past_key, present_key, V, past_value, present_value, tile_size, seqlen_k, nullptr)); } // The metadata is used to store the max and sum of each tile. @@ -420,7 +440,8 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co num_present_sequence_length_tile, tile_size, use_indirect_dispatch, present_sequence_length)); ORT_RETURN_IF_ERROR(ComputeFlashAttentionDecodeVxReduce(context, &out_split_vx, output, seqlen_k, parameters, - num_present_sequence_length_tile)); + num_total_seq_length_tile, + num_present_sequence_length_tile, use_indirect_dispatch)); return Status::OK(); } diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index a8fc06590ad4f..6481f50ab34f6 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -18,13 +18,14 @@ using namespace onnxruntime::webgpu; class CopyKVCacheProgram final : public Program { public: CopyKVCacheProgram(const std::string& kernel_name, bool has_past, bool kv_BNSH, bool past_present_share_buffer, - bool prepare_indirect_dispatch = false, uint32_t tile_size = 0, uint32_t num_heads = 0) - : Program{kernel_name}, has_past_(has_past), kv_BNSH_(kv_BNSH), past_present_share_buffer_(past_present_share_buffer), prepare_indirect_dispatch_(prepare_indirect_dispatch), tile_size_(tile_size), num_heads_(num_heads) { + bool prepare_indirect_dispatch = false) + : Program{kernel_name}, has_past_(has_past), kv_BNSH_(kv_BNSH), past_present_share_buffer_(past_present_share_buffer), prepare_indirect_dispatch_(prepare_indirect_dispatch) { } Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"copy_size", ProgramUniformVariableDataType::Uint32}, + {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, {"kv_sequence_length", ProgramUniformVariableDataType::Uint32}, {"tile_size", ProgramUniformVariableDataType::Uint32}, {"num_heads", ProgramUniformVariableDataType::Uint32}); @@ -34,8 +35,6 @@ class CopyKVCacheProgram final : public Program { bool kv_BNSH_; bool past_present_share_buffer_; bool prepare_indirect_dispatch_; - uint32_t tile_size_; - uint32_t num_heads_; }; class FlashAttentionProgram final : public Program { @@ -81,13 +80,14 @@ class FlashAttentionProgram final : public Program { class FlashAttentionDecodeQKTProgram final : public Program { public: FlashAttentionDecodeQKTProgram(const std::string& kernel_name, - bool has_attention_bias, uint32_t tile_size) - : Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size) { + bool has_attention_bias, uint32_t tile_size, bool use_indirect_dispatch) + : Program{kernel_name}, has_attention_bias_(has_attention_bias), tile_size_(tile_size), use_indirect_dispatch_(use_indirect_dispatch) { } Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"head_size_vec", ProgramUniformVariableDataType::Uint32}, + {"total_sequence_length", ProgramUniformVariableDataType::Uint32}, {"alpha", ProgramUniformVariableDataType::Float32}, {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, {"n_reps", ProgramUniformVariableDataType::Uint32}, @@ -97,17 +97,19 @@ class FlashAttentionDecodeQKTProgram final : public Program { public: - FlashAttentionDecodeSplitVxProgram(const std::string& kernel_name, uint32_t tile_size, int head_size_vec) - : Program{kernel_name}, tile_size_(tile_size), head_size_vec_(head_size_vec) { + FlashAttentionDecodeSplitVxProgram(const std::string& kernel_name, uint32_t tile_size, int head_size_vec, bool use_indirect_dispatch) + : Program{kernel_name}, tile_size_(tile_size), head_size_vec_(head_size_vec), use_indirect_dispatch_(use_indirect_dispatch) { } Status GenerateShaderCode(ShaderHelper& sh) const override; - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"head_size_vec", ProgramUniformVariableDataType::Uint32}, + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"total_sequence_length", ProgramUniformVariableDataType::Uint32}, + {"head_size_vec", ProgramUniformVariableDataType::Uint32}, {"present_sequence_length", ProgramUniformVariableDataType::Uint32}, {"n_reps", ProgramUniformVariableDataType::Uint32}, {"num_present_sequence_length_tile", ProgramUniformVariableDataType::Uint32}, @@ -116,23 +118,26 @@ class FlashAttentionDecodeSplitVxProgram final : public Program { public: - FlashAttentionDecodeVxReduceProgram(const std::string& kernel_name, uint32_t tile_size) - : Program{kernel_name}, tile_size_(tile_size) { + FlashAttentionDecodeVxReduceProgram(const std::string& kernel_name, uint32_t tile_size, bool use_indirect_dispatch) + : Program{kernel_name}, tile_size_(tile_size), use_indirect_dispatch_(use_indirect_dispatch) { } Status GenerateShaderCode(ShaderHelper& sh) const override; WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"head_size_vec", ProgramUniformVariableDataType::Uint32}, + {"num_total_seq_length_tile", ProgramUniformVariableDataType::Uint32}, {"num_present_sequence_length_tile", ProgramUniformVariableDataType::Uint32}, {"num_head_size_tile", ProgramUniformVariableDataType::Uint32}, {"num_heads", ProgramUniformVariableDataType::Uint32}); private: uint32_t tile_size_; + bool use_indirect_dispatch_; }; Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, const Tensor* attention_bias, diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkt.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkt.wgsl.template index 7b8d5b19e1029..c6f768beffa0f 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkt.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_qkt.wgsl.template @@ -5,6 +5,7 @@ #param tile_size #param tile_size_k_vec #param sub_tile_count +#param use_indirect_dispatch // Note that this shader adopts similar algorithm with dp4a generation shader. // @@ -48,7 +49,11 @@ var tile_qk: array; $MAIN { let local_row = u32(local_idx / tile_size_k_vec); let local_col = local_idx % tile_size_k_vec; +#if use_indirect_dispatch let total_sequence_length = u32(seqlens_k[0]) + 1u; +#else + let total_sequence_length = uniforms.total_sequence_length; +#endif let num_total_seq_length_tile = (total_sequence_length + tile_size - 1) / tile_size; let total_seq_offset = (workgroup_idx % num_total_seq_length_tile) * tile_size; let head_idx = u32(workgroup_idx / num_total_seq_length_tile); diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template index dd01343164499..37cf7e8f11b1f 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_split_vx.wgsl.template @@ -5,6 +5,7 @@ #param head_size_vec #param tile_size_k_vec #param sub_tile_count +#param use_indirect_dispatch // Note that this shader adopts similar algorithm with dp4a generation shader. // @@ -40,7 +41,11 @@ var qkv_values: array, $MAIN { let local_row = u32(local_idx / tile_size_k_vec); let local_col = local_idx % tile_size_k_vec; + #if use_indirect_dispatch let total_sequence_length = u32(seqlens_k[0]) + 1u; + #else + let total_sequence_length = uniforms.total_sequence_length; + #endif let num_total_seq_length_tile = (total_sequence_length + tile_size - 1) / tile_size; let total_seq_offset = (workgroup_idx % num_total_seq_length_tile) * tile_size; let head_idx = u32(workgroup_idx / num_total_seq_length_tile); diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template index f010aa9781519..08ed21810337a 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention_decode_vx_reduce.wgsl.template @@ -2,6 +2,7 @@ // Licensed under the MIT License. #param tile_size +#param use_indirect_dispatch // Inputs are splits of the GQA output, split into num_total_seq_length_tiles // rows. This shader needs to add these splits across the row dimension to @@ -23,8 +24,12 @@ $MAIN { var value = output_value_t(0); let local_row = u32(local_idx / tile_size); let local_col = local_idx % tile_size; + #if use_indirect_dispatch let total_sequence_length = u32(seqlens_k[0]) + 1u; let num_total_seq_length_tile = (total_sequence_length + 63u) / 64u; + #else + let num_total_seq_length_tile = uniforms.num_total_seq_length_tile; + #endif if (head_size_offset + local_col < uniforms.head_size_vec) { for (var r = 0u; r < num_total_seq_length_tile; r += tile_size) { From 5d66e653651f47bdcd49021e27cfd3efbd4acb3a Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 3 Sep 2025 16:12:50 +0800 Subject: [PATCH 07/16] Revert "Add int64 support for ReduceSum/Sub/Cast" This reverts commit bc4b41ed52dbd915a9b7ae71a39cee3d9f411fda. --- .../webgpu/math/binary_elementwise_ops.cc | 50 +++---------------- .../webgpu/reduction/reduction_ops.cc | 42 ++++++++-------- .../core/providers/webgpu/shader_variable.cc | 2 +- .../core/providers/webgpu/tensor/cast.cc | 39 +++------------ .../core/providers/webgpu/tensor/cast.h | 3 +- .../providers/webgpu/webgpu_supported_types.h | 13 ----- 6 files changed, 38 insertions(+), 111 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc b/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc index e52c0da282f2e..82645e30082e6 100644 --- a/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc +++ b/onnxruntime/core/providers/webgpu/math/binary_elementwise_ops.cc @@ -21,17 +21,6 @@ Status BinaryElementwiseProgram::GenerateShaderCode(ShaderHelper& shader) const shader.MainFunctionBody() << shader.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size"); - - if (c.NumComponents() == 1) { - shader.MainFunctionBody() << "let outputIndices = " << c.OffsetToIndices("global_idx") << ";\n" - << "let offset_a = " << a.BroadcastedIndicesToOffset("outputIndices", c) << ";\n" - << "let offset_b = " << b.BroadcastedIndicesToOffset("outputIndices", c) << ";\n" - << "let a = " << a.GetByOffset("offset_a") << ";\n" - << "let b = " << b.GetByOffset("offset_b") << ";\n" - << c.SetByOffset("global_idx", expression_); - return Status::OK(); - } - // check whether can use element-wise mode. // If either A or B is scalar, or A and B have the same shape, element-wise mode can be used. // In element-wise mode, no indices calculation is needed. @@ -144,38 +133,10 @@ Status BinaryElementwise::ComputeInternal(ComputeContext& context) const { return Status::OK(); } - std::string additional_impl; - if (get_additional_impl_) { - additional_impl = get_additional_impl_(lhs_tensor->GetElementType(), rhs_tensor->GetElementType()); - } - bool is_broadcast = lhs_shape != rhs_shape; bool is_lhs_scalar = lhs_shape.IsScalar(); bool is_rhs_scalar = rhs_shape.IsScalar(); - - if (output_tensor->DataType() == DataTypeImpl::GetType()) { - BinaryElementwiseProgram program{kernel_name_, - expression_, - additional_impl, - is_broadcast, - is_lhs_scalar, - is_rhs_scalar, - false, - false, - false}; - program - .SetDispatchGroupSize((size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) - .AddUniformVariables({ - {static_cast(size)}, - }) - .AddOutput({output_tensor, ProgramTensorMetadataDependency::TypeAndRank}); - program - .AddInputs({{lhs_tensor, ProgramTensorMetadataDependency::TypeAndRank}, - {rhs_tensor, ProgramTensorMetadataDependency::TypeAndRank}}); - return context.RunProgram(program); - } - // Check if either input is boolean // For boolean inputs, we need to handle them differently in the shader. This is because `bool` is not a valid type in // storage buffer. We have to use a `u32` to represent 4 boolean values. @@ -214,6 +175,11 @@ Status BinaryElementwise::ComputeInternal(ComputeContext& context) const { uint32_t vec_size = onnxruntime::narrow((size + 3) / 4); + std::string additional_impl; + if (get_additional_impl_) { + additional_impl = get_additional_impl_(lhs_tensor->GetElementType(), rhs_tensor->GetElementType()); + } + BinaryElementwiseProgram program{kernel_name_, expression_, additional_impl, @@ -346,9 +312,9 @@ WEBGPU_BINARY_VERSIONED_KERNEL(Mul, 13, 13, Mul, WebGpuSupportedNumberTypes()) WEBGPU_BINARY_KERNEL(Mul, 14, Mul, WebGpuSupportedNumberTypes()) WEBGPU_BINARY_IMPL(Sub, "a - b") -WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 7, 12, Sub, WebGpuSupportedNumberTypesExtended()) -WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 13, 13, Sub, WebGpuSupportedNumberTypesExtended()) -WEBGPU_BINARY_KERNEL(Sub, 14, Sub, WebGpuSupportedNumberTypesExtended()) +WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 7, 12, Sub, WebGpuSupportedNumberTypes()) +WEBGPU_BINARY_VERSIONED_KERNEL(Sub, 13, 13, Sub, WebGpuSupportedNumberTypes()) +WEBGPU_BINARY_KERNEL(Sub, 14, Sub, WebGpuSupportedNumberTypes()) std::string GetPowImpl(int lhs_element_type, int /* rhs_element_type */) { SS(s, 1024); diff --git a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc index bf51a40087e2b..4b23f6fc67669 100644 --- a/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/webgpu/reduction/reduction_ops.cc @@ -12,31 +12,31 @@ namespace onnxruntime { namespace webgpu { -#define REGISTER_REDUCE_VERSIONED_KERNEL(ReduceOp, begin, end) \ - ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ - ReduceOp, \ - kOnnxDomain, \ - begin, end, \ - kWebGpuExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypesExtended()), \ +#define REGISTER_REDUCE_VERSIONED_KERNEL(ReduceOp, begin, end) \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + ReduceOp, \ + kOnnxDomain, \ + begin, end, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypes()), \ ReduceOp); -#define REGISTER_REDUCE_VERSIONED_KERNEL_WITH_AXIS_IN_INPUT(ReduceOp, begin, end) \ - ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ - ReduceOp, \ - kOnnxDomain, \ - begin, end, \ - kWebGpuExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypesExtended()).InputMemoryType(OrtMemTypeCPUInput, 1), \ +#define REGISTER_REDUCE_VERSIONED_KERNEL_WITH_AXIS_IN_INPUT(ReduceOp, begin, end) \ + ONNX_OPERATOR_VERSIONED_KERNEL_EX( \ + ReduceOp, \ + kOnnxDomain, \ + begin, end, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypes()).InputMemoryType(OrtMemTypeCPUInput, 1), \ ReduceOp); -#define REGISTER_REDUCE_KERNEL(ReduceOp, version) \ - ONNX_OPERATOR_KERNEL_EX( \ - ReduceOp, \ - kOnnxDomain, \ - version, \ - kWebGpuExecutionProvider, \ - (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypesExtended()).InputMemoryType(OrtMemTypeCPUInput, 1), \ +#define REGISTER_REDUCE_KERNEL(ReduceOp, version) \ + ONNX_OPERATOR_KERNEL_EX( \ + ReduceOp, \ + kOnnxDomain, \ + version, \ + kWebGpuExecutionProvider, \ + (*KernelDefBuilder::Create()).TypeConstraint("T", WebGpuSupportedNumberTypes()).InputMemoryType(OrtMemTypeCPUInput, 1), \ ReduceOp); REGISTER_REDUCE_VERSIONED_KERNEL(ReduceMean, 1, 10); diff --git a/onnxruntime/core/providers/webgpu/shader_variable.cc b/onnxruntime/core/providers/webgpu/shader_variable.cc index ac8b3d07e6e80..c197e227e2a8c 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.cc +++ b/onnxruntime/core/providers/webgpu/shader_variable.cc @@ -308,7 +308,7 @@ std::string ShaderVariableHelper::SetByOffsetImpl(std::string_view offset, std:: ORT_THROW("Invalid type"); break; case onnxruntime::webgpu::ProgramVariableDataType::Int64: - ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, i32(" << value << ") < 0));"; + ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, " << value << " < 0));"; break; case onnxruntime::webgpu::ProgramVariableDataType::Uint64: ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), 0u);"; diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 6f6bb47a79292..313a96ba25509 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -20,8 +20,7 @@ const std::vector& CastOpTypeConstraints() { DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}; + DataTypeImpl::GetTensorType()}; return types; } } // namespace @@ -88,17 +87,12 @@ Status Cast::ComputeInternal(ComputeContext& context) const { if (size == 0) { return Status::OK(); } - bool is_from_int64 = input_tensor->DataType() == DataTypeImpl::GetType(); - const int in_components = is_from_int64 ? 1 : 4; - const int out_components = to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64 ? 1 : 4; uint32_t vec_size = onnxruntime::narrow((size + 3) / 4); - uint32_t in_vec_size = onnxruntime::narrow(in_components == 1 ? size : vec_size); - uint32_t out_vec_size = onnxruntime::narrow(out_components == 1 ? size : vec_size); - CastProgram program{to_, is_from_int64}; + CastProgram program{to_}; program - .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {in_vec_size}, in_components}) - .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {out_vec_size}, out_components}) + .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {vec_size}, 4}) .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({ {static_cast(vec_size)}, @@ -127,31 +121,12 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { case ONNX_NAMESPACE::TensorProto_DataType_BOOL: expression = "vec4(a)"; break; - case ONNX_NAMESPACE::TensorProto_DataType_INT64: - expression = "int32(a)"; - break; default: ORT_NOT_IMPLEMENTED("Cast to type ", to_, " is not supported."); } - - sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size"); - if (is_from_int64_) { - sh.MainFunctionBody() << " let a0 = " << input.GetByOffset("global_idx * 4") << ";\n" - << " let a1 = " << input.GetByOffset("global_idx * 4 + 1") << ";\n" - << " let a2 = " << input.GetByOffset("global_idx * 4 + 2") << ";\n" - << " let a3 = " << input.GetByOffset("global_idx * 4 + 3") << ";\n" - << " let a = vec4(a0, a1, a2, a3);\n"; - } else { - sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n"; - } - if (to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64) { - sh.MainFunctionBody() << output.SetByOffset("global_idx * 4", "a.x") << "\n" - << output.SetByOffset("global_idx * 4 + 1", "a.y") << "\n" - << output.SetByOffset("global_idx * 4 + 2", "a.z") << "\n" - << output.SetByOffset("global_idx * 4 + 3", "a.w") << "\n"; - } else { - sh.MainFunctionBody() << output.SetByOffset("global_idx", expression); - } + sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size") + << " let a = " << input.GetByOffset("global_idx") << ";\n " + << output.SetByOffset("global_idx", expression); return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 11e2a5b32fc7b..925cd200f0aba 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -10,7 +10,7 @@ namespace webgpu { class CastProgram final : public Program { public: - CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} + CastProgram(int32_t to) : Program{"Cast"}, to_{to} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -18,7 +18,6 @@ class CastProgram final : public Program { private: int32_t to_; - bool is_from_int64_; }; class Cast final : public WebGpuKernel { diff --git a/onnxruntime/core/providers/webgpu/webgpu_supported_types.h b/onnxruntime/core/providers/webgpu/webgpu_supported_types.h index c31c2bafee53b..ff66cd535399e 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_supported_types.h +++ b/onnxruntime/core/providers/webgpu/webgpu_supported_types.h @@ -15,14 +15,6 @@ using SupportedNumberTypes = int32_t, uint32_t>; -using SupportedNumberTypesExtended = - TypeList< - float, - MLFloat16, - int32_t, - uint32_t, - int64_t>; - using SupportedFloats = TypeList< float, @@ -33,11 +25,6 @@ inline const std::vector& WebGpuSupportedNumberTypes() { return supportedDataTypes; } -inline const std::vector& WebGpuSupportedNumberTypesExtended() { - static const std::vector supportedDataTypes = BuildKernelDefConstraintsFromTypeList(); - return supportedDataTypes; -} - inline const std::vector& WebGpuSupportedFloatTypes() { static const std::vector supportedDataTypes = BuildKernelDefConstraintsFromTypeList(); return supportedDataTypes; From 1576d1ec8a5be38e3f1082c4c8b42cee36bb6bc3 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Wed, 3 Sep 2025 16:18:07 +0800 Subject: [PATCH 08/16] Add int64 to cast --- .../core/providers/webgpu/shader_variable.cc | 2 +- .../core/providers/webgpu/tensor/cast.cc | 39 +++++++++++++++---- .../core/providers/webgpu/tensor/cast.h | 3 +- 3 files changed, 35 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/shader_variable.cc b/onnxruntime/core/providers/webgpu/shader_variable.cc index c197e227e2a8c..ac8b3d07e6e80 100644 --- a/onnxruntime/core/providers/webgpu/shader_variable.cc +++ b/onnxruntime/core/providers/webgpu/shader_variable.cc @@ -308,7 +308,7 @@ std::string ShaderVariableHelper::SetByOffsetImpl(std::string_view offset, std:: ORT_THROW("Invalid type"); break; case onnxruntime::webgpu::ProgramVariableDataType::Int64: - ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, " << value << " < 0));"; + ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), select(0u, 0xFFFFFFFFu, i32(" << value << ") < 0));"; break; case onnxruntime::webgpu::ProgramVariableDataType::Uint64: ss << name_ << "[" << offset << "]=vec2(u32(" << value << "), 0u);"; diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 313a96ba25509..6f6bb47a79292 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -20,7 +20,8 @@ const std::vector& CastOpTypeConstraints() { DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}; + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}; return types; } } // namespace @@ -87,12 +88,17 @@ Status Cast::ComputeInternal(ComputeContext& context) const { if (size == 0) { return Status::OK(); } + bool is_from_int64 = input_tensor->DataType() == DataTypeImpl::GetType(); + const int in_components = is_from_int64 ? 1 : 4; + const int out_components = to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64 ? 1 : 4; uint32_t vec_size = onnxruntime::narrow((size + 3) / 4); + uint32_t in_vec_size = onnxruntime::narrow(in_components == 1 ? size : vec_size); + uint32_t out_vec_size = onnxruntime::narrow(out_components == 1 ? size : vec_size); - CastProgram program{to_}; + CastProgram program{to_, is_from_int64}; program - .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {vec_size}, 4}) - .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {vec_size}, 4}) + .AddInput({input_tensor, ProgramTensorMetadataDependency::Type, {in_vec_size}, in_components}) + .AddOutput({output_tensor, ProgramTensorMetadataDependency::None, {out_vec_size}, out_components}) .SetDispatchGroupSize((vec_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) .AddUniformVariables({ {static_cast(vec_size)}, @@ -121,12 +127,31 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { case ONNX_NAMESPACE::TensorProto_DataType_BOOL: expression = "vec4(a)"; break; + case ONNX_NAMESPACE::TensorProto_DataType_INT64: + expression = "int32(a)"; + break; default: ORT_NOT_IMPLEMENTED("Cast to type ", to_, " is not supported."); } - sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size") - << " let a = " << input.GetByOffset("global_idx") << ";\n " - << output.SetByOffset("global_idx", expression); + + sh.MainFunctionBody() << sh.GuardAgainstOutOfBoundsWorkgroupSizes("uniforms.vec_size"); + if (is_from_int64_) { + sh.MainFunctionBody() << " let a0 = " << input.GetByOffset("global_idx * 4") << ";\n" + << " let a1 = " << input.GetByOffset("global_idx * 4 + 1") << ";\n" + << " let a2 = " << input.GetByOffset("global_idx * 4 + 2") << ";\n" + << " let a3 = " << input.GetByOffset("global_idx * 4 + 3") << ";\n" + << " let a = vec4(a0, a1, a2, a3);\n"; + } else { + sh.MainFunctionBody() << " let a = " << input.GetByOffset("global_idx") << ";\n"; + } + if (to_ == ONNX_NAMESPACE::TensorProto_DataType_INT64) { + sh.MainFunctionBody() << output.SetByOffset("global_idx * 4", "a.x") << "\n" + << output.SetByOffset("global_idx * 4 + 1", "a.y") << "\n" + << output.SetByOffset("global_idx * 4 + 2", "a.z") << "\n" + << output.SetByOffset("global_idx * 4 + 3", "a.w") << "\n"; + } else { + sh.MainFunctionBody() << output.SetByOffset("global_idx", expression); + } return Status::OK(); } diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 925cd200f0aba..11e2a5b32fc7b 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -10,7 +10,7 @@ namespace webgpu { class CastProgram final : public Program { public: - CastProgram(int32_t to) : Program{"Cast"}, to_{to} {} + CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -18,6 +18,7 @@ class CastProgram final : public Program { private: int32_t to_; + bool is_from_int64_; }; class Cast final : public WebGpuKernel { From 61c2541b83fefb9e9c56991b5c13906cd02c8c68 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 29 Sep 2025 14:28:11 +0800 Subject: [PATCH 09/16] Register conditional Cast --- .../core/providers/webgpu/tensor/cast.cc | 68 +++---------------- .../core/providers/webgpu/tensor/cast.h | 46 +++++++++++++ .../webgpu/webgpu_execution_provider.cc | 30 +++++--- 3 files changed, 77 insertions(+), 67 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 6f6bb47a79292..940962d029deb 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -10,11 +10,10 @@ namespace onnxruntime { namespace webgpu { -namespace { +// Type constraint functions const std::vector& CastOpTypeConstraints() { // currently support boolean, integer and float types that explicitly allowed in WGSL: // https://gpuweb.github.io/gpuweb/wgsl/#plain-types-section - // static std::vector types{ DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), @@ -24,62 +23,17 @@ const std::vector& CastOpTypeConstraints() { DataTypeImpl::GetTensorType()}; return types; } -} // namespace -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 6, 8, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 9, 12, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 13, 18, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 19, 20, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_VERSIONED_KERNEL_EX( - Cast, - kOnnxDomain, - 21, 22, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); -ONNX_OPERATOR_KERNEL_EX( - Cast, - kOnnxDomain, - 23, - kWebGpuExecutionProvider, - (*KernelDefBuilder::Create()) - .TypeConstraint("T1", CastOpTypeConstraints()) - .TypeConstraint("T2", CastOpTypeConstraints()), - Cast); +const std::vector& CastOpTypeConstraintsWithoutInt64() { + // Type constraints without int64_t support when graph capture is disabled + static std::vector types{ + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}; + return types; +} Status Cast::ComputeInternal(ComputeContext& context) const { const auto* input_tensor = context.Input(0); diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 11e2a5b32fc7b..fc2c93d952b29 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -3,11 +3,57 @@ #pragma once +#include "core/framework/kernel_registry.h" +#include "core/framework/op_kernel.h" #include "core/providers/webgpu/webgpu_kernel.h" namespace onnxruntime { namespace webgpu { +// Forward declaration +class Cast; + +// Type constraint functions for Cast operator +const std::vector& CastOpTypeConstraints(); +const std::vector& CastOpTypeConstraintsWithoutInt64(); + +// Create Cast kernel info with appropriate type constraints based on graph capture support +template +KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { + const auto& type_constraints = enable_graph_capture ? CastOpTypeConstraints() : CastOpTypeConstraintsWithoutInt64(); + + KernelCreateFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { + out = std::make_unique(info); + return Status::OK(); + }; + + if constexpr (StartVersion == EndVersion) { + // Non-versioned kernel + return { + KernelDefBuilder() + .SetName("Cast") + .SetDomain(kOnnxDomain) + .SinceVersion(StartVersion) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", type_constraints) + .TypeConstraint("T2", type_constraints) + .Build(), + kernel_create_fn}; + } else { + // Versioned kernel + return { + KernelDefBuilder() + .SetName("Cast") + .SetDomain(kOnnxDomain) + .SinceVersion(StartVersion, EndVersion) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", type_constraints) + .TypeConstraint("T2", type_constraints) + .Build(), + kernel_create_fn}; + } +} + class CastProgram final : public Program { public: CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index bbb3fbdd221d3..3a6d1da90179c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -28,6 +28,7 @@ #include "core/providers/webgpu/data_transfer.h" #include "core/providers/webgpu/external_data_loader.h" #include "core/providers/webgpu/webgpu_profiler.h" +#include "core/providers/webgpu/tensor/cast.h" namespace onnxruntime { @@ -417,7 +418,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxD class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 16, 17, ScatterND); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 18, ScatterND); -std::unique_ptr RegisterKernels() { +std::unique_ptr RegisterKernels(bool enable_graph_capture = true) { auto kernel_registry = std::make_unique(); static const BuildKernelCreateInfoFn function_table[] = { @@ -464,13 +465,6 @@ std::unique_ptr RegisterKernels() { KERNEL_CREATE_INFO(13, Tanh), KERNEL_CREATE_INFO(1, Not), - KERNEL_CREATE_INFO_VERSIONED(6, 8, Cast), - KERNEL_CREATE_INFO_VERSIONED(9, 12, Cast), - KERNEL_CREATE_INFO_VERSIONED(13, 18, Cast), - KERNEL_CREATE_INFO_VERSIONED(19, 20, Cast), - KERNEL_CREATE_INFO_VERSIONED(21, 22, Cast), - KERNEL_CREATE_INFO(23, Cast), - // // activations BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -771,6 +765,14 @@ std::unique_ptr RegisterKernels() { } } + // Register Cast kernels with conditional int64 support based on graph capture + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<6, 8>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<9, 12>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<13, 18>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<19, 20>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<21, 22>(enable_graph_capture))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<23>(enable_graph_capture))); + #ifndef DISABLE_CONTRIB_OPS Status status = ::onnxruntime::contrib::webgpu::RegisterWebGpuContribKernels(*kernel_registry); ORT_ENFORCE(status.IsOK(), "Failed to register WebGPU contrib kernels: " + status.ErrorMessage()); @@ -869,9 +871,17 @@ std::vector> WebGpuExecutionProvider::GetCapa } std::shared_ptr WebGpuExecutionProvider::GetKernelRegistry() const { - static std::shared_ptr registry = webgpu::RegisterKernels(); + // Use different registries based on graph capture configuration + static std::shared_ptr registry_with_int64; + static std::shared_ptr registry_without_int64; + static std::once_flag init_flag; + + std::call_once(init_flag, []() { + registry_with_int64 = webgpu::RegisterKernels(true); // with int64 support + registry_without_int64 = webgpu::RegisterKernels(false); // without int64 support + }); - return registry; + return enable_graph_capture_ ? registry_with_int64 : registry_without_int64; } std::unique_ptr WebGpuExecutionProvider::GetDataTransfer() const { From 76b059eba63741c9b41db72cf80cb73e4a9aea1f Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 29 Sep 2025 15:38:30 +0800 Subject: [PATCH 10/16] fix CI errors --- onnxruntime/core/providers/webgpu/tensor/cast.h | 3 ++- onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index fc2c93d952b29..148dd7457466d 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -23,7 +23,8 @@ KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { const auto& type_constraints = enable_graph_capture ? CastOpTypeConstraints() : CastOpTypeConstraintsWithoutInt64(); KernelCreateFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { - out = std::make_unique(info); + auto cast_kernel = std::make_unique(info); + out = std::unique_ptr(std::move(cast_kernel)); return Status::OK(); }; diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 3a6d1da90179c..52e9d887b57ac 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -418,7 +418,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxD class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 16, 17, ScatterND); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kWebGpuExecutionProvider, kOnnxDomain, 18, ScatterND); -std::unique_ptr RegisterKernels(bool enable_graph_capture = true) { +std::unique_ptr RegisterKernels(bool enable_graph_capture = false) { auto kernel_registry = std::make_unique(); static const BuildKernelCreateInfoFn function_table[] = { From 91ea82e120b2582558bfa11e2c71c3a108437b8a Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 29 Sep 2025 17:22:45 +0800 Subject: [PATCH 11/16] fix the CI errors --- .../core/providers/webgpu/tensor/cast.h | 64 +++++++++---------- 1 file changed, 30 insertions(+), 34 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 148dd7457466d..2d4ce528d3f11 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -10,8 +10,35 @@ namespace onnxruntime { namespace webgpu { -// Forward declaration -class Cast; +class CastProgram final : public Program { + public: + CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} + + Status GenerateShaderCode(ShaderHelper& sh) const override; + + WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"vec_size", ProgramUniformVariableDataType::Uint32}); + + private: + int32_t to_; + bool is_from_int64_; +}; + +class Cast final : public WebGpuKernel { + public: + Cast(const OpKernelInfo& info) : WebGpuKernel(info) { + int64_t to; + Status status = info.GetAttr("to", &to); + ORT_ENFORCE(status.IsOK(), "Attribute to is not set."); + to_ = onnxruntime::narrow(to); + + // ignore attribute 'saturate' as float8 is not supported in WebGPU + } + + Status ComputeInternal(ComputeContext& context) const override; + + private: + int32_t to_; +}; // Type constraint functions for Cast operator const std::vector& CastOpTypeConstraints(); @@ -23,8 +50,7 @@ KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { const auto& type_constraints = enable_graph_capture ? CastOpTypeConstraints() : CastOpTypeConstraintsWithoutInt64(); KernelCreateFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { - auto cast_kernel = std::make_unique(info); - out = std::unique_ptr(std::move(cast_kernel)); + out = std::unique_ptr(std::make_unique(info)); return Status::OK(); }; @@ -55,35 +81,5 @@ KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { } } -class CastProgram final : public Program { - public: - CastProgram(int32_t to, bool is_from_int64) : Program{"Cast"}, to_{to}, is_from_int64_{is_from_int64} {} - - Status GenerateShaderCode(ShaderHelper& sh) const override; - - WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES({"vec_size", ProgramUniformVariableDataType::Uint32}); - - private: - int32_t to_; - bool is_from_int64_; -}; - -class Cast final : public WebGpuKernel { - public: - Cast(const OpKernelInfo& info) : WebGpuKernel(info) { - int64_t to; - Status status = info.GetAttr("to", &to); - ORT_ENFORCE(status.IsOK(), "Attribute to is not set."); - to_ = onnxruntime::narrow(to); - - // ignore attribute 'saturate' as float8 is not supported in WebGPU - } - - Status ComputeInternal(ComputeContext& context) const override; - - private: - int32_t to_; -}; - } // namespace webgpu } // namespace onnxruntime From 39bcbe4b4cb473ca2e4b226617c7874c13cb99b2 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 30 Sep 2025 09:38:17 +0800 Subject: [PATCH 12/16] move the implementation to cc --- .../core/providers/webgpu/tensor/cast.cc | 77 +++++++++++++++---- .../core/providers/webgpu/tensor/cast.h | 39 +--------- 2 files changed, 61 insertions(+), 55 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 940962d029deb..daf4aa323c12e 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -10,30 +10,29 @@ namespace onnxruntime { namespace webgpu { -// Type constraint functions -const std::vector& CastOpTypeConstraints() { - // currently support boolean, integer and float types that explicitly allowed in WGSL: +namespace { +const std::vector& CastOpTypeConstraints(bool enable_graph_capture) { + // Base types that are always supported - boolean, integer and float types that explicitly allowed in WGSL: // https://gpuweb.github.io/gpuweb/wgsl/#plain-types-section - static std::vector types{ - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType(), - DataTypeImpl::GetTensorType()}; - return types; -} - -const std::vector& CastOpTypeConstraintsWithoutInt64() { - // Type constraints without int64_t support when graph capture is disabled - static std::vector types{ + static std::vector base_types{ DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType(), DataTypeImpl::GetTensorType()}; - return types; + + if (enable_graph_capture) { + static std::vector types_with_int64 = []() { + auto types = base_types; + types.push_back(DataTypeImpl::GetTensorType()); + return types; + }(); + return types_with_int64; + } else { + return base_types; + } } +} // namespace Status Cast::ComputeInternal(ComputeContext& context) const { const auto* input_tensor = context.Input(0); @@ -110,5 +109,49 @@ Status CastProgram::GenerateShaderCode(ShaderHelper& sh) const { return Status::OK(); } +template +KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { + const auto& type_constraints = CastOpTypeConstraints(enable_graph_capture); + + KernelCreateFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { + out = std::make_unique(info); + return Status::OK(); + }; + + if constexpr (StartVersion == EndVersion) { + // Non-versioned kernel + return { + KernelDefBuilder() + .SetName("Cast") + .SetDomain(kOnnxDomain) + .SinceVersion(StartVersion) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", type_constraints) + .TypeConstraint("T2", type_constraints) + .Build(), + kernel_create_fn}; + } else { + // Versioned kernel + return { + KernelDefBuilder() + .SetName("Cast") + .SetDomain(kOnnxDomain) + .SinceVersion(StartVersion, EndVersion) + .Provider(kWebGpuExecutionProvider) + .TypeConstraint("T1", type_constraints) + .TypeConstraint("T2", type_constraints) + .Build(), + kernel_create_fn}; + } +} + +// Explicit template instantiations +template KernelCreateInfo CreateCastKernelInfo<6, 8>(bool); +template KernelCreateInfo CreateCastKernelInfo<9, 12>(bool); +template KernelCreateInfo CreateCastKernelInfo<13, 18>(bool); +template KernelCreateInfo CreateCastKernelInfo<19, 20>(bool); +template KernelCreateInfo CreateCastKernelInfo<21, 22>(bool); +template KernelCreateInfo CreateCastKernelInfo<23>(bool); + } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.h b/onnxruntime/core/providers/webgpu/tensor/cast.h index 2d4ce528d3f11..7dfb50e3241c8 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.h +++ b/onnxruntime/core/providers/webgpu/tensor/cast.h @@ -40,46 +40,9 @@ class Cast final : public WebGpuKernel { int32_t to_; }; -// Type constraint functions for Cast operator -const std::vector& CastOpTypeConstraints(); -const std::vector& CastOpTypeConstraintsWithoutInt64(); - // Create Cast kernel info with appropriate type constraints based on graph capture support template -KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture) { - const auto& type_constraints = enable_graph_capture ? CastOpTypeConstraints() : CastOpTypeConstraintsWithoutInt64(); - - KernelCreateFn kernel_create_fn = [](FuncManager&, const OpKernelInfo& info, std::unique_ptr& out) -> Status { - out = std::unique_ptr(std::make_unique(info)); - return Status::OK(); - }; - - if constexpr (StartVersion == EndVersion) { - // Non-versioned kernel - return { - KernelDefBuilder() - .SetName("Cast") - .SetDomain(kOnnxDomain) - .SinceVersion(StartVersion) - .Provider(kWebGpuExecutionProvider) - .TypeConstraint("T1", type_constraints) - .TypeConstraint("T2", type_constraints) - .Build(), - kernel_create_fn}; - } else { - // Versioned kernel - return { - KernelDefBuilder() - .SetName("Cast") - .SetDomain(kOnnxDomain) - .SinceVersion(StartVersion, EndVersion) - .Provider(kWebGpuExecutionProvider) - .TypeConstraint("T1", type_constraints) - .TypeConstraint("T2", type_constraints) - .Build(), - kernel_create_fn}; - } -} +KernelCreateInfo CreateCastKernelInfo(bool enable_graph_capture); } // namespace webgpu } // namespace onnxruntime From dfbee482deac85b86ec6e22b80db1e22fe9762cd Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Fri, 10 Oct 2025 13:21:06 +0800 Subject: [PATCH 13/16] [webgpu] Expose webgpu ep's internal dawn test --- .../webgpu/webgpu_provider_factory.h | 50 +++++++++++++++++++ .../core/providers/webgpu/webgpu_context.cc | 49 ++++++++++++++++++ .../core/providers/webgpu/webgpu_context.h | 1 + 3 files changed, 100 insertions(+) diff --git a/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h b/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h index 0b45b847d651f..3cf14bb476c70 100644 --- a/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h +++ b/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h @@ -12,3 +12,53 @@ // The WebGPU EP can be enabled via the generic SessionOptionsAppendExecutionProvider method, so no direct usage of // the provider factory is required. + +#pragma once + +#include +#include + +// Define export macros without including onnxruntime_c_api.h to avoid conflicts +#if defined(_WIN32) + #ifdef ORT_DLL_EXPORT + #define ORT_WEBGPU_EXPORT __declspec(dllexport) + #else + #define ORT_WEBGPU_EXPORT __declspec(dllimport) + #endif + #define ORT_WEBGPU_API_CALL __stdcall +#elif defined(__APPLE__) || defined(__linux__) + #define ORT_WEBGPU_EXPORT __attribute__((visibility("default"))) + #define ORT_WEBGPU_API_CALL +#else + #define ORT_WEBGPU_EXPORT + #define ORT_WEBGPU_API_CALL +#endif + +#ifdef __cplusplus +extern "C" { +#endif + +/** + * \brief Get the Dawn proc table from WebGPU EP context + * \param context_id The WebGPU context ID (0 for default context) + * \return Pointer to the Dawn proc table, or nullptr if not available + */ +ORT_WEBGPU_EXPORT const void* ORT_WEBGPU_API_CALL OrtWebGpuGetDawnProcTable(int context_id); + +/** + * \brief Get the WebGPU instance from WebGPU EP context + * \param context_id The WebGPU context ID (0 for default context) + * \return Pointer to the WebGPU instance (WGPUInstance), or nullptr if not available + */ +ORT_WEBGPU_EXPORT void* ORT_WEBGPU_API_CALL OrtWebGpuGetInstance(int context_id); + +/** + * \brief Get the WebGPU device from WebGPU EP context + * \param context_id The WebGPU context ID (0 for default context) + * \return Pointer to the WebGPU device (WGPUDevice), or nullptr if not available + */ +ORT_WEBGPU_EXPORT void* ORT_WEBGPU_API_CALL OrtWebGpuGetDevice(int context_id); + +#ifdef __cplusplus +} +#endif diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 985fcd03f33ac..45c2c57a803e7 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -28,6 +28,13 @@ #include "core/providers/webgpu/compute_context.h" #include "core/providers/webgpu/webgpu_context.h" #include "core/providers/webgpu/buffer_manager.h" + +// Define ORT_DLL_EXPORT before including webgpu_provider_factory.h to ensure +// ORT_WEBGPU_EXPORT becomes __declspec(dllexport) instead of __declspec(dllimport) +#ifndef ORT_DLL_EXPORT +#define ORT_DLL_EXPORT +#endif +#include "core/providers/webgpu/webgpu_provider_factory.h" // For ORT_WEBGPU_EXPORT macros #include "core/providers/webgpu/webgpu_execution_provider.h" #include "core/providers/webgpu/program.h" #include "core/providers/webgpu/program_cache_key.h" @@ -959,3 +966,45 @@ WGPUDevice GetDevice(int context_id) { } // namespace webgpu } // namespace onnxruntime + +// C API functions for external access +extern "C" { + +ORT_WEBGPU_EXPORT const void* ORT_WEBGPU_API_CALL OrtWebGpuGetDawnProcTable(int context_id) { +#if !defined(__wasm__) && !defined(BUILD_DAWN_SHARED_LIBRARY) && !defined(USE_EXTERNAL_DAWN) + try { + // Ensure context is initialized + auto& context = onnxruntime::webgpu::WebGpuContextFactory::GetContext(context_id); + (void)context; // Suppress unused variable warning + return &dawn::native::GetProcs(); + } catch (...) { + return nullptr; + } +#else + return nullptr; +#endif +} + +ORT_WEBGPU_EXPORT void* ORT_WEBGPU_API_CALL OrtWebGpuGetInstance(int context_id) { +#if !defined(__wasm__) + try { + auto& context = onnxruntime::webgpu::WebGpuContextFactory::GetContext(context_id); + return context.Instance().Get(); + } catch (...) { + return nullptr; + } +#else + return nullptr; +#endif +} + +ORT_WEBGPU_EXPORT void* ORT_WEBGPU_API_CALL OrtWebGpuGetDevice(int context_id) { + try { + auto& context = onnxruntime::webgpu::WebGpuContextFactory::GetContext(context_id); + return context.Device().Get(); + } catch (...) { + return nullptr; + } +} + +} // extern "C" diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.h b/onnxruntime/core/providers/webgpu/webgpu_context.h index 0c0d116cf9394..7fefdefa6ecf0 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.h +++ b/onnxruntime/core/providers/webgpu/webgpu_context.h @@ -83,6 +83,7 @@ class WebGpuContext final { Status Wait(wgpu::Future f); + const wgpu::Instance& Instance() const { return instance_; } const wgpu::Device& Device() const { return device_; } const wgpu::AdapterInfo& AdapterInfo() const { return adapter_info_; } From f15bdbbe11e2b8da2e7d31ef848e0b905a152ef0 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Mon, 13 Oct 2025 08:53:11 +0800 Subject: [PATCH 14/16] cast 64 chagnes --- .../webgpu/webgpu_execution_provider.cc | 24 ++++++++++++------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 52e9d887b57ac..c481b41a6dc24 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -872,16 +872,22 @@ std::vector> WebGpuExecutionProvider::GetCapa std::shared_ptr WebGpuExecutionProvider::GetKernelRegistry() const { // Use different registries based on graph capture configuration - static std::shared_ptr registry_with_int64; - static std::shared_ptr registry_without_int64; - static std::once_flag init_flag; + static std::shared_ptr registry_with_graph_capture; + static std::shared_ptr registry_without_graph_capture; + static std::once_flag init_flag_with_graph_capture; + static std::once_flag init_flag_without_graph_capture; - std::call_once(init_flag, []() { - registry_with_int64 = webgpu::RegisterKernels(true); // with int64 support - registry_without_int64 = webgpu::RegisterKernels(false); // without int64 support - }); - - return enable_graph_capture_ ? registry_with_int64 : registry_without_int64; + if (enable_graph_capture_) { + std::call_once(init_flag_with_graph_capture, []() { + registry_with_graph_capture = webgpu::RegisterKernels(true); + }); + return registry_with_graph_capture; + } else { + std::call_once(init_flag_without_graph_capture, []() { + registry_without_graph_capture = webgpu::RegisterKernels(false); + }); + return registry_without_graph_capture; + } } std::unique_ptr WebGpuExecutionProvider::GetDataTransfer() const { From 35a46cf68192bf8abee98602a45b6aa7c7ece8ae Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 14 Oct 2025 12:31:21 +0800 Subject: [PATCH 15/16] fix the memory leak --- .../webgpu/webgpu_provider_factory.h | 20 +++++++++---------- .../core/providers/webgpu/webgpu_context.cc | 3 +++ 2 files changed, 13 insertions(+), 10 deletions(-) diff --git a/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h b/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h index 3cf14bb476c70..b54c8e17f46a6 100644 --- a/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h +++ b/include/onnxruntime/core/providers/webgpu/webgpu_provider_factory.h @@ -20,18 +20,18 @@ // Define export macros without including onnxruntime_c_api.h to avoid conflicts #if defined(_WIN32) - #ifdef ORT_DLL_EXPORT - #define ORT_WEBGPU_EXPORT __declspec(dllexport) - #else - #define ORT_WEBGPU_EXPORT __declspec(dllimport) - #endif - #define ORT_WEBGPU_API_CALL __stdcall +#ifdef ORT_DLL_EXPORT +#define ORT_WEBGPU_EXPORT __declspec(dllexport) +#else +#define ORT_WEBGPU_EXPORT __declspec(dllimport) +#endif +#define ORT_WEBGPU_API_CALL __stdcall #elif defined(__APPLE__) || defined(__linux__) - #define ORT_WEBGPU_EXPORT __attribute__((visibility("default"))) - #define ORT_WEBGPU_API_CALL +#define ORT_WEBGPU_EXPORT __attribute__((visibility("default"))) +#define ORT_WEBGPU_API_CALL #else - #define ORT_WEBGPU_EXPORT - #define ORT_WEBGPU_API_CALL +#define ORT_WEBGPU_EXPORT +#define ORT_WEBGPU_API_CALL #endif #ifdef __cplusplus diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 45c2c57a803e7..c496ab238402b 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -946,6 +946,9 @@ void WebGpuContextFactory::ReleaseContext(int context_id) { ORT_ENFORCE(it != contexts_.end(), "WebGPU EP context ID ", context_id, " is not found."); if (--it->second.ref_count == 0 && !it->second.context->preserve_device_) { + // TODO: Investigate why memory leak is triggered if we don't explicitly destroy the device. + // It seems that memory leak deletection is triggered before the device is destroyed. + it->second.context->Device().Destroy(); contexts_.erase(it); } } From 0171a27d34368f73504c4dc46fe2b80addea505c Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Thu, 23 Oct 2025 13:57:29 +0800 Subject: [PATCH 16/16] fix the CI errors --- onnxruntime/core/providers/webgpu/webgpu_context.cc | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/onnxruntime/core/providers/webgpu/webgpu_context.cc b/onnxruntime/core/providers/webgpu/webgpu_context.cc index 125e065096f81..de2084b74b95c 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_context.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_context.cc @@ -1016,12 +1016,9 @@ WGPUDevice GetDevice(int context_id) { // C API functions for external access extern "C" { -ORT_WEBGPU_EXPORT const void* ORT_WEBGPU_API_CALL OrtWebGpuGetDawnProcTable(int context_id) { +ORT_WEBGPU_EXPORT const void* ORT_WEBGPU_API_CALL OrtWebGpuGetDawnProcTable(int /* context_id */) { #if !defined(__wasm__) && !defined(BUILD_DAWN_SHARED_LIBRARY) && !defined(USE_EXTERNAL_DAWN) try { - // Ensure context is initialized - auto& context = onnxruntime::webgpu::WebGpuContextFactory::GetContext(context_id); - (void)context; // Suppress unused variable warning return &dawn::native::GetProcs(); } catch (...) { return nullptr;