onnxruntime
2d9f35fc - [Security] Validate CPU cache_indirection beam indices (#28486)

Commit
150 days ago
[Security] Validate CPU cache_indirection beam indices (#28486) ## Summary - reject out-of-range `cache_indirection` beam indices in the CPU beam-attention path before they are converted into past KV offsets - keep `DecoderMaskedMultiHeadAttention` beam-width handling consistent with the `cache_indirection` shape - add CPU regression tests for `MultiHeadAttention` and `DecoderMaskedMultiHeadAttention` ## Motivation `MultiHeadAttention` and `DecoderMaskedMultiHeadAttention` on the CPU provider could consume attacker-controlled `cache_indirection` values as beam indices without validating that each element stayed within `[0, beam_width)`. That let malformed models compute offsets past the past key/value buffers. This change rejects invalid indices up front and adds focused tests for the failure path. ## Key Changes - add shared CPU validation in `AttentionCPUBase::ApplyAttentionWithBeams` so the beam path fails before any past-key or past-value reads occur - report an `INVALID_ARGUMENT` error that identifies the offending beam index and its position - validate that an explicit decoder `beam_width` input matches `cache_indirection` dimension 1 when both are present - add contrib-op tests that exercise invalid cache indirection values on the CPU execution provider ## Testing - `lintrunner -a` - `cd build/Linux/Debug && make -j4 CMakeFiles/onnxruntime_providers.dir/home/tlwu/onnxruntime/onnxruntime/contrib_ops/cpu/bert/multihead_attention.cc.o CMakeFiles/onnxruntime_providers.dir/home/tlwu/onnxruntime/onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc.o CMakeFiles/onnxruntime_provider_test.dir/home/tlwu/onnxruntime/onnxruntime/test/contrib_ops/multihead_attention_op_test.cc.o CMakeFiles/onnxruntime_provider_test.dir/home/tlwu/onnxruntime/onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc.o` - full `onnxruntime_provider_test` relink/run was not completed locally --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com>
Author
Parents
Loading