[WebGPU] Support models with opset 24 ops and KV-shared decoder layers (Gemma 4)

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),

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<uint32_t>(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<Tensor*>(past_key);
+      present_value = const_cast<Tensor*>(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<uint32_t>(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"};

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc

@@ -241,14 +241,12 @@
   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 @@
   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<uint32_t>(parameters.sequence_length_ * parameters.hidden_size_);
+      const std::vector<uint32_t> q_input_output_strides{
+          batch_stride,
+          static_cast<uint32_t>(parameters.hidden_size_),
+          static_cast<uint32_t>(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 @@
                                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);