Skip to content

com.microsoft.RotaryEmbedding: validate position_ids are within cos/sin cache bounds#28206

Draft
Copilot wants to merge 1 commit intomainfrom
copilot/fix-rotary-embedding-boundary-check
Draft

com.microsoft.RotaryEmbedding: validate position_ids are within cos/sin cache bounds#28206
Copilot wants to merge 1 commit intomainfrom
copilot/fix-rotary-embedding-boundary-check

Conversation

Copy link
Copy Markdown
Contributor

Copilot AI commented Apr 23, 2026

Description

Adds explicit bounds-checking for position_ids in the CPU RotaryEmbedding kernel before indexing into the cos/sin cache. Previously, out-of-range values caused silent UB (out-of-bounds memory reads).

RunRotaryEmbedding (rotary_embedding.cc)

  • Format 0 (scalar): rejects if base_pos < 0 or base_pos + sequence_length > max_sequence_length
  • Format 1 (2D [batch, seq]): iterates all elements and rejects if any value is outside [0, max_sequence_length)
  • Returns INVALID_ARGUMENT with the offending value and its flat index

New CPU tests (rotary_embedding_op_test.cc)

  • Format 1: value exceeds max_sequence_length, negative value, OOB value in a batch
  • Format 0: base offset + sequence length exceeds cache size, negative base offset

Motivation and Context

When a caller passes position_ids with values outside the cos/sin cache range, the kernel would silently read garbage memory. The op should fail fast with a descriptive error rather than producing corrupted output or crashing.

Note: the CUDA kernel does not yet have equivalent device-side validation and is a follow-up.

Copilot AI changed the title [WIP] Fix RotaryEmbedding to handle out-of-bound position_ids com.microsoft.RotaryEmbedding: validate position_ids are within cos/sin cache bounds Apr 23, 2026
Copilot AI requested a review from xadupre April 23, 2026 12:58
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

com.microsoft.RotaryEmbedding: returns if position_ids are out of boundary

2 participants