Widen CUDA RotaryEmbedding offset arithmetic to 64-bit

[Akshay Sonawane (apsonawane)]

The kernel built tensor offsets as b * in_strides.x + s * in_strides.z + n * in_strides.y with every operand a 32-bit int. Tensors may hold more than INT32_MAX elements, so these products could wrap for large batch sizes even though each individual stride fits in an int, yielding negative offsets into the input, output, and cos/sin cache pointers.

Compute the offsets, b_s_index, and cache_offset in 64-bit, and verify the packed stride products fit in int32 before launching. Applies the same change to the contrib variant, which shares the pattern.

[Copilot]

Pull request overview

This PR updates the CUDA RotaryEmbedding kernels to avoid 32-bit overflow when computing tensor element offsets, by widening offset arithmetic (and related indices) to 64-bit and adding pre-launch stride validation. The same pattern is applied to both the core CUDA LLM implementation and the contrib BERT variant.

Changes:

• Compute input/output base offsets, cache offsets, and related indices in int64_t inside the CUDA kernels.
• Add host-side guards to ensure packed stride products fit in int before populating int4 stride structs and launching the kernel.
• Apply equivalent fixes to the contrib rotary embedding implementation.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/core/providers/cuda/llm/rotary_embedding_impl.cu	Widen kernel offset/index arithmetic to 64-bit and add stride overflow guards before launch.
onnxruntime/contrib_ops/cuda/bert/rotary_embedding_impl.cu	Mirror the 64-bit offset/index changes and stride guards for the contrib CUDA rotary embedding kernel.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.