[webgpu] Simplify the signature of CanApplyFlashAttention (microsoft#26926)

This pull request simplifies the logic for handling present key/value
tensors in the WebGPU Flash Attention implementation. The main change is
that the responsibility for creating internal present key/value tensors
is moved from the caller to the `ApplyFlashAttention` function itself.
This reduces code duplication and makes the API easier to use.
Additionally, the `CanApplyFlashAttention` function is simplified to
remove unnecessary checks for present key/value tensors.

Author: qjia7

onnxruntime/contrib_ops/webgpu/bert/attention.cc

@@ -726,16 +726,9 @@ Status Attention::ComputeInternal(onnxruntime::webgpu::ComputeContext& context)
   parameters.qkv_format_ = Q_K_V_BSNH;
 
   // Check if we can use flash attention
-  // For Attention operator, we need to create present_key and present_value tensors for flash attention
-  // even though they are not exposed as outputs
-  TensorShapeVector present_kv_shape({parameters.batch_size_, parameters.num_heads_,
-                                      parameters.total_sequence_length_, parameters.head_size_});
-  Tensor present_key = context.CreateGPUTensor(input->DataType(), present_kv_shape);
-  Tensor present_value = context.CreateGPUTensor(input->DataType(), present_kv_shape);
-
-  if (CanApplyFlashAttention(nullptr, &present_key, &present_value, parameters, context)) {
+  if (CanApplyFlashAttention(nullptr, parameters, context)) {
     // FlashAttention supports Q_K_V_BSNH format directly
-    return ApplyFlashAttention(&Q_bsd, &K_bsd, &V_bsd, attention_bias, output, nullptr, &present_key, nullptr, &present_value,
+    return ApplyFlashAttention(&Q_bsd, &K_bsd, &V_bsd, attention_bias, output, nullptr, nullptr, nullptr, nullptr,
                                parameters, context, nullptr);
   }
 

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

@@ -400,6 +400,22 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
                            const Tensor* cos_cache, const Tensor* sin_cache) {
   constexpr uint32_t tile_size = 64;
 
+  // Create present_key and present_value tensors if they are nullptr
+  Tensor internal_present_key;
+  Tensor internal_present_value;
+  if (present_key == nullptr) {
+    TensorShapeVector present_kv_shape({parameters.batch_size_, parameters.num_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_,
+                                        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]);
@@ -532,14 +548,12 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
   return Status::OK();
 }
 
-bool CanApplyFlashAttention(const Tensor* bias, const Tensor* present_key, const Tensor* present_value,
+bool CanApplyFlashAttention(const Tensor* bias,
                             const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context) {
   return !parameters.is_packed_qkv_ &&
          parameters.head_size_ == parameters.v_head_size_ &&
          bias == nullptr &&
          context.HasFeature(wgpu::FeatureName::Subgroups) &&
-         present_key != nullptr && present_value != nullptr && present_key->SizeInBytes() > 0 &&
-         present_value->SizeInBytes() > 0 &&
          ((context.AdapterInfo().vendor == std::string_view{"qualcomm"} && parameters.head_size_ % 8 == 0) || parameters.head_size_ % 4 == 0);
 }
 

onnxruntime/contrib_ops/webgpu/bert/flash_attention.h

@@ -184,7 +184,7 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
                            const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context, const Tensor* seqlen_k = nullptr,
                            const Tensor* cos_cache = nullptr, const Tensor* sin_cache = nullptr);
 
-bool CanApplyFlashAttention(const Tensor* bias, const Tensor* present_key, const Tensor* present_value,
+bool CanApplyFlashAttention(const Tensor* bias,
                             const WebgpuAttentionParameters& parameters, onnxruntime::webgpu::ComputeContext& context);
 
 // Split packed QKV with Q/K rotary embedding and copy KV cache fusion

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc

@@ -252,7 +252,7 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
     // 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(nullptr, present_key, present_value, temp_params, context);
+    will_use_flash_attention = CanApplyFlashAttention(nullptr, temp_params, context);
   }
 
   if (parameters.is_packed_qkv_ && do_rotary_) {

onnxruntime/contrib_ops/webgpu/bert/multihead_attention.cc

@@ -104,7 +104,7 @@ Status MultiHeadAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
   Tensor* output_qk = context.Output(3, output_qk_shape);
 
   if (output_qk == nullptr &&  // Flash attention does not output QK scores
-      CanApplyFlashAttention(bias, present_key, present_value, parameters, context)) {
+      CanApplyFlashAttention(bias, parameters, context)) {
     return ApplyFlashAttention(query, key, value, attention_bias, output, past_key, present_key, past_value,
                                present_value, parameters, context);
   }