Address review: fix dummy-K buffer size and add const_cast safety guard

- Size kDummy/kRotary with sequence_length (not 1) to match the kernel's
  iteration domain, preventing OOB writes during prefill (q_seq > 1)
- Add ORT_ENFORCE in flash_attention kv_empty path to guard against
  future refactors that might write through the const_cast'd pointers
- Add new test SharedKV_EmptyKV_WithPast_Rotary_Prompt_WebGPU exercising
  the rotary + kv_empty path with q_seq_len=6

Co-Authored-By: Claude Opus 4 <noreply@anthropic.com>

Author: feich-ms

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

@@ -471,12 +471,14 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
     // Use past_key/past_value directly as the present buffers for attention.
     ORT_ENFORCE(!do_rotary, "Fused SplitPackedQKVWithRotaryEmbeddingAndCopyKV should not be used with kv_sequence_length==0.");
     if (past_key != nullptr && past_value != nullptr) {
-      // Safe: flash attention kernels only read from present_key/present_value.
-      // CopyKVCache is skipped when kv_empty, so no writes through these pointers.
+      // 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);
     }
     // If past is also null, present_key/present_value were already set to internal empty tensors above.
+    ORT_ENFORCE(!parameters.past_present_share_buffer_,
+                "kv_empty path must not use past_present_share_buffer (CopyKVCache is skipped).");
   } 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.");

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc

@@ -292,8 +292,8 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
       // Apply RoPE to Q only. Use the fused kernel with a dummy K to avoid zero-element buffers.
       // K output is discarded since attention uses past_key/past_value.
       qRotary = context.CreateGPUTensor(query->DataType(), query->Shape());
-      kDummy = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, 1, parameters.kv_hidden_size_}));
-      kRotary = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, 1, parameters.kv_hidden_size_}));
+      kDummy = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.kv_hidden_size_}));
+      kRotary = context.CreateGPUTensor(query->DataType(), TensorShape({parameters.batch_size_, parameters.sequence_length_, parameters.kv_hidden_size_}));
       ORT_RETURN_IF_ERROR(RunFusedQKRotaryEmbedding(context, parameters,
                                                     query, &kDummy,
                                                     seqlen_k,

onnxruntime/test/contrib_ops/group_query_attention_op_test.cc

@@ -1884,5 +1884,38 @@ TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_WebGPU) {
   ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_Rotary_WebGPU_vs_CPU");
 }
 
+// WebGPU: shared KV with do_rotary=1 and q_seq_len > 1 (prefill phase).
+// Validates that the dummy-K rotary buffer is correctly sized to sequence_length.
+TEST(GroupQueryAttentionTest, SharedKV_EmptyKV_WithPast_Rotary_Prompt_WebGPU) {
+  auto webgpu_ep = DefaultWebGpuExecutionProvider();
+  if (!webgpu_ep) {
+    GTEST_SKIP() << "WebGPU EP not available";
+  }
+
+  constexpr int batch_size = 1;
+  constexpr int q_seq_len = 6;
+  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<float> query_data(batch_size * q_seq_len * hidden_size);
+  std::vector<float> past_key_data(batch_size * kv_num_heads * past_seq_len * head_size);
+  std::vector<float> 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<float>(i % 11 + 1);
+  for (size_t i = 0; i < past_key_data.size(); i++) past_key_data[i] = 0.2f * static_cast<float>(i % 5 + 1);
+  for (size_t i = 0; i < past_value_data.size(); i++) past_value_data[i] = 0.3f * static_cast<float>(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_Prompt_WebGPU_vs_CPU");
+}
+
 }  // namespace test
 }  // namespace onnxruntime