ONNX Attention Op — Remaining CUDA Implementation Gaps
Parent issue: #27516
Related (closed): #27880
Related (open): #28379 (PR — partial closure of item 5 via CPU oracle coverage)
After #27992 merged (3-tier dispatch, unified unfused kernel), the following gaps remain between the ONNX Attention opset 24 spec and our CUDA EP implementation.
Gap Summary
| # |
Feature |
Spec |
CUDA Status |
Code Location |
Test Coverage |
Priority |
| 1 |
qk_matmul_output_mode 1/2/3 |
Capture post-softcap, post-mask, post-softmax intermediates |
NOT_IMPLEMENTED (only mode 0 = raw QK / kQK) |
Reject: core/providers/cuda/llm/attention.cc:1432-1437 · kQK writer: contrib_ops/cuda/bert/unfused_attention.cu:440-445 · CPU reference hooks: core/providers/cpu/llm/attention.cc:500, 519, 546, 554 |
CPU: covered by spec test (mode 1/2/3 paths in cpu/llm/attention.cc). CUDA: none (the NOT_IMPLEMENTED path means tests can only assert the rejection.) |
Medium |
| 2 |
float64 / float8e4m3fn dtypes |
Type constraint allows |
Not registered for CUDA EP |
Kernel registration: core/providers/cuda/llm/attention.cc (search BuildKernelCreateInfo<...Attention...>) |
None — no registered kernel |
Low |
| 3 |
GQA + output_qk |
output_qk with num_heads > kv_num_heads |
Flash/MEA excluded by !has_output_qk; falls to unfused. Unfused supports kQK for GQA via reshape-Q trick; modes 1/2/3 share the same gap as item 1. |
Eligibility filter: attention.cc:1351, 1377 (Flash/MEA exclude on has_output_qk) · Unfused kQK writer: unfused_attention.cu:440-445 (works for GQA) · Modes 1/2/3 reject: same as item 1, attention.cc:1432-1437 |
CUDA: kQK + GQA exercised indirectly via existing unfused tests; modes 1/2/3 + GQA: none (blocked by item 1) |
Low |
| 4 |
Mask right-padding |
attn_mask last dim < total_seq_length (pad with -inf) |
Rejected at validation |
Validation: core/providers/cpu/llm/attention_helper.h:146 (ORT_ENFORCE on exact match). TODO already in source: attention_helper.h:141-145. Used by both CPU and CUDA EPs. |
None — rejected before kernel dispatch |
Low |
| 5 |
softcap + fp32 (no Flash/MEA fallback) |
Should work for all dtypes |
fp32 only has unfused path; CUDA path correctness now indirectly attested via CPU oracle parity |
fp32-GQA fall-through: attention.cc:1373-1376 (!(is_gqa && std::is_same<T, float>::value) filter) · Unfused entry: attention.cc:976-1000 (RunUnfusedAttention header) |
CPU: added in #28379 — test/python/transformers/test_onnx_attention/test_gqa.py:2892 TestONNXAttentionGQASoftcapFloat32CPU and test_mha.py:1955 TestONNXAttentionMHAUnfusedSoftcapCPU. CUDA: still implicit (covered by oracle parity at runtime when applicable). |
Low (partial) |
Details
1. qk_matmul_output_mode (Medium)
The unified unfused kernel only supports kNone (no output) and kQK (raw QK scores, scale * Q @ K^T before softcap/mask/softmax). Modes 1 (post-softcap), 2 (post-mask/bias), and 3 (post-softmax) require capturing intermediate values at different pipeline stages. Flash/MEA cannot support this — only the unfused kernel can, but it needs pipeline-stage hooks.
Proposed approach (mirrors CPU): The CPU implementation in core/providers/cpu/llm/attention.cc already handles all four modes by snapshotting the QK buffer at four well-defined stages — kQK (line 500, after scaled QK), kPostSoftCap (line 519, after softcap), kPostMaskBias (line 546, after mask + additive bias), kPostSoftMax (line 554, after softmax normalization). The CUDA unfused kernel in contrib_ops/cuda/bert/unfused_attention.cu already has the same staged structure (lines 415-470 implement softcap → mask → softmax sequentially); add three more if (output_qk != nullptr && qk_matmul_output_mode == ...) snapshot launches mirroring the CPU sites, then drop the early-reject guard at attention.cc:1432-1437. No new kernel required — reuse ScaledCopyQkKernel (already at unfused_attention.cu:63) with appropriate scale=1.0 for the post-stages.
2. float64 / float8 (Low)
The CUDA EP only registers float16, bfloat16, float32. float64 is impractical for GPU attention. float8e4m3fn would require kernel specialization (Hopper sm_90 + intrinsics).
3. GQA + output_qk (Low)
The Flash and MEA paths exclude output_qk requests via the !has_output_qk filter at attention.cc:1351 (Flash) and attention.cc:1377 (MEA), routing all output_qk cases (MHA or GQA) into the unfused path. The unfused kernel handles GQA + kQK natively via the reshape-Q trick (no K/V replication). The remaining gap is GQA + modes 1/2/3, which is blocked by item 1; closing item 1 closes this automatically.
4. Mask right-padding (Low)
Spec allows attn_mask with last_dim < total_sequence_length (remaining positions treated as -inf). The validation in cpu/llm/attention_helper.h:146 enforces exact match; this validator is shared by both CPU and CUDA EPs. A TODO note already lives at attention_helper.h:141-145 describing the fix shape: change == to <=, allocate a padded buffer, fill remainder with mask_filter_value<T>(). Workaround for users today: pad the mask to full size before calling.
5. softcap + fp32 edge case (Low — partially closed)
fp32 cannot use Flash or MEA (fp16/bf16 only). The unified unfused kernel handles softcap+fp32 correctly. PR #28379 added CPU EP twin classes (TestONNXAttentionGQASoftcapFloat32CPU and TestONNXAttentionMHAUnfusedSoftcapCPU) that exercise softcap+fp32 against the attention_ref() oracle, providing reference parity coverage. CUDA-side coverage remains implicit (any CUDA softcap+fp32 GQA/MHA invocation falls through the unfused path validated by the same oracle when CUDA-enabled torch is available — see comment #28351 (comment) for context on this dispatch).
Not Gaps (confirmed not in spec)
smoothing_factor — not an ONNX Attention attributecache_indirection — not an ONNX Attention inputsliding_window — not in ONNX Attention (GQA contrib op only)
ONNX Attention Op — Remaining CUDA Implementation Gaps
Parent issue: #27516
Related (closed): #27880
Related (open): #28379 (PR — partial closure of item 5 via CPU oracle coverage)
After #27992 merged (3-tier dispatch, unified unfused kernel), the following gaps remain between the ONNX Attention opset 24 spec and our CUDA EP implementation.
Gap Summary
qk_matmul_output_mode1/2/3core/providers/cuda/llm/attention.cc:1432-1437· kQK writer:contrib_ops/cuda/bert/unfused_attention.cu:440-445· CPU reference hooks:core/providers/cpu/llm/attention.cc:500, 519, 546, 554cpu/llm/attention.cc). CUDA: none (the NOT_IMPLEMENTED path means tests can only assert the rejection.)core/providers/cuda/llm/attention.cc(searchBuildKernelCreateInfo<...Attention...>)!has_output_qk; falls to unfused. Unfused supports kQK for GQA via reshape-Q trick; modes 1/2/3 share the same gap as item 1.attention.cc:1351, 1377(Flash/MEA exclude onhas_output_qk) · Unfused kQK writer:unfused_attention.cu:440-445(works for GQA) · Modes 1/2/3 reject: same as item 1,attention.cc:1432-1437core/providers/cpu/llm/attention_helper.h:146(ORT_ENFORCE on exact match). TODO already in source:attention_helper.h:141-145. Used by both CPU and CUDA EPs.attention.cc:1373-1376(!(is_gqa && std::is_same<T, float>::value)filter) · Unfused entry:attention.cc:976-1000(RunUnfusedAttentionheader)test/python/transformers/test_onnx_attention/test_gqa.py:2892 TestONNXAttentionGQASoftcapFloat32CPUandtest_mha.py:1955 TestONNXAttentionMHAUnfusedSoftcapCPU. CUDA: still implicit (covered by oracle parity at runtime when applicable).Details
1. qk_matmul_output_mode (Medium)
The unified unfused kernel only supports
kNone(no output) andkQK(raw QK scores, scale * Q @ K^T before softcap/mask/softmax). Modes 1 (post-softcap), 2 (post-mask/bias), and 3 (post-softmax) require capturing intermediate values at different pipeline stages. Flash/MEA cannot support this — only the unfused kernel can, but it needs pipeline-stage hooks.Proposed approach (mirrors CPU): The CPU implementation in
core/providers/cpu/llm/attention.ccalready handles all four modes by snapshotting the QK buffer at four well-defined stages — kQK (line 500, after scaled QK), kPostSoftCap (line 519, after softcap), kPostMaskBias (line 546, after mask + additive bias), kPostSoftMax (line 554, after softmax normalization). The CUDA unfused kernel incontrib_ops/cuda/bert/unfused_attention.cualready has the same staged structure (lines 415-470 implement softcap → mask → softmax sequentially); add three moreif (output_qk != nullptr && qk_matmul_output_mode == ...)snapshot launches mirroring the CPU sites, then drop the early-reject guard atattention.cc:1432-1437. No new kernel required — reuseScaledCopyQkKernel(already atunfused_attention.cu:63) with appropriate scale=1.0 for the post-stages.2. float64 / float8 (Low)
The CUDA EP only registers float16, bfloat16, float32. float64 is impractical for GPU attention. float8e4m3fn would require kernel specialization (Hopper sm_90 + intrinsics).
3. GQA + output_qk (Low)
The Flash and MEA paths exclude
output_qkrequests via the!has_output_qkfilter atattention.cc:1351(Flash) andattention.cc:1377(MEA), routing alloutput_qkcases (MHA or GQA) into the unfused path. The unfused kernel handles GQA +kQKnatively via the reshape-Q trick (no K/V replication). The remaining gap is GQA + modes 1/2/3, which is blocked by item 1; closing item 1 closes this automatically.4. Mask right-padding (Low)
Spec allows
attn_maskwithlast_dim < total_sequence_length(remaining positions treated as -inf). The validation incpu/llm/attention_helper.h:146enforces exact match; this validator is shared by both CPU and CUDA EPs. A TODO note already lives atattention_helper.h:141-145describing the fix shape: change==to<=, allocate a padded buffer, fill remainder withmask_filter_value<T>(). Workaround for users today: pad the mask to full size before calling.5. softcap + fp32 edge case (Low — partially closed)
fp32 cannot use Flash or MEA (fp16/bf16 only). The unified unfused kernel handles softcap+fp32 correctly. PR #28379 added CPU EP twin classes (
TestONNXAttentionGQASoftcapFloat32CPUandTestONNXAttentionMHAUnfusedSoftcapCPU) that exercise softcap+fp32 against theattention_ref()oracle, providing reference parity coverage. CUDA-side coverage remains implicit (any CUDA softcap+fp32 GQA/MHA invocation falls through the unfused path validated by the same oracle when CUDA-enabled torch is available — see comment #28351 (comment) for context on this dispatch).Not Gaps (confirmed not in spec)
smoothing_factor— not an ONNX Attention attributecache_indirection— not an ONNX Attention inputsliding_window— not in ONNX Attention (GQA contrib op only)