Harden contrib CPU int narrowing for attention attrs - #31648
Merged
Akshay Sonawane (apsonawane) merged 11 commits intoAug 14, 2026
Merged
Conversation
Validate int64 attributes and shape-derived values fit in int before narrowing in LinearAttention, LongformerAttentionBase, and GptSubgraph validation. Add LinearAttention regressions for oversized q_num_heads and kv_num_heads attributes.
Copilot started reviewing on behalf of
Akshay Sonawane (apsonawane)
August 4, 2026 22:54
View session
Contributor
There was a problem hiding this comment.
Pull request overview
This PR hardens contrib CPU attention-related operators against unsafe int64_t→int narrowing by validating attribute/tensor-dimension ranges and adds unit tests to ensure overflow-scale values are rejected with clear errors.
Changes:
- Add
INT_MAX-bounded validation forq_num_heads/kv_num_heads(LinearAttention) andnum_heads/window(LongformerAttentionBase) beforestatic_cast<int>(). - Add
INT_MAX-bounded validation for GPT subgraph dimension values before narrowing toint. - Add LinearAttention unit tests that exercise the new overflow-range validation for head-count attributes.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cpu/bert/linear_attention.cc | Validates q_num_heads/kv_num_heads are within [1, INT_MAX] prior to narrowing to int. |
| onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h | Validates num_heads and window are within [1, INT_MAX] prior to narrowing to int. |
| onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc | Validates key subgraph shape dims are <= INT_MAX before narrowing to int. |
| onnxruntime/test/contrib_ops/linear_attention_op_test.cc | Adds failure-mode tests for out-of-range q_num_heads/kv_num_heads. |
The error message on line 174 incorrectly referenced 'past state dimension 2' when it was actually validating the logits output dimension 2 (vocabulary size). This fix clarifies the message to correctly reference 'logits dimension 2' for better debugging experience when vocabulary size validation fails. Addresses review comment feedback from PR #31648. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Akshay Sonawane (apsonawane)
enabled auto-merge (squash)
August 5, 2026 16:58
The overflow tests for q_num_heads and kv_num_heads were failing schema validation before reaching the attribute range checks because the LinearAttention node requires two outputs. Add dummy output and present_state outputs so the invalid attribute validation is exercised as intended. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The overflow tests were failing during schema shape inference before reaching the intended q_num_heads / kv_num_heads constructor validation. Use unknown output dimensions in the test harness so the failure path exercises the attribute range checks directly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Tianlei Wu (tianleiwu)
previously approved these changes
Aug 11, 2026
Tianlei Wu (tianleiwu)
previously approved these changes
Aug 11, 2026
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Ti-Tai Wang (titaiwangms)
previously approved these changes
Aug 13, 2026
Ti-Tai Wang (titaiwangms)
approved these changes
Aug 13, 2026
Akshay Sonawane (apsonawane)
deleted the
msrc/cpu-linearattention-int-wrap-fix
branch
August 14, 2026 00:04
This was referenced Sep 10, 2026
This was referenced Sep 14, 2026
Open
Open
Closed
This was referenced Sep 22, 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.
This pull request strengthens input validation for attention-related operators and adds corresponding unit tests to ensure that large attribute values do not cause integer overflows. The changes affect the
LinearAttention,LongformerAttentionBase, and GPT subgraph implementations, improving reliability and error reporting.Input validation improvements:
onnxruntime/contrib_ops/cpu/bert/linear_attention.cc: Added checks to ensureq_num_headsandkv_num_headsattributes are within the range[1, INT_MAX], preventing integer overflows during initialization.onnxruntime/contrib_ops/cpu/bert/longformer_attention_base.h: Added validation to ensurenum_headsandwindowattributes are within[1, INT_MAX].onnxruntime/contrib_ops/cpu/transformers/subgraph_gpt.cc: Added checks to ensure certain tensor dimensions do not exceedINT_MAX, preventing unsafe casts.Testing enhancements:
onnxruntime/test/contrib_ops/linear_attention_op_test.cc: Added new tests to verify that theLinearAttentionoperator correctly rejects out-of-rangeq_num_headsandkv_num_headsvalues, ensuring the new validation logic is exercised.Code maintenance:
<limits>header where necessary to support the new validation checks. [1] [2] [3]