rope.py

"""Rotary Position Embeddings (RoPE).

RoPE encodes position in the *relationship* between query and key vectors rather than
adding it to the inputs directly. When the attention dot product Q·Kᵀ is computed, the
per-position rotations cancel to produce a score that depends only on the relative
distance between positions — not on their absolute values. This is what gives RoPE
better length generalisation than absolute learned embeddings.

Each pair of head dimensions (d, d+1) is assigned a rotation frequency
    1 / theta^(2d / head_dim)
Higher theta → slower rotation per position → position encodings remain distinguishable
further apart before wrapping. Llama 3 uses theta=500,000 as a prerequisite for
128K context support.

Supported rope types: "default" (standard unscaled RoPE), "linear", and "yarn".
HuggingFace's ROPE_INIT_FUNCTIONS handles inv_freq computation for linear and yarn;
the default case is not in that registry and is computed directly here.
"""

import torch
import torch.nn as nn
from transformers import PretrainedConfig
from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS

_SUPPORTED_ROPE_TYPES = {"default", "linear", "yarn"}


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Apply the 90° rotation used in the RoPE update formula.

    Splits the last dimension into two halves [x1, x2] and returns [-x2, x1].
    Combined with ``x * cos + rotate_half(x) * sin``, this implements a 2D rotation
    on each consecutive pair of dimensions.
    """
    d = x.shape[-1] // 2
    x1, x2 = x[..., :d], x[..., d:]
    return torch.cat([-x2, x1], dim=-1)


class RotaryEmbedding(nn.Module):
    """Rotary Position Embeddings as an nn.Module.

    Computes position-dependent rotation frequencies from the model config, maintains
    a lazily-extended cos/sin cache, and applies the rotations to query and key tensors.

    The cos/sin cache grows automatically at runtime when a sequence longer than the
    current cache is encountered. ``config.max_position_embeddings`` records the
    training context length (required by HF's scaling computations) but does not cap
    inference length.

    Args:
        config: Model config. Must expose ``rope_theta``, ``rope_parameters`` (set by
            HF's RotaryEmbeddingConfigMixin), and ``head_dim``.
        device: Optional device for initial buffer placement. Buffers move with the
            model on ``.to()`` / ``.cuda()`` calls.

    Raises:
        NotImplementedError: If ``config.rope_parameters`` specifies an unsupported
            rope type. Supported types: "default", "linear", "yarn".
    """

    def __init__(self, config: PretrainedConfig, device: torch.device | None = None) -> None:
        super().__init__()
        self.config = config

        # rope_parameters is None when no rope_scaling was passed to the config.
        rope_params = config.rope_parameters
        self.rope_type = (
            rope_params.get("rope_type", "default") if rope_params is not None else "default"
        )

        if self.rope_type not in _SUPPORTED_ROPE_TYPES:
            raise NotImplementedError(
                f"rope_type '{self.rope_type}' is not supported. "
                f"Supported types: {sorted(_SUPPORTED_ROPE_TYPES)}"
            )

        if self.rope_type == "default":
            # Standard RoPE: inv_freq = 1 / theta^(2i / head_dim).
            # Not in ROPE_INIT_FUNCTIONS, so computed directly.
            inv_freq = 1.0 / (
                config.rope_theta
                ** (torch.arange(0, config.head_dim, 2, dtype=torch.float32, device=device) / config.head_dim)
            )
            self.attention_scaling: float = 1.0
        else:
            inv_freq, self.attention_scaling = ROPE_INIT_FUNCTIONS[self.rope_type](config, device)

        self.register_buffer("inv_freq", inv_freq, persistent=False)
        # Initialised as None; built on first forward call and extended lazily thereafter.
        # Registered as buffers so they move with the model across devices.
        self.register_buffer("_cos_cached", None, persistent=False)
        self.register_buffer("_sin_cached", None, persistent=False)

    def _extend_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype) -> None:
        """Build the cos/sin table to cover positions [0, seq_len).

        Registered as buffers so subsequent calls to ``.to()`` / ``.cuda()`` will
        move them to the correct device. Rebuilds whenever the sequence grows or
        the dtype changes (e.g. switching between fp32 and bf16).
        """
        positions = torch.arange(seq_len, device=device, dtype=torch.float32)
        # outer product → (seq_len, head_dim // 2); duplicate → (seq_len, head_dim)
        freqs = torch.outer(positions, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("_cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("_sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        position_ids: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, float]:
        """Apply rotary embeddings to query and key tensors.

        The cos/sin cache is extended lazily when position_ids reference positions
        beyond its current length.

        Args:
            q: Query tensor of shape (batch, num_heads, seq_len, head_dim).
            k: Key tensor of shape (batch, num_kv_heads, seq_len, head_dim).
            position_ids: Integer positions of shape (batch, seq_len).

        Returns:
            Tuple of (q_rotated, k_rotated, attention_scaling). attention_scaling is
            1.0 for default and linear; YaRN returns a value != 1.0 that callers must
            apply to attention logits to correct for frequency magnitude changes.
        """
        seq_len = int(position_ids.max().item()) + 1

        if self._cos_cached is None or seq_len > self._cos_cached.shape[0] or self._cos_cached.dtype != q.dtype:
            self._extend_cache(seq_len, device=q.device, dtype=q.dtype)

        # Gather cos/sin for the given positions → (batch, seq_len, head_dim),
        # then unsqueeze the head axis for broadcast over all heads.
        cos = self._cos_cached[position_ids].unsqueeze(1)
        sin = self._sin_cached[position_ids].unsqueeze(1)

        q_rotated = q * cos + _rotate_half(q) * sin
        k_rotated = k * cos + _rotate_half(k) * sin

        return q_rotated, k_rotated, self.attention_scaling