library_name: kernels
license: apache-2.0

mrope

Interleaved multimodal rotary position embedding, applied in place, loadable through kernels. The reference baseline is transformers' Qwen3_5TextRotaryEmbedding, matched to 3.1e-7 in fp32, with an fp64 rotation as the accuracy oracle.

A multimodal model needs three positions per token rather than one: an image patch has a place in time, in height, and in width, so the rotary frequency band is partitioned across those axes. The interleaved layout cycles them, [THWTHW...], so each axis keeps a full spread of wavelengths. This kernel applies all of it in one pass, one block per token with the angles computed once into shared memory and reused across every head, and evaluates angles in double precision, which is what keeps a 262,144-token context accurate where the eager fp32 path has lost four orders of magnitude.

The kernel's own axis map probed live: 32 frequency dials colored time, height, width in the interleaved pattern, each phase advancing one axis and spinning only that axis's dials. At position 262,143 the kernel's angle errs 9e-8 radians against an exact reference; the fp32 eager path errs 8e-5.

Usage

import torch
from kernels import get_kernel

mrope = get_kernel("phanerozoic/mrope", version=1, trust_remote_code=True)

# q is [B, S, Hq, D], k is [B, S, Hk, D], pos is [3, B, S] int32 holding the
# time, height and width position of every token. Both rotate in place.
mrope.apply_mrope(q, k, pos, head_dim=256, partial_rotary_factor=0.25,
                  mrope_section=(11, 11, 10), rope_theta=1e7)

version selects the release branch; trust_remote_code is required by kernels for publishers without the trusted-publisher mark.

API

Method

Two details make this its own kernel rather than a case of ordinary RoPE. Partial rotary: only the first head_dim * partial_rotary_factor channels rotate, the rest pass through untouched; at head dimension 256 and factor 0.25 that is 64 rotated channels of 256. Interleaved sections: frequency j uses the time position unless j % 3 == 1 and j < 3 * section[1] (height) or j % 3 == 2 and j < 3 * section[2] (width); with sections [11, 11, 10] over 32 frequencies that reduces exactly to j % 3. One block per (batch, token) computes the 32 angles once into shared memory and reuses them across every head of that token.

Measured

Against the eager transformers implementation, head dimension 256, 24 query and 4 key heads:

The eager path builds a [3, B, S, 32] frequency tensor, slices it three ways, concatenates, takes a cosine and a sine, and runs four elementwise passes; this kernel touches each element once.

Accuracy at long positions, against an independent fp64 rotation:

Angles are evaluated in double before being reduced to float; the eager reference computes them in float32, which loses the fractional part of a cycle at large positions.

Correctness

Verified against transformers 5.13:

• The frequency-to-axis assignment is torch.equal to the reference slice construction, and per-axis counts equal the configured sections.
• Output matches to 3.1e-7 relative in fp32 and 1.0e-3 in bf16, across batch 1 to 3, 8 to 512 tokens, grouped and ungrouped head counts, with distinct positions on all three axes.
• Moving only the height position changes exactly the frequencies assigned to height, checked elementwise; each rotation pair keeps its magnitude to 1e-3; the tail beyond rotary_dim is bitwise unchanged; position zero is the exact identity.

Requirements and limits

• NVIDIA GPU with compute capability 8.0+.
• q and k contiguous [B, S, H, D], bf16 or fp32, sharing a dtype; [B, H, S, D] tensors need a permute first.
• rotary_dim / 2 at most 128; positions int32 (contexts to 2^31 tokens).
• Forward only: RoPE is its own inverse up to a sign, so backward is a second call with negated positions.

References

Su et al., "RoFormer" (2021); Wang et al., "Qwen2-VL" (2024) for multimodal rotary sections; the interleaved layout as implemented in transformers' Qwen3_5TextRotaryEmbedding.

License

Apache-2.0.