WebGPU: Add indirect dispatch for flash attention graph capture to support Gemma4

onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc

@@ -52,6 +52,16 @@ fn populate_indirect_dispatch_buffer(x: u32, y: u32, z: u32) {
 }
 )";
 
+// Emits kPopulateIndirectDispatchBufferFn into AdditionalImplementation and appends the two
+// WGSL lines that compute num_total_seq_length_tile and call populate_indirect_dispatch_buffer.
+// `total_seq_len_expr` is the WGSL expression that evaluates to the total sequence length.
+static void AppendPopulateIndirectDispatchShader(ShaderHelper& shader, std::string_view total_seq_len_expr) {
+  shader.AdditionalImplementation() << kPopulateIndirectDispatchBufferFn;
+  shader.MainFunctionBody()
+      << "  let num_total_seq_length_tile = (" << total_seq_len_expr << " + uniforms.tile_size - 1u) / uniforms.tile_size;\n"
+      << "  populate_indirect_dispatch_buffer(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size);\n";
+}
+
 Status SplitPackedQKVWithRotaryEmbeddingAndCopyKVProgram::GenerateShaderCode(ShaderHelper& sh) const {
   const auto& packed_qkv = sh.AddInput("packed_qkv", ShaderUsage::UseUniform);
   const auto& seqlens = sh.AddInput("seqlens", ShaderUsage::UseUniform);
@@ -159,6 +169,14 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const {
   return Status::OK();
 }
 
+Status PrepareIndirectDispatchProgram::GenerateShaderCode(ShaderHelper& shader) const {
+  shader.AddInput("seqlen_k", ShaderUsage::None);
+  shader.AddOutput("indirect_buffer", ShaderUsage::None);
+  shader.MainFunctionBody() << "  let total_seq_length = u32(seqlen_k[0u]) + 1u;\n";
+  AppendPopulateIndirectDispatchShader(shader, "total_seq_length");
+  return Status::OK();
+}
+
 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,
@@ -511,11 +529,14 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
   Tensor* indirect_buffer_ptr = nullptr;
   Tensor indirect_buffer;
 
-  // Prepare indirect dispatch buffer for split-reduce path with static KV cache.
-  // When graph capture is enabled, total_sequence_length_ may be 0 (GPU-based
-  // seqlen_k), so the indirect buffer computes dispatch sizes on GPU.
-  // Static KV cache (past_present_share_buffer_) is guaranteed by GQA's
-  // ORT_ENFORCE when graph capture is enabled.
+  // Prepare indirect dispatch buffer for split-reduce path when CPU-side
+  // total_sequence_length is unusable for dispatch sizing:
+  //   - past_present_share_buffer layers: when graph capture is enabled,
+  //     total_sequence_length_ may be 0 (GPU-based seqlen_k), so the indirect
+  //     buffer computes dispatch sizes on GPU. Static KV cache (past_present_share_buffer_)
+  //     is guaranteed by GQA's ORT_ENFORCE when graph capture is enabled.
+  //   - kv_empty layers: no KV tokens are appended, so total_sequence_length_ is always 0;
+  //     using it directly would produce dispatch(0) and skip attention entirely.
   const bool use_indirect_dispatch = seqlen_k != nullptr &&
                                      total_seqlen != nullptr &&
                                      context.IsGraphCaptureEnabled();
@@ -547,6 +568,21 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co
       present_key = const_cast<Tensor*>(past_key);
       present_value = const_cast<Tensor*>(past_value);
     }
+
+    // CopyKVCache normally prepares the indirect buffer. Since we skip CopyKVCache
+    // for kv_empty layers, we must prepare it here.
+    if (use_indirect_dispatch) {
+      PrepareIndirectDispatchProgram program;
+      program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None});
+      program.AddOutput({indirect_buffer_ptr, ProgramTensorMetadataDependency::None});
+      program.SetDispatchGroupSize(1)
+          .SetWorkgroupSize(1)
+          .AddUniformVariables({{tile_size},
+                                {static_cast<uint32_t>(parameters.num_heads_)},
+                                {num_q_tiles},
+                                {static_cast<uint32_t>(parameters.batch_size_)}});
+      ORT_RETURN_IF_ERROR(context.RunProgram(program));
+    }
   }
 
   // When TurboQuant is active, create u32 tensor views over present/past KV cache buffers.

onnxruntime/contrib_ops/webgpu/bert/flash_attention.h

@@ -70,6 +70,20 @@ class CopyKVCacheProgram final : public Program<CopyKVCacheProgram> {
   bool use_seqlen_k_;
 };
 
+class PrepareIndirectDispatchProgram final : public Program<PrepareIndirectDispatchProgram> {
+ public:
+  PrepareIndirectDispatchProgram()
+      : Program{"PrepareIndirectDispatch"} {}
+
+  Status GenerateShaderCode(ShaderHelper& sh) const override;
+
+  WEBGPU_PROGRAM_DEFINE_UNIFORM_VARIABLES(
+      {"tile_size", ProgramUniformVariableDataType::Uint32},
+      {"num_heads", ProgramUniformVariableDataType::Uint32},
+      {"num_q_tiles", ProgramUniformVariableDataType::Uint32},
+      {"batch_size", ProgramUniformVariableDataType::Uint32});
+};
+
 class FlashAttentionProgram final : public Program<FlashAttentionProgram> {
  public:
   FlashAttentionProgram(const std::string& kernel_name,

onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc

@@ -353,7 +353,13 @@ Status GroupQueryAttention::ComputeInternal(onnxruntime::webgpu::ComputeContext&
                                           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.");
-  ORT_ENFORCE(!context.IsGraphCaptureEnabled() || parameters.past_present_share_buffer_,
+  // 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);
+  // kv_empty layers (e.g. Gemma4 layers 15-34) reuse KV from another layer so
+  // past/present cannot share the same buffer — exempt them from this check.
+  ORT_ENFORCE(!context.IsGraphCaptureEnabled() || kv_empty || parameters.past_present_share_buffer_,
               "Graph capture requires past/present KV cache to share the same buffer (static KV cache).");
 
   Tensor qSplit;
@@ -363,11 +369,6 @@ 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;

onnxruntime/test/contrib_ops/group_query_attention_op_test.cc

@@ -3160,6 +3160,134 @@ TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_RejectsNonPowerOf2HeadSize) {
              {}, nullptr, &execution_providers);
 }
 
+// ---------------------------------------------------------------------------
+// kv_empty WebGPU tests
+//
+// kv_empty is triggered when key/value have sequence_length=0 (Gemma4 shared-KV
+// layers where KV is reused from another layer's cache). The WebGPU kernel
+// aliases present to past and runs flash attention directly against past_key/value.
+// Correctness is verified by cross-checking against the CPU reference path.
+// ---------------------------------------------------------------------------
+
+// WebGPU: kv_empty decode (q_seq=1, past_seq=8).
+// Exercises the kv_empty branch in ApplyFlashAttention for a short past.
+TEST(GroupQueryAttentionTest, WebGPU_SharedKV_KvEmpty_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;
+  constexpr int kv_hidden_size = kv_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 % 7 + 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);
+
+  // WebGPU run: kv_empty path — present is aliased to past, flash attention
+  // runs directly against past_key/value.
+  OpTester webgpu_tester("GroupQueryAttention", 1, onnxruntime::kMSDomain);
+  webgpu_tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(num_heads));
+  webgpu_tester.AddAttribute<int64_t>("kv_num_heads", static_cast<int64_t>(kv_num_heads));
+  webgpu_tester.AddInput<float>("query", {batch_size, q_seq_len, hidden_size}, query_data);
+  webgpu_tester.AddInput<float>("key", {batch_size, 0, kv_hidden_size}, {});    // kv_empty
+  webgpu_tester.AddInput<float>("value", {batch_size, 0, kv_hidden_size}, {});  // kv_empty
+  webgpu_tester.AddInput<float>("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data);
+  webgpu_tester.AddInput<float>("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data);
+  webgpu_tester.AddInput<int32_t>("seqlens_k", {batch_size}, {static_cast<int32_t>(past_seq_len - 1)});
+  webgpu_tester.AddInput<int32_t>("total_sequence_length", {1}, {past_seq_len});
+  webgpu_tester.AddOptionalInputEdge<float>();    // cos_cache
+  webgpu_tester.AddOptionalInputEdge<float>();    // sin_cache
+  webgpu_tester.AddOptionalInputEdge<int64_t>();  // position_ids
+  webgpu_tester.AddOptionalInputEdge<float>();    // attention_bias
+  webgpu_tester.AddOptionalInputEdge<float>();    // head_sink
+  const int output_size = batch_size * q_seq_len * hidden_size;
+  const int present_size = batch_size * kv_num_heads * past_seq_len * head_size;
+  webgpu_tester.AddOutput<float>("output", {batch_size, q_seq_len, hidden_size}, std::vector<float>(output_size, 0.0f));
+  webgpu_tester.AddOutput<float>("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector<float>(present_size, 0.0f));
+  webgpu_tester.AddOutput<float>("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector<float>(present_size, 0.0f));
+  webgpu_tester.SetOutputTolerance(1e6f);
+  std::vector<std::unique_ptr<IExecutionProvider>> webgpu_eps;
+  webgpu_eps.push_back(DefaultWebGpuExecutionProvider());
+  webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &webgpu_eps);
+  const float* webgpu_out = webgpu_tester.GetFetches()[0].Get<Tensor>().Data<float>();
+  std::vector<float> webgpu_output(webgpu_out, webgpu_out + output_size);
+
+  // CPU reference: use real total_sequence_length so CPU path is correct.
+  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, GqaTargetEp::kCpu);
+
+  ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_KvEmpty_Decode_WebGPU_vs_CPU");
+}
+
+// WebGPU: kv_empty decode with larger past (past_seq=32, exercises tile boundary).
+// past_seq_len=32 crosses the tile size threshold, exercising multi-tile flash attention.
+TEST(GroupQueryAttentionTest, WebGPU_SharedKV_KvEmpty_LargerPast) {
+  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 = 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;
+
+  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 % 7 + 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);
+
+  OpTester webgpu_tester("GroupQueryAttention", 1, onnxruntime::kMSDomain);
+  webgpu_tester.AddAttribute<int64_t>("num_heads", static_cast<int64_t>(num_heads));
+  webgpu_tester.AddAttribute<int64_t>("kv_num_heads", static_cast<int64_t>(kv_num_heads));
+  webgpu_tester.AddInput<float>("query", {batch_size, q_seq_len, hidden_size}, query_data);
+  webgpu_tester.AddInput<float>("key", {batch_size, 0, kv_hidden_size}, {});    // kv_empty
+  webgpu_tester.AddInput<float>("value", {batch_size, 0, kv_hidden_size}, {});  // kv_empty
+  webgpu_tester.AddInput<float>("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data);
+  webgpu_tester.AddInput<float>("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data);
+  webgpu_tester.AddInput<int32_t>("seqlens_k", {batch_size}, {static_cast<int32_t>(past_seq_len - 1)});
+  webgpu_tester.AddInput<int32_t>("total_sequence_length", {1}, {past_seq_len});
+  webgpu_tester.AddOptionalInputEdge<float>();    // cos_cache
+  webgpu_tester.AddOptionalInputEdge<float>();    // sin_cache
+  webgpu_tester.AddOptionalInputEdge<int64_t>();  // position_ids
+  webgpu_tester.AddOptionalInputEdge<float>();    // attention_bias
+  webgpu_tester.AddOptionalInputEdge<float>();    // head_sink
+  const int output_size = batch_size * q_seq_len * hidden_size;
+  const int present_size = batch_size * kv_num_heads * past_seq_len * head_size;
+  webgpu_tester.AddOutput<float>("output", {batch_size, q_seq_len, hidden_size}, std::vector<float>(output_size, 0.0f));
+  webgpu_tester.AddOutput<float>("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector<float>(present_size, 0.0f));
+  webgpu_tester.AddOutput<float>("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector<float>(present_size, 0.0f));
+  webgpu_tester.SetOutputTolerance(1e6f);
+  std::vector<std::unique_ptr<IExecutionProvider>> webgpu_eps;
+  webgpu_eps.push_back(DefaultWebGpuExecutionProvider());
+  webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &webgpu_eps);
+  const float* webgpu_out = webgpu_tester.GetFetches()[0].Get<Tensor>().Data<float>();
+  std::vector<float> webgpu_output(webgpu_out, webgpu_out + output_size);
+
+  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, GqaTargetEp::kCpu);
+
+  ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_KvEmpty_LargerPast_WebGPU_vs_CPU");
+}
+
 // --- Success paths: TurboQuant with flash attention at various K sizes ---
 // K=1 (decode with minimal past), K=24 (moderate), K=128 (large)
 // These exercise the split-reduce decode path (QKV + VxReduce kernels) for seq_len=1,