Contrib Op Shape Inference Memory Safety

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 than N. 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.

1. The pattern

// onnxruntime/core/graph/contrib_ops/bert_defs.cc  (before)
propagateElemTypeFromInputToOutput(ctx, 0, 0);
if (ctx.getNumOutputs() > 1) {                       // guard says "> 1"
  propagateElemTypeFromInputToOutput(ctx, 0, 1);
  propagateElemTypeFromInputToOutput(ctx, 0, 2);     // but writes index 2
}

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::Optional lower 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.

2. The sink (why the write is not caught)

// onnxruntime/core/graph/graph.cc  -  InferenceContextImpl
const TypeProto* getInputType(size_t index) const override {
  return node_.InputDefs().at(index)->TypeAsProto();   // .at()  -> bounds-checked
}
TypeProto* getOutputType(size_t index) override {
  return &node_output_types_[index];                   // operator[]  -> NOT bounds-checked
}

• 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 *thrown*

fail_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.

git grep -n 'getNumOutputs' -- \
  onnxruntime/core/graph/contrib_ops/*.cc \
  onnxruntime/core/graph/contrib_ops/shape_inference_functions.cc

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.

// before
if (ctx.getNumOutputs() > 1) { ... writes idx 2 ... }
// after
if (ctx.getNumOutputs() > 2) {  // both present_key (idx 1) AND present_value (idx 2)
  ...
}

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::getOutputType
TypeProto* 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.

• Sibling skill: ort-test (test executables, --gtest_filter, contrib-op test layout);

ort-lint (lintrunner -a); ort-build (build flags, ASan).