Fix rotary embedding oob issue

Merged
AP
4mo ago
+85-2

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.h to ensure that the effective rotary embedding dimension (cos_cache width × 2) does not exceed hidden_size when rotary_embedding_dim is 0, returning an error if the condition is violated.

Unit test additions:

  • Added ContribRotaryEmbedding_RejectsCosCacheExceedsHiddenSize test in rotary_embedding_op_test.cc to verify that an invalid configuration is correctly rejected in the contrib op.
  • Added RotaryEmbedding_RejectsCosCacheExceedsHiddenSize test in rotary_embedding_op_test.cc (providers/cpu/llm) to verify the same validation in the mainline op.