Validate CUDA FMHA seqstart/seqlen tensor values - #31641
Open
apsonawane wants to merge 3 commits into
Open
Conversation
- validate MASK_1D_KEY_SEQ_LEN_START contents before memory-efficient attention launch - validate GQA seqlens_k bounds before right-padding FMHA path - add CUDA MultiHeadAttention regression for invalid seqstart values
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds host-side validation for CUDA attention inputs (mask seqstart/seqlen and GQA seqlens) to fail fast with clearer errors before launching fused/MEM-efficient attention kernels, and adds a regression test to ensure invalid seqstart values are rejected.
Changes:
- Added host-side validation for
MASK_1D_KEY_SEQ_LEN_STARTmask contents in CUDA memory-efficient attention (MultiHeadAttention path). - Added host-side validation for GQA
seqlens_kvalues prior to launching CUDA memory-efficient attention. - Added a CUDA-only unit test that asserts invalid
seqstart_qvalues are rejected with an error.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cuda/bert/attention_impl.cu | Adds host-side validation for 1D mask (seqlen + cumulative start arrays) before MEA kernel launch. |
| onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu | Adds host-side validation for GQA seqlens_k values and wires it into the MEA path. |
| onnxruntime/test/contrib_ops/multihead_attention_op_test.cc | Adds a CUDA test that expects failure for invalid seqstart_q in the 1D mask encoding. |
apsonawane
enabled auto-merge (squash)
August 5, 2026 17:49
Gate the new host-side seqstart/seqlen validation behind a shared ORT_CUDA_ATTENTION_VALIDATE_SEQ_LENS environment variable so the release path avoids per-call device-to-host copies and stream synchronization in the hot attention execution path. Also switch the validation copies to explicit cudaMemcpyDeviceToHost and update the CUDA tests to enable the validation env var when they expect failures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
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.
This pull request adds validation logic for attention mask and sequence length inputs in CUDA attention kernels, improving input robustness and error reporting. It introduces host-side checks to catch invalid sequence length and start values before launching GPU kernels, and expands test coverage to ensure these checks work as intended.
Input Validation Enhancements:
onnxruntime/contrib_ops/cuda/bert/attention_impl.cu: AddedValidateMask1DKeySeqLenStartValues, which copies the mask to the host and checks for out-of-range or inconsistent sequence length and start values in the 1D mask input. This validation is now called inEfficientAttentionbefore using the mask. [1] [2] [3]onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu: AddedValidateGqaSeqLensValuesto check that all sequence lengths are within valid bounds, and invoked it inEfficientAttentionfor the relevant sequence length input. [1] [2] [3]Test Coverage:
onnxruntime/test/contrib_ops/multihead_attention_op_test.cc: Added a new test case,CudaMask1DKeySeqLenStartRejectsInvalidSeqstartValues, to verify that invalid mask start values are correctly rejected with an error.These changes improve error detection for invalid input data, making debugging easier and preventing undefined behavior in CUDA attention operations.