Rotary Position Embeddings (RoPE)#

Rotary position embedding (RoPE) modules and primitives.

Overview#

Rotary Position Embedding (RoPE) encodes token position by rotating query and key vectors before the attention dot-product. Because the dot-product of a rotated query and a rotated key depends only on the relative angle between them, RoPE gives attention position-awareness without adding any learned parameters. No positional vectors are added to the token features — instead, the position is woven into the rotation of each head’s Q/K projections.

This module exposes two levels of API:

Shared table-provider modules (owners of the cos/sin tables):

RotaryEmbedding2DTables — owns axial 2D RoPE cos/sin tables for an \(h \times w\) token grid, in explicit (h, w, head_dim) layout.
RotaryEmbedding1DTables — owns standard 1D sequence RoPE cos/sin tables of shape (max_seq_len, head_dim).

A provider holds no projections and applies no rotation itself: its forward simply returns the (cos, sin) tables. The intended pattern is that a top-level, multi-block model constructs a single provider and passes the returned tables into every attention block’s forward (which rotates Q/K with the functional apply_rotary_pos_emb()), so the tables are built, stored, and — under domain parallelism — sharded exactly once instead of once per block. See RopeNatten2DSelfAttention and DiT for a reference wiring.

Low-level functional helpers (build_axial_rope_cos_sin_2d(), build_rope_cos_sin_1d(), apply_rotary_pos_emb()):

Used internally by the providers above and by attention implementations that need direct control over the table layout (e.g. NATTEN windowed attention, which keeps explicit spatial (h, w) dimensions, or domain-parallel paths that shard the tables across GPUs).

Choosing the right API#

Building a multi-block transformer (2D grid or 1D sequence)? Construct one RotaryEmbedding2DTables / RotaryEmbedding1DTables at the top level of the model and share its tables across every block, applying them with apply_rotary_pos_emb().
Implementing a single attention block or need full control over the table layout? Call the functional helpers directly and apply them with apply_rotary_pos_emb().

Math (axial 2D RoPE)#

head_dim is split in half: the first half rotates by row index, the second by column index. Each axis has head_dim/4 rotation pairs sharing a frequency \(\theta_k = \text{base}^{-2k/(head\_dim/2)}\) for \(k = 0 \ldots head\_dim/4 - 1\). For an adjacent channel pair \((x_a, x_b)\) at angle \(\phi\), the rotation is \((x_a \cos\phi - x_b \sin\phi,\ x_a \sin\phi + x_b \cos\phi)\).

class physicsnemo.nn.module.rope.RotaryEmbedding1DTables(*args, **kwargs)[source]#

Bases: Module

Shared owner of standard 1D RoPE cos/sin tables for a token sequence.

This module owns the cos/sin tables and nothing else: it holds no projections and applies no rotation. Its forward() returns the (cos, sin) tables of shape \((seq\_len, head\_dim)\), which a consumer applies to its query/key with apply_rotary_pos_emb(). This is the same RoPE variant used by most autoregressive and encoder transformer architectures (LLaMA, GPT-NeoX, etc.).

A top-level, multi-block transformer constructs a single instance and shares its tables across all blocks, so the tables are built and stored once instead of once per block. Sequences shorter than max_seq_len are served by returning the leading positions of the precomputed table, so one instance covers any length up to max_seq_len without rebuilding.

The tables are stored as persistent=False buffers (they are deterministically reconstructed from (max_seq_len, head_dim, theta) and do not need to be saved with the model weights).

Parameters:

head_dim (int) – Per-head channel dimension. Must be even (rotation acts on adjacent channel pairs).
max_seq_len (int) – Maximum sequence length for which to precompute tables.
theta (float, optional, default=10000.0) – Base used for the RoPE frequency schedule.

Forward:

seq_len (int, optional) – Number of leading positions to return. If None, the full max_seq_len table is returned.

Outputs:

Tuple[torch.Tensor, torch.Tensor] – (cos, sin), each of shape \((seq\_len, head\_dim)\).

Examples

>>> import torch
>>> from physicsnemo.nn.module.rope import (
...     RotaryEmbedding1DTables,
...     apply_rotary_pos_emb,
... )
>>> rope = RotaryEmbedding1DTables(head_dim=16, max_seq_len=128)
>>> cos, sin = rope(seq_len=100)
>>> cos.shape
torch.Size([100, 16])
>>> q = torch.randn(2, 8, 100, 16)  # (B, heads, seq, head_dim)
>>> q_rot = apply_rotary_pos_emb(q, cos, sin)
>>> q_rot.shape
torch.Size([2, 8, 100, 16])

forward( seq_len: int | None = None, ) → Tuple[Float[Tensor, 'seq head_dim'], Float[Tensor, 'seq head_dim']][source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class physicsnemo.nn.module.rope.RotaryEmbedding2DTables(*args, **kwargs)[source]#

Bases: Module

Shared owner of axial 2D RoPE cos/sin tables for an \(h \times w\) grid.

This module owns the cos/sin tables and nothing else: it holds no projections and applies no rotation. Its forward() returns the (cos, sin) tables in explicit (h, w, head_dim) spatial layout, which a consumer applies to its query/key with apply_rotary_pos_emb().

The intended pattern is that a top-level, multi-block model constructs a single instance and passes the returned tables into every attention block’s forward (see RopeNatten2DSelfAttention and DiT). Building, storing, and — under domain parallelism — sharding the tables then happens exactly once for the whole model instead of once per block.

The tables are stored as persistent=False buffers named rope_cos / rope_sin: they are deterministically reconstructed from (latent_hw, head_dim, theta) and do not need to be saved with the model weights. The names and the height-first (h, w, head_dim) layout are chosen so that domain-parallel sharding along dimension 0 (height) gives each rank globally-correct rows with no explicit rank offset in model code.

Parameters:

head_dim (int) – Per-head channel dimension. Must be divisible by 4 (half per spatial axis, then adjacent channel pairs within each half).
latent_hw (Tuple[int, int]) – Spatial size \((h, w)\) of the token grid.
theta (float, optional, default=10000.0) – Base used for the RoPE frequency schedule.

Forward:

latent_hw (Tuple[int, int], optional) – Override the spatial grid size at call time. If given and different from the current grid, the cos/sin tables are rebuilt in place before being returned (off the torch.compile fast path). Under domain parallelism the in-place rebuild replaces the sharded buffers with plain tensors, so it is only appropriate for single-device variable-resolution inference.

Outputs:

Tuple[torch.Tensor, torch.Tensor] – (rope_cos, rope_sin), each of shape \((h, w, head\_dim)\).

Examples

>>> import torch
>>> from physicsnemo.nn.module.rope import (
...     RotaryEmbedding2DTables,
...     apply_rotary_pos_emb,
... )
>>> rope = RotaryEmbedding2DTables(head_dim=16, latent_hw=(4, 4))
>>> cos, sin = rope()
>>> cos.shape
torch.Size([4, 4, 16])
>>> # q reshaped to spatial layout (B, heads, h, w, head_dim)
>>> q = torch.randn(2, 8, 4, 4, 16)
>>> q_rot = apply_rotary_pos_emb(q, cos.unsqueeze(0).unsqueeze(0), sin.unsqueeze(0).unsqueeze(0))
>>> q_rot.shape
torch.Size([2, 8, 4, 4, 16])

forward( latent_hw: Tuple[int, int] | None = None, ) → Tuple[Float[Tensor, 'h w head_dim'], Float[Tensor, 'h w head_dim']][source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

physicsnemo.nn.module.rope.apply_rotary_pos_emb( x: Float[Tensor, '...'], cos: Float[Tensor, '...'], sin: Float[Tensor, '...'], ) → Float[Tensor, '...'][source]#

Apply precomputed RoPE cos/sin tables to a query or key tensor.

Rotates each adjacent channel pair \((x_a, x_b)\) in x by the angle encoded in the corresponding position of cos/sin:

\[(x_a,\, x_b) \;\mapsto\; (x_a \cos\phi - x_b \sin\phi,\;\; x_a \sin\phi + x_b \cos\phi)\]

This is the standard rotate-half formulation x * cos + rotate_half(x) * sin. The arithmetic is promoted to fp32 regardless of x’s dtype (the sign-flipped term accumulates error in half precision) and cast back before returning.

Call this directly when you manage the cos/sin tables yourself — for example, inside a custom NATTEN or domain-parallel attention block where you obtain the tables from a RotaryEmbedding2DTables / RotaryEmbedding1DTables provider (or build them with build_axial_rope_cos_sin_2d() / build_rope_cos_sin_1d()) and need to apply them independently to queries and keys.

Parameters:

x (torch.Tensor) – Query or key tensor of shape \((\ldots, \text{positions}, head\_dim)\).
cos (torch.Tensor) – Rotation tables broadcastable to x over the trailing (positions, head_dim) dimensions (e.g. shape \((\text{positions}, head\_dim)\)), as produced by build_axial_rope_cos_sin_2d() or build_rope_cos_sin_1d().
sin (torch.Tensor) – Rotation tables broadcastable to x over the trailing (positions, head_dim) dimensions (e.g. shape \((\text{positions}, head\_dim)\)), as produced by build_axial_rope_cos_sin_2d() or build_rope_cos_sin_1d().

Returns:

Rotated tensor of the same shape and dtype as x.

Return type:

torch.Tensor

physicsnemo.nn.module.rope.build_axial_rope_cos_sin_2d( h: int, w: int, head_dim: int, theta: float = 10000.0, device: device | None = None, ) → Tuple[Tensor, Tensor][source]#

Precompute axial 2D RoPE cos/sin tables for an \(h \times w\) token grid.

The first head_dim/2 channels are rotated by the row index, the last head_dim/2 by the column index. Within each axis-half, frequency \(\theta_k = \text{theta}^{-2k/(head\_dim/2)}\) drives the adjacent channel pair (2k, 2k+1).

Parameters:

h (int) – Token grid height.
w (int) – Token grid width.
head_dim (int) – Per-head channel dimension. Must be divisible by 4 (half per axis, then adjacent pairs within each half).
theta (float, optional, default=10000.0) – Base used for the RoPE frequency schedule.
device (torch.device, optional) – Device for the generated tables.

Returns:

(cos, sin), each of shape \((h, w, head\_dim)\) in fp32.

Return type:

Tuple[torch.Tensor, torch.Tensor]

physicsnemo.nn.module.rope.build_rope_cos_sin_1d( seq_len: int, head_dim: int, theta: float = 10000.0, device: device | None = None, ) → Tuple[Tensor, Tensor][source]#

Precompute 1D RoPE cos/sin tables for a length-seq_len sequence.

The standard sequence RoPE: every channel rotates by the token position, with head_dim/2 frequencies \(\theta_k = \text{theta}^{-2k/head\_dim}\) for \(k = 0 \ldots head\_dim/2 - 1\), each driving the adjacent channel pair (2k, 2k+1).

Parameters:

seq_len (int) – Number of positions in the sequence.
head_dim (int) – Per-head channel dimension. Must be even (rotation acts on adjacent channel pairs).
theta (float, optional, default=10000.0) – Base used for the RoPE frequency schedule.
device (torch.device, optional) – Device for the generated tables.

Returns:

(cos, sin), each of shape \((seq\_len, head\_dim)\) in fp32.

Return type:

Tuple[torch.Tensor, torch.Tensor]