Fix CUDA ONNX Attention: min_bias_align crash on SM<80 and MEA NaN for fully-masked batches#27831
Merged
Conversation
yuslepukhin
requested review from
Copilot and
tianleiwu
and removed request for
Copilot
yuslepukhin
reviewed
Mar 24, 2026
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
titaiwangms
added a commit
that referenced
this pull request
Mar 24, 2026
…igned dims PR #27831 fell back to CUBLAS_DEFAULT_MATH which still uses TF32 on Ampere+ GPUs (SM>=80) since cuBLAS 11.0. Changed to CUBLAS_PEDANTIC_MATH when dimensions are not 4-aligned to guarantee no tensor core usage, preventing CUDA error 716 (misaligned address) on CUDA 12.9+. Three-way logic in all three float cuBLAS GEMM helpers: - TF32 requested + dimensions aligned: CUBLAS_TF32_TENSOR_OP_MATH - TF32 requested + dimensions NOT aligned: CUBLAS_PEDANTIC_MATH - TF32 not requested: CUBLAS_DEFAULT_MATH (unchanged) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
titaiwangms
force-pushed
the
titaiwang/fix_cuda
branch
2 times, most recently
from
38fef79 to
0796037
Compare
titaiwangms
changed the title
Contributor
Author
|
Now we improve from crash to 3 tests mismatched expected results. |
titaiwangms
force-pushed
the
titaiwang/fix_cuda
branch
2 times, most recently
from
9a1bf88 to
2cccef2
Compare
titaiwangms
commented
Mar 27, 2026
titaiwangms
force-pushed
the
titaiwang/fix_cuda
branch
from
2cccef2 to
678c840
Compare
titaiwangms
changed the title
Contributor
There was a problem hiding this comment.
Pull request overview
This PR fixes correctness and stability issues in the CUDA EP implementation of the ONNX Attention op, focusing on decode-time bool masking with past KV, and NaNs in CUTLASS Memory Efficient Attention (MEA) for fully-masked batches.
Changes:
- Aligns MEA bias-stride eligibility logic in
attention.ccand adds post-MEA output zeroing for fully-masked batches. - Fixes unfused decode behavior for bool masks with past KV by using variable-length KV concat consistent with Flash layout semantics (and zero-inits present buffers).
- Adds C++ and Python tests covering decode bool-mask edge cases (partial masks, divergent per-batch seqlens).
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/python/transformers/test_onnx_attention/test_mha.py | Adds CUDA-only graph-level tests for bool-mask decode with past KV (partial mask + divergent seqlens). |
| onnxruntime/test/providers/cpu/llm/attention_op_test.cc | Adds CUDA execution tests forcing unfused path and verifying variable-length concat correctness in decode. |
| onnxruntime/core/providers/cuda/llm/attention_mask_impl.h | Declares a CUDA helper to zero outputs for fully-masked batches (seqlens_k==0). |
| onnxruntime/core/providers/cuda/llm/attention_mask_impl.cu | Implements and instantiates the ZeroOutputForFullyMaskedBatches kernel and launcher. |
| onnxruntime/core/providers/cuda/llm/attention.cc | Applies fully-masked output zeroing for MEA and unfused nonpad paths; fixes unfused bool-mask decode concat semantics; updates MEA eligibility alignment check. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
titaiwangms
force-pushed
the
titaiwang/fix_cuda
branch
from
678c840 to
ab43783
Compare
Contributor
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 5 out of 5 changed files in this pull request and generated 3 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Contributor
Author
This was referenced Jun 1, 2026
This was referenced Jun 15, 2026
This was referenced Jun 22, 2026
This was referenced Jun 29, 2026
This was referenced Jul 6, 2026
This was referenced Jul 13, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description:
Summary
Fixes three issues in the CUDA ONNX Attention operator and improves spec compliance:
min_bias_align crash on SM<80: The alignment check for Memory Efficient Attention (MEA) bias used 4*sizeof(T) (bytes), but the check is against element counts. Fixed to 4 elements, matching CUTLASS kMinimumAlignment. This prevented valid MEA dispatch on SM<80.
MEA NaN for fully-masked batches: When nonpad_kv_seqlen=0, CUTLASS MEA computes 1/s_prime where s_prime=0, producing NaN. Added ZeroOutputForFullyMaskedBatches kernel (MEA path only) to zero output for these batches. Uses int64_t for element count to prevent overflow at large context lengths.
Flash rejects attn_mask for spec compliance: Flash Attention's paged KV cache produces spec-divergent present_key/present_value layout when used with attn_mask + past_key. Flash now requires attn_mask == nullptr — cases with bool mask + past_key fall to the unfused runner which handles them spec-correctly. Removed ~137 lines of dead code (ConvertMaskToSeqlensKernel, LaunchConvertMaskToFlashSeqlensK) no longer needed after this change.
Known limitation
Related
attn_maskas padding mask, diverges from ONNX spec #27885: Flash Attention bool attn_mask semantic divergence (root cause documented)