Eliminate Legacy MHA Unfused path from ONNX Attention; unify on 3-tier dispatch with causal alignment fix

Merged
TI
6mo ago
+2695-772

Motivation

Eliminate the legacy MHA Unfused path (QkvToContext in attention_impl.cu) from the ONNX standard Attention op, simplifying the CUDA dispatch to a clean 3-tier cascade.

Design

Flash Attention → Memory-Efficient Attention (MEA) → Unified Unfused Attention
  • Flash: Handles fp16/bf16 with head_size ≤ 256, no explicit attn_mask. Fastest path.
  • MEA (CUTLASS): Handles cases Flash cannot (explicit masks, softcap+mask combos). Requires head_size % 8 == 0.
  • Unified Unfused: Fallback for everything else — fp32, small heads, H≠H_v, output_qk. Handles both MHA and GQA via FP32 QK accumulation.

The legacy RunUnfusedAttention wrapper (which called contrib ops QkvToContext) is deleted. The contrib MHA op is unaffected.

Key Behavior Changes

  • Unified unfused kernel replaces separate GQA-only and MHA-only unfused paths
  • Causal alignment: lower-right when past_key is present, upper-left otherwise (per ONNX spec)
  • H≠H_v + past KV now supported (separate K/V concat calls)
  • output_qk (mode 0) supported in unified kernel via ScaledCopyQkKernel
  • 29 ONNX backend test filters removed — tests now pass natively

Testing

All existing tests pass (40 C++ attention tests, 215 Python parametrized cases) plus new coverage for causal alignment on CPU EP and softcap ordering verification.

Closes #27880. Related: #27516, #28198.