Fix CPU Attention softcap/attn_mask ordering (onnx#7867, #7913) + consolidate CUDA spec coverage tests from #28371

.agents/skills/cuda-attention-kernel-patterns/SKILL.md

@@ -36,7 +36,7 @@ Unified Unfused      → RunUnfusedAttention()
 
 **Flash eligibility**: fp16/bf16 only, SM≥8.0 (Ampere+), `head_size == v_head_size`, `head_size <= 256`, no `output_qk`, `attn_mask == nullptr`. Uses `mha_fwd` / `mha_fwd_kvcache`.
 
-**MEA eligibility**: SM50+/53+/80+ by dtype, `head_size <= 1024` and divisible by 8, no `output_qk`. Decode requires `head_size == v_head_size` (for `LaunchConcatNewToPastKV`). Bias stride must satisfy `total_sequence_length % 4 == 0`. GQA with FP32 is excluded (LaunchUngroup only has fp16/bf16 instantiations). Supports `softcap + attn_mask` — CUTLASS applies softcap before bias in kernel tiles, matching ONNX spec ordering (onnx/onnx#7865).
+**MEA eligibility**: SM50+/53+/80+ by dtype, `head_size <= 1024` and divisible by 8 (enforced by `has_memory_efficient_attention`), no `output_qk`. GQA additionally requires `head_size == v_head_size` (for `LaunchUngroup`); decode also requires it (for `LaunchConcatNewToPastKV`). Bias stride must satisfy `total_sequence_length % 4 == 0`. GQA with FP32 is excluded (LaunchUngroup only has fp16/bf16 instantiations). Supports `softcap + attn_mask` — CUTLASS applies softcap before bias in kernel tiles, matching ONNX spec ordering (onnx/onnx#7867, supersedes the now-closed onnx/onnx#7865 issue).
 
 **Unified Unfused Attention**: Always available as the final fallback. Handles both MHA (`num_heads == kv_num_heads`, group=1) and GQA (`num_heads != kv_num_heads`, group>1) via a reshape-Q trick with stride-based cuBLAS batched GEMM (no K/V head replication). Uses FP32 QK scratch for precision. Supports all features:
 - softcap + attn_mask (spec-correct ordering)
@@ -78,15 +78,45 @@ if (parameters.total_sequence_length % min_bias_align != 0) {
 
 ## 4. Softcap Ordering
 
-ONNX spec ordering (onnx/onnx#7865): `QK → scale → softcap → add mask/bias → softmax`
+ONNX Attention opset 23/24 spec ordering (per onnx/onnx#7867, which superseded
+the now-closed onnx/onnx#7865 issue, and onnx/onnx#7913 which swapped
+`qk_matmul_output_mode` values 1 and 2 to align with the corrected pipeline):
 
-- **MEA (CUTLASS)**: Fuses softcap before bias in kernel tile loop (`kernel_forward.h`). Matches spec ordering.
-- **Flash**: Handles softcap natively in `mha_fwd`/`mha_fwd_kvcache` but rejects `attn_mask`, so ordering with mask is moot.
-- **Unfused**: Handles spec-correct ordering in the fused softmax kernel: `QK → scale → softcap → add bias → softmax`.
-
-All three paths apply softcap BEFORE mask/bias. If softcap were applied after masking, `tanh(-inf/sc) = -sc` (finite), leaking probability to masked positions.
+```
+scale * (Q @ K^T)        # stage 0: raw scaled QK
+    |
+softcap (if > 0)         # stage 1: tanh(qk / softcap) * softcap
+    |
++ attn_bias / + attn_mask # stage 2: additive (mask -inf survives to stage 3)
+    |
+softmax                  # stage 3
+    |
+@ V
+```
 
-The unfused path does: `QK → scale → softcap → add bias → softmax` (all fused in `UnfusedSoftmaxKernel`).
+`qk_matmul_output_mode` integer values follow pipeline stage order:
+0 = raw scale*QK, 1 = post-softcap (pre-mask), 2 = post-mask/bias (pre-softmax),
+3 = post-softmax.
+
+CUDA implementation status (all spec-correct):
+- **MEA (CUTLASS)**: `kernel_forward.h` applies softcap inside the score-compute
+  tile loop BEFORE `attn_bias` is added.
+- **Flash**: `mha_fwd` / `mha_fwd_kvcache` handle softcap natively; reject
+  explicit `attn_mask`, so ordering with float mask is moot for this path.
+- **Unfused**: `UnfusedSoftmaxKernel` does `QK -> scale -> softcap -> add bias -> softmax`
+  (all fused).
+
+CPU implementation status: `core/providers/cpu/llm/attention.cc::ComputeAttentionProbs<T>`
+applies softcap BEFORE the mask add (post-fix; pre-fix it inverted the order
+and leaked probability through masked positions).
+
+Why this ordering matters: a -inf in `attn_mask` must survive to softmax. If
+softcap were applied AFTER the mask-add, then `tanh(-inf/softcap) * softcap = -softcap`
+(a finite value), and softmax would assign non-zero weight to the masked
+position — leaking poison V values into the output. The CUDA-side guard tests
+at `test_onnx_attention/test_gqa.py:1501` and `:1761`, and the CPU-side guards
+at `TestONNXAttentionCPUSoftcapMaskOrdering` in the same file, exercise this
+property by combining small softcap, a -inf mask entry, and a poison V value.
 
 ## 5. Grid-Stride Loops for CUDA Kernels
 

cmake/onnxruntime_python.cmake

@@ -476,6 +476,9 @@ if (onnxruntime_BUILD_UNIT_TESTS)
   file(GLOB onnxruntime_python_transformers_test_srcs CONFIGURE_DEPENDS
       "${ONNXRUNTIME_ROOT}/test/python/transformers/*.py"
   )
+  file(GLOB onnxruntime_python_transformers_test_onnx_attention_srcs CONFIGURE_DEPENDS
+      "${ONNXRUNTIME_ROOT}/test/python/transformers/test_onnx_attention/*.py"
+  )
   file(GLOB onnxruntime_python_transformers_testdata_srcs CONFIGURE_DEPENDS
       "${ONNXRUNTIME_ROOT}/test/python/transformers/test_data/models/*.onnx"
   )
@@ -629,6 +632,7 @@ add_custom_command(
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/quantization/neural_compressor
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/quantization
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers
+  COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_onnx_attention
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_data/models
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_data/models/whisper
   COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/eager_test
@@ -829,6 +833,9 @@ if (onnxruntime_BUILD_UNIT_TESTS)
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_transformers_test_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/transformers/
+    COMMAND ${CMAKE_COMMAND} -E copy
+        ${onnxruntime_python_transformers_test_onnx_attention_srcs}
+        $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_onnx_attention/
     COMMAND ${CMAKE_COMMAND} -E copy
         ${onnxruntime_python_transformers_testdata_srcs}
         $<TARGET_FILE_DIR:${build_output_target}>/transformers/test_data/models/

onnxruntime/core/providers/cpu/llm/attention.cc

@@ -94,6 +94,27 @@ inline void ComputeAttentionSoftcapInplace(MLFloat16* scores, int sequence_lengt
   }
 }
 
+// In-place elementwise add: scores[i] += addend[i].
+//
+// Used to apply attn_mask / attn_bias to the QK scores after softcap, per the
+// ONNX Attention v23/24 spec ordering (onnx/onnx#7867 + #7913). For float,
+// delegates to MLAS. For MLFloat16, uses a portable scalar fallback because
+// MlasEltwiseAdd<MLAS_FP16> requires the per-platform EltwiseDispatch->Add_Fp16
+// kernel slot to be populated, and only the ARM NEON build provides it
+// (see onnxruntime/core/mlas/lib/eltwise.cpp:92-103); x86 and other targets
+// would throw at runtime.
+template <typename T>
+inline void AddInPlace(T* scores, const T* addend, size_t count) {
+  MlasEltwiseAdd<T>(addend, scores, scores, count);
+}
+
+template <>
+inline void AddInPlace<MLFloat16>(MLFloat16* scores, const MLFloat16* addend, size_t count) {
+  for (size_t i = 0; i < count; ++i) {
+    scores[i] = MLFloat16(scores[i].ToFloat() + addend[i].ToFloat());
+  }
+}
+
 // Dispatches a GEMM operation across float and MLFloat16 types.
 //   C = alpha * op(A) * op(B) + beta * C
 //
@@ -199,10 +220,10 @@ Attention<T>::Attention(const OpKernelInfo& info) : AttentionBase<T>(info) {
                                : attention_helper::QKMatMulOutputMode::kNone;
   ORT_ENFORCE(qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kNone ||
                   qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kQK ||
-                  qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kQKMask ||
-                  qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kQKSoftCap ||
-                  qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kQKSoftMax,
-              "qk_matmul_output_mode must be 0, 1, 2, or 3.");
+                  qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kPostSoftCap ||
+                  qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kPostMaskBias ||
+                  qk_matmul_output_mode_ == attention_helper::QKMatMulOutputMode::kPostSoftMax,
+              "qk_matmul_output_mode must be -1 (absent), 0, 1, 2, or 3.");
   // The default scale depends on the input dimensions. It is set to nan to indicate that it should be computed.
   scale_ = info.GetAttrOrDefault<float>("scale", std::numeric_limits<T>::quiet_NaN());
   softcap_ = info.GetAttrOrDefault<float>("softcap", 0.0f);
@@ -402,16 +423,6 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
 
       T* output = attention_probs + output_offset;
       T* out_qk = output_qk == nullptr ? nullptr : output_qk + output_offset;
-      float beta;
-
-      if (mask_data != nullptr &&
-          (out_qk == nullptr || parameters.qk_matmul_output_mode != attention_helper::QKMatMulOutputMode::kQK)) {
-        // Broadcast mask data: SxT -> SxT
-        memcpy(output, mask_data + mask_data_offset, probs_matrix_bytes);
-        beta = 1;
-      } else {
-        beta = 0;
-      }
 
       // handling GQA
       std::ptrdiff_t head_ki = head_i * parameters.kv_num_heads / parameters.q_num_heads;
@@ -431,10 +442,41 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
       }
 
       // Compute Q*K' + AttentionMask
+      //
+      // ONNX Attention v23/24 (per onnx/onnx#7867 + #7913) requires the mask/bias
+      // to be applied AFTER softcap, otherwise -inf mask values get squashed by
+      // tanh into -c, leaking probability through softmax onto masked positions.
+      //
+      // When softcap is disabled, mask add commutes with the (no-op) softcap, so we
+      // follow the original code path verbatim:  fold the mask add into the GEMM as
+      // `beta = 1` (preload mask, accumulate via FMA).  This preserves the FMA-fused
+      // numerics that pre-spec-fix tests were calibrated against.
+      // When softcap is active, we run GEMM with `beta = 0`, apply softcap inplace,
+      // then add the mask explicitly via AddInPlace.
+      //
+      // The fold is also skipped when the caller wants a kQK / kPostSoftCap snapshot,
+      // so the snapshot reflects the raw / post-softcap QK without the mask folded in.
+      //
       //                     original                 transposed             each iteration
       // A: Q                (B x N x) S x H          (B x N x) S x H        S x H
       // B: K'               (B x N x) T x H          (B x N x) H x T        H x T
       // C: attention_probs  (B x N x) S x T          (B x N x) S x T        S x T
+      const bool softcap_active = (parameters.softcap > 0.0f);
+      const bool snapshot_needs_pre_mask =
+          out_qk != nullptr &&
+          (parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQK ||
+           parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kPostSoftCap);
+      const bool fold_mask_into_gemm =
+          (mask_data != nullptr) && !softcap_active && !snapshot_needs_pre_mask;
+      float beta;
+      if (fold_mask_into_gemm) {
+        // Broadcast mask data: SxT -> SxT
+        memcpy(output, mask_data + mask_data_offset, probs_matrix_bytes);
+        beta = 1;
+      } else {
+        beta = 0;
+      }
+
       const T* q_ptr = parameters.transpose_output
                            ? Q + q_input_chunk_length * parameters.q_num_heads * batch_i + head_i * parameters.head_size
                            : Q + q_input_chunk_length * i;
@@ -452,17 +494,44 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
                     parameters.q_sequence_length, parameters.total_sequence_length, parameters.head_size,
                     alpha, q_ptr, q_lda, k_ptr, k_ldb, beta, output, parameters.total_sequence_length,
                     &mlas_backend_kernel_selector_config_);
+
+      // Snapshot kQK (raw scale*Q*K^T): only reachable when fold path was skipped.
       if (out_qk != nullptr &&
-          (parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQKMask ||
-           parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQK)) {
+          parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQK) {
         memcpy(out_qk, output, SafeInt<size_t>(probs_matrix_size) * sizeof(T));
-        if (mask_data != nullptr && parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQK) {
-          // We need to add the bias we could not add because out_qk was requested without the mask.
-          // This can be optimized with vectorized add using MlasAddFloat32x4.
-          MlasEltwiseAdd(output, mask_data + mask_data_offset, output, probs_matrix_size);
+      }
+
+      if (softcap_active) {
+        // Softcap path (mask was NOT folded into GEMM since beta=0 above).
+        if constexpr (std::is_same<T, float>::value) {
+          ComputeAttentionSoftcapInplace(output, static_cast<int>(probs_matrix_size), parameters.softcap);
+        } else if constexpr (std::is_same<T, MLFloat16>::value) {
+          ComputeAttentionSoftcapInplace(output, static_cast<int>(probs_matrix_size), MLFloat16(parameters.softcap));
+        } else {
+          ORT_THROW("Unsupported data type for ComputeAttentionSoftcapInplace: ",
+                    DataTypeImpl::ToString(DataTypeImpl::GetType<T>()));
         }
       }
+
+      // Snapshot kPostSoftCap (post-softcap, pre-mask/bias).  When softcap is disabled
+      // this equals raw scale*Q*K^T (kQK).  Reachable only when fold was skipped above.
+      if (out_qk != nullptr &&
+          parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kPostSoftCap) {
+        memcpy(out_qk, output, SafeInt<size_t>(probs_matrix_size) * sizeof(T));
+      }
+
+      // Add mask explicitly when it wasn't folded into GEMM (single source of truth:
+      // a non-zero `beta` is exactly the case where the mask was preloaded into C).
+      if (mask_data != nullptr && !fold_mask_into_gemm) {
+        AddInPlace(output, mask_data + mask_data_offset, probs_matrix_size);
+      }
+
       // Apply nonpad_kv_seqlen masking (Opset 24+): mask out KV positions >= valid length per batch.
+      // Done AFTER softcap+mask so the masked positions hold the `mask_filter_value<T>()` sentinel
+      // (`std::numeric_limits<T>::lowest()` for floats, `MLFloat16::MinValue` for fp16 — see
+      // `onnxruntime/core/providers/cpu/llm/attention.h`). The CPU softmax uses this finite sentinel
+      // (not IEEE -inf) because MLAS' softmax kernel expects only finite inputs; the value is small
+      // enough relative to any softcap-saturated score that the corresponding softmax weight is 0.
       if (parameters.has_nonpad_kv_seqlen) {
         int valid_kv_len = static_cast<int>(parameters.nonpad_kv_seqlen_data[batch_i]);
         for (int s = 0; s < parameters.q_sequence_length; ++s) {
@@ -471,24 +540,19 @@ void AttentionBase<T>::ComputeAttentionProbs(T* attention_probs,
                     mask_filter_value<T>());
         }
       }
-      if (parameters.softcap > 0.0f) {
-        if constexpr (std::is_same<T, float>::value) {
-          ComputeAttentionSoftcapInplace(output, static_cast<int>(probs_matrix_size), parameters.softcap);
-        } else if constexpr (std::is_same<T, MLFloat16>::value) {
-          ComputeAttentionSoftcapInplace(output, static_cast<int>(probs_matrix_size), MLFloat16(parameters.softcap));
-        } else {
-          ORT_THROW("Unsupported data type for ComputeAttentionSoftcapInplace: ",
-                    DataTypeImpl::ToString(DataTypeImpl::GetType<T>()));
-        }
-      }
-      if (out_qk != nullptr && parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQKSoftCap) {
+
+      // Snapshot kPostMaskBias (post-mask/bias, pre-softmax).
+      if (out_qk != nullptr &&
+          parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kPostMaskBias) {
         memcpy(out_qk, output, SafeInt<size_t>(probs_matrix_size) * sizeof(T));
       }
+
       ComputeAttentionSoftmaxInplace(output, parameters.q_sequence_length, parameters.total_sequence_length, nullptr, allocator);
 
-      if (output_qk != nullptr && parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kQKSoftMax) {
-        memcpy(output_qk + output_offset, output,
-               SafeInt<size_t>(parameters.q_sequence_length) * parameters.total_sequence_length * sizeof(T));
+      // Snapshot kPostSoftMax (post-softmax).
+      if (out_qk != nullptr &&
+          parameters.qk_matmul_output_mode == attention_helper::QKMatMulOutputMode::kPostSoftMax) {
+        memcpy(out_qk, output, SafeInt<size_t>(probs_matrix_size) * sizeof(T));
       }
     }
   });

onnxruntime/core/providers/cpu/llm/attention_parameters.h

@@ -9,13 +9,41 @@ namespace onnxruntime {
 // Declares enum QKMatMulOutputMode and struct AttentionParameters inside namespace onnxruntime::attention_helper.
 namespace attention_helper {
 
-// enum equivalent to the onnx defintion of qk_matmul_output_mode
+// enum equivalent to the ONNX definition of qk_matmul_output_mode.
+//
+// IMPORTANT — ORT intentionally LEADS the bundled ONNX submodule on this
+// numbering. The bundled ONNX (cmake/external/onnx) is currently v1.21.0,
+// which was tagged BEFORE onnx/onnx#7913 was merged upstream; that bundled
+// schema therefore still reflects the OLD value mapping (1 = post-mask/bias,
+// 2 = post-softcap). This implementation already follows the corrected
+// post-#7913 numbering so that ORT will be spec-correct as soon as the next
+// ONNX release (v1.22) is bundled, with no behavior change required at
+// that point. As a side effect, the as-shipped opset-23 ONNX backend node
+// tests under cmake/external/onnx that pin the OLD numbering will fail
+// against this implementation until the submodule is bumped — this is
+// expected; see PR description for details.
+//
+// Mode integer numbering follows the ONNX Attention v23/24 pipeline stage
+// order, per onnx/onnx#7867 (which corrected the ordering to apply softcap
+// BEFORE bias/mask add) and onnx/onnx#7913 (which swapped the integer
+// values of modes 1 and 2 to align with the corrected pipeline):
+//
+//   stage 0: scale * (Q @ K^T)
+//   stage 1: softcap (if > 0)
+//   stage 2: + attn_bias / + attn_mask
+//   stage 3: softmax
+//
+// TODO(onnx-v1.22): when cmake/external/onnx is bumped to v1.22+ which
+// includes ONNX PRs #7867 + #7913, drop the "ORT leads ONNX" caveat above
+// and re-enable the corresponding ONNX backend node tests by removing the
+// skip blocks in onnxruntime/test/onnx/TestCase.cc::GetBrokenTests() and
+// onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc.
 enum QKMatMulOutputMode {
-  kNone = -1,      // No output Q*K
-  kQK = 0,         // Output Q*K
-  kQKMask = 1,     // Output Q*K + Mask
-  kQKSoftCap = 2,  // Output SoftCap(Q*K + Mask)
-  kQKSoftMax = 3,  // Output SoftMax(SoftCap(Q*K + Mask))
+  kNone = -1,         // No optional output.
+  kQK = 0,            // Raw scale * Q @ K^T (pre-softcap).
+  kPostSoftCap = 1,   // Post-softcap, pre-mask/bias.
+  kPostMaskBias = 2,  // Post-mask/bias, pre-softmax.
+  kPostSoftMax = 3,   // Post-softmax.
 };
 
 // Parameters deduced from node attributes and inputs/outputs.