diff --git a/onnxruntime/contrib_ops/webgpu/bert/attention_common.h b/onnxruntime/contrib_ops/webgpu/bert/attention_common.h index fb237d8cb9e9a..5955b588e119e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/attention_common.h +++ b/onnxruntime/contrib_ops/webgpu/bert/attention_common.h @@ -42,7 +42,7 @@ struct WebgpuAttentionParameters { explicit WebgpuAttentionParameters(onnxruntime::contrib::GroupQueryAttentionParameters parameters) : is_gqa_(true), batch_size_(parameters.batch_size), sequence_length_(parameters.sequence_length), - kv_sequence_length_(parameters.sequence_length), + kv_sequence_length_(parameters.kv_sequence_length), past_sequence_length_(parameters.seqlen_past_kv_cache), total_sequence_length_(parameters.total_sequence_length), hidden_size_(parameters.hidden_size), diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 8217a07448266..684e050f0201a 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -422,26 +422,25 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co const Tensor* cos_cache, const Tensor* sin_cache, const Tensor* head_sink) { constexpr uint32_t tile_size = 64; - // Create present_key and present_value tensors if they are nullptr + // Create present_key and present_value tensors if they are nullptr. + // Skip allocation for kv_empty — present will be aliased to past below. Tensor internal_present_key; Tensor internal_present_value; - if (present_key == nullptr) { - TensorShapeVector present_kv_shape({parameters.batch_size_, parameters.num_heads_, + const int present_kv_heads = parameters.is_gqa_ ? parameters.kv_num_heads_ : parameters.num_heads_; + const bool kv_empty = (parameters.kv_sequence_length_ == 0); + if (!kv_empty && present_key == nullptr) { + TensorShapeVector present_kv_shape({parameters.batch_size_, present_kv_heads, parameters.total_sequence_length_, parameters.head_size_}); internal_present_key = context.CreateGPUTensor(Q->DataType(), TensorShape(present_kv_shape)); present_key = &internal_present_key; } - if (present_value == nullptr) { - TensorShapeVector present_kv_shape({parameters.batch_size_, parameters.num_heads_, + if (!kv_empty && present_value == nullptr) { + TensorShapeVector present_kv_shape({parameters.batch_size_, present_kv_heads, parameters.total_sequence_length_, parameters.head_size_}); internal_present_value = context.CreateGPUTensor(Q->DataType(), TensorShape(present_kv_shape)); present_value = &internal_present_value; } - // Extract present_sequence_length directly from present_key tensor shape: - // (batch_size, num_heads, total_sequence_length/max_sequence_length, head_size) - const uint32_t present_sequence_length = static_cast(present_key->Shape()[2]); - const bool use_seqlen_k = seqlen_k != nullptr && context.IsGraphCaptureEnabled(); // Declare query_output at function scope to ensure it persists throughout the function @@ -452,7 +451,8 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co Tensor indirect_buffer; // Prepare indirect dispatch buffer for decode path with static KV cache - const bool use_indirect_dispatch = parameters.sequence_length_ == 1 && + const bool use_indirect_dispatch = !kv_empty && + parameters.sequence_length_ == 1 && parameters.past_present_share_buffer_ && seqlen_k != nullptr && context.IsGraphCaptureEnabled(); @@ -464,7 +464,24 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co const bool do_rotary = (cos_cache != nullptr && sin_cache != nullptr); - if (do_rotary) { + if (kv_empty) { + // kv_sequence_length==0: K/V inputs are empty (shared KV layer). + // Skip CopyKVCache and fused split+rotary+copyKV. + // Use past_key/past_value directly as the present buffers for attention. + // Note: do_rotary is always false here because GQA passes cos_cache=nullptr, sin_cache=nullptr + // for kv_empty layers (rotary is applied to Q separately in GQA before calling ApplyFlashAttention). + ORT_ENFORCE(!do_rotary, "Fused SplitPackedQKVWithRotaryEmbeddingAndCopyKV should not be used with kv_sequence_length==0."); + ORT_ENFORCE(past_key != nullptr && past_value != nullptr, + "kv_empty path requires past KV context (KV-shared layers reuse another layer's cache)."); + // When past_present_share_buffer_ is true (MayInplace optimization), present already + // shares the past buffer. No aliasing needed — the data is already in place. + if (!parameters.past_present_share_buffer_) { + // Alias past as present — flash attention only reads present_key/present_value, + // and CopyKVCache is skipped when kv_empty, so no writes occur through these pointers. + present_key = const_cast(past_key); + present_value = const_cast(past_value); + } + } else if (do_rotary) { ORT_ENFORCE(parameters.is_packed_qkv_, "Fused SplitPackedQKVWithRotaryEmbeddingAndCopyKV requires packed QKV input."); ORT_ENFORCE(parameters.past_present_share_buffer_, "Fused SplitPackedQKVWithRotaryEmbeddingAndCopyKV requires static KV cache."); @@ -481,6 +498,11 @@ 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, tile_size, use_seqlen_k ? seqlen_k : nullptr, indirect_buffer_ptr)); } + // Extract present_sequence_length directly from present_key tensor shape + // after kv_empty aliasing ensures present_key is valid: + // (batch_size, num_heads, total_sequence_length/max_sequence_length, head_size) + const uint32_t present_sequence_length = static_cast(present_key->Shape()[2]); + if (parameters.sequence_length_ > 1) { bool has_attention_bias = attention_bias != nullptr; bool is_qualcomm = context.AdapterInfo().vendor == std::string_view{"qualcomm"}; diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index e3b91bdbb82f4..6b39733b212eb 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc @@ -241,14 +241,12 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& Tensor* present_key = context.Output(1, present_kv_shape); Tensor* present_value = context.Output(2, present_kv_shape); - // WebGPU flash attention requires present_key/present_value as working KV buffers. - if (present_key == nullptr || present_value == nullptr) { - return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, - "WebGPU GroupQueryAttention requires present_key and present_value outputs. " - "Optional present outputs are supported on CPU and CUDA EPs only."); - } - - parameters.past_present_share_buffer_ = past_key != nullptr && past_value != nullptr && past_key->DataRaw() == present_key->DataRaw() && past_value->DataRaw() == present_value->DataRaw(); + // When present_key/present_value outputs are not requested (nullptr), this is a + // KV-shared layer. Flash attention will create internal GPU buffers as needed. + parameters.past_present_share_buffer_ = present_key != nullptr && present_value != nullptr && + past_key != nullptr && past_value != nullptr && + past_key->DataRaw() == present_key->DataRaw() && + past_value->DataRaw() == present_value->DataRaw(); ORT_ENFORCE(parameters.total_sequence_length_ <= parameters.seqlen_present_kv_cache_, "Total sequence length cannot be greater than the existing KV cache length."); @@ -259,17 +257,47 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& Tensor qRotary; Tensor kRotary; + // kv_sequence_length==0 fast path: K/V inputs are empty (shared KV layer). + // Skip all K/V processing; only apply RoPE to Q if needed. + // Use past_key/past_value directly as the KV context. + const bool kv_empty = (parameters.kv_sequence_length_ == 0); + // Use a sliding window if the total sequence exceeds the window's length. bool use_sliding_window = (local_window_size_ != -1 && local_window_size_ < parameters.total_sequence_length_); bool will_use_flash_attention = false; - if (!use_smooth_softmax_ && !use_sliding_window) { + // For kv_empty layers (shared KV), sliding window is irrelevant — there's no new KV to window + // over, the layer reuses another layer's already-computed KV cache. Flash attention is required + // for these layers, so we bypass the sliding window check to allow it. + if (!use_smooth_softmax_ && (!use_sliding_window || kv_empty)) { // Create a temporary parameters copy with is_packed_qkv_ set to false to check if flash attention can be applied after unpacking WebgpuAttentionParameters temp_params = parameters; temp_params.is_packed_qkv_ = false; will_use_flash_attention = CanApplyFlashAttention(temp_params, context); } - if (parameters.is_packed_qkv_ && do_rotary_) { + if (kv_empty) { + // KV inputs are empty - shared KV layer. Only need to optionally apply RoPE to Q. + ORT_ENFORCE(!parameters.is_packed_qkv_, "Packed QKV is not supported with kv_sequence_length==0 (shared KV layers)."); + if (do_rotary_) { + // Apply RoPE to Q only — K doesn't need rotation since we reuse another layer's already-rotated KV cache. + qRotary = context.CreateGPUTensor(query->DataType(), query->Shape()); + // Query is BSD (3 dims): [batch, sequence, hidden]. Strides for bsnh layout: + // {batch_stride, hidden_size, head_size, 1}. + const auto batch_stride = static_cast(parameters.sequence_length_ * parameters.hidden_size_); + const std::vector q_input_output_strides{ + batch_stride, + static_cast(parameters.hidden_size_), + static_cast(parameters.head_size_), + 1u}; + ORT_RETURN_IF_ERROR(RunRotaryEmbedding(context, + query, seqlen_k, cos_cache, sin_cache, &qRotary, + parameters.batch_size_, parameters.sequence_length_, + parameters.hidden_size_, parameters.head_size_, + parameters.scale_, parameters.rotary_interleaved_, + /*use_seqlens_for_position=*/true, q_input_output_strides)); + query = &qRotary; + } + } else if (parameters.is_packed_qkv_ && do_rotary_) { // Use the ultimate fused operation when FlashAttention and static KV cache is enabled. if (will_use_flash_attention && parameters.past_present_share_buffer_) { // Directly call ApplyFlashAttention with fused split/rotary/copyKV enabled @@ -322,6 +350,13 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext& present_value, parameters, context, seqlen_k, nullptr, nullptr, head_sink); } + // Non-flash attention path does not support kv_sequence_length==0 (shared KV layers). + if (kv_empty) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "WebGPU non-flash attention path does not support kv_sequence_length==0 (shared KV layers). " + "Flash attention is required for KV-shared decoder layers."); + } + TensorShapeVector q_new_dims({parameters.batch_size_, parameters.num_heads_, parameters.sequence_length_, parameters.head_size_}); TensorShape q_new_shape(q_new_dims); diff --git a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc index 69d2db391ce3c..2963cefc05b1a 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.cc @@ -22,45 +22,84 @@ ONNX_OPERATOR_KERNEL_EX( Status RotaryEmbeddingProgram::GenerateShaderCode(ShaderHelper& shader) const { const auto& input = shader.AddInput("input", ShaderUsage::UseUniform); - const auto& position_ids = shader.AddInput("position_ids", ShaderUsage::UseUniform); + // The second input is either seqlens (use_seqlens_for_position_) or position_ids (legacy path). + // Declared here so the input order matches the caller's AddInputs order: + // [input, seqlens|position_ids, cos_cache, sin_cache]. + const auto& position_ids_or_seqlens = use_seqlens_for_position_ + ? shader.AddInput("seqlens", ShaderUsage::UseUniform) + : shader.AddInput("position_ids", ShaderUsage::UseUniform); const auto& cos_cache = shader.AddInput("cos_cache", ShaderUsage::UseUniform); const auto& sin_cache = shader.AddInput("sin_cache", ShaderUsage::UseUniform); const auto& output = shader.AddOutput("output", ShaderUsage::UseUniform); - // TODO: remove output_indices. - const auto& output_indices = shader.AddIndices("output_indices", ShaderUsage::None); const auto interleaved_str = interleaved_ ? "true" : "false"; - shader.MainFunctionBody() << " let half_rotary_emb_dim = uniforms.cos_cache_shape[1];\n" - " let bsnh = global_idx / uniforms.global_stride % uniforms.global_shape;\n" - " let size = uniforms.global_shape[0] * uniforms.global_stride[0];\n" - " if (global_idx >= size) { return; }\n" - " if (bsnh[3] < half_rotary_emb_dim) {\n" - << " let position_ids_idx = " << position_ids.BroadcastedIndicesToOffset("bsnh.xy", output_indices) << ";\n" - << " let raw_pos = " << position_ids.GetByOffset("position_ids_idx") << ";\n" - << " let i = dot(bsnh, uniforms.input_output_stride) + select(0, bsnh[3], " << interleaved_str << ");\n" - << " let j = i + select(half_rotary_emb_dim, 1, " << interleaved_str << ");\n" - " let max_position = uniforms.cos_cache_shape[0];\n" - // Bounds check: raw_pos < 0 catches negative position_ids (i32 from truncated int64). - // After u32 conversion + offset, check >= max_position catches too-large values. - // On OOB, pass through input unchanged (same as CUDA kernel behavior). - " if (raw_pos < 0) {\n" - << " " << output.SetByOffset("i", input.GetByOffset("i")) << "\n" - << " " << output.SetByOffset("j", input.GetByOffset("j")) << "\n" - " } else {\n" - " let position_id = u32(raw_pos) + select(0, bsnh[1], position_ids_idx == 0);\n" - " if (position_id >= max_position) {\n" - << " " << output.SetByOffset("i", input.GetByOffset("i")) << "\n" - << " " << output.SetByOffset("j", input.GetByOffset("j")) << "\n" - " } else {\n" - << " let re = " << input.GetByOffset("i") << " * " << cos_cache.GetByIndices("vec2(position_id, bsnh[3])") << " - " << input.GetByOffset("j") << " * " << sin_cache.GetByIndices("vec2(position_id, bsnh[3])") << ";\n" - << " " << output.SetByOffset("i", "re") << "\n" - << " let im = " << input.GetByOffset("i") << " * " << sin_cache.GetByIndices("vec2(position_id, bsnh[3])") << " + " << input.GetByOffset("j") << " * " << cos_cache.GetByIndices("vec2(position_id, bsnh[3])") << ";\n" - << " " << output.SetByOffset("j", "im") << "\n" - " }\n" + if (use_seqlens_for_position_) { + // Seqlens path (GQA): inputs are [input, seqlens, cos_cache, sin_cache]. + // Compute per-batch past_seqlen from seqlens[batch_idx] = total_seqlen - 1. + shader.MainFunctionBody() << " let half_rotary_emb_dim = uniforms.cos_cache_shape[1];\n" + " let bsnh = global_idx / uniforms.global_stride % uniforms.global_shape;\n" + " let size = uniforms.global_shape[0] * uniforms.global_stride[0];\n" + " if (global_idx >= size) { return; }\n" + " if (bsnh[3] < half_rotary_emb_dim) {\n" + " let batch_idx = bsnh[0];\n" + << " let seqlen_i = " << position_ids_or_seqlens.GetByOffset("batch_idx") << ";\n" + << " let seqlen = u32(seqlen_i);\n" + " let total_seqlen = seqlen + 1u;\n" + " let past_seqlen = total_seqlen - uniforms.global_shape[1];\n" + " let position_id = past_seqlen + bsnh[1];\n" + << " let i = dot(bsnh, uniforms.input_output_stride) + select(0u, bsnh[3], " << interleaved_str << ");\n" + << " let j = i + select(half_rotary_emb_dim, 1u, " << interleaved_str << ");\n" + " let max_position = uniforms.cos_cache_shape[0];\n" + " if (position_id >= max_position) {\n" + << " " << output.SetByOffset("i", input.GetByOffset("i")) << "\n" + << " " << output.SetByOffset("j", input.GetByOffset("j")) << "\n" + " } else {\n" + << " let re = " << input.GetByOffset("i") << " * " << cos_cache.GetByIndices("vec2(position_id, bsnh[3])") << " - " << input.GetByOffset("j") << " * " << sin_cache.GetByIndices("vec2(position_id, bsnh[3])") << ";\n" + << " " << output.SetByOffset("i", "re") << "\n" + << " let im = " << input.GetByOffset("i") << " * " << sin_cache.GetByIndices("vec2(position_id, bsnh[3])") << " + " << input.GetByOffset("j") << " * " << cos_cache.GetByIndices("vec2(position_id, bsnh[3])") << ";\n" + << " " << output.SetByOffset("j", "im") << "\n" " }\n" - << " } else { \n" - " let k = dot(bsnh, uniforms.input_output_stride) + half_rotary_emb_dim;\n" - << " " << output.SetByOffset("k", input.GetByOffset("k")) << "\n" - << " }"; + << " } else {\n" + " let k = dot(bsnh, uniforms.input_output_stride) + half_rotary_emb_dim;\n" + << " " << output.SetByOffset("k", input.GetByOffset("k")) << "\n" + << " }"; + } else { + // Original path: inputs are [input, position_ids, cos_cache, sin_cache]. + const auto& position_ids = position_ids_or_seqlens; + // TODO: remove output_indices. + const auto& output_indices = shader.AddIndices("output_indices", ShaderUsage::None); + shader.MainFunctionBody() << " let half_rotary_emb_dim = uniforms.cos_cache_shape[1];\n" + " let bsnh = global_idx / uniforms.global_stride % uniforms.global_shape;\n" + " let size = uniforms.global_shape[0] * uniforms.global_stride[0];\n" + " if (global_idx >= size) { return; }\n" + " if (bsnh[3] < half_rotary_emb_dim) {\n" + << " let position_ids_idx = " << position_ids.BroadcastedIndicesToOffset("bsnh.xy", output_indices) << ";\n" + << " let raw_pos = " << position_ids.GetByOffset("position_ids_idx") << ";\n" + << " let i = dot(bsnh, uniforms.input_output_stride) + select(0, bsnh[3], " << interleaved_str << ");\n" + << " let j = i + select(half_rotary_emb_dim, 1, " << interleaved_str << ");\n" + " let max_position = uniforms.cos_cache_shape[0];\n" + // Bounds check: raw_pos < 0 catches negative position_ids (i32 from truncated int64). + // After u32 conversion + offset, check >= max_position catches too-large values. + // On OOB, pass through input unchanged (same as CUDA kernel behavior). + " if (raw_pos < 0) {\n" + << " " << output.SetByOffset("i", input.GetByOffset("i")) << "\n" + << " " << output.SetByOffset("j", input.GetByOffset("j")) << "\n" + " } else {\n" + " let position_id = u32(raw_pos) + select(0, bsnh[1], position_ids_idx == 0);\n" + " if (position_id >= max_position) {\n" + << " " << output.SetByOffset("i", input.GetByOffset("i")) << "\n" + << " " << output.SetByOffset("j", input.GetByOffset("j")) << "\n" + " } else {\n" + << " let re = " << input.GetByOffset("i") << " * " << cos_cache.GetByIndices("vec2(position_id, bsnh[3])") << " - " << input.GetByOffset("j") << " * " << sin_cache.GetByIndices("vec2(position_id, bsnh[3])") << ";\n" + << " " << output.SetByOffset("i", "re") << "\n" + << " let im = " << input.GetByOffset("i") << " * " << sin_cache.GetByIndices("vec2(position_id, bsnh[3])") << " + " << input.GetByOffset("j") << " * " << cos_cache.GetByIndices("vec2(position_id, bsnh[3])") << ";\n" + << " " << output.SetByOffset("j", "im") << "\n" + " }\n" + " }\n" + << " } else { \n" + " let k = dot(bsnh, uniforms.input_output_stride) + half_rotary_emb_dim;\n" + << " " << output.SetByOffset("k", input.GetByOffset("k")) << "\n" + << " }"; + } return Status::OK(); } @@ -142,6 +181,57 @@ RotaryEmbedding::RotaryEmbedding(const OpKernelInfo& info) : WebGpuKernel(info) is_packed_batching_ = (info.GetAttrOrDefault("is_packed_batching", 0) == 1); } +Status RunRotaryEmbedding(onnxruntime::webgpu::ComputeContext& context, + const Tensor* input, + const Tensor* position_ids_or_seqlens, + const Tensor* cos_cache, + const Tensor* sin_cache, + Tensor* output, + int batch_size, + int sequence_length, + int hidden_size, + int head_size, + float scale, + bool rotary_interleaved, + bool use_seqlens_for_position, + const std::vector& input_output_strides) { + const auto half_rotary_embedding_dim = onnxruntime::narrow(cos_cache->Shape()[1]); + const auto num_heads = hidden_size / head_size; + + // Rotary embeddings are calculated in a pair-wise fashion. Use the shape + // [batch, sequence, heads, half_rotary_dim_complement] to unfold the global index in shader. + const TensorShape global_shape({static_cast(batch_size), + static_cast(sequence_length), + static_cast(num_heads), + static_cast(head_size - half_rotary_embedding_dim)}); + const auto rank = global_shape.NumDimensions(); + std::vector global_dims(rank); + std::vector global_strides(rank); + for (size_t j = 0; j < rank; ++j) { + global_dims[j] = onnxruntime::narrow(global_shape[j]); + global_strides[j] = onnxruntime::narrow(global_shape.SizeFromDimension(j + 1)); + } + const auto output_size = onnxruntime::narrow(global_shape.Size()); + + RotaryEmbeddingProgram program(rotary_interleaved, use_seqlens_for_position); + program + .CacheHint(rotary_interleaved, use_seqlens_for_position) + .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank}, + {position_ids_or_seqlens, ProgramTensorMetadataDependency::TypeAndRank}, + {cos_cache, ProgramTensorMetadataDependency::Rank}, + {sin_cache, ProgramTensorMetadataDependency::Rank}}) + .AddOutput({output, ProgramTensorMetadataDependency::None}) + .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) + .AddUniformVariables({{scale}, + {gsl::make_span(global_dims)}, + {gsl::make_span(global_strides)}, + {gsl::make_span(input_output_strides)}}); + if (!use_seqlens_for_position) { + program.AddIndices(TensorShape{1, 1}); + } + return context.RunProgram(program); +} + Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& context) const { const auto* input = context.Input(0); const auto input_shape = input->Shape(); @@ -162,24 +252,6 @@ Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& con // because WebGPU program inputs must be GPU buffers (InputMemoryType(OrtMemTypeCPUInput) is // incompatible with AddInputs). - // Rotary embeddings will be calculated in a pair-wise fashion. In accordance, use the shape - // [batch size, sequence length, num of heads, num of pairs to rotate + num of dims to copy] - // to unfold the global index in shader. - const TensorShape global_shape({batch_size, - sequence_length, - hidden_size / head_size, - head_size - half_rotary_embedding_dim}); - - const auto rank = global_shape.NumDimensions(); - std::vector global_dims(rank); - std::vector global_strides(rank); - for (size_t j = 0; j < rank; ++j) { - global_dims[j] = onnxruntime::narrow(global_shape[j]); - global_strides[j] = onnxruntime::narrow(global_shape.SizeFromDimension(j + 1)); - } - - const auto output_size = onnxruntime::narrow(global_shape.Size()); - RotaryEmbeddingProgram program{interleaved_}; const auto input_output_strides = input_shape.NumDimensions() == 3 ? std::vector({batch_stride, hidden_size, head_size, 1}) @@ -187,20 +259,10 @@ Status RotaryEmbedding::ComputeInternal(onnxruntime::webgpu::ComputeContext& con ? std::vector({batch_stride, head_size, sequence_length * head_size, 1}) : std::vector({})); - program - .CacheHint(interleaved_) - .AddInputs({{input, ProgramTensorMetadataDependency::TypeAndRank}, - {position_ids, ProgramTensorMetadataDependency::Rank}, - {cos_cache, ProgramTensorMetadataDependency::Rank}, - {sin_cache, ProgramTensorMetadataDependency::Rank}}) - .AddOutput({output, ProgramTensorMetadataDependency::None}) - .SetDispatchGroupSize((output_size + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE) - .AddUniformVariables({{scale_}, - {gsl::make_span(global_dims)}, - {gsl::make_span(global_strides)}, - {gsl::make_span(input_output_strides)}}) - .AddIndices(TensorShape{1, 1}); - return context.RunProgram(program); + return RunRotaryEmbedding(context, input, position_ids, cos_cache, sin_cache, output, + static_cast(batch_size), static_cast(sequence_length), + static_cast(hidden_size), static_cast(head_size), + scale_, interleaved_, /*use_seqlens_for_position=*/false, input_output_strides); } } // namespace webgpu diff --git a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.h b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.h index e3dc4468cb3ed..a1482376179ae 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.h +++ b/onnxruntime/contrib_ops/webgpu/bert/rotary_embedding.h @@ -15,8 +15,8 @@ using onnxruntime::webgpu::ComputeContext; class RotaryEmbeddingProgram final : public Program { public: - RotaryEmbeddingProgram(bool interleaved) : Program{"RotaryEmbedding"}, interleaved_{interleaved} { - } + RotaryEmbeddingProgram(bool interleaved, bool use_seqlens_for_position = false) + : Program{"RotaryEmbedding"}, interleaved_{interleaved}, use_seqlens_for_position_{use_seqlens_for_position} {} Status GenerateShaderCode(ShaderHelper& sh) const override; @@ -27,6 +27,7 @@ class RotaryEmbeddingProgram final : public Program { private: const bool interleaved_; + const bool use_seqlens_for_position_; }; class FusedQKRotaryEmbeddingProgram final : public Program { @@ -63,6 +64,32 @@ class RotaryEmbedding final : public WebGpuKernel { bool is_packed_batching_; }; +// Apply rotary embedding to a single tensor using RotaryEmbeddingProgram. +// +// If use_seqlens_for_position is true, `position_ids_or_seqlens` must be the seqlens tensor (shape +// [batch_size], containing per-batch seqlen_k values where +// seqlen_k = past_sequence_length + kv_sequence_length - 1). The shader derives position_id +// per batch as: past_seqlen + sequence_index, where +// past_seqlen = (seqlens[batch] + 1) - global_shape[1]. +// +// If use_seqlens_for_position is false, `position_ids_or_seqlens` must be the position_ids tensor +// (shape [batch, seq] or [1, 1] for broadcast). The shader reads position from this tensor +// directly. +Status RunRotaryEmbedding(ComputeContext& context, + const Tensor* input, + const Tensor* position_ids_or_seqlens, + const Tensor* cos_cache, + const Tensor* sin_cache, + Tensor* output, + int batch_size, + int sequence_length, + int hidden_size, + int head_size, + float scale, + bool rotary_interleaved, + bool use_seqlens_for_position, + const std::vector& input_output_strides); + } // namespace webgpu } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/tensor/cast.cc b/onnxruntime/core/providers/webgpu/tensor/cast.cc index 2695c2800d37a..a7c1791a86f31 100644 --- a/onnxruntime/core/providers/webgpu/tensor/cast.cc +++ b/onnxruntime/core/providers/webgpu/tensor/cast.cc @@ -128,7 +128,8 @@ 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); +template KernelCreateInfo CreateCastKernelInfo<23, 23>(bool); +template KernelCreateInfo CreateCastKernelInfo<24>(bool); } // namespace webgpu } // namespace onnxruntime diff --git a/onnxruntime/core/providers/webgpu/tensor/shape_op.cc b/onnxruntime/core/providers/webgpu/tensor/shape_op.cc index 09194aa9f4dbb..de3547b8840f0 100644 --- a/onnxruntime/core/providers/webgpu/tensor/shape_op.cc +++ b/onnxruntime/core/providers/webgpu/tensor/shape_op.cc @@ -125,10 +125,21 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .TypeConstraint("T1", DataTypeImpl::GetTensorType()), Shape); +ONNX_OPERATOR_VERSIONED_KERNEL_EX( + Shape, + kOnnxDomain, + 23, 23, + kWebGpuExecutionProvider, + (*KernelDefBuilder::Create()) + .OutputMemoryType(OrtMemTypeCPU, 0) + .TypeConstraint("T", WebGpuSupportedNumberTypes()) + .TypeConstraint("T1", DataTypeImpl::GetTensorType()), + Shape); + ONNX_OPERATOR_KERNEL_EX( Shape, kOnnxDomain, - 23, + 24, kWebGpuExecutionProvider, (*KernelDefBuilder::Create()) .OutputMemoryType(OrtMemTypeCPU, 0) diff --git a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc index 7e11ddf6b13a0..a8eb772f71c46 100644 --- a/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc +++ b/onnxruntime/core/providers/webgpu/webgpu_execution_provider.cc @@ -192,7 +192,8 @@ static const BuildKernelCreateInfoFn build_kernel_create_info_function_table[] = BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -472,7 +473,8 @@ std::unique_ptr RegisterKernels(bool enable_graph_capture, bool ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<13, 18>(enable_int64))); ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<19, 20>(enable_int64))); ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<21, 22>(enable_int64))); - ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<23>(enable_int64))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<23, 23>(enable_int64))); + ORT_THROW_IF_ERROR(kernel_registry->Register(CreateCastKernelInfo<24>(enable_int64))); // Register Range kernels with conditional int64 support RegisterRangeKernels(*kernel_registry, enable_int64); diff --git a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc index 112d6f1eecc72..90fc623f77985 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -369,7 +369,8 @@ static std::vector RunGQASharedKV( int num_heads, int kv_num_heads, int head_size, - bool use_cuda = false) { + bool use_cuda = false, + bool use_webgpu = false) { const int hidden_size = num_heads * head_size; const int total_seq_len = past_seq_len; // all KV data is in past @@ -414,6 +415,8 @@ static std::vector RunGQASharedKV( std::vector> execution_providers; if (use_cuda) { execution_providers.push_back(DefaultCudaExecutionProvider()); + } else if (use_webgpu) { + execution_providers.push_back(DefaultWebGpuExecutionProvider()); } else { execution_providers.push_back(DefaultCpuExecutionProvider()); } @@ -828,7 +831,8 @@ static std::vector RunGQASharedKVWithRotary( int num_heads, int kv_num_heads, int head_size, - bool use_cuda = false) { + bool use_cuda = false, + bool use_webgpu = false) { const int hidden_size = num_heads * head_size; const int total_seq_len = past_seq_len; const int rotary_dim = head_size; // full rotary @@ -896,6 +900,8 @@ static std::vector RunGQASharedKVWithRotary( std::vector> execution_providers; if (use_cuda) { execution_providers.push_back(DefaultCudaExecutionProvider()); + } else if (use_webgpu) { + execution_providers.push_back(DefaultWebGpuExecutionProvider()); } else { execution_providers.push_back(DefaultCpuExecutionProvider()); } @@ -1747,5 +1753,235 @@ TEST(GroupQueryAttentionTest, SeqlensKExceedsCosCache_MultiBatch) { {}, nullptr, &execution_providers); } +// --------------------------------------------------------------------------- +// WebGPU: shared KV tests (Gemma4 kv_sequence_length=0 pattern) +// Each test cross-checks WebGPU against CPU for correctness. +// --------------------------------------------------------------------------- + +// WebGPU: kv_sequence_length=0 with past, decode (q_seq=1). +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_Decode) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto webgpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/true); + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_WebGPU_vs_CPU"); +} + +// WebGPU: kv_sequence_length=0 with past, prompt phase (q_seq_len > 1). +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_Prefill) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 8; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto webgpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/true); + auto cpu_output = RunGQASharedKV( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_Prompt_WebGPU_vs_CPU"); +} + +// WebGPU: kv_sequence_length=0 with past and do_rotary=1 (Q-only RoPE path). +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_Rotary) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto webgpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/true); + auto cpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_Rotary_WebGPU_vs_CPU"); +} + +// WebGPU: kv_sequence_length=0 with do_rotary=1 and q_seq_len > 1 (prefill). +// Validates position_offset + bsnh[1] arithmetic for multiple sequence positions. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_Rotary_Prefill) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 4; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto webgpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/true); + auto cpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_Rotary_Prefill_WebGPU_vs_CPU"); +} + +// WebGPU: kv_sequence_length=0 with do_rotary=1 and batch_size > 1. +// Validates batch stride calculations in the rotary embedding path. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_Rotary_MultiBatch) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 2; + constexpr int q_seq_len = 1; + constexpr int past_seq_len = 8; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 16; + constexpr int hidden_size = num_heads * head_size; + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + auto webgpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/true); + auto cpu_output = RunGQASharedKVWithRotary( + batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + num_heads, kv_num_heads, head_size, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_Rotary_MultiBatch_WebGPU_vs_CPU"); +} + +// WebGPU: kv_sequence_length=0 with sliding window active (total_seq > local_window_size). +// Regression test: sliding window must not block flash attention for kv_empty layers. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_SlidingWindow) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + + constexpr int batch_size = 1; + constexpr int q_seq_len = 4; + constexpr int past_seq_len = 32; + constexpr int num_heads = 2; + constexpr int kv_num_heads = 1; + constexpr int head_size = 8; + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int local_window_size = 16; // < past_seq_len to trigger sliding window + constexpr int total_seq_len = past_seq_len; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("local_window_size", static_cast(local_window_size)); + + std::vector query_data(batch_size * q_seq_len * hidden_size); + std::vector past_key_data(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data.size(); i++) query_data[i] = 0.1f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast(i % 5 + 1); + for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast(i % 3 + 1); + + tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); + tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); + + std::vector seqlens_k_data(batch_size, static_cast(total_seq_len - 1)); + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); + tester.AddInput("total_sequence_length", {1}, {static_cast(total_seq_len)}); + + tester.AddOptionalInputEdge(); // cos_cache + tester.AddOptionalInputEdge(); // sin_cache + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * q_seq_len * hidden_size; + tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, + std::vector(output_size, 0.0f)); + const int present_size = batch_size * kv_num_heads * past_seq_len * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, + std::vector(present_size, 0.0f)); + + tester.SetOutputTolerance(1e6f); + + std::vector> execution_providers; + execution_providers.push_back(DefaultWebGpuExecutionProvider()); + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); +} + } // namespace test } // namespace onnxruntime