[Security] Validate CPU cache_indirection beam indices - #28486
#28486[Security] Validate CPU cache_indirection beam indices#28486
Conversation
…ests Both tests crashed with null pointer dereferences because present_key and present_value were marked as optional output edges. The DMMHA and MHA compute functions unconditionally dereference present_key/present_value when past_present_share_buffer is enabled, causing SEGV/SEH exceptions before the cache indirection validation could run. Provide actual output tensors so the code reaches ValidateCacheIndirectionValues.
…ember DecoderMaskedMultiHeadAttentionParameters shadows AttentionParameters::beam_width with a default of 1. CheckInputs sets the base class member via reinterpret_cast, but the Compute function reads the derived class's shadowed member which stays at 1. Read the beam width from cache_indirection dim 1 directly instead.
There was a problem hiding this comment.
Pull request overview
This PR hardens the CPU-provider beam-attention path by validating cache_indirection beam indices before they are used to compute past-K/V offsets, and adds regression coverage for invalid indices.
Changes:
- Add CPU-side validation for
cache_indirectionindices in the beam attention path (AttentionCPUBase::ApplyAttentionWithBeams). - Keep
DecoderMaskedMultiHeadAttentionbeam width handling consistent withcache_indirection’s shape (and error out on mismatch). - Add CPU regression tests for invalid
cache_indirectionvalues inMultiHeadAttentionandDecoderMaskedMultiHeadAttention.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cpu/bert/attention_cpu_base.h | Adds shared validation for cache_indirection indices in the CPU beam path. |
| onnxruntime/contrib_ops/cpu/bert/decoder_masked_multihead_attention.cc | Aligns decoder beam width handling with cache_indirection shape and validates mismatch. |
| onnxruntime/test/contrib_ops/multihead_attention_op_test.cc | Adds a CPU regression test expecting failure on out-of-range cache_indirection values. |
| onnxruntime/test/contrib_ops/decoder_masked_multihead_attention_op_test.cc | Adds a CPU regression test for invalid cache_indirection indices in decoder MHA. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
|
@copilot apply changes based on the comments in this thread |
… test Agent-Logs-Url: https://github.com/microsoft/onnxruntime/sessions/45ea6801-f9bc-4978-9f4f-ba63c70fbde2 Co-authored-by: tianleiwu <30328909+tianleiwu@users.noreply.github.com>
Implemented in 18c947b. I updated CPU cache_indirection validation to also check beam_width==1 (only index 0 allowed), and added a new MultiHeadAttention CPU regression test for cache_indirection shape with beam_width==1 containing an invalid index. I also rebuilt and ran the targeted regression tests, which now pass. |
|
There is a known parity issue between the CPU and CUDA implementations for
If I recall correctly, the CPU unit tests are not correct as well and the Whisper CPU model with Does this PR fix the issue? |
No, this PR is for security. That parity issue can be addressed in another PR. |
Summary
cache_indirectionbeam indices in the CPU beam-attention path before they are converted into past KV offsetsDecoderMaskedMultiHeadAttentionbeam-width handling consistent with thecache_indirectionshapeMultiHeadAttentionandDecoderMaskedMultiHeadAttentionMotivation
MultiHeadAttentionandDecoderMaskedMultiHeadAttentionon the CPU provider could consume attacker-controlledcache_indirectionvalues 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
AttentionCPUBase::ApplyAttentionWithBeamsso the beam path fails before any past-key or past-value reads occurINVALID_ARGUMENTerror that identifies the offending beam index and its positionbeam_widthinput matchescache_indirectiondimension 1 when both are presentTesting
lintrunner -acd 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.oonnxruntime_provider_testrelink/run was not completed locally