Skip to content

Positional Embeddings

About

The d9d.module.block.positional package implements rotary position embedding (RoPE) from RoFormer.

Rotary Position Embedding

RoPE uses two classes:

  • RotaryEmbeddingProvider usually belongs to the model. It caches the (cos, sin) embedding tensors and returns them for given position IDs.
  • RotaryEmbeddingApplicator usually belongs to the attention layer. It rotates the query and key states with these tensors.

Layout Styles

RotaryEmbeddingStyle selects how features are paired for rotation:

  • HALF pairs the first half of the head dimension with the second half.
  • INTERLEAVED pairs adjacent features.

The provider and the applicator must use the same style.

Scaling

RotaryEmbeddingProvider takes an optional RopeScaling strategy to extend the context length:

  • NoRopeScaling is the default and applies no scaling.
  • LinearRopeScaling divides all frequencies by a constant factor.
  • NtkRopeScaling implements NTK-aware scaling.
  • YarnRopeScaling implements YaRN.

Usage

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import torch

from d9d.module.block.positional import (
    RotaryEmbeddingApplicator,
    RotaryEmbeddingProvider,
    RotaryEmbeddingStyle,
    YarnRopeScaling,
)

provider = RotaryEmbeddingProvider(
    rope_base=10_000,
    head_dim=128,
    max_position_ids=16384,
    style=RotaryEmbeddingStyle.HALF,
    rope_scaling=YarnRopeScaling(
        factor=4.0, beta_fast=32.0, beta_slow=1.0, original_max_position_embeddings=4096
    ),
)
provider.reset_parameters()
applicator = RotaryEmbeddingApplicator(style=RotaryEmbeddingStyle.HALF)

position_ids = torch.arange(16).expand(2, -1)
cos, sin = provider(position_ids)  # (2, 16, 128) each

# Shapes: (batch, seq_len, num_heads, head_dim) and (batch, seq_len, num_kv_heads, head_dim).
query_states = torch.randn(2, 16, 32, 128)
key_states = torch.randn(2, 16, 4, 128)
query_states, key_states = applicator(query_states, key_states, cos, sin)

API Reference

d9d.module.block.positional

Position embeddings, such as rotary position embeddings (RoPE).

LinearRopeScaling

Bases: RopeScaling

Linear scaling strategy for rotary position embeddings.

attention_mscale property

The attention scale multiplier (mscale) applied to the cosine and sine embeddings.

__init__(factor)

Constructs the LinearRopeScaling object.

Parameters:

Name Type Description Default
factor float

Linear scaling factor. The inverse frequencies are divided by it.

 required

NoRopeScaling

Bases: RopeScaling

Strategy that applies no scaling to rotary position embeddings.

attention_mscale property

The attention scale multiplier (mscale) applied to the cosine and sine embeddings.

NtkRopeScaling

Bases: RopeScaling

NTK-Aware (Neural Tangent Kernel) scaling strategy for position embeddings.

References

NTK-Aware Scaled RoPE

attention_mscale property

The attention scale multiplier (mscale) applied to the cosine and sine embeddings.

__init__(factor)

Constructs the NtkRopeScaling object.

Parameters:

Name Type Description Default
factor float

Sequence length expansion factor.

 required

RopeScaling

Bases: ABC

Abstract base class for rotary position embedding (RoPE) scaling strategies.

attention_mscale property

The attention scale multiplier (mscale) applied to the cosine and sine embeddings.

inverse_frequencies(rope_base, head_dim) abstractmethod

Computes the RoPE inverse frequencies for this scaling strategy.

Parameters:

Name Type Description Default
rope_base int

Base of the geometric progression of RoPE frequencies.

 required
head_dim int

Dimensionality of the attention head.

 required

Returns:

Type Description
Tensor

The inverse frequencies. Shape: (head_dim // 2,).

RotaryEmbeddingApplicator

Bases: Module

Layer that applies rotary position embeddings (RoPE) to Q and K projections.

__init__(style)

Constructs the RotaryEmbeddingApplicator object.

Parameters:

Name Type Description Default
style RotaryEmbeddingStyle

RoPE layout style.

 required
forward(query_states, key_states, position_embedding_cos, position_embedding_sin)

Rotates query and key states using provided cosine and sine embeddings.

Parameters:

Name Type Description Default
query_states Tensor

Query tensor. Shape: (batch, seq_len, num_heads, head_dim).

 required
key_states Tensor

Key tensor. Shape: (batch, seq_len, num_kv_heads, head_dim).

 required
position_embedding_cos Tensor

Cosine values for the positions. Shape: (batch, seq_len, head_dim).

 required
position_embedding_sin Tensor

Sine values for the positions. Shape: (batch, seq_len, head_dim).

 required

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple of the rotated query and key tensors, with the input shapes.

RotaryEmbeddingProvider

Bases: Module, ModuleLateInit

Module that caches rotary position embeddings and returns them for given positions.

__init__(rope_base, head_dim, max_position_ids, style, rope_scaling=None)

Constructs the RotaryEmbeddingProvider object.

Parameters:

Name Type Description Default
rope_base int

Base of the geometric progression of RoPE frequencies.

 required
head_dim int

Dimensionality of the attention head.

 required
max_position_ids int

Number of cached positions. Position indices must be smaller than this value.

 required
style RotaryEmbeddingStyle

RoPE layout style.

 required
rope_scaling RopeScaling | None

Optional scaling strategy for extended context lengths. If None, NoRopeScaling is used.

 None
forward(position_ids)

Returns the cached cosine and sine embeddings for the given positions.

Parameters:

Name Type Description Default
position_ids Tensor

Position indices, usually with shape (batch, seq_len).

 required

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple of (cos, sin) tensors. Shape of each: (*position_ids.shape, head_dim).
reset_parameters()

Recomputes the cached cosine and sine buffers.

RotaryEmbeddingStyle

Bases: StrEnum

Supported rotary position embedding (RoPE) layout styles.

Attributes:

Name Type Description
HALF

Rotates pairs formed by the first and second halves of the feature dimension.
INTERLEAVED

Rotates pairs of adjacent feature elements.

YarnRopeScaling

Bases: RopeScaling

YaRN (Yet another RoPE extensioN) scaling strategy for position embeddings.

References

YaRN: Efficient Context Window Extension of Large Language Models

__init__(factor, beta_fast, beta_slow, original_max_position_embeddings)

Constructs the YarnRopeScaling object.

Parameters:

Name Type Description Default
factor float

Context extension factor.

 required
beta_fast float

Fast boundary (upper bound) of the frequency ramp, in rotations.

 required
beta_slow float

Slow boundary (lower bound) of the frequency ramp, in rotations.

 required
original_max_position_embeddings int

Original context length of the base model.

 required

Raises:

Type Description
ValueError

If beta_fast is less than or equal to beta_slow.