Fix softcap/attn_mask ordering in CUDA Attention to match ONNX v24 spec and CPU EP

.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, no `output_qk`. MEA requires `head_size == v_head_size` in two internal paths: `LaunchConcatNewToPastKV` (decode with past_key) and `LaunchUngroup` (GQA head expansion). Both share the constraint because each uses a single `head_size` parameter for the K and V tiles. 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 bias before softcap in kernel tiles, matching the ONNX opset 24 reference function (`cmake/external/onnx/onnx/defs/nn/defs.cc` lines 3657-3675: bias-add then softcap then softmax).
 
 **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,38 @@ 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`
-
-- **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.
-
-The unfused path does: `QK → scale → softcap → add bias → softmax` (all fused in `UnfusedSoftmaxKernel`).
+ONNX opset 24 Attention spec ordering: `QK → scale → +bias/mask → softcap → softmax`.
+The reference function in `cmake/external/onnx/onnx/defs/nn/defs.cc` (lines 3322-3341
+ASCII pipeline; lines 3657-3675 reference graph) is unambiguous: `Add(QK, AttnBias)`
+runs *before* the `softcap` block (`Div`/`Tanh`/`Mul`). The CPU EP at
+`core/providers/cpu/llm/attention.cc` matches this. All CUDA kernels must too.
+
+- **Unfused** (`unfused_attention.cu`): the fused `UnfusedSoftmaxKernel` does
+  `qk*scale → +bias → softcap → softmax`. Three identical orderings across the
+  max/sum/normalize passes.
+- **MEA (CUTLASS)** (`cutlass_fmha/kernel_forward.h`): in the score-compute block
+  the order is `scale → bias-add → softcap`. Bias is loaded gmem→smem and added
+  into the accumulator register fragment **before** the `1/softcap`,`fast_tanh`,
+  `*softcap` triple.
+- **Flash**: handles softcap natively in `mha_fwd`/`mha_fwd_kvcache` but rejects
+  `attn_mask` at eligibility, so bias/softcap ordering with mask never executes
+  on this path.
+
+Why this ordering matters: softcap saturates scores into `[-cap, +cap]`. Adding
+bias *after* saturation lets bias values escape the cap (a `-10` additive mask
+combined with `cap=5` becomes `≈ -5 + (-10) = -15` instead of the spec's
+`tanh((scaled_qk - 10)/5)*5 ≈ -5`). For hard masks (`bias = lowest()`) both
+orderings drive masked positions to ~0 softmax weight, but for moderate float
+biases (e.g. ALiBi-style position biases combined with softcap, as in some
+Gemma-family models) the two orderings produce materially different softmax
+distributions.
+
+Historical note: an earlier version of this skill cited `onnx/onnx#7865` and
+asserted softcap-then-bias was correct. That claim was wrong relative to the
+v24 spec text bundled in this repo and was inverted in late 2025 to match the
+ONNX v24 spec (`cmake/external/onnx/onnx/defs/nn/defs.cc` lines 3657-3675).
+If you find any remaining "softcap before bias" references in ORT, treat them
+as bugs.
 
 ## 5. Grid-Stride Loops for CUDA Kernels
 
@@ -176,7 +199,7 @@ Run tests with `disable_cpu=false` to always validate against CPU. The C++ test
 
 | File | Purpose |
 |------|---------|
-| `contrib_ops/cuda/bert/unfused_attention.cu` | Unified unfused attention: QK GEMM (FP32), fused softmax kernel (scale+softcap+bias+causal), V GEMM. Handles MHA and GQA. |
+| `contrib_ops/cuda/bert/unfused_attention.cu` | Unified unfused attention: QK GEMM (FP32), fused softmax kernel (scale+bias+softcap+causal, in spec order), V GEMM. Handles MHA and GQA. |
 | `contrib_ops/cuda/bert/unfused_attention.h` | `UnfusedAttentionParams`, `LaunchUnfusedAttention`, workspace size |
 | `contrib_ops/cuda/bert/attention_impl.cu` | Legacy unfused `QkvToContext` (contrib MHA only). Also `ApplySoftcap`, `ConcatPastToPresent` |
 | `contrib_ops/cuda/bert/attention_softmax.h` | CUDA softmax kernels (`ComputeSoftmax`, `ComputeSoftmaxWithRawMask`) — used by legacy contrib path |

onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h

@@ -859,16 +859,12 @@ struct AttentionKernel {
             cutlass::multiplies<typename MM0::Mma::FragmentC>()(p.scale, accum);
       }
 
-      // apply softcap if applicable
-      if (p.softcap > 0.0) {
-        accum = cutlass::multiplies<typename MM0::Mma::FragmentC>()(1.0 / p.softcap, accum);
-        for (int i = 0; i < accum.size(); ++i) {
-          accum[i] = cutlass::fast_tanh(accum[i]);
-        }
-        accum = cutlass::multiplies<typename MM0::Mma::FragmentC>()(p.softcap, accum);
-      }
-
-      // apply attention bias if applicable
+      // apply attention bias if applicable (BEFORE softcap, per ONNX opset 24
+      // Attention spec: scale*Q@K^T + bias -> softcap -> softmax. See
+      // cmake/external/onnx/onnx/defs/nn/defs.cc reference function lines
+      // 3657-3675.) Bias is loaded from gmem to smem and added into the
+      // accumulator register fragment before softcap saturates the combined
+      // pre-softmax scores.
       if (kSupportsBias && p.attn_bias_ptr != nullptr) {
         // load bias tile Bij into shared memory
         typename MM0::BiasLoader::GmemTileIterator bias_iter(
@@ -899,6 +895,15 @@ struct AttentionKernel {
             [&](int accum_m) {});
       }
 
+      // apply softcap if applicable (AFTER bias, per ONNX opset 24 spec)
+      if (p.softcap > 0.0) {
+        accum = cutlass::multiplies<typename MM0::Mma::FragmentC>()(1.0 / p.softcap, accum);
+        for (int i = 0; i < accum.size(); ++i) {
+          accum[i] = cutlass::fast_tanh(accum[i]);
+        }
+        accum = cutlass::multiplies<typename MM0::Mma::FragmentC>()(p.softcap, accum);
+      }
+
       // Mask out last if causal
       // This is only needed if upper-right corner of current query / key block
       // intersects the mask Coordinates of upper-right corner of current block

onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu

@@ -73,10 +73,13 @@ __global__ void ScaledCopyQkKernel(
 // ---------------------------------------------------------------------------
 // Softmax kernel: reads FP32 QK scores, writes T softmax output.
 //
-// Applies (in this order):
+// Applies (in this order, matching the ONNX opset 24 Attention reference function:
+// QKAttnWeight = MatMul(Q, K^T); QKAttnWeightWithBias = Add(QKAttnWeight, AttnBias);
+// then softcap; then softmax). See cmake/external/onnx/onnx/defs/nn/defs.cc lines
+// 3322-3341 (pipeline diagram) and 3657-3675 (reference function).
 //   1. scale: x = scale * qk
-//   2. softcap (if > 0):            x = softcap * tanh(x / softcap)
-//   3. attn_bias (if provided):     x += bias
+//   2. attn_bias (if provided):     x += bias        (BEFORE softcap, per spec)
+//   3. softcap (if > 0):            x = softcap * tanh(x / softcap)
 //   4. mask (causal + sliding window + per-batch seqlens_k)
 //   5. stable softmax across [start, end) for each row
 //
@@ -152,16 +155,18 @@ __global__ void UnfusedSoftmaxKernel(
   __shared__ float s_inv_sum;
 
   // Pass 1: compute max of masked values.
+  // Score order: scale → +bias (if any) → softcap (if any). bias-then-softcap
+  // matches the ONNX opset 24 spec — see header doc and SKILL.md §4.
   float thread_max = -CUDART_INF_F;
   for (int i = threadIdx.x; i < total_kv_length; i += TPB) {
     if (i < start || i >= end) continue;
     float x = qk_in[row_offset + i] * scale;
-    if (softcap > 0.f) {
-      x = softcap * tanhf(x / softcap);
-    }
     if (has_bias) {
       x += ToFloat(attn_bias[bias_row_offset + i]);
     }
+    if (softcap > 0.f) {
+      x = softcap * tanhf(x / softcap);
+    }
     if (x > thread_max) thread_max = x;
   }
   float block_max;
@@ -182,16 +187,17 @@ __global__ void UnfusedSoftmaxKernel(
   }
 
   // Pass 2: compute sum of exp.
+  // Score order: scale → +bias → softcap (bias-then-softcap, see SKILL.md §4).
   float thread_sum = 0.f;
   for (int i = threadIdx.x; i < total_kv_length; i += TPB) {
     if (i < start || i >= end) continue;
     float x = qk_in[row_offset + i] * scale;
-    if (softcap > 0.f) {
-      x = softcap * tanhf(x / softcap);
-    }
     if (has_bias) {
       x += ToFloat(attn_bias[bias_row_offset + i]);
     }
+    if (softcap > 0.f) {
+      x = softcap * tanhf(x / softcap);
+    }
     thread_sum += expf(x - s_max);
   }
   float block_sum;
@@ -204,16 +210,17 @@ __global__ void UnfusedSoftmaxKernel(
   __syncthreads();
 
   // Pass 3: write softmax output in type T.
+  // Score order: scale → +bias → softcap (bias-then-softcap, see SKILL.md §4).
   for (int i = threadIdx.x; i < total_kv_length; i += TPB) {
     float y = 0.f;
     if (i >= start && i < end) {
       float x = qk_in[row_offset + i] * scale;
-      if (softcap > 0.f) {
-        x = softcap * tanhf(x / softcap);
-      }
       if (has_bias) {
         x += ToFloat(attn_bias[bias_row_offset + i]);
       }
+      if (softcap > 0.f) {
+        x = softcap * tanhf(x / softcap);
+      }
       y = expf(x - s_max) * s_inv_sum;
     }
     softmax_out[row_offset + i] = T(y);

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

@@ -543,9 +543,11 @@ Status Attention<T>::RunFlashAttention(
 //   Eligibility: see has_memory_efficient_attention() (SM50+/53+/80+ by dtype,
 //                head_size <= 1024, head_size divisible by 8), plus: no output_qk, bias stride alignment.
 //   Note: softcap is forwarded to the MEA kernel via p.softcap. CUTLASS applies
-//   softcap before bias (fused in kernel tiles), matching ONNX spec ordering
-//   (onnx/onnx#7865): QK → softcap → mask/bias → softmax. softmax_precision
-//   is inherently satisfied (cutlass FMHA accumulates softmax in FP32).
+//   bias before softcap (fused in kernel tiles), matching ONNX opset 24
+//   Attention spec ordering: QK -> scale -> +bias -> softcap -> softmax.
+//   See cmake/external/onnx/onnx/defs/nn/defs.cc reference function lines
+//   3657-3675.
+//   softmax_precision is inherently satisfied (cutlass FMHA accumulates softmax in FP32).
 //
 template <typename T>
 Status Attention<T>::RunMemoryEfficientAttention(

onnxruntime/test/providers/cpu/llm/attention_op_test.cc

@@ -2975,5 +2975,146 @@ TEST(AttentionTest, Attention4DCausalSquareNoPast) {
   );
 }
 
+// Regression test for the softcap / attn_mask ordering bug — spec ordering is
+// scale*Q@K^T + bias -> softcap -> softmax (see ONNX opset 24 Attention
+// reference function in cmake/external/onnx/onnx/defs/nn/defs.cc lines
+// 3322-3341 and 3657-3675). Pre-fix CUDA kernels applied softcap *before*
+// bias, which let large negative bias values escape the softcap saturation.
+//
+// This fp32 / head_size=8 / kv_seq=2 config has total_sequence_length=2 which
+// fails MEA's bias-stride %4 alignment, so it routes to the unified unfused
+// kernel on CUDA. The CPU EP runs in the same RunTest4D loop. Both must
+// produce the spec-correct output.
+//
+// Inputs are designed so the two orderings produce materially different Y:
+//   spec:  Y[i] ≈ 1.0320  (softcap saturates the masked combined score to ≈ -5)
+//   buggy: Y[i] ≈ 1.0002  (bias of -10 added after softcap dominates softmax)
+// A pre-fix run fails the 3e-5f tolerance; the post-fix run passes.
+TEST(AttentionTest, Attention4DSoftcapFloatBias_Unfused_FP32) {
+  // CUDA-EP regression: skip on CPU-only builds so the test does not silently
+  // pass via CPU alone. CPU is left enabled below for parity verification when
+  // CUDA is available.
+  if (!HasCudaEnvironment(0)) {
+    return;
+  }
+  int batch_size = 1;
+  int q_num_heads = 1;
+  int q_sequence_length = 1;
+  int head_size = 8;
+  int kv_sequence_length = 2;  // total_sequence_length=2, %4!=0 → MEA rejected → unfused
+  int kv_num_heads = 1;
+  int v_head_size = 8;
+  int past_sequence_length = 0;
+
+  std::vector<float> q(static_cast<size_t>(batch_size * q_num_heads * q_sequence_length * head_size), 0.01f);
+  std::vector<float> k(static_cast<size_t>(batch_size * kv_num_heads * kv_sequence_length * head_size), 0.01f);
+  std::vector<float> v(static_cast<size_t>(batch_size * kv_num_heads * kv_sequence_length * v_head_size));
+  // V row 0 = 1.0, V row 1 = 5.0 (across all v_head_size dims).
+  for (int s = 0; s < kv_sequence_length; ++s) {
+    float val = (s == 0) ? 1.0f : 5.0f;
+    for (int h = 0; h < v_head_size; ++h) {
+      v[static_cast<size_t>(s * v_head_size + h)] = val;
+    }
+  }
+
+  // Float attn_mask = additive bias. Position 1 has -10 bias.
+  std::vector<float> attn_mask = {0.0f, -10.0f};
+
+  // Expected Y derivation (spec ordering, per ONNX opset 24 reference function:
+  // QK -> scale -> +bias -> softcap -> softmax -> @V):
+  //   scale       = 1/sqrt(8) ≈ 0.353553
+  //   Q[i]*K[j]   = 8 * (0.01 * 0.01) = 8e-4   for both kv positions j=0,1
+  //   scaled      = 0.353553 * 8e-4 ≈ 2.828427e-4
+  //   +bias       = [2.828427e-4, 2.828427e-4 + (-10)] ≈ [2.828e-4, -9.99972]
+  //   softcap=5   = 5*tanh(x/5) → [2.828e-4 (≈linear here), -4.99977 (saturated)]
+  //   softmax     ≈ [0.99326, 0.00674]
+  //   Y           = 0.99326 * 1.0 + 0.00674 * 5.0 ≈ 1.0319962
+  // Pre-fix (softcap-then-bias) collapses softcap before bias dominates the
+  // softmax, yielding Y ≈ 1.0002 — delta 0.032, well outside 3e-5.
+  std::vector<float> y(static_cast<size_t>(batch_size * q_num_heads * q_sequence_length * v_head_size), 1.0319962f);
+
+  // disable_dml=true: DML EP does not implement opset-24 Attention with float
+  // additive bias + softcap; excluding it keeps this test focused on
+  // CPU<->CUDA parity for the spec ordering fix.
+  RunTest4D(batch_size, q_num_heads, q_sequence_length, head_size, kv_sequence_length, kv_num_heads, v_head_size, past_sequence_length,
+            q, k, v, attn_mask, std::initializer_list<bool>(), std::vector<float>(), std::vector<float>(),
+            -1, -1, std::numeric_limits<float>::quiet_NaN(), 5.0f, -1, TensorType::kFloat,  // is_causal, qk_matmul_output_mode, scale, softcap, softmax_precision, tensor_type
+            y, std::vector<float>(), std::vector<float>(), std::vector<float>(),
+            false, false, true  // disable_cpu, disable_cuda, disable_dml
+  );
+}
+
+// Companion test forcing the CUTLASS MEA path: fp16, kv_seq=4 so
+// total_sequence_length=4 satisfies MEA's bias-stride alignment, and Flash is
+// disabled via env var so the dispatcher must pick MEA. Verifies the CUTLASS
+// score-compute loop also applies bias before softcap (per spec).
+TEST(AttentionTest, Attention4DSoftcapFloatBias_MEA_FP16) {
+  if (!HasCudaEnvironment(530)) {
+    return;
+  }
+  // Force MEA: disable Flash. (MEA is preferred over unfused when eligible.)
+  ScopedEnvironmentVariables scoped_env_vars{
+      EnvVarMap{{onnxruntime::contrib::attention::kDisableFlashAttention, "1"}}};
+
+  int batch_size = 1;
+  int q_num_heads = 1;
+  int q_sequence_length = 1;
+  int head_size = 8;
+  int kv_sequence_length = 4;  // total=4, %4==0 → MEA bias-stride OK
+  int kv_num_heads = 1;
+  int v_head_size = 8;
+
+  OpTester test("Attention", 24, onnxruntime::kOnnxDomain);
+  test.AddAttribute<float>("softcap", 5.0f);
+
+  std::vector<float> q(static_cast<size_t>(batch_size * q_num_heads * q_sequence_length * head_size), 0.01f);
+  std::vector<float> k(static_cast<size_t>(batch_size * kv_num_heads * kv_sequence_length * head_size), 0.01f);
+  std::vector<float> v(static_cast<size_t>(batch_size * kv_num_heads * kv_sequence_length * v_head_size));
+  // V[s] = (1.0, 5.0, 3.0, 2.0) for s = 0..3, replicated across head dim.
+  const float v_vals[4] = {1.0f, 5.0f, 3.0f, 2.0f};
+  for (int s = 0; s < kv_sequence_length; ++s) {
+    for (int h = 0; h < v_head_size; ++h) {
+      v[static_cast<size_t>(s * v_head_size + h)] = v_vals[s];
+    }
+  }
+  // attn_mask: positions 1 and 2 strongly biased.
+  std::vector<float> attn_mask = {0.0f, -10.0f, -10.0f, 0.0f};
+
+  // Expected Y derivation (spec ordering, fp16; mirrors the FP32 test):
+  //   scale       = 1/sqrt(8) ≈ 0.353553
+  //   Q[i]*K[j]   = 8 * (0.01 * 0.01) = 8e-4   for all 4 kv positions
+  //   scaled      ≈ 2.828e-4
+  //   +bias       ≈ [2.828e-4, -9.99972, -9.99972, 2.828e-4]
+  //   softcap=5   ≈ [2.828e-4, -4.99977, -4.99977, 2.828e-4]
+  //   softmax     ≈ [0.49831, 0.00169, 0.00169, 0.49831]
+  //   Y           = 0.49831*1 + 0.00169*5 + 0.00169*3 + 0.49831*2 ≈ 1.5202
+  //   (rounded to 1.5200 within fp16 + 5e-3 tolerance)
+  // Pre-fix (softcap-then-bias) yields Y ≈ 1.5001 — delta ≈ 0.020, outside 5e-3.
+
+  test.AddInput<MLFloat16>("Q", {batch_size, q_num_heads, q_sequence_length, head_size}, ToFloat16(q));
+  test.AddInput<MLFloat16>("K", {batch_size, kv_num_heads, kv_sequence_length, head_size}, ToFloat16(k));
+  test.AddInput<MLFloat16>("V", {batch_size, kv_num_heads, kv_sequence_length, v_head_size}, ToFloat16(v));
+  test.AddInput<MLFloat16>("attn_mask", {q_sequence_length, kv_sequence_length}, ToFloat16(attn_mask));
+  test.AddOptionalInputEdge<MLFloat16>();  // past_key
+  test.AddOptionalInputEdge<MLFloat16>();  // past_value
+
+  // Spec-ordered expected: ≈ 1.5200. Pre-fix CUDA produced ≈ 1.5001 — a delta
+  // of ≈ 0.02 that exceeds the 5e-3 tolerance below.
+  std::vector<float> expected_y(
+      static_cast<size_t>(batch_size * q_num_heads * q_sequence_length * v_head_size), 1.5200f);
+  test.AddOutput<MLFloat16>("Y", {batch_size, q_num_heads, q_sequence_length, v_head_size},
+                            ToFloat16(expected_y), false, 0, 5e-3f);
+  test.AddOptionalOutputEdge<MLFloat16>();  // present_key
+  test.AddOptionalOutputEdge<MLFloat16>();  // present_value
+
+  std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
+  execution_providers.push_back(DefaultCudaExecutionProvider());
+  // TODO(pr-followup): Assert MEA is the actually-dispatched kernel (not a
+  // silent fall-through to unfused). Current dispatcher does not expose a
+  // public counter; the env var + shape selection is the strongest indirect
+  // guarantee available without adding test-only instrumentation.
+  test.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers);
+}
+
 }  // namespace test
 }  // namespace onnxruntime