Harden shape inference for contrib ops - #32607
Conversation
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
|
Dwayne Robinson (@fdwr) PTAL, thanks! |
There was a problem hiding this comment.
🟡 Changes recommended
Invalid ndim values remain accepted when input shape metadata is absent, and the new guards lack regression tests.
Get a fresh assessment by requesting another Copilot review.
Pull request overview
Hardens contrib-op shape inference against malformed model metadata.
Changes:
- Validates CausalConvWithState ranks and
ndim. - Guards GatedDeltaNet width arithmetic against overflow.
- Rejects empty attention sequence-length initializers.
File summaries
| File | Description |
|---|---|
onnxruntime/core/graph/contrib_ops/bert_defs.cc |
Adds shape, attribute, initializer, and overflow validation. |
Review details
- Files reviewed: 1/1 changed files
- Comments generated: 4
- Review effort level: Balanced
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
There was a problem hiding this comment.
🟡 Changes recommended
SparseAttention validates the wrong input, GatedDeltaNet validation is output-dependent, and regression tests are missing.
Get a fresh assessment by requesting another Copilot review.
Review details
Suppressed comments (3)
Previously missed (2) — in code that hasn't changed since the last review.
onnxruntime/core/graph/contrib_ops/bert_defs.cc:309
ctx.getInputData(6)is thetotal_sequence_lengthinput for GroupQueryAttention, but input 6 isblock_col_indicesfor SparseAttention; SparseAttention's scalar is input 7 (bert_defs.cc:1861-1868). This therefore validates the wrong initializer for SparseAttention. Parameterize this input index and pass 6 for GroupQueryAttention and 7 for SparseAttention.
This issue also appears on line 307 of the same file.
onnxruntime/core/graph/contrib_ops/bert_defs.cc:3576
- The positivity check is nested under
getNumOutputs() > 2, so a GatedDeltaNet node requesting onlyoutput(oroutputplusfinal_state) can still resolve with a statically zero head count/size. Move validation of known head dimensions before the optionalstate_updatebranch; only the width-overflow calculation should depend on output 2 being present.
onnxruntime/core/graph/contrib_ops/bert_defs.cc:309
- This condition does not enforce the diagnostic's “single element” contract: a multi-element initializer is accepted and only
data[0]is used, while the runtime requires a scalar or one-element vector. Reject every size other than one.
if (data.empty()) {
fail_shape_inference("total_sequence_length input must contain a single element");
}
- Files reviewed: 1/1 changed files
- Comments generated: 1
- Review effort level: Balanced
|
Akshay Sonawane (@apsonawane) Addressed your comments and added some regression tests, please take antoher look, thanks! |
There was a problem hiding this comment.
🟡 Changes recommended
The new GatedDeltaNet test reads dangling initializer-list storage, and validation remains conditional on an optional output.
Get a fresh assessment by requesting another Copilot review.
Review details
- Files reviewed: 4/4 changed files
- Comments generated: 6
- Review effort level: Balanced
|
shiyi (@shiyi9801) can you review the copilot comments |
|
Akshay Sonawane (@apsonawane) Addressed the copilot comments, please take another look. |
There was a problem hiding this comment.
Copilot review overview
Review effort: Balanced
Findings: None
Resolved since last review (6)
fail_shape_inference, which aborts rather than returning a…Graph::Resolve()to catch afail_shape_inferenceexception, but no-exceptions…fail_shape_inference; in anORT_NO_EXCEPTIONS…state_updateoutput branch, so a node declaring…Casestores non-owninginitializer_listviews, but the backing arrays belong to temporaries…
bca9d15
into
microsoft:main
Summary
Strengthens attribute/input validation in three contrib-op shape-inference functions in bert_defs.cc so malformed or malicious models are rejected at
Graph::Resolveinstead of triggering out-of-bounds reads or signed-integer overflow during shape inference.Changes
CausalConvWithState: validate the ndim attribute is in [1, 3] and cross-check tensor ranks against it (weight == ndim+2, channels-first input == ndim+2, channels-last input >= 3) before the spatial-dim loop indexes input.dim(2+i). Previously only a rank >= 2 guard existed, so ndim=2/3 with a low-rank input read past the shape's dimensions.GatedDeltaNet: require the head counts/sizes to be positive and add step-by-step overflow guards before computing the state_update capsule width, preventing signed int64 overflow (UB) in state_update_capacity * (num_heads_v + num_heads_khead_size_qk + num_heads_vhead_size_v).GroupQueryAttention/SparseAttention: check the parsedtotal_sequence_lengthinitializer is non-empty before indexing data[0].