Rotary Embedding Torch

https://github.com/lucidrains/rotary-embedding-torch
Admin
A PyTorch library that implements rotary embeddings for transformers, providing state-of-the-art...
Tokens:33,114
Snippets:175
Trust Score:10
License:MIT
Update:1 month ago
Show doc for...
Tokens:
Raw
### Example: rotate_queries_and_keys usage

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example usage of rotate_queries_and_keys with use_xpos=True. Rotates queries and keys, applying XPos decay.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(
    dim = 32,
    use_xpos = True   # set to True to extrapolate beyond training sequence length
)

q = torch.randn(1, 8, 1024, 64)
k = torch.randn(1, 8, 1024, 64)

q, k = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### Install rotary-embedding-torch via pip

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Install the library using pip. This is the standard installation command for the package.

```bash
$ pip install rotary-embedding-torch
```

--------------------------------

### Example: get_scale usage

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example usage of get_scale to compute the XPos decay scale and apply it to queries and keys with apply_rotary_emb.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb

rotary_emb = RotaryEmbedding(dim = 32, use_xpos = True)
seq = rotary_emb.get_seq_pos(1024)
freqs = rotary_emb(seq)
scale = rotary_emb.get_scale(seq, seq_len = 1024)   # (1024, 32)

q = torch.randn(1, 8, 1024, 64)
k = torch.randn(1, 8, 1024, 64)
q = apply_rotary_emb(freqs, q, scale = scale)
k = apply_rotary_emb(freqs, k, scale = scale ** -1)
```

--------------------------------

### forward method usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example of computing frequency tensors for a sequence of 1024 positions and for an offset query. The first call returns shape (1024, 32) with cos/sin pairs per position.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(dim = 32)
freqs = rotary_emb(torch.arange(1024))   # (1024, 32) — cos/sin pairs per position

# with offset (query positions beyond cached region)
freqs_offset = rotary_emb(torch.arange(8), seq_len = 8, offset = 1024)
```

--------------------------------

### device property usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Shows how to access the device property after moving the model to CUDA. Output is 'cuda:0'.

```python
rotary_emb = RotaryEmbedding(dim = 32).to('cuda')
print(rotary_emb.device)  # cuda:0
```

--------------------------------

### Usage with XPos enabled

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example with use_xpos=True, showing that the query is scaled up and keys scaled down. Includes a try/except block to handle the AssertionError when the query length exceeds the key length.

```python
rotary_emb = RotaryEmbedding(dim = 32, use_xpos = True)

q = torch.randn(1, 8, 1, 64)
k = torch.randn(1, 8, 1024, 64)
try:
    q, k = rotary_emb.rotate_queries_with_cached_keys(q, torch.randn(1, 8, 0, 64))  # empty keys
    print("rotated", q.shape, k.shape)
except AssertionError as e:
    print("q_len must be <= k_len")
```

--------------------------------

### RotaryEmbedding usage examples

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Examples showing basic usage of RotaryEmbedding for language models and the XPos variant. Includes an assertion check for interpolate_factor.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

# basic language-model rotary embedding
rotary_emb = RotaryEmbedding(dim = 32)
q = torch.randn(1, 8, 1024, 64)
q = rotary_emb.rotate_queries_or_keys(q)

# XPos variant (assertion-safe: interpolate_factor must be >= 1.)
rotary_emb = RotaryEmbedding(dim = 32, use_xpos = True, xpos_scale_base = 512)
try:
    RotaryEmbedding(dim = 32, interpolate_factor = 0.5)
except AssertionError as e:
    print("interpolate_factor must be at least 1.")
```

--------------------------------

### rotate_queries_or_keys usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example usage of rotate_queries_or_keys. Creates a RotaryEmbedding instance, a random query tensor, rotates it, and shows a manual KV-cache offset example.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(dim = 32)
q = torch.randn(1, 8, 1024, 64)

q = rotary_emb.rotate_queries_or_keys(q)

# manual KV-cache offset (see rotate_queries_with_cached_keys)
q_new = rotary_emb.rotate_queries_or_keys(q, offset = 1024 - q.shape[-2])
```

--------------------------------

### get_seq_pos usage examples

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Two usage examples for get_seq_pos. The first computes positions for 1024 positions, the second computes a single position at offset 512.

```python
seq = rotary_emb.get_seq_pos(1024)                    # arange(1024) / interpolate_factor
seq = rotary_emb.get_seq_pos(1, offset = 512)         # single position at 512 + offset
```

--------------------------------

### get_axial_freqs with per-axis offsets

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example of using get_axial_freqs with an offsets tuple to shift the starting positions of each axis. The offsets length must equal the number of dimensions.

```python
freqs = pos_emb.get_axial_freqs(8, 64, 32, offsets = (0, 4, 4))  # start axes at shifted positions
```

--------------------------------

### Basic usage of rotate_queries_with_cached_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example showing how to rotate a single query against a key cache. The query is automatically offset to align with the tail of the keys.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(dim = 32)

q = torch.randn(1, 8, 1, 64)     # one query at a time
k = torch.randn(1, 8, 1024, 64)  # keys with cache concatenated

q, k = rotary_emb.rotate_queries_with_cached_keys(q, k)
```

--------------------------------

### broadcat usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding_torch.md

Example of using broadcat to broadcast two tensors and concatenate them along dimension -2, resulting in a tensor of shape (2, 4, 4).

```python
import torch
from rotary_embedding_torch import broadcat

a = torch.randn(2, 1, 4)
b = torch.randn(1, 3, 4)
out = broadcat((a, b), dim = -2)   # (2, 4, 4) — broadcast to (2, 3, 4), cat on dim -2
```

--------------------------------

### Example: flash_attn_with_rotary with rotary position embeddings and indices

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Shows how to use flash_attn_with_rotary with rotary_pos_emb and rotary_pos_emb_indices for causal attention. The example uses a RotaryEmbedding instance to generate frequencies, and includes extra tokens (e.g., CLS, register) that are skipped via pos_indices. The output shape is (1, 8, 1026, 64).

```python
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))   # (1024, 32)

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

pos_indices = torch.arange(1024).cuda() + 2  # skip the leading non-rotary tokens

out = flash_attn_with_rotary(
    q, k, v,
    rotary_pos_emb = freqs,
    rotary_pos_emb_indices = pos_indices,
    is_causal = True
)
# out.shape == (1, 8, 1026, 64)
```

--------------------------------

### reference_flash_attention usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Applies the reference attention with GQA (8 query heads, 4 kv heads) and is_causal=False. The output shape is (1, 8, 64, 64).

```python
import torch
from rotary_embedding_torch.flash_attn_with_rotary import reference_flash_attention

q = torch.randn(1, 8, 64, 64)
k = torch.randn(1, 4, 64, 64)   # GQA: 8 query heads, 4 kv heads
v = torch.randn(1, 4, 64, 64)

out = reference_flash_attention(q, k, v, is_causal = False)
```

--------------------------------

### Apply rotary embeddings to query and key tensors

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding_torch.md

Shows how to create a RotaryEmbedding instance, generate frequencies for a sequence, and apply them to query and key tensors using apply_rotary_emb. The example uses dim=32, batch=2, heads=8, seq=1024, and dim_head=64.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb

rotary_emb = RotaryEmbedding(dim = 32)
q = torch.randn(2, 8, 1024, 64)   # (batch, heads, seq, dim_head)
k = torch.randn(2, 8, 1024, 64)

seq = torch.arange(1024)
freqs = rotary_emb(seq)            # (1024, 32)

q = apply_rotary_emb(freqs, q)
k = apply_rotary_emb(freqs, k)
```

--------------------------------

### get_axial_freqs usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Example of using get_axial_freqs for a video with 8 frames and a 64x32 image. The method returns frequencies of shape (8, 64, 32, 48) for 'pixel' mode with dim=16, which are then applied to query and key tensors using apply_rotary_emb.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding, apply_rotary_emb

pos_emb = RotaryEmbedding(
    dim = 16,
    freqs_for = 'pixel',
    max_freq = 256
)

# video: 8 frames, 64x32 image, feature dimension last
q = torch.randn(1, 8, 64, 32, 64)
k = torch.randn(1, 8, 64, 32, 64)

freqs = pos_emb.get_axial_freqs(8, 64, 32)   # (8, 64, 32, 48) — partial rotary (16 -> 48)

q = apply_rotary_emb(freqs, q)
k = apply_rotary_emb(freqs, k)
```

--------------------------------

### get_flash_attention_fused usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Creates a CPU-safe reference attention callable and applies it to random tensors with is_causal=False. The output shape is (1, 8, 128, 64).

```python
import torch
from rotary_embedding_torch.flash_attn_with_rotary import get_flash_attention_fused

flash_attn_ref = get_flash_attention_fused(force_reference = True)   # CPU-safe implementation

q = torch.randn(1, 8, 128, 64)
k = torch.randn(1, 8, 128, 64)
v = torch.randn(1, 8, 128, 64)

out = flash_attn_ref(q, k, v, is_causal = False)   # plain attention, no rotary
# out.shape == (1, 8, 128, 64)
```

--------------------------------

### FlashAttentionFused usage example

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Shows how to use FlashAttentionFused.apply with rotary embeddings. It creates a RotaryEmbedding, computes cosine and sine frequencies, reshapes them to the expected layout, and applies the fused attention. The output shape is (1, 8, 128, 64).

```python
import torch
from rotary_embedding_torch import RotaryEmbedding
from rotary_embedding_torch.flash_attn_with_rotary import FlashAttentionFused

rotary_emb = RotaryEmbedding(dim = 32)
freqs = rotary_emb(torch.arange(128)).cuda()
cos = freqs.cos().reshape(128, 1, 1, 32)   # kernel expects '... 1 1 n d' layout
sin = freqs.sin().reshape(128, 1, 1, 32)

q = torch.randn(1, 8, 128, 64).cuda()
k = torch.randn(1, 8, 128, 64).cuda()
v = torch.randn(1, 8, 128, 64).cuda()

out = FlashAttentionFused.apply(q, k, v, cos, sin, None, None, True) if TRITON_AVAILABLE else None
# out.shape == (1, 8, 128, 64)
```

--------------------------------

### Example: Non-causal masked attention with flash_attn_with_rotary

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Shows how to use flash_attn_with_rotary with a non-causal attention mask. The mask is a boolean tensor of shape (1, 1024) where the last quarter is set to False, masking out those positions. The is_causal parameter is set to False.

```python
attn_mask = torch.ones(1, 1024, dtype = torch.bool).cuda()
attn_mask[0, 768:] = False   # mask out the last quarter

out = flash_attn_with_rotary(
    q, k, v,
    rotary_pos_emb = freqs,
    attn_mask = attn_mask,
    is_causal = False
)
```

--------------------------------

### Import flash_attn_with_rotary module

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Import the main API symbols from the module. Requires the optional `triton` package for the fused path; otherwise falls back to a PyTorch reference implementation.

```python
from rotary_embedding_torch.flash_attn_with_rotary import (
    flash_attn_with_rotary,
    get_flash_attention_fused,
    reference_flash_attention,
    FlashAttentionFused,
    TRITON_AVAILABLE
)
```

--------------------------------

### Use XPos for length extrapolation with rotate_queries_and_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Set use_xpos=True to enable XPos, which improves extrapolation to longer sequences. Only for autoregressive transformers. Use rotate_queries_and_keys instead of rotate_queries_or_keys.

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

# instantiate the positional embedding in your transformer and pass to all your attention layers

rotary_emb = RotaryEmbedding(
    dim = 32,
    use_xpos = True   # set this to True to make rotary embeddings extrapolate better to sequence lengths greater than the one used at training time
)

# mock queries and keys - dimensions should end with (seq_len, feature dimension), and any number of preceding dimensions (batch, heads, etc)

q = torch.randn(1, 8, 1024, 64) # queries - (batch, heads, seq len, dimension of head)
k = torch.randn(1, 8, 1024, 64) # keys

# apply the rotations to your queries and keys after the heads have been split out, but prior to the dot product and subsequent softmax (attention)

# instead of using `rotate_queries_or_keys`, you will use `rotate_queries_and_keys`, the rest is taken care of

q, k = rotary_emb.rotate_queries_and_keys(q, k)
```

--------------------------------

### rotate_queries_and_keys method signature

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Method signature for rotate_queries_and_keys. Requires use_xpos=True, otherwise raises AssertionError.

```python
def rotate_queries_and_keys(self, q, k, seq_dim = None):
```

--------------------------------

### device (property)

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Returns the device of the dummy buffer, used to anchor creation of position tensors.

```APIDOC
## device (property)

### Description
Returns the device of the `dummy` buffer, used to anchor creation of position tensors.

### Return Type
`torch.device`

### Example
```python
rotary_emb = RotaryEmbedding(dim = 32).to('cuda')
print(rotary_emb.device)  # cuda:0
```
```

--------------------------------

### rotate_queries_and_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Rotates queries and keys together with XPos length-extrapolatable decay applied. Queries are scaled by the XPos scale, keys by its inverse. Requires use_xpos=True.

```APIDOC
## rotate_queries_and_keys

### Description
Rotates queries and keys together with XPos length-extrapolatable decay applied. The queries are scaled by the XPos scale, and the keys by its inverse (`scale ** -1`). Requires `use_xpos = True`.

### Method
rotate_queries_and_keys

### Signature
```python
def rotate_queries_and_keys(self, q, k, seq_dim = None):
```

### Parameters
- **q** (torch.Tensor) - Required - Query tensor, feature dimension last.
- **k** (torch.Tensor) - Required - Key tensor, feature dimension last.
- **seq_dim** (int) - Optional - Default: `self.default_seq_dim` - Sequence dimension of both tensors.

### Return Type
`tuple[torch.Tensor, torch.Tensor]` — `(rotated_q, rotated_k)`, cast back to `q.dtype`.

### Throws
- **AssertionError** - If `use_xpos` is `False` — the method is only valid for length-extrapolatable (XPos) embeddings.

### Example
```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(
    dim = 32,
    use_xpos = True   # set to True to extrapolate beyond training sequence length
)

q = torch.randn(1, 8, 1024, 64)
k = torch.randn(1, 8, 1024, 64)

q, k = rotary_emb.rotate_queries_and_keys(q, k)
```
```

--------------------------------

### RotaryEmbedding.__init__

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Constructor for the RotaryEmbedding class. Initializes the rotary embedding with configurable frequency schedules, optional XPos decay, learned frequencies, and caching.

```APIDOC
## RotaryEmbedding.__init__

### Description
Initializes a RotaryEmbedding instance with configurable frequency schedules, optional XPos decay, learned frequencies, and caching.

### Method
Constructor

### Signature
```python
def __init__(
    self,
    dim,
    custom_freqs: Tensor | None = None,
    freqs_for:  Literal['lang', 'pixel', 'constant'] = 'lang',
    theta = 10000,
    max_freq = 10,
    num_freqs = 1,
    learned_freq = False,
    use_xpos = False,
    xpos_scale_base = 512,
    interpolate_factor = 1.,
    theta_rescale_factor = 1.,
    seq_before_head_dim = False,
    cache_if_possible = True,
    cache_max_seq_len = 8192
):
```

### Parameters
- **dim** (int) - Required - Feature dimension over which frequencies are generated (`dim // 2` frequencies in `'lang'`/`'pixel'` modes).
- **custom_freqs** (torch.Tensor | None) - Optional - If provided, used directly as the frequency vector (shape `(f,)`), overriding `freqs_for`.
- **freqs_for** (Literal['lang', 'pixel', 'constant']) - Optional - Frequency schedule. Default: `'lang'`.
- **theta** (float / int) - Optional - Base of the inverse geometric progression for `'lang'` mode. Default: `10000`.
- **max_freq** (float) - Optional - Maximum frequency for `'pixel'` mode (sample points `torch.linspace(1., max_freq / 2, dim // 2) * pi`). Default: `10`.
- **num_freqs** (int) - Optional - Number of frequencies in `'constant'` mode (`torch.ones(num_freqs)`). Default: `1`.
- **learned_freq** (bool) - Optional - If `True`, the frequency vector becomes an `nn.Parameter(freqs, requires_grad = True)`. Default: `False`.
- **use_xpos** (bool) - Optional - Enables XPos length-extrapolatable decay (autoregressive only). Default: `False`.
- **xpos_scale_base** (int) - Optional - Base of the XPos decay power. Default: `512`.
- **interpolate_factor** (float) - Optional - Divides sequence positions by this factor (Position Interpolation). Default: `1.`.
- **theta_rescale_factor** (float) - Optional - NTK-aware rescaling: `theta *= theta_rescale_factor ** (dim / (dim - 2))`. Default: `1.`.
- **seq_before_head_dim** (bool) - Optional - If `True`, sequence dimension precedes head dimension, making `default_seq_dim = -3`. Default: `False`.
- **cache_if_possible** (bool) - Optional - Cache computed frequencies (and scales) in registered buffers. Default: `True`.
- **cache_max_seq_len** (int) - Optional - Size of the `cached_freqs` (and `cached_scales`) buffers. Default: `8192`.

### Frequency Construction
- `'lang'`: `freqs = 1. / (theta ** (torch.arange(0, dim, 2)[:(dim // 2)].float() / dim))`
- `'pixel'`: `freqs = torch.linspace(1., max_freq / 2, dim // 2) * pi`
- `'constant'`: `freqs = torch.ones(num_freqs).float()`

### Registered Buffers / Parameters
- **cached_freqs** (buffer, persistent = False) - Shape `(cache_max_seq_len, dim)` - Frequency cache.
- **cached_scales** (buffer, persistent = False) - Shape `(cache_max_seq_len, dim)` - XPos scale cache (only when `use_xpos = True`).
- **dummy** (buffer, persistent = False) - Shape `()` - Device anchor for the `device` property.
- **freqs** (parameter) - Shape `(f,)` - Frequency vector; `requires_grad = learned_freq`.
- **scale** (buffer, persistent = False) - Shape `(dim // 2,)` - XPos base scales `(torch.arange(0, dim, 2) + 0.4 * dim) / (1.4 * dim)` (only when `use_xpos = True`).

### Throws
- **AssertionError** - If `interpolate_factor < 1.`.

### Additional Attributes
- `freqs_for`, `cache_if_possible`, `cache_max_seq_len`, `cached_freqs_seq_len` (int, starts `0`), `learned_freq`, `seq_before_head_dim`, `default_seq_dim` (`-3` if `seq_before_head_dim` else `-2`), `interpolate_factor`, `use_xpos`. When `use_xpos = True`, `self.apply_rotary_emb = staticmethod(apply_rotary_emb)` is also attached.

### Example
```python
import torch
from rotary_embedding_torch import RotaryEmbedding

# basic language-model rotary embedding
rotary_emb = RotaryEmbedding(dim = 32)
q = torch.randn(1, 8, 1024, 64)
q = rotary_emb.rotate_queries_or_keys(q)

# XPos variant (assertion-safe: interpolate_factor must be >= 1.)
rotary_emb = RotaryEmbedding(dim = 32, use_xpos = True, xpos_scale_base = 512)
try:
    RotaryEmbedding(dim = 32, interpolate_factor = 0.5)
except AssertionError as e:
    print("interpolate_factor must be at least 1.")
```
```

--------------------------------

### apply_rotary_emb function signature

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding_torch.md

Defines the apply_rotary_emb function with parameters for frequency tensor, input tensor, start index, scale, and sequence dimensions. Decorated with @autocast('cuda', enabled = False) to ensure full precision rotation math.

```python
@autocast('cuda', enabled = False)
def apply_rotary_emb(
    freqs,
    t,
    start_index = 0,
    scale = 1.,
    seq_dim = -2,
    freqs_seq_dim = None
):
```

--------------------------------

### device property implementation

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Returns the device of the dummy buffer, used to anchor creation of position tensors. Source: rotary_embedding_torch/rotary_embedding_torch.py:168.

```python
@property
def device(self):
    return self.dummy.device
```

--------------------------------

### Fused Flash Attention with rotary embeddings using flash_attn_with_rotary

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Computes attention with rotary embeddings in a single fused kernel, falling back to a reference PyTorch implementation if Triton is unavailable. Supports ignoring tokens (e.g., CLS or register) by passing rotary_pos_emb_indices. The example uses CUDA tensors and a causal mask.

```python
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 = torch.randn(1, 8, 1024 + 2, 64).cuda() # 2 extra tokens (e.g. CLS, register)
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 (CLS / register)
pos_indices = torch.arange(1024).cuda() + 2

# fused flash attention with rotary
out = flash_attn_with_rotary(
    q, k, v,
    rotary_pos_emb = freqs,
    rotary_pos_emb_indices = pos_indices,
    is_causal = True
)
```

--------------------------------

### get_scale method signature

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Method signature for get_scale. Requires use_xpos=True, otherwise raises AssertionError.

```python
def get_scale(
    self,
    t: Tensor,
    seq_len: int | None = None,
    offset = 0
):
```

--------------------------------

### Usage: Apply rotary embeddings to queries and keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Shows how to instantiate RotaryEmbedding, create mock queries and keys, and apply the rotation before attention. The dimensions should end with (seq_len, feature dimension) and can have any number of preceding dimensions (batch, heads, etc.).

```python
import torch
from rotary_embedding_torch import RotaryEmbedding

# instantiate the positional embedding in your transformer and pass to all your attention layers

rotary_emb = RotaryEmbedding(dim = 32)

# mock queries and keys - dimensions should end with (seq_len, feature dimension), and any number of preceding dimensions (batch, heads, etc)

q = torch.randn(1, 8, 1024, 64) # queries - (batch, heads, seq len, dimension of head)
k = torch.randn(1, 8, 1024, 64) # keys

# apply the rotations to your queries and keys after the heads have been split out, but prior to the dot product and subsequent softmax (attention)

q = rotary_emb.rotate_queries_or_keys(q)
k = rotary_emb.rotate_queries_or_keys(k)

# then do your attention with your queries (q) and keys (k) as usual
```

--------------------------------

### reference_flash_attention(q, k, v, cos = None, sin = None, pos_mask = None, attn_mask = None, is_causal = True)

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

The pure-PyTorch, non-fused reference implementation of attention with optional rotary embeddings, used as fallback and ground truth. Supports GQA (keys/values repeated `g = heads_q // heads_kv` times via `repeat(k, 'b h n d -> b (h g) n d')`).

```APIDOC
## reference_flash_attention(q, k, v, cos = None, sin = None, pos_mask = None, attn_mask = None, is_causal = True)

### Description
The pure-PyTorch, non-fused reference implementation of attention with optional rotary embeddings, used as fallback and ground truth. Supports GQA (keys/values repeated `g = heads_q // heads_kv` times via `repeat(k, 'b h n d -> b (h g) n d')`).

### Signature
```python
def reference_flash_attention(
    q, k, v,
    cos = None,
    sin = None,
    pos_mask = None,
    attn_mask = None,
    is_causal = True
):
```

### Parameters
- **q** (torch.Tensor) - Required - `(batch, heads_q, seq_len_q, dim)`
- **k** (torch.Tensor) - Required - `(batch, heads_kv, seq_len_k, dim)`
- **v** (torch.Tensor) - Required - `(batch, heads_kv, seq_len_k, dim)`
- **cos** (torch.Tensor) - Optional, default `None` - Cosine frequencies (rot dim `cos.shape[-1]`).
- **sin** (torch.Tensor) - Optional, default `None` - Sine frequencies.
- **pos_mask** (torch.Tensor) - Optional, default `None` - `(seq,)` boolean mask; `sim = torch.where(is_pos_2d, sim_rot, sim)` where `is_pos_2d = pos_mask[:, None] & pos_mask[None, :]`.
- **attn_mask** (torch.Tensor) - Optional, default `None` - Additive mask added to logits (`sim = sim + attn_mask`).
- **is_causal** (bool) - Optional, default `True` - Applies a `torch.ones(seq_len_q, seq_len_k).tril()` mask; unmasked entries become `-inf`.

### Return Type
`torch.Tensor` — `torch.einsum('b h i j, b h j d -> b h i d', attn, v)` where `attn = sim.softmax(dim = -1)` and `sim = sim * (dim ** -0.5)`.

### Example
```python
import torch
from rotary_embedding_torch.flash_attn_with_rotary import reference_flash_attention

q = torch.randn(1, 8, 64, 64)
k = torch.randn(1, 4, 64, 64)   # GQA: 8 query heads, 4 kv heads
v = torch.randn(1, 4, 64, 64)

out = reference_flash_attention(q, k, v, is_causal = False)
```
```

--------------------------------

### Partial rotation with start_index and assertion error handling

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding_torch.md

Demonstrates partial rotation by specifying start_index=32, rotating only features [32, 80) of a 128-feature tensor. Also shows catching the AssertionError when the rotation dimension exceeds the feature dimension.

```python
import torch
from rotary_embedding_torch import apply_rotary_emb

freqs = torch.randn(64, 48)   # only 48 of 64 features will be rotated

t = torch.randn(2, 4, 64, 128)  # (batch, heads, seq, features=128)
out = apply_rotary_emb(freqs, t, start_index = 32, seq_dim = -2)
# features [32, 80) rotated, the rest untouched — out.shape == (2, 4, 64, 128)

try:
    apply_rotary_emb(torch.randn(64, 96), torch.randn(64, 64))
except AssertionError as e:
    print(e)  # 'feature dimension 64 is not of sufficient size to rotate in all the positions 96'
```

--------------------------------

### rotate_queries_with_cached_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Rotates a single query against a full key/value cache during inference. The query is automatically offset by k_len - q_len so its positions align with the tail of the keys. When use_xpos is True, the query is scaled by the XPos scale of its aligned positions and the keys by the inverse scale.

```APIDOC
## rotate_queries_with_cached_keys

### Description
Rotates a single query against a full key/value cache during inference. The query is automatically offset by `k_len - q_len` so its positions align with the tail of the keys. When `use_xpos = True`, the query is scaled by the XPos scale of its aligned positions and the keys by the inverse scale.

### Method
rotate_queries_with_cached_keys

### Signature
```python
def rotate_queries_with_cached_keys(self, q, k, seq_dim = None, offset = 0):
```

### Parameters
- `q` (torch.Tensor) - Required - Query tensor, feature dimension last.
- `k` (torch.Tensor) - Required - Key tensor (concatenated cache), feature dimension last.
- `seq_dim` (int) - Optional - Sequence dimension (must be `-2` or `-3` consistently for both tensors). Default: `self.default_seq_dim`.
- `offset` (int) - Optional - Additional offset on top of `k_len - q_len`. Default: `0`.

### Return Type
`tuple[torch.Tensor, torch.Tensor]` — `(rotated_q, rotated_k)`, both cast back to `q.dtype` via `.type(q.dtype)`.

### Throws
- `AssertionError` - If `q_len > k_len` (the query sequence length must not exceed the key sequence length).

### Examples
```python
import torch
from rotary_embedding_torch import RotaryEmbedding

rotary_emb = RotaryEmbedding(dim = 32)

q = torch.randn(1, 8, 1, 64)     # one query at a time
k = torch.randn(1, 8, 1024, 64)  # keys with cache concatenated

q, k = rotary_emb.rotate_queries_with_cached_keys(q, k)
```

With XPos enabled (query scaled up, keys scaled down):

```python
rotary_emb = RotaryEmbedding(dim = 32, use_xpos = True)

q = torch.randn(1, 8, 1, 64)
k = torch.randn(1, 8, 1024, 64)
try:
    q, k = rotary_emb.rotate_queries_with_cached_keys(q, torch.randn(1, 8, 0, 64))  # empty keys
    print("rotated", q.shape, k.shape)
except AssertionError as e:
    print("q_len must be <= k_len")
```
```

--------------------------------

### FlashAttentionFused backward method signature

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Defines the backward static method of FlashAttentionFused. It takes the context and gradient of the output, allocates zero gradients for query, key, and value, launches the backward Triton kernel, and returns gradients for the tensor inputs and None for non-tensor arguments.

```python
    @staticmethod
    def backward(ctx, dout):
```

--------------------------------

### Import RotaryEmbedding

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Import the RotaryEmbedding class from the rotary_embedding_torch package.

```python
from rotary_embedding_torch import RotaryEmbedding
```

--------------------------------

### Manually rotate queries with cached keys using rotate_queries_or_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/README.md

Alternative to rotate_queries_with_cached_keys. Manually specify the offset as the difference between key/value sequence length and query sequence length.

```python
q = rotary_emb.rotate_queries_or_keys(q, offset = k.shape[-2] - q.shape[-2])
```

--------------------------------

### FlashAttentionFused.backward

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Static method that launches the fused backward Triton kernel to compute gradients for query, key, and value tensors.

```APIDOC
## FlashAttentionFused.backward

### Description
Static method that launches the fused backward Triton kernel to compute gradients for the query, key, and value tensors given the gradient of the output.

### Method
Static method (typically called via autograd)

### Signature
```python
@staticmethod
def backward(ctx, dout)
```

### Parameters
- **ctx** - Context object used by autograd.
- **dout** (torch.Tensor) - Gradient of the output.

### Return Type
- **tuple** - Gradients `(dq, dk, dv, None, None, None, None, None)` for the input tensors and non-tensor arguments.

### Example
(No direct example provided; backward is typically invoked automatically by autograd.)
```

--------------------------------

### apply_learned_rotations with learned frequency ranges

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding_torch.md

Shows how to use apply_learned_rotations with a freq_ranges tensor. The rotations are expanded by the frequency ranges via einsum before being paired into cos/sin components.

```python
import torch
from rotary_embedding_torch import apply_learned_rotations

t = torch.randn(2, 8, 128, 64)
rotations = torch.randn(2, 8, 128, 4, 8)   # (..., r, f) form
freq_ranges = torch.randn(8)               # per-frequency factor

out = apply_learned_rotations(rotations, t, freq_ranges = freq_ranges)
# rotations expanded to (..., 4, 8) -> (..., 32), then paired to 64 features
```

--------------------------------

### RotaryEmbedding __init__ signature

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

The constructor signature for RotaryEmbedding, showing all parameters and their defaults. It supports custom frequencies, different frequency schedules, XPos, interpolation, and caching.

```python
def __init__(
    self,
    dim,
    custom_freqs: Tensor | None = None,
    freqs_for:  Literal['lang', 'pixel', 'constant'] = 'lang',
    theta = 10000,
    max_freq = 10,
    num_freqs = 1,
    learned_freq = False,
    use_xpos = False,
    xpos_scale_base = 512,
    interpolate_factor = 1.,
    theta_rescale_factor = 1.,
    seq_before_head_dim = False,
    cache_if_possible = True,
    cache_max_seq_len = 8192
):
```

--------------------------------

### FlashAttentionFused.forward

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Static method that launches the fused forward Triton kernel. Accepts query, key, value tensors, cosine and sine frequency tensors, and optional position and attention masks. Returns the attention output tensor.

```APIDOC
## FlashAttentionFused.forward

### Description
Static method that launches the fused forward Triton kernel for flash attention with rotary embeddings. It processes query, key, value tensors along with cosine and sine frequency tensors, and optional position and attention masks, returning the attention output.

### Method
Static method (typically called via `FlashAttentionFused.apply(...)`)

### Signature
```python
@staticmethod
def forward(
    ctx,
    q, k, v,
    cos, sin,
    pos_mask = None,
    attn_mask = None,
    is_causal = True
)
```

### Parameters
- **ctx** - Context object used by autograd.
- **q** (torch.Tensor) - Query tensor of shape `(b, h, n, d)`.
- **k** (torch.Tensor) - Key tensor of shape `(b, h, n, d)`.
- **v** (torch.Tensor) - Value tensor of shape `(b, h, n, d)`.
- **cos** (torch.Tensor) - Cosine frequency tensor, shape ending `... 1 1 n d` after rearrangement by the caller.
- **sin** (torch.Tensor) - Sine frequency tensor, shape ending `... 1 1 n d` after rearrangement by the caller.
- **pos_mask** (optional) - Position mask.
- **attn_mask** (optional) - Attention mask.
- **is_causal** (bool, default True) - Whether to use causal masking.

### Return Type
- **torch.Tensor** - Output tensor of shape `(batch, heads_q, seq_len_q, dim)`, filled by the kernel.

### Throws
- Relies on the caller's assertions; directly calling `forward` without `cos`/`sin` colocated on the same device as `q` will raise a device-mismatch `RuntimeError` from Triton/PyTorch.

### Example
```python
import torch
from rotary_embedding_torch import RotaryEmbedding
from rotary_embedding_torch.flash_attn_with_rotary import FlashAttentionFused

rotary_emb = RotaryEmbedding(dim = 32)
freqs = rotary_emb(torch.arange(128)).cuda()
cos = freqs.cos().reshape(128, 1, 1, 32)   # kernel expects '... 1 1 n d' layout
sin = freqs.sin().reshape(128, 1, 1, 32)

q = torch.randn(1, 8, 128, 64).cuda()
k = torch.randn(1, 8, 128, 64).cuda()
v = torch.randn(1, 8, 128, 64).cuda()

out = FlashAttentionFused.apply(q, k, v, cos, sin, None, None, True) if TRITON_AVAILABLE else None
# out.shape == (1, 8, 128, 64)
```
```

--------------------------------

### Method signature for rotate_queries_with_cached_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

The method signature as defined in the source file. It takes query and key tensors, an optional sequence dimension, and an offset.

```python
def rotate_queries_with_cached_keys(self, q, k, seq_dim = None, offset = 0):
```

--------------------------------

### get_flash_attention_fused(force_reference = False)

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Factory function that returns an attention callable. Passing force_reference = True yields a callable that always runs the pure-PyTorch reference implementation, useful for testing and CPU execution.

```APIDOC
## get_flash_attention_fused(force_reference = False)

### Description
Factory function that returns an attention callable. Passing `force_reference = True` yields a callable that always runs the pure-PyTorch reference implementation (useful for testing and CPU execution).

### Signature
```python
def get_flash_attention_fused(force_reference = False):
```

### Parameters
- **force_reference** (bool) - Optional, default `False` - If `True`, the returned inner function always dispatches to `reference_flash_attention`.

### Return Type
`Callable[[...], torch.Tensor]` — the `inner` function with signature `(q, k, v, cos = None, sin = None, pos_mask = None, attn_mask = None, is_causal = True, rotary_pos_emb = None, rotary_pos_emb_indices = None)`.

### Example
```python
import torch
from rotary_embedding_torch.flash_attn_with_rotary import get_flash_attention_fused

flash_attn_ref = get_flash_attention_fused(force_reference = True)   # CPU-safe implementation

q = torch.randn(1, 8, 128, 64)
k = torch.randn(1, 8, 128, 64)
v = torch.randn(1, 8, 128, 64)

out = flash_attn_ref(q, k, v, is_causal = False)   # plain attention, no rotary
# out.shape == (1, 8, 128, 64)
```
```

--------------------------------

### rotate_queries_and_keys

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding.md

Rotates both queries and keys.

```APIDOC
## rotate_queries_and_keys

### Description
Rotates both queries and keys.

### Method
`rotate_queries_and_keys` (line 215)
```

--------------------------------

### FlashAttentionFused forward method signature

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Defines the forward static method of FlashAttentionFused. It takes query, key, value tensors, cosine and sine frequency tensors, optional position and attention masks, and a causal flag. The method makes tensors contiguous, allocates output and log-sum-exp buffers, and launches the forward Triton kernel.

```python
class FlashAttentionFused(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        q, k, v,
        cos, sin,
        pos_mask = None,
        attn_mask = None,
        is_causal = True
    ):
```

--------------------------------

### reference_flash_attention

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/flash_attn_with_rotary.md

Reference implementation of Flash Attention with rotary embeddings in pure PyTorch. Used as fallback when Triton is not available.

```APIDOC
## reference_flash_attention

### Description
Reference implementation of Flash Attention with rotary embeddings in pure PyTorch. Used as fallback when Triton is not available.

### Import
```python
from rotary_embedding_torch.flash_attn_with_rotary import reference_flash_attention
```

### Parameters
- **q** (Tensor) - Query tensor.
- **k** (Tensor) - Key tensor.
- **v** (Tensor) - Value tensor.
- **rotary_pos_emb** (Tensor) - Rotary position embeddings.
- **rotary_pos_emb_indices** (Tensor, optional) - Indices to ignore tokens from rotation.
- **causal** (bool, optional) - Whether to apply causal masking.
- **attn_mask** (Tensor, optional) - Additive attention mask.

### Returns
- **output** (Tensor) - Attention output with rotary embeddings applied.

### Notes
- Runs on CPU or GPU.
```

--------------------------------

### Basic usage of apply_learned_rotations

Source: https://github.com/lucidrains/rotary-embedding-torch/blob/main/_autodocs/api-reference/rotary_embedding_torch.md

Shows how to apply learned rotations to a tensor. The rotations tensor has one learned rotation per position, which is repeated pairwise to match the feature dimension of the input tensor.

```python
import torch
from rotary_embedding_torch import apply_learned_rotations

t = torch.randn(2, 8, 128, 64)          # (batch, heads, seq, features)
rotations = torch.randn(2, 8, 128, 32)  # one learned rotation per position

out = apply_learned_rotations(rotations, t)
# rotations repeated to 64 features, so all 64 are rotated
```
Tokens:
Raw