Fix rotary embedding oob issue
Merged
AP
4mo ago+85+85+01234567890123456789-2-2-0123456789
This pull request improves the validation logic for the RotaryEmbedding operator to prevent out-of-bounds reads when the rotary embedding dimension derived from cos_cache exceeds the input tensor's hidden_size. It also adds dedicated unit tests to verify that this validation triggers as expected.
Validation improvements:
- Added a check in
rotary_embedding_helper.hto ensure that the effective rotary embedding dimension (cos_cachewidth × 2) does not exceedhidden_sizewhenrotary_embedding_dimis 0, returning an error if the condition is violated.
Unit test additions:
- Added
ContribRotaryEmbedding_RejectsCosCacheExceedsHiddenSizetest inrotary_embedding_op_test.ccto verify that an invalid configuration is correctly rejected in the contrib op. - Added
RotaryEmbedding_RejectsCosCacheExceedsHiddenSizetest inrotary_embedding_op_test.cc(providers/cpu/llm) to verify the same validation in the mainline op.