Skip to content

Positional Embeddings

About

The d9d.module.block.positional package manages positional encoding logic.

Features

Rotary Positional Encoding

Rotary Positional Encoding from RoFormer.

See RotaryEmbeddingProvider and RotaryEmbeddingApplicator classes.

First one is typically bound to a model class and is used for providing (cos, sin) embedding tensors for specified position IDs.

Second one is typically bound to attention module implementation and is used for modifying query and key states in runtime.

Embedding Layout Styles

The package supports multiple internal memory layouts for RoPE operations via the RotaryEmbeddingStyle enumeration. It is critical that both the provider and applicator share the identical style configuration:

d9d.module.block.positional

Provides modules for positional embeddings, such as Rotary Positional Embeddings.

LinearRopeScaling

Bases: RopeScaling

Linear scaling strategy for Rotary Position Embeddings.

attention_mscale property

Calculates the attention multiplier scale.

Returns:

Type Description
float

The attention multiplier scale.

__init__(factor)

Constructs a linear RoPE scaling object.

Parameters:

Name Type Description Default
factor float

The linear scaling factor to apply.

 required

NoRopeScaling

Bases: RopeScaling

Strategy that applies no scaling to Rotary Position Embeddings.

attention_mscale property

Calculates the attention multiplier scale.

Returns:

Type Description
float

The attention multiplier scale.

NtkRopeScaling

Bases: RopeScaling

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

References

https://www.reddit.com/r/LocalLLaMA/comments/14lz7j5/ntkaware_scaled_rope_allows_llama_models_to_have/

attention_mscale property

Calculates the attention multiplier scale.

Returns:

Type Description
float

The attention multiplier scale.

__init__(factor)

Constructs an NTK-Aware RoPE scaling object.

Parameters:

Name Type Description Default
factor float

The sequence length expansion factor.

 required

RopeScaling

Bases: ABC

Abstract base class for Rotary Position Embedding (RoPE) scaling strategies.

attention_mscale property

Calculates the attention multiplier scale.

Returns:

Type Description
float

The attention multiplier scale.

inverse_frequencies(rope_base, head_dim) abstractmethod

Calculates the inverse frequencies for the given RoPE scaling strategy.

Parameters:

Name Type Description Default
rope_base int

The base value used for calculating frequencies.

 required
head_dim int

The dimension of the attention head.

 required

Returns:

Type Description
Tensor

The computed inverse frequencies tensor.

RotaryEmbeddingApplicator

Bases: Module

Applies Rotary Positional Embeddings (RoPE) to Q and K projections.

__init__(style)

Constructs RotaryEmbeddingApplicator object.

Parameters:

Name Type Description Default
style RotaryEmbeddingStyle

Rotary embedding layout style alignment.

 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, n_heads, seq_len, head_dim).

 required
key_states Tensor

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

 required
position_embedding_cos Tensor

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

 required
position_embedding_sin Tensor

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

 required

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple containing the rotated query and key tensors.

RotaryEmbeddingProvider

Bases: Module, ModuleLateInit

Module that manages and provides Rotary Positional Embeddings.

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

Constructs the RotaryEmbeddingProvider.

Parameters:

Name Type Description Default
rope_base int

Base geometrical progression period for RoPE.

 required
head_dim int

Dimensionality of the attention head.

 required
max_position_ids int

Maximum supported sequence length for caching.

 required
style RotaryEmbeddingStyle

Embedding layout alignment.

 required
rope_scaling RopeScaling | None

Optional scaling configuration for extended context lengths. When None (default), NoRopeScaling is used — no scaling applied.

 None

forward(position_ids)

Retrieves cached cosine and sine embeddings for specific positions.

Parameters:

Name Type Description Default
position_ids Tensor

Tensor of position indices.

 required

Returns:

Type Description
tuple[Tensor, Tensor]

A tuple of (cos, sin) tensors aligned with the input positions.

reset_parameters()

Resets module buffer populated values.

RotaryEmbeddingStyle

Bases: StrEnum

Supported Rotary Positional Embedding (RoPE) layout styles.

Attributes:

Name Type Description
HALF

Applies transformations by splitting the feature dimension into two halves.
INTERLEAVED

Applies transformations by treating adjacent feature elements as pairs.

YarnRopeScaling

Bases: RopeScaling

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

References

https://arxiv.org/abs/2309.00071

__init__(factor, beta_fast, beta_slow, original_max_position_embeddings)

Constructs a YaRN RoPE scaling object.

Parameters:

Name Type Description Default
factor float

The context scaling extension factor.

 required
beta_fast float

The fast boundary (upper bound) frequency multiplier.

 required
beta_slow float

The slow boundary (lower bound) frequency multiplier.

 required
original_max_position_embeddings int

The original context limit of the base model.

 required

Raises:

Type Description
ValueError

If beta_fast is less than or equal to beta_slow.