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 { 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 84d3b18de73fe..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]; @@ -2586,6 +2598,10 @@ static void RunBatchedRightPaddedRotaryPrefillForEP(GqaTargetEp target_ep) { } } +TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CPU) { + RunBatchedRightPaddedRotaryPrefillForEP(GqaTargetEp::kCpu); +} + TEST(GroupQueryAttentionTest, BatchedRightPaddedRotaryPrefill_CUDA) { auto cuda_ep = DefaultCudaExecutionProvider(); if (!cuda_ep) {