com.microsoft.RotaryEmbedding

com.microsoft · ONNX Runtime contrib operator · contrib since_version 1

Description

Rotary positional embedding (RoPE): each head's embedding vector is rotated with the cos_cache and sin_cache rows selected by position_ids, which is either a single base offset (token s reads row position_ids[0] + s) or a (batch_size, sequence_length) table. input is rank 3 (batch_size, sequence_length, hidden_size) or rank 4 (batch_size, num_heads, sequence_length, head_size). rotary_embedding_dim rotates a prefix and copies the tail unchanged. Rotation arithmetic is float32 with one narrowing store. Bfloat16 and non-default scale are not implemented.

See the ONNX Runtime RotaryEmbedding contrib-operator spec for the reference semantics.

Inputs

Name	Upstream name	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`x`	`input`	`T`	same as logical dtype	—	—	Input token embeddings, shaped `(batch_size, sequence_length, hidden_size)` at rank 3 or `(batch_size, num_heads, sequence_length, head_size)` at rank 4. At rank 3 the head size comes from `num_heads` when that attribute is positive and from `2 * cos_cache.shape[1]` otherwise; at rank 4 both the head count and head size are read from the shape.	required
`positionIds`	`position_ids`	`M`	`uint32`	—	—	Logical int64 cache-row selector in either upstream format: a scalar or one-element vector holding a base offset, so token `s` reads row `position_ids[0] + s`; or a `(batch_size, sequence_length)` table read per token. Valid positions are non-negative rows of the caches and use uint32 WebGPU storage, so a negative value is rejected at the host boundary.	required
`cos`	`cos_cache`	`T`	same as logical dtype	`2`	—	Precomputed cosine values of shape `(max_sequence_length, rotary_dim / 2)`, where `rotary_dim` is `rotary_embedding_dim` when that is positive and the head size otherwise.	required
`sin`	`sin_cache`	`T`	same as logical dtype	`2`	—	Precomputed sine values with the same shape and type as `cos_cache`.	required

Outputs

Name	Upstream name	Logical dtype	Rank	Shape	Description	Presence
`y`	`output`	`T`	same as `x`	same as `x`	Rotary-position-encoded tensor with the same shape and type as `input`.	required

Attributes

Default values (overridable per request):

Attribute	Default	Description
`interleaved`	`0`	Set to 1 to pair adjacent even/odd elements, or 0 to pair each element of the first half of the rotary window with the matching element of the second half. Default is 0.
`is_packed_batching`	`0`	Ragged (packed) batch inputs. Its only upstream effect is to lift the `sequence_length <= max_sequence_length` bound, and this implementation never imposes that bound: every gathered row is required to index the caches whatever the sequence length. Default is 0.
`num_heads`	`0`	Number of attention heads. Default is 0, which asks the rank-3 path to take the head size from the cache width instead; a positive value is required whenever `rotary_embedding_dim` is nonzero. At rank 4 the head count comes from the input shape and this attribute is not consulted.
`rotary_embedding_dim`	`0`	Positive even count of leading head-dimension elements to rotate; the remaining tail is copied unchanged. Default is 0, meaning the whole head dimension, which must then be even. An odd value is rejected.
`scale`	`1`	Declared scale for the gathered rotation. No ONNX Runtime provider applies it, so the default and only accepted value is 1.0 and a caller's other value is rejected rather than silently discarded.

Type constraints

Variable	Allowed dtypes
`T`	`float32`, `float16`
`M`	`int64`

Files

metadata.json — kernel metadata (id, digests, per-variant templates, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark cases
rotary-embedding-slices.wgsl.jinja

Use with `@huggingface/kernels`

npm install --save-exact @huggingface/kernels@0.0.1-preview.3

Required output shapes and logical data types are inferred from the supplied inputs and attributes; result tensors are allocated automatically.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model version. It follows the v1 branch as fixes land. To pin exact artifact bytes, pass a 40-character commit revision instead of version.

Replace each *Data placeholder with a typed array containing the corresponding input data.

import { getKernel } from "@huggingface/kernels";

const kernel = await getKernel("webgpu-kernels/com.microsoft.RotaryEmbedding", { version: 1 });
const { y } = await kernel({
  x: { data: xData, shape: [1, 2, 18] },
  positionIds: { data: positionIdsData, shape: [1, 2] },
  cos: { data: cosData, shape: [4, 3] },
  sin: { data: sinData, shape: [4, 3] },
});

Downloads last month: -

kernel

webgpu

wgsl

apache-2.0

WebGPU

Requires WebGPU support. See the compatibility table.