diff --git a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md b/.agents/skills/cuda-attention-kernel-patterns/SKILL.md index 5325a1bf22bdc..6476f4ef2f1ef 100644 --- a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md +++ b/.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 | diff --git a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h index af7a202ef559f..8ba9c5afd79bb 100644 --- a/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h +++ b/onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/kernel_forward.h @@ -859,16 +859,12 @@ struct AttentionKernel { cutlass::multiplies()(p.scale, accum); } - // apply softcap if applicable - if (p.softcap > 0.0) { - accum = cutlass::multiplies()(1.0 / p.softcap, accum); - for (int i = 0; i < accum.size(); ++i) { - accum[i] = cutlass::fast_tanh(accum[i]); - } - accum = cutlass::multiplies()(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()(1.0 / p.softcap, accum); + for (int i = 0; i < accum.size(); ++i) { + accum[i] = cutlass::fast_tanh(accum[i]); + } + accum = cutlass::multiplies()(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 diff --git a/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu b/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu index a0c9d4666cae3..0c9cbcc83593d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/unfused_attention.cu +++ b/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); diff --git a/onnxruntime/core/providers/cuda/llm/attention.cc b/onnxruntime/core/providers/cuda/llm/attention.cc index 15f9dcbf8e7f2..b6fbd9d8a24bc 100644 --- a/onnxruntime/core/providers/cuda/llm/attention.cc +++ b/onnxruntime/core/providers/cuda/llm/attention.cc @@ -543,9 +543,11 @@ Status Attention::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 Status Attention::RunMemoryEfficientAttention( diff --git a/onnxruntime/test/providers/cpu/llm/attention_op_test.cc b/onnxruntime/test/providers/cpu/llm/attention_op_test.cc index 40c45db2dfd66..43ed4d4c447e9 100644 --- a/onnxruntime/test/providers/cpu/llm/attention_op_test.cc +++ b/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 q(static_cast(batch_size * q_num_heads * q_sequence_length * head_size), 0.01f); + std::vector k(static_cast(batch_size * kv_num_heads * kv_sequence_length * head_size), 0.01f); + std::vector v(static_cast(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(s * v_head_size + h)] = val; + } + } + + // Float attn_mask = additive bias. Position 1 has -10 bias. + std::vector 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 y(static_cast(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(), std::vector(), std::vector(), + -1, -1, std::numeric_limits::quiet_NaN(), 5.0f, -1, TensorType::kFloat, // is_causal, qk_matmul_output_mode, scale, softcap, softmax_precision, tensor_type + y, std::vector(), std::vector(), std::vector(), + 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("softcap", 5.0f); + + std::vector q(static_cast(batch_size * q_num_heads * q_sequence_length * head_size), 0.01f); + std::vector k(static_cast(batch_size * kv_num_heads * kv_sequence_length * head_size), 0.01f); + std::vector v(static_cast(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(s * v_head_size + h)] = v_vals[s]; + } + } + // attn_mask: positions 1 and 2 strongly biased. + std::vector 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("Q", {batch_size, q_num_heads, q_sequence_length, head_size}, ToFloat16(q)); + test.AddInput("K", {batch_size, kv_num_heads, kv_sequence_length, head_size}, ToFloat16(k)); + test.AddInput("V", {batch_size, kv_num_heads, kv_sequence_length, v_head_size}, ToFloat16(v)); + test.AddInput("attn_mask", {q_sequence_length, kv_sequence_length}, ToFloat16(attn_mask)); + test.AddOptionalInputEdge(); // past_key + test.AddOptionalInputEdge(); // 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 expected_y( + static_cast(batch_size * q_num_heads * q_sequence_length * v_head_size), 1.5200f); + test.AddOutput("Y", {batch_size, q_num_heads, q_sequence_length, v_head_size}, + ToFloat16(expected_y), false, 0, 5e-3f); + test.AddOptionalOutputEdge(); // present_key + test.AddOptionalOutputEdge(); // present_value + + std::vector> 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