rotary-embedding-torch

In standard implementations, rotary embeddings are applied to $Q$ and $K$ as a separate step: $Q_{rot} = ext{rotate}(Q)$ and $K_{rot} = ext{rotate}(K)$.

flash_attn_with_rotary fuses this operation into the attention kernel. During the kernel execution, the rotary transformation is applied to the loaded blocks of $Q$ and $K$ on-the-fly using the provided cos and sin values. This reduces memory bandwidth overhead by avoiding the need to write the intermediate rotated $Q$ and $K$ tensors back to global memory.

Key Features:

• Fused Kernel: Uses Triton to perform rotation and attention in a single pass.
• Support for GQA/MQA: Automatically handles cases where heads_q != heads_kv by repeating $K$ and $V$ heads.
• Positional Masking: Supports a pos_mask to selectively apply rotary embeddings to specific positions.
• Causal Masking: Built-in support for causal attention via is_causal.

The flash_attn_with_rotary function provides a fused kernel that computes attention with rotary embeddings in a single pass. It uses Triton if available and falls back to a PyTorch implementation otherwise. It supports skipping specific tokens (like CLS or register tokens) via the rotary_pos_emb_indices argument.

import torch
from rotary_embedding_torch import RotaryEmbedding
from rotary_embedding_torch.flash_attn_with_rotary import flash_attn_with_rotary

rotary_emb = RotaryEmbedding(dim = 32)
freqs = rotary_emb(torch.arange(1024))

# q, k, v with extra tokens (e.g. CLS/register) at the start
q = torch.randn(1, 8, 1024 + 2, 64).cuda()
k = torch.randn(1, 8, 1024 + 2, 64).cuda()
v = torch.randn(1, 8, 1024 + 2, 64).cuda()

# Indices for the 1024 rotary positions, skipping the first 2 tokens
pos_indices = torch.arange(1024).cuda() + 2

out = flash_attn_with_rotary(
    q, k, v,
    rotary_pos_emb = freqs,
    rotary_pos_emb_indices = pos_indices,
    is_causal = True
)

The flash_attn_with_rotary function provides a high-performance, fused implementation of attention that integrates rotary positional embeddings directly into the attention kernel. This is significantly more efficient than applying rotary embeddings to queries and keys separately before calling a standard attention function.

If triton is installed, it uses optimized Triton kernels. If triton is unavailable or if force_reference=True is passed to the factory, it falls back to a standard PyTorch implementation (reference_flash_attention).

Input Tensor Shapes

• q: (batch, heads_q, seq_len_q, dim)
• k: (batch, heads_kv, seq_len_k, dim)
• v: (batch, heads_kv, seq_len_k, dim)
• rotary_pos_emb: (seq_len, rotary_dim) or similar, representing the frequencies to be rotated.
• rotary_pos_emb_indices: (seq_len,) used for indexing into a padded frequency buffer (useful for non-contiguous or sparse positions).
• attn_mask: (batch, seq_len_k) or (batch, seq_len_q, seq_len_k). If boolean, it is treated as a mask where False values are replaced with -inf.
• pos_mask: (seq_len_q,) used to define which positions in the query sequence should receive rotary embeddings (in conjunction with rotary_pos_emb_indices).

from rotary_embedding_torch.flash_attn_with_rotary import flash_attn_with_rotary

# Assuming you have q, k, v, and rotary_pos_emb prepared
# rotary_pos_emb should be the frequencies (e.g., from a RotaryEmbedding object)

output = flash_attn_with_rotary(
    q = q, 
    k = k, 
    v = v, 
    rotary_pos_emb = rotary_pos_emb, 
    is_causal = True
)

The RotaryEmbedding class is the primary interface for applying rotary positional embeddings to queries and keys in a transformer. It supports various frequency generation modes, XPOS for length extrapolation, and learned frequencies.

Key Parameters

• dim: The dimension of the rotary embedding (must match the feature dimension of the part of the tensor being rotated).
• freqs_for: Determines how frequencies are generated. Options are 'lang' (default, standard RoPE), 'pixel' (for images/video), or 'constant'.
• theta: The base for frequency calculation (default 10000).
• learned_freq: If True, frequencies become trainable parameters.
• use_xpos: Enables XPOS (Extrapolatable Positional Encoding) for better length extrapolation.
• interpolate_factor: Used for scaling sequence positions (default 1.0).
• seq_before_head_dim: If True, assumes the sequence dimension is at index -3 instead of -2 (e.g., [batch, seq, heads, dim]).
• cache_if_possible: Enables caching of computed frequencies to improve performance.

from rotary_embedding_torch import RotaryEmbedding
import torch

# Setup
dim = 64
rotary = RotaryEmbedding(dim = dim)

# Mock queries and keys: (batch, heads, seq_len, dim)
q = torch.randn(1, 8, 128, dim)
k = torch.randn(1, 8, 128, dim)

# Apply rotation
# By default, it expects seq_dim = -2
q, k = rotary.rotate_queries_or_keys(q, k)