ONNX Runtime CUDA Attention Patterns

ONNX Domain Attention (Opset 23/24) CUDA Patterns

Reusable knowledge from ONNX Attention CUDA development in ORT.

Scope: This skill covers the ONNX domain Attention operator (opset 23/24) implemented at core/providers/cuda/llm/attention.cc. This is separate from the contrib domain MultiHeadAttention / GroupQueryAttention at contrib_ops/cuda/bert/. They share some underlying kernels (CUTLASS FMHA, Flash Attention) and infrastructure (attention_softmax.h) but have different dispatch logic, parameter structs, and eligibility checks.

• Shared infrastructure: CUTLASS FMHA kernel, Flash kernel, unified unfused kernel (unfused_attention.cu), attention_softmax.h, attention_impl.cu (contrib only)
• ONNX-specific: Dispatch cascade in attention.cc, ConvertAttnMaskToBias, mask_filter_value cap, parameter bridge to contrib structs, attention_mask_impl.cu
• Contrib-specific: Own dispatch in contrib MHA/GQA ops, uses contrib::AttentionParameters directly, has XQA kernel, past-present buffer sharing

1. Runner Dispatch Cascade

CUDA attention dispatches in priority order: Flash → MEA (Memory Efficient) → Unified Unfused Attention.

// onnxruntime/core/providers/cuda/llm/attention.cc — ComputeInternal()
Flash eligible?      → RunFlashAttention()
  ↓ no
MEA eligible?        → RunMemoryEfficientAttention()
  ↓ no
Unified Unfused      → RunUnfusedAttention()
  (handles both MHA and GQA via reshape-Q trick)

Eligibility anchors (symbols are stable; line numbers as of cc34d0b914)

head_size=512 IS routed to MEA, but its MEA kernel is NOT portably launchable — so it is not a robust test probe. By the predicate, 512 > 256 fails Flash and 512 ≤ 1024 with 512 & 7 == 0 passes the MEA predicate, so dispatch selects MEA — but the MEA eligibility check (memory_efficient_attention.h:68-73) gates only on SM + head&7==0 + head<=1024, with no shared-memory check. For head_size=512 FP16 the CUTLASS MEA SharedStorage exceeds the dynamic-smem opt-in cap on capacity-limited arches (sm86 ~99KB, sm80 ~163KB, sm90 ~227KB — non-monotonic, no clean SM-version guard). fmha_launch_template.h calls cudaFuncSetAttribute(..., cudaFuncAttributeMaxDynamicSharedMemorySize, ...) but ignores its return value and launches anyway, so on sm86 the kernel dies at launch with CUDA failure 1: invalid argument — there is no fallback to unfused (live bug #28388; its fix PR #28383 was never merged). So head_size=512's MEA kernel launches only on large-smem arches like sm90/H100.

To force the MEA path portably in a test, use ORT_DISABLE_FLASH_ATTENTION=1 with a small head_size (e.g. 64) whose SharedStorage fits every target arch — NOT head_size=512. Also guard with SKIP_IF_MEA_NOT_COMPILED (see §7) so a MEA-OFF build SKIPs rather than false-greens via the (correct) unfused fallback.

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.

QUICK_BUILD caveat (false hypothesis trap). (General principle — build flags can silently reroute kernel dispatch — lives in the ort-build skill, "Agent tips". The attention-specific instance:) With onnxruntime_QUICK_BUILD=ON (-DORT_QUICK_BUILD), Flash is compiled for head_dim 128 only: flash_api.h:147 is_supported<T> returns false for head_size != 128, and static_switch.h:80 HEADDIM_SWITCH only instantiates kHeadDim=128. So under QUICK_BUILD nearly every shape routes to MEA, not FlashAttention-2. If a head_size!=128 test "fails only on some SM", suspect MEA (CUTLASS, arch-independent), NOT a Flash/FA2 hardware bug. head_size=512 is routed to MEA in all MEA-enabled builds (Flash caps at 256), but its MEA kernel fails to launch on small-smem GPUs — see the head_size=512 caveat above (#28388).

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)
• output_qk (kQK mode: copies raw QK before softcap/mask mutations)
• past_key + past_value with head_size != v_head_size (separate K/V concat)
• causal masking, nonpad_kv_seqlen, all dtypes (fp16/bf16/fp32)

2. CUTLASS kLog2e Overflow

CUTLASS iterative_softmax multiplies all attention scores by kLog2e ≈ 1.4427 internally (for exp2f instead of expf). For float/bf16:

mask_filter_value = std::numeric_limits<float>::lowest() ≈ -3.40e+38
-3.40e+38 × 1.4427 ≈ -4.91e+38  →  overflows fp32  →  -inf

When all values become -inf, CUTLASS's special-case path produces s_prime=0 → 1/s_prime=inf → 0 × inf = NaN.

Fix: Cap mask_filter_value to -1.0e+30f in ConvertAttnMaskToBias. This value is safe: 1e30 × 1.4427 ≈ 1.4e30 << FLT_MAX, and exp(-1e30) ≈ 0 (effectively masked).

fp16 is NOT affected: lowest() = -65504, and -65504 × 1.4427 ≈ -94500 stays within fp32 range.

This cap is ONLY applied in MEA paths. The unfused path uses lowest() directly (its softmax subtracts max first, avoiding overflow).

Subtlety: When bias is present (kSupportsBias=true), CUTLASS pre-applies p.scale to QK (line 858) and uses scaling=1.0f in the softmax loop (line 981). So the full kLog2e multiplier hits the bias-dominated values — the overflow is head_size-independent. Without bias, scaling = p.scale * kLog2e = kLog2e/sqrt(head_size), which is much smaller.

3. Bias Alignment

CUTLASS FMHA requires the attention bias row stride to satisfy minimum alignment. The bias has shape [B, H, S, T] where T = total_sequence_length is the row stride.

constexpr int min_bias_align = 4;  // elements, not bytes
if (parameters.total_sequence_length % min_bias_align != 0) {
    mea_eligible = false;  // fall through to unfused
}

Impact on tests: If a test uses total_sequence_length not divisible by 4 (e.g., past=5 + new=6 = 11), MEA is rejected and unfused handles it. To test MEA with bias, ensure total_sequence_length % 4 == 0.

4. Softcap Ordering

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

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

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

Always cap grid size to prevent exceeding gridDim.x limits, and use grid-stride loops for large workloads:

constexpr int64_t kMaxGridDimX = 65535;
int threads = static_cast<int>(std::min(static_cast<int64_t>(max_threads_per_block), total));
int64_t blocks = (total + threads - 1) / threads;
unsigned int grid_size = static_cast<unsigned int>(std::min(blocks, kMaxGridDimX));

MyKernel<<<grid_size, threads, 0, stream>>>(...);

// Inside the kernel:
for (int64_t idx = blockIdx.x * blockDim.x + threadIdx.x;
     idx < total;
     idx += static_cast<int64_t>(gridDim.x) * blockDim.x) {
    // work
}

Never cast int64_t block count directly to unsigned int without capping — it silently truncates.

Always call CUDA_CALL(cudaGetLastError()) after kernel launches in standalone helper functions. This is the established pattern in the file (see ConcatPastToPresent, PastPresentBufferShare).

6. Fully-Masked Rows and Batches

All-false bool masks, an all--inf attn_mask row, or a causal/nonpad frontier with no allowed key produce NaN in CUTLASS MEA (the uniform/empty softmax degenerates: s_prime=0 → 1/s_prime=inf → 0 × inf = NaN). Per onnx/onnx#8068 (Bug-2), a fully-masked query row — one with no key allowed by the composed causal + nonpad + mask constraints — must output a zero row (Y = 0), not mean-of-V.

This Y = 0 behavior is now consistent on BOTH EPs (the earlier mean(V)-vs-zero cross-EP divergence is RESOLVED — there is no longer an open TODO here):

• CUDA: ZeroFullyMaskedRowsKernel (in attention_mask_impl.cu) runs after the MEA/CUTLASS output and zeros each fully-masked row with a select (not multiply, so 0 @ V = 0 even when V is poisoned). It detects a fully-masked row with an exact per-key predicate (within the causal/nonpad frontier AND the additive-bias slot is above the mask sentinel), matching the onnx#8068 isneginf-of-row-max reference. A finite (even very negative) user bias is not the sentinel, so its key stays unmasked and the row is left untouched.
• CPU: core/providers/cpu/llm/attention.cc applies the same Bug-2 guard — after softmax it zeros any row whose composed frontier admitted no unmasked key.

Additive-bias path (bool mask converted to bias): mask_filter_value is capped to -1e+30f (see section 2) so CUTLASS does not overflow to NaN; a row that is nonetheless fully masked is then zeroed by the per-row guard above.

Whole-batch empty (seqlens_k[b] == 0): the structural case where an entire batch has zero valid keys is additionally handled by ZeroOutputForFullyMaskedBatches, which zeros that batch's output. (The per-row guard covers the finer-grained case where only some query rows are fully masked.)

qk_matmul_output_mode (mode 3 / post-softmax debug output): for a fully-masked row the mode-3 snapshot is mandated to be 0 (zero row), consistent with Y = 0, per the onnx#8068 SIG decision (this superseded the earlier "unspecified" proposal). The CPU post-softmax snapshot is taken after the row-zeroing guard — matching the onnx reference and the v23/v24 function bodies, where the guard runs before the mode-3 capture — so the debug tensor reflects the same zero row as the output. Note this mode-3=0 behavior is served by the CPU path: CUDA qk_matmul_output_mode beyond kNone/kQK (i.e. kPostSoftCap/kPostMaskBias/kPostSoftMax) returns NOT_IMPLEMENTED (attention.cc), so an agent must not assume CUDA produces mode-3=0.

7. Test Runner Targeting

Use ScopedEnvironmentVariables to force specific CUDA runners:

// Force MEA (disable Flash)
ScopedEnvironmentVariables scoped_env({
    {"ORT_DISABLE_FLASH_ATTENTION", "1"},
});

// Force Unfused (disable both Flash and MEA)
ScopedEnvironmentVariables scoped_env({
    {"ORT_DISABLE_FLASH_ATTENTION", "1"},
    {"ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION", "1"},
});

Always verify which runner a test actually hits. A test designed for MEA may silently fall to unfused if:

• total_sequence_length % 4 != 0 (bias alignment)
• head_size != v_head_size (decode path)
• fp32 dtype with GQA (LaunchUngroup fp16/bf16 only)
• fp32 dtype on SM < 80

Enable verbose logging to confirm: LOGS_DEFAULT(VERBOSE) << "ONNX Attention: using ...".

SKIP_IF_MEA_NOT_COMPILED is a local gtest macro (defined in test/providers/cpu/llm/attention_op_test.cc) that GTEST_SKIPs — rather than silently passes — when USE_MEMORY_EFFICIENT_ATTENTION is OFF, so an MEA-targeted test cannot false-green via the (correct) unfused fallback. Use it in any test that must prove the MEA path ran (see the ort-test skill → "Verify which path/kernel actually executed").

8. Cross-EP Consistency

CPU is the spec reference implementation. CUDA outputs should match CPU for all valid inputs.

• CPU uses mask_filter_value = std::numeric_limits<T>::lowest() (finite, not -inf)
• CPU softmax: subtract-max-first → works correctly with extreme finite values
• CPU zeros fully-masked query rows (onnx#8068 Bug-2 guard) — output Y = 0, matching CUDA's ZeroFullyMaskedRowsKernel. (Earlier docs claimed CPU produced mean(V) here; that divergence is resolved — both EPs now emit a zero row.)

Run tests with disable_cpu=false to always validate against CPU. The C++ test framework (RunTest4D) supports disable_cpu, disable_cuda, disable_dml flags.