onnxruntime
a83b776e - Fix optional-output guard in DecoderAttention/MultiHeadAttention shape inference (#29268)

Commit
13 days ago
Fix optional-output guard in DecoderAttention/MultiHeadAttention shape inference (#29268) ### Description The `DecoderAttention` and `MultiHeadAttention` shape-inference functions guarded population of their optional `present_key` (output 1) and `present_value` (output 2) outputs with `getNumOutputs() > 1`, but then write output index 2. `present_key` and `present_value` are produced as a both-or-neither pair, so this requires all three outputs (`> 2`) to be present before populating them — matching the existing `BaseGroupQueryAttention` (`>= 3`) and `EmbedLayerNorm` guards. It also adds an output-index range check in `InferenceContextImpl::getOutputType` so an output index beyond the declared output count fails inference cleanly instead of indexing past the end of the outputs container, mirroring the existing `DataPropagationContextImpl::getOutputType` and `getInputType` behavior. ### Motivation and Context A model that declares fewer outputs than the optional present outputs could previously drive shape inference to access an output index that was not declared. This makes the guard consistent with the other attention-family contrib ops. ### Changes - `onnxruntime/core/graph/contrib_ops/bert_defs.cc` — require all present outputs before populating `present_key`/`present_value` in `DecoderAttention` and `MultiHeadAttention`. - `onnxruntime/core/graph/graph.cc` — add an output-index range check in `InferenceContextImpl::getOutputType`. - `onnxruntime/test/contrib_ops/attention_optional_outputs_shape_inference_test.cc` — regression tests covering omitted optional present outputs, the 3-output positive cases, and the MHA/DMMHA two-output cases. - Adds a contrib-op shape-inference output-index safety skill doc plus a one-line coding-convention note. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --------- Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Author
Parents
Loading