Xenova HF Staff

sync 2e7068faf55e

5ce7237 verified 5 days ago

5.43 kB

library_name: kernels
license: apache-2.0
tags:
  - kernel
  - webgpu
  - wgsl

com.microsoft.MRotaryEmbedding

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

Description

Multimodal rotary position embedding (M-RoPE) for Qwen models. Each token has temporal, height, and width position streams; mrope_section partitions the half-rotary axis and mrope_layout assigns them. Text-only tokens set all streams equal, reducing the op to RotaryEmbedding. The effective rotary dimension must be positive and even; an odd head size is supported with a smaller even rotary_embedding_dim. This package supports float16/float32 and non-packed mode; bfloat16 and packed batching are not implemented. Position ids must be valid non-negative cache-row indices.

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

Inputs

Name	Bind key	Logical dtype	WebGPU storage	Rank	Shape	Description	Presence
`input`	`x`	`T`	same as logical dtype	—	—	Input token embeddings. Shape is `(batch_size, sequence_length, hidden_size)` for rank 3 or `(batch_size, num_heads, sequence_length, head_size)` for rank 4. The effective rotary dimension must be even, and `num_heads` is required for rank-3 input.	required
`position_ids`	`positionIds`	`M`	`uint32`	`3`	—	Logical int64 position indices of shape `(3, batch_size, sequence_length)`, holding the temporal, height and width streams in that order along the first axis. Valid positions are non-negative cache-row indices and use uint32 WebGPU storage.	required
`cos_cache`	`cos`	`T`	same as logical dtype	`2`	—	Precomputed cosine values of shape `(max_sequence_length, rotary_dim/2)`, shared by all three position streams.	required
`sin_cache`	`sin`	`T`	same as logical dtype	`2`	—	Precomputed sine values with the same shape and type as `cos_cache`.	required

Outputs

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

Attributes

Attributes and default values (overridable per request):

Attribute	Default	Description
`interleaved`	`0`	Set to 1 to rotate using an interleaved pattern (even/odd elements), or 0 to split the head dimension into two contiguous halves. Default is 0. This is the rotation pairing and is independent of `mrope_layout`.
`is_packed_batching`	`0`	Whether `position_ids` uses packed-batch metadata. The default and only supported value is 0; packed batching (1) is not implemented.
`mrope_layout`	`0`	How the three sections are combined into one per-token cos/sin vector: `0` for the sectioned/chunked layout (Qwen2-VL, Qwen2.5-VL) or `1` for the interleaved layout (Qwen3-VL, Qwen3.5). Default is 0.
`num_heads`	`0`	Number of attention heads. The schema default is 0. A positive value is required for rank-3 `input` and whenever `rotary_embedding_dim` is nonzero; rank-4 execution otherwise infers the head count from `input`.
`rotary_embedding_dim`	`0`	Positive even number of head-dimension elements to rotate; `0` means the full head dimension, which must then be even. A smaller even value permits an odd head size and copies the remaining tail unchanged.
`scale`	`1`	Scale applied to the gathered cosine and sine values before the rotation. Default is 1.0.
`mrope_section`	—	Three non-negative integers `[section_t, section_h, section_w]` dividing the half-rotary axis among the temporal, height and width streams. They must sum to `rotary_embedding_dim / 2`, or to `head_size / 2` when `rotary_embedding_dim` is 0. Required.

Type constraints

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

Files

metadata.json — kernel metadata (id, digests, provenance)
manifest.json — the op contract (source of truth)
test.json — correctness cases
bench.json — benchmark + tuning cases
mrotary-embedding.wgsl.jinja

Use with `@huggingface/kernels`

The loader derives every required output's shape and logical dtype from the manifest contract and this call. It then allocates the result tensors automatically.

The version: 1 option selects the published kernel contract; it is independent of any operator opset, contrib since_version, or model 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.MRotaryEmbedding", { version: 1 });
const { y } = await kernel({
  x: { data: xData, shape: [1, 4, 16] },
  positionIds: { data: positionIdsData, shape: [3, 1, 4] },
  cos: { data: cosData, shape: [8, 4] },
  sin: { data: sinData, shape: [8, 4] },
}, {
  attrs: { num_heads: 2, mrope_section: [2, 1, 1] },
});