Fix CPU GQA NaN output for right-padded batched prompts with rotary embeddings

onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h

@@ -429,13 +429,19 @@
           for (size_t seq = 0; seq < static_cast<size_t>(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<size_t>(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<size_t>(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 @@
               ApplyAttentionBias(sm + start_off, attn_bias + start_off, static_cast<int>(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 @@
         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<size_t>(local_window_size_);
+                                                 effective_causal_length > static_cast<size_t>(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<U, float>::value) {
                 output_softmax[total_seq_id] = 0.f;
               } else {
@@ -1120,8 +1132,8 @@
             }
           }
 
-          // 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<U, float>::value) {
               output_softmax[total_seq_id] = 0.f;
             } else {

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) {