From ae30cf3e20f01a8f0a243ae44d6e157292a77b21 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 16 Jun 2026 07:20:35 +0000 Subject: [PATCH 1/5] Initial plan From 48e7455f4d0072182f09de605f5130ca1e96c6c4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 16 Jun 2026 07:45:25 +0000 Subject: [PATCH 2/5] Initial plan: Fix CPU GQA NaN for right-padded batched prefill --- .../group_query_attention_op_test.cc | 212 ++++++++++++++++++ 1 file changed, 212 insertions(+) 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 821f43971848a..7ce5eab1ee6b2 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2392,5 +2392,217 @@ TEST(GroupQueryAttentionTest, WebGPU_SharedKV_SlidingWindow) { tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); } +// --------------------------------------------------------------------------- +// Batched right-padded packed-QKV prefill with do_rotary. +// +// In a multi-batch prefill where individual prompts have different real lengths, +// GenAI right-pads short prompts up to the max sequence_length and reports each +// batch's real length via seqlens_k[b] = real_len[b] - 1. The property under +// test: each batch's real-last-token output (the one used to predict the next +// token) must equal what we get from running that prompt singly as a batch=1 +// prefill. This is a generic correctness check that any GQA-supporting EP +// should satisfy. +// --------------------------------------------------------------------------- + +// Builds a packed QKV tensor with deterministic values at real positions and +// zeros at right-padded positions. Layout per token: [Q(hidden), K(kv), V(kv)]. +// Uses values of order ~1.0 (well above the 5e-3 mismatch tolerance) so the +// rotated-vs-unrotated divergence is unambiguously detectable. +static void FillBatchedRightPaddedPackedQKV(int batch_size, + int sequence_length, + int num_heads, + int kv_num_heads, + int head_size, + const std::vector& real_lens, + std::vector& packed_out) { + const int hidden_size = num_heads * head_size; + const int kv_hidden_size = kv_num_heads * head_size; + const int token_size = hidden_size + 2 * kv_hidden_size; + packed_out.assign(batch_size * sequence_length * token_size, 0.0f); + for (int b = 0; b < batch_size; ++b) { + const int real_len = real_lens[b]; + for (int s = 0; s < real_len; ++s) { + float* token = &packed_out[(b * sequence_length + s) * token_size]; + for (int c = 0; c < hidden_size; ++c) { + token[c] = 0.1f + 0.3f * static_cast(((b * 7 + s * 3 + c) % 13) + 1); + } + for (int c = 0; c < kv_hidden_size; ++c) { + token[hidden_size + c] = + 0.1f + 0.25f * static_cast(((b * 5 + s * 2 + c) % 11) + 1); + token[hidden_size + kv_hidden_size + c] = + 0.1f + 0.2f * static_cast(((b * 3 + s + c) % 9) + 1); + } + } + } +} + +// Runs a packed-QKV GQA prefill with do_rotary=1 and the given per-batch +// seqlens_k. Returns the output tensor [batch_size, sequence_length, hidden_size]. +static std::vector RunGQAPackedQKVRotaryPrefill( + int batch_size, + int sequence_length, + int num_heads, + int kv_num_heads, + int head_size, + const std::vector& seqlens_k_data, + const std::vector& packed_qkv_data, + bool use_cuda = false, + bool use_webgpu = false) { + const int hidden_size = num_heads * head_size; + const int kv_hidden_size = kv_num_heads * head_size; + const int qkv_hidden = hidden_size + 2 * kv_hidden_size; + const int total_sequence_length = sequence_length; // prefill: no past + const int half_rotary = head_size / 2; + const int max_seq_len = sequence_length + 8; + + OpTester tester("GroupQueryAttention", 1, onnxruntime::kMSDomain); + tester.AddAttribute("num_heads", static_cast(num_heads)); + tester.AddAttribute("kv_num_heads", static_cast(kv_num_heads)); + tester.AddAttribute("do_rotary", static_cast(1)); + + // Packed QKV: pass through `query` input, leave key/value as optional edges. + tester.AddInput("query", {batch_size, sequence_length, qkv_hidden}, packed_qkv_data); + tester.AddOptionalInputEdge(); // key (signals packed) + tester.AddOptionalInputEdge(); // value (signals packed) + + tester.AddOptionalInputEdge(); // past_key + tester.AddOptionalInputEdge(); // past_value + + tester.AddInput("seqlens_k", {batch_size}, seqlens_k_data); + tester.AddInput("total_sequence_length", {1}, {total_sequence_length}, + /*is_initializer=*/true); + + std::vector cos_cache(max_seq_len * half_rotary); + std::vector sin_cache(max_seq_len * half_rotary); + for (int pos = 0; pos < max_seq_len; ++pos) { + for (int d = 0; d < half_rotary; ++d) { + const float freq = 1.0f / std::pow(10000.0f, 2.0f * static_cast(d) / + static_cast(head_size)); + cos_cache[pos * half_rotary + d] = std::cos(static_cast(pos) * freq); + sin_cache[pos * half_rotary + d] = std::sin(static_cast(pos) * freq); + } + } + tester.AddInput("cos_cache", {max_seq_len, half_rotary}, cos_cache); + tester.AddInput("sin_cache", {max_seq_len, half_rotary}, sin_cache); + + tester.AddOptionalInputEdge(); // position_ids + tester.AddOptionalInputEdge(); // attention_bias + tester.AddOptionalInputEdge(); // head_sink + + const int output_size = batch_size * sequence_length * hidden_size; + tester.AddOutput("output", {batch_size, sequence_length, hidden_size}, + std::vector(output_size, 0.0f)); + const int present_size = batch_size * kv_num_heads * total_sequence_length * head_size; + tester.AddOutput("present_key", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_size, 0.0f)); + tester.AddOutput("present_value", {batch_size, kv_num_heads, total_sequence_length, head_size}, + std::vector(present_size, 0.0f)); + + tester.SetOutputTolerance(1e6f); // We fetch and compare outputs ourselves. + + std::vector> execution_providers; + if (use_cuda) { + execution_providers.push_back(DefaultCudaExecutionProvider()); + } else if (use_webgpu) { + execution_providers.push_back(DefaultWebGpuExecutionProvider()); + } else { + execution_providers.push_back(DefaultCpuExecutionProvider()); + } + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers); + + auto fetches = tester.GetFetches(); + const float* out_data = fetches[0].Get().Data(); + return std::vector(out_data, out_data + output_size); +} + +// Inner helper: builds packed-QKV inputs, computes per-prompt references, runs +// the right-padded batched prefill, and asserts each batch's real-last-token +// output matches its single-prompt reference. Both reference and batched runs +// go through the same EP, so this validates per-batch consistency within each +// EP rather than cross-EP equivalence. +static void RunBatchedRightPaddedRotaryPrefillForEP(bool use_cuda, bool use_webgpu) { + constexpr int batch_size = 3; + constexpr int num_heads = 4; + constexpr int kv_num_heads = 2; + constexpr int head_size = 16; // multiple of 4 for FlashAttention gate; rotary half = 8 + constexpr int hidden_size = num_heads * head_size; + constexpr int kv_hidden_size = kv_num_heads * head_size; + constexpr int qkv_hidden = hidden_size + 2 * kv_hidden_size; + + // Real prompt lengths per batch; max = sequence_length (right-padding extends + // shorter batches up to this length). The bug only manifests when at least + // one batch is shorter than sequence_length. + const std::vector real_lens = {4, 2, 6}; + const int sequence_length = *std::max_element(real_lens.begin(), real_lens.end()); + + std::vector packed_batched; + FillBatchedRightPaddedPackedQKV(batch_size, sequence_length, num_heads, kv_num_heads, + head_size, real_lens, packed_batched); + + // Build single-prompt references by extracting each batch's real-len slice + // and running it as a batch_size=1 prefill (which is known correct). + std::vector> ref_outputs(batch_size); + for (int b = 0; b < batch_size; ++b) { + const int real_len = real_lens[b]; + std::vector packed_single(real_len * qkv_hidden); + for (int s = 0; s < real_len; ++s) { + std::copy_n(&packed_batched[(b * sequence_length + s) * qkv_hidden], qkv_hidden, + &packed_single[s * qkv_hidden]); + } + ref_outputs[b] = RunGQAPackedQKVRotaryPrefill( + /*batch_size=*/1, /*sequence_length=*/real_len, + num_heads, kv_num_heads, head_size, + /*seqlens_k_data=*/{static_cast(real_len - 1)}, + packed_single, use_cuda, use_webgpu); + } + + // Now run all batches together with right-padding. + std::vector seqlens_k_data(batch_size); + for (int b = 0; b < batch_size; ++b) { + seqlens_k_data[b] = static_cast(real_lens[b] - 1); + } + const auto batched_output = RunGQAPackedQKVRotaryPrefill( + batch_size, sequence_length, num_heads, kv_num_heads, head_size, + seqlens_k_data, packed_batched, use_cuda, use_webgpu); + + // Each batch's real-last-token output (used to predict next token) must match + // its single-prompt reference. The tolerance is loose enough for fp16 rounding + // while still catching the underflow bug (which produces values that differ + // by orders of magnitude or are NaN/Inf). + constexpr float tolerance = 5e-3f; + for (int b = 0; b < batch_size; ++b) { + const int real_len = real_lens[b]; + const int q_last = real_len - 1; + const float* batched_last = + batched_output.data() + (b * sequence_length + q_last) * hidden_size; + const float* ref_last = ref_outputs[b].data() + q_last * hidden_size; + for (int c = 0; c < hidden_size; ++c) { + EXPECT_NEAR(batched_last[c], ref_last[c], tolerance) + << "batch " << b << " real_len=" << real_len + << " channel " << c << " mismatch"; + } + } +} + +TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CPU) { + RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/false, /*use_webgpu=*/false); +} + +TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CUDA) { + auto cuda_ep = DefaultCudaExecutionProvider(); + if (!cuda_ep) { + GTEST_SKIP() << "CUDA EP not available"; + } + RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/true, /*use_webgpu=*/false); +} + +TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_WebGPU) { + auto webgpu_ep = DefaultWebGpuExecutionProvider(); + if (!webgpu_ep) { + GTEST_SKIP() << "WebGPU EP not available"; + } + RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/false, /*use_webgpu=*/true); +} + } // namespace test } // namespace onnxruntime From 17ceadcaa08ac84159eea5b5ea27ec96f2cb7ca0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 16 Jun 2026 08:18:02 +0000 Subject: [PATCH 3/5] Fix CPU GQA NaN softmax bug for right-padded batched prompts --- .../contrib_ops/cpu/bert/gqa_attention_base.h | 34 +++++++++++++------ 1 file changed, 23 insertions(+), 11 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h index 12f61cddea18c..59313cf527c91 100644 --- a/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h +++ b/onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h @@ -429,13 +429,19 @@ class GQAAttentionBase { for (size_t seq = 0; seq < static_cast(sequence_length); seq++) { size_t seq_causal_length = causal_past_seqlen + seq + 1; + // Cap effective causal length at total_seqlen so the softmax window stays within + // the region filled by the QK GEMM. For right-padded batched prompts, padding + // positions have seq_causal_length > total_seqlen; without this cap the softmax + // would read uninitialized memory and produce NaN. + const size_t effective_causal_length = std::min(seq_causal_length, total_seqlen); + const bool apply_local = local_window_size_ >= 0 && - seq_causal_length > static_cast(local_window_size_); - const size_t start_off = apply_local ? seq_causal_length - local_window_size_ : 0; - const size_t win_size = apply_local ? local_window_size_ : seq_causal_length; + effective_causal_length > static_cast(local_window_size_); + const size_t start_off = apply_local ? effective_causal_length - local_window_size_ : 0; + const size_t win_size = apply_local ? local_window_size_ : effective_causal_length; if (apply_local) { - for (size_t t = 0; t < seq_causal_length - local_window_size_; t++) { + for (size_t t = 0; t < effective_causal_length - local_window_size_; t++) { sm[t] = 0.f; } } @@ -448,7 +454,7 @@ class GQAAttentionBase { ApplyAttentionBias(sm + start_off, attn_bias + start_off, static_cast(win_size)); } - for (size_t t = seq_causal_length; t < total_seqlen; t++) { + for (size_t t = effective_causal_length; t < total_seqlen; t++) { sm[t] = 0.f; } @@ -1084,15 +1090,21 @@ class GQAAttentionBase { for (size_t seq = 0; seq < sequence_length; seq++) { size_t seq_causal_length = causal_past_seqlen + seq + 1; + // For right-padded batched prompts, padding positions have seq_causal_length > total_seqlen. + // The GEMM only fills columns [0, total_seqlen); beyond that the buffer is uninitialized. + // Cap the effective causal length so the softmax window stays within the filled region, + // preventing NaN from uninitialized memory propagating into the output. + const size_t effective_causal_length = std::min(seq_causal_length, total_seqlen); + const bool should_apply_local_window = local_window_size_ >= 0 && - seq_causal_length > static_cast(local_window_size_); + effective_causal_length > static_cast(local_window_size_); - const size_t start_offset = should_apply_local_window ? seq_causal_length - local_window_size_ : 0; - const size_t window_size = should_apply_local_window ? local_window_size_ : seq_causal_length; + const size_t start_offset = should_apply_local_window ? effective_causal_length - local_window_size_ : 0; + const size_t window_size = should_apply_local_window ? local_window_size_ : effective_causal_length; // Mask everything before local window, if local window should be applied if (should_apply_local_window) { - for (size_t total_seq_id = 0; total_seq_id < seq_causal_length - local_window_size_; total_seq_id++) { + for (size_t total_seq_id = 0; total_seq_id < effective_causal_length - local_window_size_; total_seq_id++) { if constexpr (std::is_same::value) { output_softmax[total_seq_id] = 0.f; } else { @@ -1120,8 +1132,8 @@ class GQAAttentionBase { } } - // set causal [seq_causal_length, total_seqlen) to 0.f - for (size_t total_seq_id = seq_causal_length; total_seq_id < total_seqlen; total_seq_id++) { + // set causal [effective_causal_length, total_seqlen) to 0.f + for (size_t total_seq_id = effective_causal_length; total_seq_id < total_seqlen; total_seq_id++) { if constexpr (std::is_same::value) { output_softmax[total_seq_id] = 0.f; } else { From 9c692b339e042d098de4050d1b5367ddb06066ed Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Tue, 16 Jun 2026 14:02:28 +0000 Subject: [PATCH 4/5] Remove BatchedRightPaddedRotaryPrefill_WebGPU test per review feedback --- .../group_query_attention_op_test.cc | 23 +++++-------------- 1 file changed, 6 insertions(+), 17 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 7ce5eab1ee6b2..4c45cea3bf15e 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2446,8 +2446,7 @@ static std::vector RunGQAPackedQKVRotaryPrefill( int head_size, const std::vector& seqlens_k_data, const std::vector& packed_qkv_data, - bool use_cuda = false, - bool use_webgpu = false) { + bool use_cuda = false) { const int hidden_size = num_heads * head_size; const int kv_hidden_size = kv_num_heads * head_size; const int qkv_hidden = hidden_size + 2 * kv_hidden_size; @@ -2503,8 +2502,6 @@ static std::vector RunGQAPackedQKVRotaryPrefill( std::vector> execution_providers; if (use_cuda) { execution_providers.push_back(DefaultCudaExecutionProvider()); - } else if (use_webgpu) { - execution_providers.push_back(DefaultWebGpuExecutionProvider()); } else { execution_providers.push_back(DefaultCpuExecutionProvider()); } @@ -2520,7 +2517,7 @@ static std::vector RunGQAPackedQKVRotaryPrefill( // output matches its single-prompt reference. Both reference and batched runs // go through the same EP, so this validates per-batch consistency within each // EP rather than cross-EP equivalence. -static void RunBatchedRightPaddedRotaryPrefillForEP(bool use_cuda, bool use_webgpu) { +static void RunBatchedRightPaddedRotaryPrefillForEP(bool use_cuda) { constexpr int batch_size = 3; constexpr int num_heads = 4; constexpr int kv_num_heads = 2; @@ -2553,7 +2550,7 @@ static void RunBatchedRightPaddedRotaryPrefillForEP(bool use_cuda, bool use_webg /*batch_size=*/1, /*sequence_length=*/real_len, num_heads, kv_num_heads, head_size, /*seqlens_k_data=*/{static_cast(real_len - 1)}, - packed_single, use_cuda, use_webgpu); + packed_single, use_cuda); } // Now run all batches together with right-padding. @@ -2563,7 +2560,7 @@ static void RunBatchedRightPaddedRotaryPrefillForEP(bool use_cuda, bool use_webg } const auto batched_output = RunGQAPackedQKVRotaryPrefill( batch_size, sequence_length, num_heads, kv_num_heads, head_size, - seqlens_k_data, packed_batched, use_cuda, use_webgpu); + seqlens_k_data, packed_batched, use_cuda); // Each batch's real-last-token output (used to predict next token) must match // its single-prompt reference. The tolerance is loose enough for fp16 rounding @@ -2585,7 +2582,7 @@ static void RunBatchedRightPaddedRotaryPrefillForEP(bool use_cuda, bool use_webg } TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CPU) { - RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/false, /*use_webgpu=*/false); + RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/false); } TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CUDA) { @@ -2593,15 +2590,7 @@ TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CUDA) { if (!cuda_ep) { GTEST_SKIP() << "CUDA EP not available"; } - RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/true, /*use_webgpu=*/false); -} - -TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_WebGPU) { - auto webgpu_ep = DefaultWebGpuExecutionProvider(); - if (!webgpu_ep) { - GTEST_SKIP() << "WebGPU EP not available"; - } - RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/false, /*use_webgpu=*/true); + RunBatchedRightPaddedRotaryPrefillForEP(/*use_cuda=*/true); } } // namespace test From 84d937047da8cd36c4ec0322e2eb2e5721150446 Mon Sep 17 00:00:00 2001 From: Jiajia Qin Date: Tue, 23 Jun 2026 11:30:13 +0800 Subject: [PATCH 5/5] Address review feedback on batched right-padded rotary prefill test - Reword the tolerance-justification comment to name the actual failure modes per EP (CPU: uninitialized attention-probs reads; WebGPU: u32 underflow on rotary past_seqlen, see PR #29002) instead of calling it an "underflow bug" generically. - Add an explicit std::isfinite check over the full batched output so the regression is caught deterministically regardless of whether the allocator returns zeroed pages. --- .../group_query_attention_op_test.cc | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 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 01d043cbbe607..8a78199ccc54f 100644 --- a/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc +++ b/onnxruntime/test/contrib_ops/group_query_attention_op_test.cc @@ -2567,10 +2567,22 @@ static void RunBatchedRightPaddedRotaryPrefillForEP(GqaTargetEp target_ep) { batch_size, sequence_length, num_heads, kv_num_heads, head_size, seqlens_k_data, packed_batched, target_ep); + // Guard the regression deterministically: every element of the batched output + // (including padding rows) must be finite. The CPU root cause is uninitialized + // attention-probs memory, so a NaN/Inf at any padding position would otherwise + // depend on the allocator returning non-zero pages. + for (size_t i = 0; i < batched_output.size(); ++i) { + ASSERT_TRUE(std::isfinite(batched_output[i])) + << "non-finite value at index " << i << " in batched GQA output"; + } + // Each batch's real-last-token output (used to predict next token) must match - // its single-prompt reference. The tolerance is loose enough for fp16 rounding - // while still catching the underflow bug (which produces values that differ - // by orders of magnitude or are NaN/Inf). + // its single-prompt reference. Tolerance is loose enough for fp16 rounding, + // tight enough to catch the right-padding regressions across EPs: + // - CPU: uninitialized attention-probs reads at padding positions -> NaN. + // - WebGPU: u32 underflow on rotary past_seqlen -> out-of-range cos/sin + // index -> garbage Q/K (see PR #29002). + // Both manifest as NaN/Inf or values differing by orders of magnitude. constexpr float tolerance = 5e-3f; for (int b = 0; b < batch_size; ++b) { const int real_len = real_lens[b];