ONNX Attention CUDA: Coverage Gaps in Runner Fallback Paths

[titaiwangms]

ONNX Attention Op — CUDA Implementation Gap Tracking

Parent issue: #27516

Status: All Gaps Closed ✅

All gaps addressed by #28198 (merged) and #27992 (in review).

Dispatch Cascade (Final)

Flash Attention → MEA (CUTLASS) → Unified Unfused Attention

The Legacy MHA Unfused path (QkvToContext from attention_impl.cu) has been eliminated from the ONNX Attention op. The unified kernel handles both MHA and GQA.

Support Matrix (Unified Unfused Attention)

Feature	Flash	MEA	Unified Unfused
MHA (num_heads == kv_num_heads)	✅	✅	✅
GQA (num_heads > kv_num_heads)	✅	✅	✅
fp16 / bf16	✅	✅	✅
fp32	❌	❌ (GQA)	✅
Softcap	✅	✅	✅
Padding mask (seqlens_k)	✅	✅	✅
Explicit attn_mask	❌	✅	✅
Softcap + mask	❌ (no mask)	✅	✅
Past KV	✅	✅	✅
Past KV + H≠H_v	❌	❌	✅
Causal	✅	✅	✅
output_qk (mode 0)	❌	❌	✅

Previously Open Gaps — Now Closed

• output_qk: Only supported in Legacy path → Added ScaledCopyQkKernel to unified kernel (Eliminate Legacy MHA Unfused path from ONNX Attention; unify on 3-tier dispatch with causal alignment fix #27992)
• H≠H_v + past KV: Not supported in GQA Unfused → Separate K/V concat calls (Eliminate Legacy MHA Unfused path from ONNX Attention; unify on 3-tier dispatch with causal alignment fix #27992)
• MHA in unfused path: Required Legacy wrapper → Unified kernel handles MHA (group_size=1) (Eliminate Legacy MHA Unfused path from ONNX Attention; unify on 3-tier dispatch with causal alignment fix #27992)

PRs

• GQA unfused attention with FP32 QK accumulation (fixes #28195) #28198 — GQA unfused attention with FP32 QK accumulation ✅ Merged
• Eliminate Legacy MHA Unfused path from ONNX Attention; unify on 3-tier dispatch with causal alignment fix #27992 — Fix softcap ordering, unify unfused kernel, eliminate Legacy path 🔄 In Review

ONNX Spec

• Spec bug: Softcap in Attention op onnx/onnx#7865 (softcap ordering)
• Spec fix: Fix Attention op softcap ordering: apply before mask/bias onnx/onnx#7867

Remaining Follow-ups (not blocking)

• [ONNX Attention] Follow-up: causal+softcap ordering, defensive assertion, and backend test filter cleanup #28215 — Causal + softcap interaction (spec divergence)

This issue will be closed when #27992 merges.

[titaiwangms]

#27831 (review)

[titaiwangms]

#28198

[titaiwangms]

Update: PR #28198 Merged — GQA Unfused Attention Path

PR #28198 (GQA unfused attention with FP32 QK accumulation) has been merged. This adds a new dispatch path to the cascade:

Flash → MEA → GQA Unfused (NEW) → MHA Unfused (legacy)

What #28198 Added

• GQA unfused attention kernel (gqa_unfused_attention.cu) — 3-stage pipeline: QK GEMM (FP32 output) → softmax → AV GEMM
• FP32 QK accumulation — prevents fp16 overflow at head_size > 256 (fixes Gemma 4 NaN at head_dim=512, CUDA EP: Unfused Attention runner produces NaN for fp16 with head_dim > 256 #28195)
• Softcap support with correct ordering per Softcap in Attention op onnx/onnx#7865 (softcap → mask → softmax)
• Sliding window support
• Large head_size support (256, 512) — MAX_HEAD_SIZE raised to 512
• "Reshape-Q trick" for native GQA — no K/V replication needed

Gap Table Update (Post-#28198)

#	Gap	Previous Status	New Status	Notes
R1	GQA + fp32	🔴 OPEN	✅ CLOSED	GQA unfused kernel uses FP32 QK GEMM; accepts fp32 inputs
R2	GQA + decode + h ≠ v_h	🔴 OPEN	✅ CLOSED	GQA unfused supports different Q/V head sizes
R3	GQA + decode + output_qk	🟡 OPEN	🟡 REMAINS	GQA unfused does not implement output_qk modes
R4	GQA + softcap + attn_mask	🟡 OPEN	✅ CLOSED	GQA unfused supports softcap + bias with correct ordering
R5	output_qk beyond kQK	🟢 OPEN	🟢 REMAINS	Not addressed

3 gaps closed by #28198: R1, R2, R4
2 gaps remaining: R3 (GQA + output_qk), R5 (output_qk advanced modes) — both low/medium priority

Updated Dispatch Cascade

ComputeInternal() — core/providers/cuda/llm/attention.cc
│
├─ Flash eligible? ──YES──→ RunFlashAttention()
│   [fp16/bf16, SM≥8.0, no mask, h==v_h, h≤256]
│
├─ MEA eligible?   ──YES──→ RunMemoryEfficientAttention()
│   [fp16/bf16, decode when h==v_h, no softcap+mask]
│
├─ GQA Unfused eligible? ──YES──→ RunGqaUnfusedAttention()  ← NEW
│   [all dtypes, GQA+MHA, softcap, masks, sliding window, h≤512]
│
└─ MHA Unfused fallback
    [MHA only, all dtypes, all masks, softcap (spec-correct)]

Overall Status

PR	Status	Gaps Closed
#27714	✅ Merged	softcap + softmax_precision on Flash/MEA
#27831	✅ Merged	Flash rejects mask, min_bias_align fix, MEA NaN fix
#27992	🔄 Open	MEA decode + unfused softcap — closes original gaps 2, 5, 6, 8
#28198	✅ Merged	GQA unfused path — closes R1, R2, R4

After #27992 merges, only 2 low-priority gaps remain (both output_qk related). The ONNX Attention CUDA operator will have comprehensive coverage for all practical workloads.