Kernels:

---

webgpu-kernels

/

com.microsoft.RotaryEmbedding

Like 0

Follow

WebGPU Kernels Community 38

---
library_name: kernels
license: apache-2.0
tags:
- kernel
- webgpu
- wgsl
---
# 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](https://github.com/microsoft/onnxruntime/blob/main/docs/ContribOperators.md#com.microsoft.RotaryEmbedding) 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`](build/webgpu/metadata.json) — kernel metadata (id, digests, per-variant templates, provenance)
- [`manifest.json`](build/webgpu/manifest.json) — the op contract (source of truth)
- [`test.json`](build/webgpu/test.json) — correctness cases
- [`bench.json`](build/webgpu/bench.json) — benchmark cases
- [`rotary-embedding-slices.wgsl.jinja`](build/webgpu/rotary-embedding-slices.wgsl.jinja)

## Use with `@huggingface/kernels`

```sh
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.

```js
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] },
});
```