[Feature][Relax][ONNX] Support RotaryEmbedding (RoPE) operator in Relax ONNX frontend #20967

Feature Description

Support the ONNX RotaryEmbedding (Rotary Position Embedding / RoPE) operator in the Apache TVM Relax ONNX frontend (python/tvm/relax/frontend/onnx/onnx_frontend.py).

Motivation

Rotary Position Embedding (RoPE) has become the de facto standard positional encoding technique for modern Large Language Models (LLMs), including LLaMA (1, 2, 3), Mistral, Gemma, Qwen 2/2.5, and DeepSeek.

When exporting modern LLMs from PyTorch / Hugging Face Transformers to ONNX, attention modules frequently emit the RotaryEmbedding operator (both in standard ONNX opset 23+ and the common com.microsoft domain). Attempting to import these models into TVM Relax currently fails with:

NotImplementedError: The following operators are not supported: ['RotaryEmbedding']

Adding native converter support for RotaryEmbedding in the Relax ONNX frontend will enable direct compilation of state-of-the-art generative LLMs into Relax IR without requiring manual graph surgery or decomposition passes.

Operator Specification (ONNX Opset 23+ / com.microsoft)

• Inputs:
  1. X: Input tensor of shape (batch_size, sequence_length, num_heads, head_size) or (batch_size, num_heads, sequence_length, head_size).
  2. pos_ids: 1D or 2D tensor specifying token position IDs for each sequence.
  3. cos_cache: Precomputed cosine frequency cache matrix.
  4. sin_cache: Precomputed sine frequency cache matrix.
• Attributes:
  • interleaved: int (default 0) - whether real and imaginary components are interleaved (e.g. [-x2, x1, -x4, x3]) or separated along half the head dimension (e.g. [-x_half2, x_half1]).
  • num_heads: int (optional) - number of attention heads.
  • rotary_embedding_dim: int (optional) - dimension to rotate (can be smaller than head_size for partial rotary embeddings).
• Outputs:
  1. Y: Output tensor with the same shape and dtype as X with rotary embeddings applied.

Expected Behavior

The Relax ONNX frontend should implement RotaryEmbedding by slicing the rotary dimension, rotating components according to the interleaved flag, multiplying with corresponding cos and sin tables, and concatenating any pass-through non-rotated dimensions.