RotaryEmbedding (CPU)#

  • Domain: com.microsoft

  • Device: CPU

  • Kernel: onnx_light_cpu::RotaryEmbedding

  • Supported types: FLOAT, FLOAT16, BFLOAT16

  • Opset bounds: since opset 1

Operator support#

Capability

Implementation

Shape inference

onnx_light_cpu::ComputeShapeRotaryEmbedding()

Peak memory

onnx_light_cpu::ComputePeakMemoryRotaryEmbedding()

Fusion patterns

None

Gradient

Not provided

LightOpSchema#

Version: 1

Rotates the first rotary_embedding_dim values of each head using the indexed cosine and sine caches. Supports scalar start offsets or per-token position ids; packed batching and dynamic cache extension are not supported.

Inputs

  • input (T): Rank-3 (batch, sequence, hidden) or rank-4 (batch, heads, sequence, head).

  • position_ids (M): INT64 scalar offset or (batch, sequence) positions.

  • cos_cache (T): Rank-2 cosine table with rotary dimension / 2 or head size / 2 columns.

  • sin_cache (T): Matching sine table.

Outputs

  • output (T): Rotated tensor with input shape and type.

Attributes

  • interleaved (INT): Interleaved layout if 1. (optional, default: 0)

  • is_packed_batching (INT): Only 0 is supported. (optional, default: 0)

  • num_heads (INT): Head count for rank-3 input. (optional, default: 0)

  • rotary_embedding_dim (INT): Zero rotates the full head. (optional, default: 0)

  • scale (FLOAT): Used only by upstream dynamic cache extension. (optional, default: 1)

Type constraints

  • T: Matching input and cache types. Allowed types: tensor(bfloat16), tensor(float), tensor(float16).

  • M: Position indices. Allowed types: tensor(int64).