Audit and fix out-of-range output writes in ONNX Runtime operator shape-inference functions. Use when reviewing or fixing a contrib (or standard) op TypeAndShapeInference where a getNumOutputs() guard precedes a write to a higher output index - optional trailing outputs make a smaller output count schema-valid, so getOutputType(index) can run one past the declared outputs at Graph::Resolve.
Audit and fix out-of-range output writes in ONNX Runtime operator shape-inference functions. Use when reviewing or fixing a contrib (or standard) op TypeAndShapeInference where a getNumOutputs() guard precedes a write to a higher output index - optional trailing outputs make a smaller output count schema-valid, so getOutputType(index) can run one past the declared outputs at Graph::Resolve.
Contrib-Op Shape-Inference Output-Index Safety
Reusable method for finding and fixing the bug class where an operator's
TypeAndShapeInference function guards an output write with getNumOutputs() > N but then
writes an output index greater thanN. For a node that declares fewer outputs, the
written index is past the end of the inference context's output vector.
Scope: schema-level shape inference in onnxruntime/core/graph/contrib_ops/*.cc and
shape_inference_functions.cc. This runs once during Graph::Resolve (model-load time),
EP-agnostic - there is no per-EP (CPU/CUDA/ROCm) kernel duplicate of this code to
chase. Op kernels allocate outputs via the bounds-safe OpKernelContext::Output(index)
and are a separate concern.
The guard getNumOutputs() > 1 admits a node with exactly 2 outputs (indices 0, 1), yet
the body writes index 2. The implication "> 1 ⇒ index 2 exists" is false: > 1 only
guarantees indices 0 and 1.
Why a smaller output count is valid
Trailing outputs declared OpSchema::Optionallower min_output. ONNX derives
min_output = number of required outputs, max_output = total declared. The model checker
(checker::check_node) only enforces min_output <= N <= max_output.
Op
Output decls
min / max
A 2-output node?
DecoderAttention
out (req), new_key_cache (Opt), new_value_cache (Opt)
1 / 3
passes checker
MultiHeadAttention
out (req), present_key (Opt), present_value (Opt), qk (Opt)
1 / 4
passes checker
DecoderMaskedMultiHeadAttention
out (req) + 3 Optional
1 / 4
passes checker
So a node with output=['out','present_key'] is schema-valid, passes the checker, and then
reaches the index-2 write. A passing checker is not a guarantee the index is in range.
node_output_types_ is sized to node.OutputDefs().size() in the InferenceContextImpl
ctor, so for a 2-output node it has 2 elements; getOutputType(2) returns one past the end.
getInputType uses .at() (would throw on a bad index); getOutputType uses raw
operator[] (no check) - the asymmetry is the root cause.
The call runs at Graph::Resolve → InferAndVerifyTypeMatch → RunInferencing. The
surrounding ORT_TRY/ORT_CATCH(const std::exception&) only catches thrownfail_shape_inference; a raw out-of-range operator[] does not throw, so the catch does
not help.
Because this is schema-level inference, it is EP-independent - no CUDA/ROCm copy.
3. Audit technique — always sweep siblings
Do not stop at the reported function. Grep every shape-inference guard and compare its
threshold against the highest output index written before the next guard.
For each if (ctx.getNumOutputs() > N) block, find the largest index passed to
propagateElemTypeFromInputToOutput(ctx, _, index) / updateOutputShape(ctx, index, _) /
getOutputType(index) inside it. Rule: the guard must require strictly more outputs than
the highest index written (write index k ⇒ guard must ensure getNumOutputs() > k).
Correct exemplars already in the tree to copy:
Exemplar
Pattern
Why it is safe
BaseGroupQueryAttention...
if (getNumOutputs() >= 3) then writes idx 2
guard covers highest index
PagedAttention...
nested > 1 + inner if (getNumOutputs() != 3) fail_shape_inference
fails before any write
EmbedLayerNormalizationShapeInference
> 2 then writes idx 2
fixed by PR #28176 (precedent)
SkipLayerNormalizationShapeInference
each idx k guarded by > k
per-index guard
Gotcha — conditional writes can hide a vacuous audit. A write may sit behind an inner
condition (e.g. hasInputShape(past_key_index) before writing index 2). The site is still
a bug, but you can only observe it when that inner condition is also satisfied. Keep this
in mind both for the audit and for tests (§5).
4. Fix patterns
Point fix (required): raise the guard to cover the highest index written.
Justify the threshold with the op's output semantics. For these attention ops the two trailing
outputs - present_key (idx 1) and present_value (idx 2) for MultiHeadAttention,
new_key_cache / new_value_cache for DecoderAttention (see the §1 table for each op's
exact output names) - are a both-or-neither pair: there is no valid configuration that
emits one without the other, so requiring all three outputs before populating indices 1 and 2
is behavior-preserving. (PagedAttention encodes the same invariant via its nested != 3
check.)
Defense-in-depth (recommended): bound the sink so a future author cannot reintroduce the
class.
// onnxruntime/core/graph/graph.cc - InferenceContextImpl::getOutputTypeTypeProto* getOutputType(size_t index)override{
if (index >= node_output_types_.size()) {
fail_type_inference("output index ", index, " is out of range; node has ",
node_output_types_.size(), " outputs");
}
return &node_output_types_[index];
}
This mirrors getInputType's .at() and the existing bounds checks in the sibling
DataPropagationContextImpl. Placing it at the base layer transitively protects the NHWC and
quantization wrapper contexts. After the point fix this branch is unreachable through a normal
model (the guard already prevents the out-of-range index), so it is pure defense-in-depth. Its
failure mode is build-dependent: with exceptions enabled, fail_type_inference raises
InferenceError (a std::exception), caught by the existing ORT_CATCH(const std::exception&)
around RunInferencing and surfaced as a clean load-time error; under ORT_NO_EXCEPTIONS it is
not compiled out - ONNX's no-exceptions path prints the message to std::cerr and calls
abort(), a deterministic fail-fast (consistent with getInputType's .at(), which likewise
terminates under no-exceptions). Either way the result is a controlled failure rather than an
out-of-range write.
5. Test recipe
Tests live in onnxruntime/test/contrib_ops/*.cc and are auto-globbed into the
onnxruntime_provider_test target by cmake/onnxruntime_unittests.cmake
(test/contrib_ops/*.cc pattern) - no cmake edit needed for a new file. See the
ort-test skill for the executable taxonomy (onnxruntime_provider_test vs
onnxruntime_test_all).
Rules that make the regression test actually guard the fix:
Drive through Model + Graph::Resolve, not ONNX's standalone TestShapeInference.
Only the full resolve path constructs the real InferenceContextImpl and hits the
getOutputType sink described in §2. A standalone ONNX shape-inference helper uses a
different context and bypasses the sink, so it cannot reproduce the bug.
Negative tests must be NON-VACUOUS - they must actually enter the write branch on
pre-fix source. If a write is gated by an inner condition (§3 gotcha), satisfy it: e.g. for
MultiHeadAttention/DecoderMaskedMultiHeadAttention, supply a shaped past_key
(and past_sequence_length / past_present_share_buffer as the op requires) so the
index-2 block runs. A negative test that only supplies query skips the block and passes
even on pre-fix source - regression-proof in name only.
Add positive (all-outputs) cases: a node with every output present must still infer the
trailing output types - proves the tightened guard did not over-restrict.
Keep tests throw-free post-fix so they are valid under ORT_NO_EXCEPTIONS. Any case
that is expected to fail_shape_inference (throws) must be excluded with
#ifndef ORT_NO_EXCEPTIONS. The "2 outputs must not go out of range" case is throw-free
after the point fix and is safe in all builds.
Verify the negative test is non-vacuous (sanitizer A/B) - the most reliable way to prove a
negative test enters the previously-out-of-range branch: build the test at the pre-fix
commit with AddressSanitizer and confirm it flags the out-of-range output access; then
confirm it is clean after the fix.
# Functional run (any Debug build):
cmake --build build/Linux/Debug --target onnxruntime_provider_test -j"$(nproc)"
./build/Linux/Debug/onnxruntime_provider_test \
--gtest_filter='AttentionOptionalOutputsShapeInferenceTest.*'# A/B proof (isolated worktree at the pre-fix commit, CPU-only Debug + sanitizer):
git worktree add --detach ../ort-prefix-check <fix_commit>~1
# copy the new test file in, then:
python3 tools/ci_build/build.py --build_dir build/asan --config Debug --parallel \
--skip_tests --enable_address_sanitizer --skip_submodule_sync \
--cmake_generator Ninja --target onnxruntime_provider_test
# Pre-fix: the negative tests fail (the sanitizer flags the out-of-range output access).# Post-fix (cherry-pick the guard fix): all tests pass, no sanitizer report.
6. Process / wording conventions
Run lintrunner -a before pushing so the CLANGFORMAT / Python-format gate passes. See
the ort-lint skill.
Use correctness/robustness framing in code, comments, commit messages, and the PR body
describe the change as fixing an optional-output guard, not as a security fix. This
matches repo convention (compare python-kwargs-setattr-security) and keeps the PR neutral.
7. Audit checklist (per-operator review)
When reviewing or hardening any operator implementation or its shape inference:
Read the op's spec - ONNX standard op page, or for a contrib op its OpSchema
registration (.Input/.Output/.Attr, and Optional/Variadic markers). A local ONNX
checkout has the standard-op spec pages; contrib ops are defined only in ORT.
Enumerate all inputs, attributes, and outputs, noting which are optional and the
resulting min/max input and output counts.
Validate every input/attribute before indexing into it, to avoid out-of-range reads
(which can cascade into worse failures). Match each output-index write to a guard that
guarantees the index is in range (§3 rule).
Prefer ORT_RETURN_IF / ORT_RETURN_IF_NOT for validation; use ORT_ENFORCE in
constructors. In shape inference use fail_shape_inference / fail_type_inference.
Use SafeInt<> / narrow<>() for index and size arithmetic and casts to avoid overflow
or truncation that yields a wrong index. See core/common/safeint.h and
docs/Coding_Conventions_and_Standards.md.
Ensure tests build and pass under no-exceptions builds; #ifndef ORT_NO_EXCEPTIONS
around any case expected to throw.
Exclude EPs known not to support the op, with a comment explaining why.
Check whether other EPs (notably CUDA/ROCm) implement the same op and whether the
same issue exists there. (For shape inference specifically, the logic is EP-agnostic
and single-source - confirm there is no kernel-side analogue.)
References
PR #28176 - "Fix ... in EmbedLayerNormalizationShapeInference": the precedent that fixed
the identical > 1 → > 2 primitive in one site; the sibling attention sites were missed,
motivating the sweep in §3.
PR #29268 - this fix: guards corrected in DecoderAttention / MultiHeadAttention /
DecoderMaskedMultiHeadAttention shape inference, plus the getOutputType bounds check and
non-vacuous regression tests.
