From c5f03453d633f1117bcd282f616dadd2c011d5c4 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Tue, 23 Jun 2026 17:01:42 +0800 Subject: [PATCH 01/16] add indirect dispatch support for graph capture --- .../webgpu/bert/flash_attention.cc | 36 ++++++++++++++++--- .../contrib_ops/webgpu/bert/flash_attention.h | 14 ++++++++ 2 files changed, 45 insertions(+), 5 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index ff6938c0e9bf4..3f4502dc5ec18 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -159,6 +159,16 @@ 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.AdditionalImplementation() << kNormalizeDispatchGroupSizeFn; + shader.MainFunctionBody() << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n" + << " let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n" + << " normalize_dispatch_group_size(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size);\n"; + 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 +521,12 @@ 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 the split-reduce path when the CPU-side + // total_sequence_length is unusable for dispatch sizing: + // - past_present_share_buffer layers: graph capture keeps seqlen_k GPU-resident, so + // total_sequence_length_ is 0 on CPU and the true value must be read from the GPU tensor. + // - 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 +558,21 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co present_key = const_cast(past_key); present_value = const_cast(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(parameters.num_heads_)}, + {num_q_tiles}, + {static_cast(parameters.batch_size_)}}); + ORT_RETURN_IF_ERROR(context.RunProgram(program)); + } } // When TurboQuant is active, create u32 tensor views over present/past KV cache buffers. diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h index 02f25a753fff8..f917f50c7fd72 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.h @@ -70,6 +70,20 @@ class CopyKVCacheProgram final : public Program { bool use_seqlen_k_; }; +class PrepareIndirectDispatchProgram final : public Program { + 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 { public: FlashAttentionProgram(const std::string& kernel_name, From 2f70dec349b1d63a5d4611371868e27935a4080e Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Fri, 26 Jun 2026 13:09:14 +0800 Subject: [PATCH 02/16] webgpu: refactor indirect dispatch normalization in flash attention - Extract AppendNormalizeDispatchShader() helper so CopyKVCacheProgram and PrepareIndirectDispatchProgram share one copy of the tile-count and normalize_dispatch_group_size call instead of duplicating it. - Express use_indirect_dispatch as use_seqlen_k && (share_buffer || kv_empty) to make the subset relationship explicit and eliminate the repeated condition. - Tighten use_seqlen_k / use_indirect_dispatch guards from == 0 to <= 0 to handle a negative total_sequence_length_ defensively. - Add comment on the WGSL template's normalize call pointing back to the C++ helper so the two stay in sync. Co-Authored-By: Claude --- .../contrib_ops/webgpu/bert/flash_attention.cc | 16 ++++++++++++---- 1 file changed, 12 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 3f4502dc5ec18..fc3b2abd3698f 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -52,6 +52,16 @@ fn populate_indirect_dispatch_buffer(x: u32, y: u32, z: u32) { } )"; +// Emits kNormalizeDispatchGroupSizeFn into AdditionalImplementation and appends the two +// WGSL lines that compute num_total_seq_length_tile and call normalize_dispatch_group_size. +// `total_seq_len_expr` is the WGSL expression that evaluates to the total sequence length. +static void AppendNormalizeDispatchShader(ShaderHelper& shader, std::string_view total_seq_len_expr) { + shader.AdditionalImplementation() << kNormalizeDispatchGroupSizeFn; + shader.MainFunctionBody() + << " let num_total_seq_length_tile = (" << total_seq_len_expr << " + uniforms.tile_size - 1u) / uniforms.tile_size;\n" + << " normalize_dispatch_group_size(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); @@ -162,10 +172,8 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { Status PrepareIndirectDispatchProgram::GenerateShaderCode(ShaderHelper& shader) const { shader.AddInput("seqlen_k", ShaderUsage::None); shader.AddOutput("indirect_buffer", ShaderUsage::None); - shader.AdditionalImplementation() << kNormalizeDispatchGroupSizeFn; - shader.MainFunctionBody() << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n" - << " let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n" - << " normalize_dispatch_group_size(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_size);\n"; + shader.MainFunctionBody() << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n"; + AppendNormalizeDispatchShader(shader, "total_seq_length"); return Status::OK(); } From faee205e8cfc82833614f2b84a9a6ff66ff43ffb Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Fri, 26 Jun 2026 16:11:04 +0800 Subject: [PATCH 03/16] webgpu: add indirect dispatch path tests for flash attention Add two WebGPU GQA tests that exercise PrepareIndirectDispatchProgram: - WebGPU_SharedKV_IndirectDispatch_Decode: kv_empty + total_sequence_length=0 (decode, past_seq=8), triggers use_seqlen_k=true and use_indirect_dispatch=true via the kv_empty path, cross-checked against CPU reference. - WebGPU_SharedKV_IndirectDispatch_LargerPast: same path with past_seq=32 to exercise num_total_seq_length_tile > 1 in the tile count arithmetic. Co-Authored-By: Claude --- .../group_query_attention_op_test.cc | 135 +++++++++++++++++- 1 file changed, 133 insertions(+), 2 deletions(-) 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 ef7a04bc67d4f..7e662a7ebe38e 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2822,8 +2822,7 @@ TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CUDA) { RunBatchedRightPaddedRotaryPrefillForEP(GqaTargetEp::kCuda); } -TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_WebGPU) { - auto webgpu_ep = DefaultWebGpuExecutionProvider(); +TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_WebGPU) { auto webgpu_ep = DefaultWebGpuExecutionProvider(); if (!webgpu_ep) { GTEST_SKIP() << "WebGPU EP not available"; } @@ -3160,6 +3159,138 @@ TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_RejectsNonPowerOf2HeadSize) { {}, nullptr, &execution_providers); } +// --------------------------------------------------------------------------- +// Indirect dispatch tests (PrepareIndirectDispatchProgram) +// +// These tests pass total_sequence_length=0 to trigger the kv_empty indirect +// dispatch path: use_seqlen_k=true so seqlen_k tensor is used on the GPU, and +// use_indirect_dispatch=true so PrepareIndirectDispatchProgram sizes the flash- +// attention dispatch from the GPU-resident seqlen_k value instead of the zero +// CPU value. Correctness is verified by cross-checking against CPU with the +// real total. +// --------------------------------------------------------------------------- + +// WebGPU: kv_empty + total_sequence_length=0, decode (q_seq=1). +// Exercises PrepareIndirectDispatchProgram for the kv_empty indirect dispatch path. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_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 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); + + // WebGPU run: total_sequence_length=0 forces use_seqlen_k=true and + // use_indirect_dispatch=true; PrepareIndirectDispatchProgram sizes the + // dispatch from seqlen_k[0] on the GPU rather than the zero CPU value. + OpTester webgpu_tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + webgpu_tester.AddAttribute("num_heads", static_cast(num_heads)); + webgpu_tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + webgpu_tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + webgpu_tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); + webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); + webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); + webgpu_tester.AddInput("total_sequence_length", {1}, {0}); // 0 → indirect dispatch path + webgpu_tester.AddOptionalInputEdge(); // cos_cache + webgpu_tester.AddOptionalInputEdge(); // sin_cache + webgpu_tester.AddOptionalInputEdge(); // position_ids + webgpu_tester.AddOptionalInputEdge(); // attention_bias + webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); + webgpu_tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); + webgpu_tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); + webgpu_tester.SetOutputTolerance(1e6f); + std::vector> 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().Data(); + std::vector 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, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatch_Decode_WebGPU_vs_CPU"); +} + +// WebGPU: kv_empty + total_sequence_length=0, larger past (exercises tile boundary). +// past_seq_len=32 ensures num_total_seq_length_tile > 1, validating the tile +// count arithmetic in PrepareIndirectDispatchProgram more thoroughly. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_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 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); + + OpTester webgpu_tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + webgpu_tester.AddAttribute("num_heads", static_cast(num_heads)); + webgpu_tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + webgpu_tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); + webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); + webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + webgpu_tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); + webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); + webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); + webgpu_tester.AddInput("total_sequence_length", {1}, {0}); // 0 → indirect dispatch path + webgpu_tester.AddOptionalInputEdge(); // cos_cache + webgpu_tester.AddOptionalInputEdge(); // sin_cache + webgpu_tester.AddOptionalInputEdge(); // position_ids + webgpu_tester.AddOptionalInputEdge(); // attention_bias + webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); + webgpu_tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); + webgpu_tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); + webgpu_tester.SetOutputTolerance(1e6f); + std::vector> 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().Data(); + std::vector 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, /*use_cuda=*/false, /*use_webgpu=*/false); + + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatch_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, From e5f62cde0314a7c57adce73604598e7a6c9e5b6e Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Fri, 26 Jun 2026 16:55:58 +0800 Subject: [PATCH 04/16] webgpu: clang-format and remove verbose comment in flash attention Co-Authored-By: Claude --- .../webgpu/bert/flash_attention.cc | 8 ++++--- .../group_query_attention_op_test.cc | 23 ++++++++++--------- 2 files changed, 17 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index fc3b2abd3698f..4baa982641d7e 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -529,10 +529,12 @@ 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 the split-reduce path when the CPU-side + // Prepare indirect dispatch buffer for split-reduce path when CPU-side // total_sequence_length is unusable for dispatch sizing: - // - past_present_share_buffer layers: graph capture keeps seqlen_k GPU-resident, so - // total_sequence_length_ is 0 on CPU and the true value must be read from the GPU tensor. + // - 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 && 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 7e662a7ebe38e..83e4fccf16f83 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2822,7 +2822,8 @@ TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CUDA) { RunBatchedRightPaddedRotaryPrefillForEP(GqaTargetEp::kCuda); } -TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_WebGPU) { auto webgpu_ep = DefaultWebGpuExecutionProvider(); +TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_WebGPU) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); if (!webgpu_ep) { GTEST_SKIP() << "WebGPU EP not available"; } @@ -3207,11 +3208,11 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_Decode) { webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); webgpu_tester.AddInput("total_sequence_length", {1}, {0}); // 0 → indirect dispatch path - webgpu_tester.AddOptionalInputEdge(); // cos_cache - webgpu_tester.AddOptionalInputEdge(); // sin_cache - webgpu_tester.AddOptionalInputEdge(); // position_ids - webgpu_tester.AddOptionalInputEdge(); // attention_bias - webgpu_tester.AddOptionalInputEdge(); // head_sink + webgpu_tester.AddOptionalInputEdge(); // cos_cache + webgpu_tester.AddOptionalInputEdge(); // sin_cache + webgpu_tester.AddOptionalInputEdge(); // position_ids + webgpu_tester.AddOptionalInputEdge(); // attention_bias + webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); @@ -3267,11 +3268,11 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_LargerPast) { webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); webgpu_tester.AddInput("total_sequence_length", {1}, {0}); // 0 → indirect dispatch path - webgpu_tester.AddOptionalInputEdge(); // cos_cache - webgpu_tester.AddOptionalInputEdge(); // sin_cache - webgpu_tester.AddOptionalInputEdge(); // position_ids - webgpu_tester.AddOptionalInputEdge(); // attention_bias - webgpu_tester.AddOptionalInputEdge(); // head_sink + webgpu_tester.AddOptionalInputEdge(); // cos_cache + webgpu_tester.AddOptionalInputEdge(); // sin_cache + webgpu_tester.AddOptionalInputEdge(); // position_ids + webgpu_tester.AddOptionalInputEdge(); // attention_bias + webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); From 4ee58bb088229358ddf08d76697505f216f962df Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Fri, 26 Jun 2026 17:38:56 +0800 Subject: [PATCH 05/16] webgpu: fix RunGQASharedKV call sites for GqaTargetEp API change Replace deprecated bool use_cuda/use_webgpu params with GqaTargetEp::kCpu. Co-Authored-By: Claude --- onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 83e4fccf16f83..cf56ad74739af 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -3228,7 +3228,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_Decode) { // 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, /*use_cuda=*/false, /*use_webgpu=*/false); + num_heads, kv_num_heads, head_size, GqaTargetEp::kCpu); ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatch_Decode_WebGPU_vs_CPU"); } @@ -3287,7 +3287,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_LargerPast) { 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); + num_heads, kv_num_heads, head_size, GqaTargetEp::kCpu); ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatch_LargerPast_WebGPU_vs_CPU"); } From 9272ceb4ffa5e6091f0a1530d54d9d416df594be Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Fri, 26 Jun 2026 18:24:18 +0800 Subject: [PATCH 06/16] webgpu: fix GQA kv_empty tests to use valid total_sequence_length OpTester cannot enable graph capture, so use_indirect_dispatch is never triggered. Rewrite the tests to exercise the kv_empty path directly with a real positive total_sequence_length instead of 0. Co-Authored-By: Claude --- .../group_query_attention_op_test.cc | 66 +++++++++---------- 1 file changed, 31 insertions(+), 35 deletions(-) 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 cf56ad74739af..ecf91ce30d308 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -3161,19 +3161,17 @@ TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_RejectsNonPowerOf2HeadSize) { } // --------------------------------------------------------------------------- -// Indirect dispatch tests (PrepareIndirectDispatchProgram) +// kv_empty WebGPU tests // -// These tests pass total_sequence_length=0 to trigger the kv_empty indirect -// dispatch path: use_seqlen_k=true so seqlen_k tensor is used on the GPU, and -// use_indirect_dispatch=true so PrepareIndirectDispatchProgram sizes the flash- -// attention dispatch from the GPU-resident seqlen_k value instead of the zero -// CPU value. Correctness is verified by cross-checking against CPU with the -// real total. +// 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 + total_sequence_length=0, decode (q_seq=1). -// Exercises PrepareIndirectDispatchProgram for the kv_empty indirect dispatch path. -TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_Decode) { +// 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"; @@ -3195,24 +3193,23 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_Decode) { 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); - // WebGPU run: total_sequence_length=0 forces use_seqlen_k=true and - // use_indirect_dispatch=true; PrepareIndirectDispatchProgram sizes the - // dispatch from seqlen_k[0] on the GPU rather than the zero CPU value. + // 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("num_heads", static_cast(num_heads)); webgpu_tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); webgpu_tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); - webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); - webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); // kv_empty + webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); // kv_empty webgpu_tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); - webgpu_tester.AddInput("total_sequence_length", {1}, {0}); // 0 → indirect dispatch path - webgpu_tester.AddOptionalInputEdge(); // cos_cache - webgpu_tester.AddOptionalInputEdge(); // sin_cache - webgpu_tester.AddOptionalInputEdge(); // position_ids - webgpu_tester.AddOptionalInputEdge(); // attention_bias - webgpu_tester.AddOptionalInputEdge(); // head_sink + webgpu_tester.AddInput("total_sequence_length", {1}, {past_seq_len}); + webgpu_tester.AddOptionalInputEdge(); // cos_cache + webgpu_tester.AddOptionalInputEdge(); // sin_cache + webgpu_tester.AddOptionalInputEdge(); // position_ids + webgpu_tester.AddOptionalInputEdge(); // attention_bias + webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); @@ -3230,13 +3227,12 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_Decode) { 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_IndirectDispatch_Decode_WebGPU_vs_CPU"); + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_KvEmpty_Decode_WebGPU_vs_CPU"); } -// WebGPU: kv_empty + total_sequence_length=0, larger past (exercises tile boundary). -// past_seq_len=32 ensures num_total_seq_length_tile > 1, validating the tile -// count arithmetic in PrepareIndirectDispatchProgram more thoroughly. -TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_LargerPast) { +// 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"; @@ -3262,17 +3258,17 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_LargerPast) { webgpu_tester.AddAttribute("num_heads", static_cast(num_heads)); webgpu_tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); webgpu_tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); - webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); - webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); + webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); // kv_empty + webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); // kv_empty webgpu_tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); - webgpu_tester.AddInput("total_sequence_length", {1}, {0}); // 0 → indirect dispatch path - webgpu_tester.AddOptionalInputEdge(); // cos_cache - webgpu_tester.AddOptionalInputEdge(); // sin_cache - webgpu_tester.AddOptionalInputEdge(); // position_ids - webgpu_tester.AddOptionalInputEdge(); // attention_bias - webgpu_tester.AddOptionalInputEdge(); // head_sink + webgpu_tester.AddInput("total_sequence_length", {1}, {past_seq_len}); + webgpu_tester.AddOptionalInputEdge(); // cos_cache + webgpu_tester.AddOptionalInputEdge(); // sin_cache + webgpu_tester.AddOptionalInputEdge(); // position_ids + webgpu_tester.AddOptionalInputEdge(); // attention_bias + webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); @@ -3289,7 +3285,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatch_LargerPast) { 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_IndirectDispatch_LargerPast_WebGPU_vs_CPU"); + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_KvEmpty_LargerPast_WebGPU_vs_CPU"); } // --- Success paths: TurboQuant with flash attention at various K sizes --- From aa60c436ec291d8450b2fdec5611ae9aed589d4f Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Mon, 29 Jun 2026 12:39:21 +0800 Subject: [PATCH 07/16] webgpu: exempt kv_empty layers from graph capture share_buffer assertion Gemma4 layers 15-34 are KV-sharing layers (kv_sequence_length==0): they reuse KV from a different layer's output, so past/present cannot share the same GPU buffer even with past_present_share_buffer=true in the config. The assertion added in 92b4c66349 correctly enforces that graph capture requires static KV cache (shared buffers), but must exempt kv_empty layers since their cross-layer KV sharing is expected and valid. Move kv_empty definition before the ORT_ENFORCE so it can be used directly in the condition rather than inlining kv_sequence_length_ == 0. Co-Authored-By: Claude --- .../webgpu/bert/group_query_attention.cc | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc index 879a63a2482ca..7cc00555b6789 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc +++ b/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; From 49936038964146aa1f81031bbad0e9b9a6aa19d8 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Tue, 30 Jun 2026 11:49:00 +0800 Subject: [PATCH 08/16] webgpu: fix stale kNormalizeDispatchGroupSizeFn reference after rebase AppendNormalizeDispatchShader referenced the old constant and WGSL function names from before main renamed them to kPopulateIndirectDispatchBufferFn / populate_indirect_dispatch_buffer. Co-Authored-By: Claude --- onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 4baa982641d7e..28fef778aad21 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -52,14 +52,14 @@ fn populate_indirect_dispatch_buffer(x: u32, y: u32, z: u32) { } )"; -// Emits kNormalizeDispatchGroupSizeFn into AdditionalImplementation and appends the two -// WGSL lines that compute num_total_seq_length_tile and call normalize_dispatch_group_size. +// 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 AppendNormalizeDispatchShader(ShaderHelper& shader, std::string_view total_seq_len_expr) { - shader.AdditionalImplementation() << kNormalizeDispatchGroupSizeFn; + shader.AdditionalImplementation() << kPopulateIndirectDispatchBufferFn; shader.MainFunctionBody() << " let num_total_seq_length_tile = (" << total_seq_len_expr << " + uniforms.tile_size - 1u) / uniforms.tile_size;\n" - << " normalize_dispatch_group_size(num_total_seq_length_tile, uniforms.num_heads * uniforms.num_q_tiles, uniforms.batch_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 { From 9cefb106ebcec47b6ffb78dd79579185ea55f194 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Tue, 30 Jun 2026 14:45:09 +0800 Subject: [PATCH 09/16] webgpu: rename AppendNormalizeDispatchShader to AppendPopulateIndirectDispatchShader Aligns the helper name with the underlying WGSL function and constant it emits (populate_indirect_dispatch_buffer / kPopulateIndirectDispatchBufferFn) after main renamed these from the old normalize_dispatch_group_size naming. Co-Authored-By: Claude --- onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 28fef778aad21..4e5baba7c9ab9 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -55,7 +55,7 @@ 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 AppendNormalizeDispatchShader(ShaderHelper& shader, std::string_view total_seq_len_expr) { +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" @@ -173,7 +173,7 @@ Status PrepareIndirectDispatchProgram::GenerateShaderCode(ShaderHelper& shader) shader.AddInput("seqlen_k", ShaderUsage::None); shader.AddOutput("indirect_buffer", ShaderUsage::None); shader.MainFunctionBody() << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n"; - AppendNormalizeDispatchShader(shader, "total_seq_length"); + AppendPopulateIndirectDispatchShader(shader, "total_seq_length"); return Status::OK(); } From 495aac70f5fe71ec69080b56d8f04f085961b2e9 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Tue, 30 Jun 2026 15:15:45 +0800 Subject: [PATCH 10/16] webgpu: inline AppendPopulateIndirectDispatchShader into its only call site Single-use static helpers add indirection without benefit. Co-Authored-By: Claude --- .../contrib_ops/webgpu/bert/flash_attention.cc | 17 +++++------------ 1 file changed, 5 insertions(+), 12 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 4e5baba7c9ab9..7b86528a6e6e9 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -52,16 +52,6 @@ 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); @@ -172,8 +162,11 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { 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"); + shader.AdditionalImplementation() << kPopulateIndirectDispatchBufferFn; + shader.MainFunctionBody() + << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n" + << " let num_total_seq_length_tile = (total_seq_length + 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"; return Status::OK(); } From bc84e8d33fa35edf7535d1dcdab77ca04d5b21e2 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Tue, 30 Jun 2026 16:19:20 +0800 Subject: [PATCH 11/16] webgpu: switch PrepareIndirectDispatchProgram to use total_seqlen input Addresses reviewer feedback: use total_sequence_length_input (the global max, batch-safe) instead of seqlen_k[0] (per-batch index 0, unsafe for batch > 1). Aligns naming and logic with CopyKVCacheProgram which reads the same input the same way. Also clarifies comments to make explicit that indirect dispatch is only needed under graph capture, when total_seqlen is GPU-resident and CPU-side dispatch sizing is unavailable. Co-Authored-By: Claude --- .../webgpu/bert/flash_attention.cc | 27 +++++++++---------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc index 7b86528a6e6e9..ce4018956f5c0 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc +++ b/onnxruntime/contrib_ops/webgpu/bert/flash_attention.cc @@ -160,12 +160,12 @@ Status CopyKVCacheProgram::GenerateShaderCode(ShaderHelper& shader) const { } Status PrepareIndirectDispatchProgram::GenerateShaderCode(ShaderHelper& shader) const { - shader.AddInput("seqlen_k", ShaderUsage::None); + shader.AddInput("total_sequence_length_input", ShaderUsage::None); shader.AddOutput("indirect_buffer", ShaderUsage::None); shader.AdditionalImplementation() << kPopulateIndirectDispatchBufferFn; shader.MainFunctionBody() - << " let total_seq_length = u32(seqlen_k[0u]) + 1u;\n" - << " let num_total_seq_length_tile = (total_seq_length + uniforms.tile_size - 1u) / uniforms.tile_size;\n" + << " let global_total_seq_length = u32(total_sequence_length_input[0]);\n" + << " let num_total_seq_length_tile = (global_total_seq_length + 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"; return Status::OK(); } @@ -522,14 +522,11 @@ 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 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. + // 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. const bool use_indirect_dispatch = seqlen_k != nullptr && total_seqlen != nullptr && context.IsGraphCaptureEnabled(); @@ -562,11 +559,13 @@ Status ApplyFlashAttention(const Tensor* Q, const Tensor* K, const Tensor* V, co present_value = const_cast(past_value); } - // CopyKVCache normally prepares the indirect buffer. Since we skip CopyKVCache - // for kv_empty layers, we must prepare it here. + // CopyKVCache normally prepares the indirect dispatch buffer. For kv_empty layers + // CopyKVCache is skipped, so we prepare it here. Only needed under graph capture + // because that is when total_seqlen is GPU-resident and CPU-side dispatch sizing + // is unavailable. if (use_indirect_dispatch) { PrepareIndirectDispatchProgram program; - program.AddInput({seqlen_k, ProgramTensorMetadataDependency::None}); + program.AddInput({total_seqlen, ProgramTensorMetadataDependency::None}); program.AddOutput({indirect_buffer_ptr, ProgramTensorMetadataDependency::None}); program.SetDispatchGroupSize(1) .SetWorkgroupSize(1) From 99c41ee356cdfdc827099bb4df412ae03e95788a Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 30 Jun 2026 14:02:30 +0000 Subject: [PATCH 12/16] Fix kv_empty WebGPU tests: reuse EP, drop unvalidated present outputs --- .../test/contrib_ops/group_query_attention_op_test.cc | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) 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 ecf91ce30d308..7a1dd9c8dcf86 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -3211,13 +3211,10 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_KvEmpty_Decode) { webgpu_tester.AddOptionalInputEdge(); // attention_bias webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); - webgpu_tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); - webgpu_tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); webgpu_tester.SetOutputTolerance(1e6f); std::vector> webgpu_eps; - webgpu_eps.push_back(DefaultWebGpuExecutionProvider()); + webgpu_eps.push_back(std::move(webgpu_ep)); webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &webgpu_eps); const float* webgpu_out = webgpu_tester.GetFetches()[0].Get().Data(); std::vector webgpu_output(webgpu_out, webgpu_out + output_size); @@ -3270,13 +3267,10 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_KvEmpty_LargerPast) { webgpu_tester.AddOptionalInputEdge(); // attention_bias webgpu_tester.AddOptionalInputEdge(); // 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("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); - webgpu_tester.AddOutput("present_key", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); - webgpu_tester.AddOutput("present_value", {batch_size, kv_num_heads, past_seq_len, head_size}, std::vector(present_size, 0.0f)); webgpu_tester.SetOutputTolerance(1e6f); std::vector> webgpu_eps; - webgpu_eps.push_back(DefaultWebGpuExecutionProvider()); + webgpu_eps.push_back(std::move(webgpu_ep)); webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &webgpu_eps); const float* webgpu_out = webgpu_tester.GetFetches()[0].Get().Data(); std::vector webgpu_output(webgpu_out, webgpu_out + output_size); From dabdf99b135706fadfab105e0298e09d69b7093b Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Wed, 1 Jul 2026 17:28:43 +0800 Subject: [PATCH 13/16] webgpu: add graph capture test for kv_empty indirect dispatch Adds WebGPU_SharedKVLayers_IndirectDispatchForGraphCapture to verify PrepareIndirectDispatchProgram correctly computes GPU-side dispatch sizes for kv_empty (shared-KV) layers when graph capture is enabled and total_seqlen is GPU-resident. Uses gpu_graph_id=-1 to keep IsGraphCaptureEnabled()=true without triggering actual Dawn capture/replay. Co-Authored-By: Claude --- .../group_query_attention_op_test.cc | 87 ++++--------------- 1 file changed, 15 insertions(+), 72 deletions(-) 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 7a1dd9c8dcf86..592d10dfa64a1 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -15,6 +15,7 @@ #include "test/util/include/default_providers.h" #ifdef USE_WEBGPU #include "core/providers/webgpu/webgpu_provider_options.h" +#include "core/session/onnxruntime_run_options_config_keys.h" #endif namespace onnxruntime { @@ -3160,77 +3161,15 @@ 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 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); - - // 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("num_heads", static_cast(num_heads)); - webgpu_tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - webgpu_tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); - webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); // kv_empty - webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); // kv_empty - webgpu_tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); - webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); - webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); - webgpu_tester.AddInput("total_sequence_length", {1}, {past_seq_len}); - webgpu_tester.AddOptionalInputEdge(); // cos_cache - webgpu_tester.AddOptionalInputEdge(); // sin_cache - webgpu_tester.AddOptionalInputEdge(); // position_ids - webgpu_tester.AddOptionalInputEdge(); // attention_bias - webgpu_tester.AddOptionalInputEdge(); // head_sink - const int output_size = batch_size * q_seq_len * hidden_size; - webgpu_tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); - webgpu_tester.SetOutputTolerance(1e6f); - std::vector> webgpu_eps; - webgpu_eps.push_back(std::move(webgpu_ep)); - webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &webgpu_eps); - const float* webgpu_out = webgpu_tester.GetFetches()[0].Get().Data(); - std::vector 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(); +// layers where KV is reused from another layer's cache). Under graph capture, +// total_seqlen is GPU-resident so dispatch sizes must be computed on GPU via +// PrepareIndirectDispatchProgram. +TEST(GroupQueryAttentionTest, WebGPU_SharedKVLayers_IndirectDispatchForGraphCapture) { + ConfigOptions config_options{}; + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto webgpu_ep = WebGpuExecutionProviderWithOptions(config_options); if (!webgpu_ep) { GTEST_SKIP() << "WebGPU EP not available"; } @@ -3269,9 +3208,13 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_KvEmpty_LargerPast) { const int output_size = batch_size * q_seq_len * hidden_size; webgpu_tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); webgpu_tester.SetOutputTolerance(1e6f); + // gpu_graph_id=-1 keeps IsGraphCaptureEnabled()=true (exercising the indirect dispatch + // code path) without triggering actual Dawn command capture/replay. + RunOptions run_options; + ORT_THROW_IF_ERROR(run_options.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, "-1")); std::vector> webgpu_eps; webgpu_eps.push_back(std::move(webgpu_ep)); - webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &webgpu_eps); + webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, &run_options, &webgpu_eps); const float* webgpu_out = webgpu_tester.GetFetches()[0].Get().Data(); std::vector webgpu_output(webgpu_out, webgpu_out + output_size); @@ -3279,7 +3222,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_KvEmpty_LargerPast) { 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"); + ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKVLayers_IndirectDispatchForGraphCapture_vs_CPU"); } // --- Success paths: TurboQuant with flash attention at various K sizes --- From f99046c20b619b08d6e80a32f7fcd45d7f5ba3c6 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Wed, 1 Jul 2026 20:53:24 +0800 Subject: [PATCH 14/16] webgpu: rewrite graph capture GQA test to exercise real ORT GC path Build the model via Graph API (opset 17) instead of OpTester::BuildModel so graph inputs are properly declared after graph.Resolve(). Bind all inputs as GPU-resident OrtValues via IOBinding: PrepareIndirectDispatchProgram now reads total_seqlen from a real WGPUBuffer instead of a CPU pointer cast, exercising the full capture-then-replay path as the reviewer requested. Co-Authored-By: Claude --- .../group_query_attention_op_test.cc | 239 +++++++++++++----- 1 file changed, 174 insertions(+), 65 deletions(-) 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 592d10dfa64a1..019f236f1db90 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -14,8 +14,12 @@ #include "test/providers/provider_test_utils.h" #include "test/util/include/default_providers.h" #ifdef USE_WEBGPU +#include "core/graph/model.h" #include "core/providers/webgpu/webgpu_provider_options.h" -#include "core/session/onnxruntime_run_options_config_keys.h" +#include "core/session/inference_session.h" +#include "core/session/IOBinding.h" +#include "test/test_environment.h" +#include "test/unittest_util/framework_test_utils.h" #endif namespace onnxruntime { @@ -2563,6 +2567,175 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_SlidingWindow) { tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// WebGPU graph capture test for kv_empty (Gemma4 shared-KV) layers. +// +// When graph capture is enabled, total_seqlen is GPU-resident and +// PrepareIndirectDispatchProgram must compute dispatch sizes on GPU. +// This test exercises the full ORT graph capture path by allocating all +// inputs as GPU tensors via IOBinding, running capture then replay, and +// verifying the replay output matches the CPU reference. +TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { +#ifdef USE_WEBGPU + 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; + + // Build a GQA model directly using the Graph API so graph inputs are + // properly declared (OpTester::BuildModel doesn't call graph.Resolve or + // SetInputs, producing a proto that InferenceSession::Load rejects). + std::unique_ptr p_model; + { + std::unordered_map domain_to_version; + domain_to_version[kOnnxDomain] = 17; + domain_to_version[kMSDomain] = 1; + p_model = std::make_unique( + "gqa_gc_test", true, ModelMetaData(), PathString(), + IOnnxRuntimeOpSchemaRegistryList(), domain_to_version, + std::vector{}, + DefaultLoggingManager().DefaultLogger(), + ModelOptions(true, true)); + onnxruntime::Graph& graph = p_model->MainGraph(); + + ONNX_NAMESPACE::TypeProto tp_float, tp_int32; + tp_float.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_FLOAT); + tp_int32.mutable_tensor_type()->set_elem_type(ONNX_NAMESPACE::TensorProto_DataType_INT32); + + std::vector in_defs = { + &graph.GetOrCreateNodeArg("query", &tp_float), + &graph.GetOrCreateNodeArg("key", &tp_float), + &graph.GetOrCreateNodeArg("value", &tp_float), + &graph.GetOrCreateNodeArg("past_key", &tp_float), + &graph.GetOrCreateNodeArg("past_value", &tp_float), + &graph.GetOrCreateNodeArg("seqlens_k", &tp_int32), + &graph.GetOrCreateNodeArg("total_sequence_length", &tp_int32), + // optional inputs 8-12 are absent — use empty NodeArgs + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + &graph.GetOrCreateNodeArg("", nullptr), + }; + std::vector out_defs = { + &graph.GetOrCreateNodeArg("output", &tp_float), + }; + + auto& node = graph.AddNode("gqa", "GroupQueryAttention", "GQA", + in_defs, out_defs, nullptr, kMSDomain); + node.AddAttribute("num_heads", static_cast(num_heads)); + node.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + + ORT_THROW_IF_ERROR(graph.Resolve()); + } + + std::string model_data; + p_model->ToProto().SerializeToString(&model_data); + + // Create InferenceSession with graph capture EP. + SessionOptions so; + InferenceSession session{so, GetEnvironment()}; + ConfigOptions config_options{}; + ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kEnableGraphCapture, + webgpu::options::kEnableGraphCapture_ON)); + auto webgpu_ep = WebGpuExecutionProviderWithOptions(config_options); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + IExecutionProvider* ep_ptr = webgpu_ep.get(); + ORT_THROW_IF_ERROR(session.RegisterExecutionProvider(std::move(webgpu_ep))); + std::istringstream model_stream(model_data); + ORT_THROW_IF_ERROR(session.Load(model_stream)); + ORT_THROW_IF_ERROR(session.Initialize()); + + // Get GPU allocator from session. + OrtMemoryInfo gpu_mem_info(WEBGPU_BUFFER, OrtAllocatorType::OrtDeviceAllocator, + OrtDevice(OrtDevice::GPU, OrtDevice::MemType::DEFAULT, OrtDevice::VendorIds::NONE, 0)); + auto gpu_alloc = session.GetAllocator(gpu_mem_info); + AllocatorPtr cpu_alloc = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; + + // Prepare test data. + 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); + + // Helper: create GPU OrtValue by copying from CPU data (data may be nullptr for zero-size tensors). + auto make_gpu_value = [&](const void* data, MLDataType dtype, TensorShape shape) -> OrtValue { + Tensor gpu_tensor(dtype, shape, gpu_alloc); + if (data != nullptr && shape.Size() > 0) { + Tensor cpu_tensor(dtype, shape, const_cast(data), cpu_alloc->Info()); + ORT_THROW_IF_ERROR(ep_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, gpu_tensor)); + } + OrtValue val; + Tensor::InitOrtValue(std::move(gpu_tensor), val); + return val; + }; + + auto q_val = make_gpu_value(query_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, q_seq_len, hidden_size}); + auto key_val = make_gpu_value(nullptr, DataTypeImpl::GetType(), + TensorShape{batch_size, 0, kv_hidden_size}); + auto value_val = make_gpu_value(nullptr, DataTypeImpl::GetType(), + TensorShape{batch_size, 0, kv_hidden_size}); + auto pk_val = make_gpu_value(past_key_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + auto pv_val = make_gpu_value(past_value_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + std::vector seqlens_k_data = {past_seq_len - 1}; + auto sk_val = make_gpu_value(seqlens_k_data.data(), DataTypeImpl::GetType(), + TensorShape{batch_size}); + std::vector total_seqlen_data = {past_seq_len}; + auto ts_val = make_gpu_value(total_seqlen_data.data(), DataTypeImpl::GetType(), + TensorShape{1}); + + // Preallocate GPU output. + Tensor gpu_out_tensor(DataTypeImpl::GetType(), + TensorShape{batch_size, q_seq_len, hidden_size}, gpu_alloc); + OrtValue out_val; + Tensor::InitOrtValue(std::move(gpu_out_tensor), out_val); + + // Bind inputs and output. + std::unique_ptr io_binding; + ORT_THROW_IF_ERROR(session.NewIOBinding(&io_binding)); + ORT_THROW_IF_ERROR(io_binding->BindInput("query", q_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("key", key_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("value", value_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("past_key", pk_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("past_value", pv_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("seqlens_k", sk_val)); + ORT_THROW_IF_ERROR(io_binding->BindInput("total_sequence_length", ts_val)); + ORT_THROW_IF_ERROR(io_binding->BindOutput("output", out_val)); + ORT_THROW_IF_ERROR(io_binding->SynchronizeInputs()); + + // Run 1: capture. + RunOptions run_options; + ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); + + // Run 2: replay. + ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); + + // Copy output GPU -> CPU and compare to CPU reference. + const int output_size = batch_size * q_seq_len * hidden_size; + auto& gpu_out = io_binding->GetOutputs()[0].Get(); + Tensor cpu_out_tensor(DataTypeImpl::GetType(), gpu_out.Shape(), cpu_alloc); + ORT_THROW_IF_ERROR(ep_ptr->GetDataTransfer()->CopyTensor(gpu_out, cpu_out_tensor)); + std::vector webgpu_output(cpu_out_tensor.Data(), + cpu_out_tensor.Data() + 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_IndirectDispatchForGraphCapture_vs_CPU"); +#endif +} + // --------------------------------------------------------------------------- // Batched right-padded packed-QKV prefill with do_rotary. // @@ -3161,70 +3334,6 @@ TEST(GroupQueryAttentionTest, WebGPU_TurboQuant_RejectsNonPowerOf2HeadSize) { {}, nullptr, &execution_providers); } -// kv_empty is triggered when key/value have sequence_length=0 (Gemma4 shared-KV -// layers where KV is reused from another layer's cache). Under graph capture, -// total_seqlen is GPU-resident so dispatch sizes must be computed on GPU via -// PrepareIndirectDispatchProgram. -TEST(GroupQueryAttentionTest, WebGPU_SharedKVLayers_IndirectDispatchForGraphCapture) { - ConfigOptions config_options{}; - ORT_THROW_IF_ERROR(config_options.AddConfigEntry(webgpu::options::kEnableGraphCapture, - webgpu::options::kEnableGraphCapture_ON)); - auto webgpu_ep = WebGpuExecutionProviderWithOptions(config_options); - 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 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); - - OpTester webgpu_tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); - webgpu_tester.AddAttribute("num_heads", static_cast(num_heads)); - webgpu_tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); - webgpu_tester.AddInput("query", {batch_size, q_seq_len, hidden_size}, query_data); - webgpu_tester.AddInput("key", {batch_size, 0, kv_hidden_size}, {}); // kv_empty - webgpu_tester.AddInput("value", {batch_size, 0, kv_hidden_size}, {}); // kv_empty - webgpu_tester.AddInput("past_key", {batch_size, kv_num_heads, past_seq_len, head_size}, past_key_data); - webgpu_tester.AddInput("past_value", {batch_size, kv_num_heads, past_seq_len, head_size}, past_value_data); - webgpu_tester.AddInput("seqlens_k", {batch_size}, {static_cast(past_seq_len - 1)}); - webgpu_tester.AddInput("total_sequence_length", {1}, {past_seq_len}); - webgpu_tester.AddOptionalInputEdge(); // cos_cache - webgpu_tester.AddOptionalInputEdge(); // sin_cache - webgpu_tester.AddOptionalInputEdge(); // position_ids - webgpu_tester.AddOptionalInputEdge(); // attention_bias - webgpu_tester.AddOptionalInputEdge(); // head_sink - const int output_size = batch_size * q_seq_len * hidden_size; - webgpu_tester.AddOutput("output", {batch_size, q_seq_len, hidden_size}, std::vector(output_size, 0.0f)); - webgpu_tester.SetOutputTolerance(1e6f); - // gpu_graph_id=-1 keeps IsGraphCaptureEnabled()=true (exercising the indirect dispatch - // code path) without triggering actual Dawn command capture/replay. - RunOptions run_options; - ORT_THROW_IF_ERROR(run_options.config_options.AddConfigEntry(kOrtRunOptionsConfigCudaGraphAnnotation, "-1")); - std::vector> webgpu_eps; - webgpu_eps.push_back(std::move(webgpu_ep)); - webgpu_tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, &run_options, &webgpu_eps); - const float* webgpu_out = webgpu_tester.GetFetches()[0].Get().Data(); - std::vector 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, "SharedKVLayers_IndirectDispatchForGraphCapture_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, From ec1d50c5424384c4f07e50f672f80a21f4e43016 Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Thu, 2 Jul 2026 11:06:14 +0800 Subject: [PATCH 15/16] webgpu: replay graph capture test with different inputs to verify correctness MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add a second input set (set B) with distinct values. For the replay run, copy set B into the same GPU buffers via CopyTensor — rebinding is not allowed since WGPUBuffer pointers are baked in at capture time. Verify the replay output matches the CPU reference for set B, not set A, proving the captured graph actually executes with the updated data. Co-Authored-By: Claude --- .../group_query_attention_op_test.cc | 35 ++++++++++++++++--- 1 file changed, 30 insertions(+), 5 deletions(-) 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 019f236f1db90..4a03a997d6715 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2657,7 +2657,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { auto gpu_alloc = session.GetAllocator(gpu_mem_info); AllocatorPtr cpu_alloc = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; - // Prepare test data. + // Capture-run input data (set A). 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); @@ -2665,7 +2665,15 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { 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); - // Helper: create GPU OrtValue by copying from CPU data (data may be nullptr for zero-size tensors). + // Replay-run input data (set B) — different values to verify replay actually uses new data. + std::vector query_data2(batch_size * q_seq_len * hidden_size); + std::vector past_key_data2(batch_size * kv_num_heads * past_seq_len * head_size); + std::vector past_value_data2(batch_size * kv_num_heads * past_seq_len * head_size); + for (size_t i = 0; i < query_data2.size(); i++) query_data2[i] = 0.5f * static_cast(i % 11 + 1); + for (size_t i = 0; i < past_key_data2.size(); i++) past_key_data2[i] = 0.4f * static_cast(i % 7 + 1); + for (size_t i = 0; i < past_value_data2.size(); i++) past_value_data2[i] = 0.6f * static_cast(i % 4 + 1); + + // Helper: create GPU OrtValue by copying from CPU data (nullptr = zero-size tensor). auto make_gpu_value = [&](const void* data, MLDataType dtype, TensorShape shape) -> OrtValue { Tensor gpu_tensor(dtype, shape, gpu_alloc); if (data != nullptr && shape.Size() > 0) { @@ -2677,6 +2685,15 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { return val; }; + // Helper: overwrite an existing GPU tensor's data in-place from CPU (same buffer, new values). + // Under graph capture the WGPUBuffer pointer is baked in; rebinding is not allowed. + auto update_gpu_value = [&](const OrtValue& gpu_val, const void* data, MLDataType dtype, TensorShape shape) { + if (shape.Size() > 0) { + Tensor cpu_tensor(dtype, shape, const_cast(data), cpu_alloc->Info()); + ORT_THROW_IF_ERROR(ep_ptr->GetDataTransfer()->CopyTensor(cpu_tensor, const_cast(gpu_val.Get()))); + } + }; + auto q_val = make_gpu_value(query_data.data(), DataTypeImpl::GetType(), TensorShape{batch_size, q_seq_len, hidden_size}); auto key_val = make_gpu_value(nullptr, DataTypeImpl::GetType(), @@ -2717,10 +2734,18 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { RunOptions run_options; ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); - // Run 2: replay. + // Run 2: replay with different inputs (set B) written into the same GPU buffers. + // Rebinding is not allowed under graph capture — the WGPUBuffer pointers are baked + // in at capture time, so new data must be copied into the existing buffers. + update_gpu_value(q_val, query_data2.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, q_seq_len, hidden_size}); + update_gpu_value(pk_val, past_key_data2.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); + update_gpu_value(pv_val, past_value_data2.data(), DataTypeImpl::GetType(), + TensorShape{batch_size, kv_num_heads, past_seq_len, head_size}); ORT_THROW_IF_ERROR(session.Run(run_options, *io_binding)); - // Copy output GPU -> CPU and compare to CPU reference. + // Copy replay output GPU -> CPU and compare against CPU reference for set B. const int output_size = batch_size * q_seq_len * hidden_size; auto& gpu_out = io_binding->GetOutputs()[0].Get(); Tensor cpu_out_tensor(DataTypeImpl::GetType(), gpu_out.Shape(), cpu_alloc); @@ -2729,7 +2754,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { cpu_out_tensor.Data() + output_size); auto cpu_output = RunGQASharedKV( - batch_size, q_seq_len, past_seq_len, query_data, past_key_data, past_value_data, + batch_size, q_seq_len, past_seq_len, query_data2, past_key_data2, past_value_data2, num_heads, kv_num_heads, head_size, GqaTargetEp::kCpu); ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatchForGraphCapture_vs_CPU"); From 053c2b98a64eedb5037f410ee4209ddcf079cd3b Mon Sep 17 00:00:00 2001 From: Fei Chen Date: Thu, 2 Jul 2026 11:18:21 +0800 Subject: [PATCH 16/16] webgpu: move #ifdef USE_WEBGPU to wrap the whole GC test, not just body Co-Authored-By: Claude --- onnxruntime/test/contrib_ops/group_query_attention_op_test.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 4a03a997d6715..16af4a4e489aa 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2567,6 +2567,7 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_SlidingWindow) { tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +#ifdef USE_WEBGPU // WebGPU graph capture test for kv_empty (Gemma4 shared-KV) layers. // // When graph capture is enabled, total_seqlen is GPU-resident and @@ -2575,7 +2576,6 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_SlidingWindow) { // inputs as GPU tensors via IOBinding, running capture then replay, and // verifying the replay output matches the CPU reference. TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { -#ifdef USE_WEBGPU constexpr int batch_size = 1; constexpr int q_seq_len = 1; constexpr int past_seq_len = 32; @@ -2758,8 +2758,8 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_IndirectDispatchForGraphCapture) { num_heads, kv_num_heads, head_size, GqaTargetEp::kCpu); ExpectOutputsMatch(webgpu_output, cpu_output, 0.05f, "SharedKV_IndirectDispatchForGraphCapture_vs_CPU"); -#endif } +#endif // USE_WEBGPU // --------------------------------------------------------------------------- // Batched right-padded packed-QKV prefill with do_rotary.