[Security] Validate CPU cache_indirection beam indices

Merged
TI
5mo ago
+124-0

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