rotary-embedding-torch: RoPE as a Standalone Pytorch Module

The gap it fills: RoPE without copying a model repository

Rotary positional embeddings appear in most open transformer implementations, but usually as a private function buried inside a model file. If you are building a transformer outside an existing training framework, your choice has been to copy that function, reimplement the rotation math, or pull in an entire model repository for one operation. This package exists to remove that choice. It is a standalone library for adding rotary embeddings to transformers in Pytorch, in the README's phrasing, and its scope stops at the embedding. It does not ship a transformer, a training loop, or a tokenizer. The audience is narrow and specific: researchers and engineers who already have an attention implementation and want the positional encoding as a dependency rather than as pasted code. The README also frames it as making rotation into any axis of a tensor easy, which is the thread that leads to the axial frequency support for video later in the document.

What rotate_queries_or_keys does to your q and k tensors

The mechanism is a rotation applied to the query and key tensors, not to the token embeddings. You instantiate RotaryEmbedding with a dim, pass the instance to your attention layers, and call rotate_queries_or_keys on q and k. The README is explicit about placement: apply the rotations after the heads have been split out, but prior to the dot product and subsequent softmax. The shape contract follows from that. Tensors end with (seq_len, feature dimension) and may carry any number of preceding dimensions for batch, heads, and so on; the example uses q of shape (1, 8, 1024, 64), meaning batch 1, 8 heads, 1024 tokens, head dimension 64. Because the rotation is applied to q and k separately and before the dot product, the relative position information survives into the attention scores. The README states that if you do the steps correctly, you should see a dramatic improvement during training. That is the author's claim, not a measured result, and the README supplies no benchmark to support it. What can be said from the code path is narrower: the embedding is applied per attention layer, so a model with N layers holds N references to the same rotary_emb instance in the documented pattern.

Inference with a KV cache: the offset problem and its two fixes

Autoregressive decoding breaks naive RoPE. When you generate one token at a time, the query has length 1 while the cached keys have length 1024, and the query position must be offset by key_value_seq_length - query_seq_length. The README states this directly and offers two routes. The convenience method is rotate_queries_with_cached_keys, which takes q of shape (1, 8, 1, 64) and k of shape (1, 8, 1024, 64) and returns both after handling the offset. The manual route is to pass offset explicitly: rotate_queries_or_keys(q, offset = k.shape[-2] - q.shape[-2]). Having both is a reasonable design. The explicit offset form is the one to reach for when your cache layout differs from the documented one, for example when you interleave multiple sequences in a batch and each needs its own position count. Note that the README does not describe how the method behaves for a batch of sequences with different cache lengths. If your serving setup packs sequences, that is the case to test before trusting the convenience wrapper.

XPos, axial frequencies, and interpolation: three optional behaviours

Three optional paths sit behind initialization flags. Setting use_xpos = True applies the XPos technique from the cited paper, which gives the embedding a decay similar to ALiBi so it extrapolates to sequence lengths longer than training. The README states a hard constraint on this: it can only be used for autoregressive transformers, and you must call rotate_queries_and_keys instead of rotate_queries_or_keys, because the method handles both tensors together. The axial path targets video. With freqs_for = 'pixel' and a max_freq, you call get_axial_freqs with the frame count and spatial dimensions, and the README notes it will automatically do partial rotary when the frequency tensor is wider than the head dimension; the example produces a frequency tensor of shape (8, 64, 32, 48) for a head dimension of 16 across three axes. You then apply it with the module-level apply_rotary_emb function rather than a method. The third option, interpolate_factor, follows the cited MetaAI paper on fine-tuning with interpolated positions. The README's own update is the notable part: someone in the community has reported that it does not work well, and the author asks to be emailed with either a positive or negative result. Treat that flag as unproven.

The fused Flash Attention path and its Triton dependency

The newest surface in the README is flash_attn_with_rotary, imported from rotary_embedding_torch.flash_attn_with_rotary. It computes attention with rotary embeddings in one pass. The README states that it automatically falls back to a reference Pytorch implementation on CPU or GPU if Triton is not available, which means the import does not hard-fail on a machine without Triton; you get correct results through a slower path. The calling convention differs from the rest of the library. You first materialise frequencies by calling the embedding on a position range, rotary_emb(torch.arange(1024)), then pass q, k, v along with rotary_pos_emb = freqs and rotary_pos_emb_indices = pos_indices. That indices argument is the feature worth noting: it lets you exclude tokens from rotation, which the README illustrates with two extra tokens for CLS or register positions, using pos_indices = torch.arange(1024).cuda() + 2 to skip the first two. The example also sets is_causal = True. This is the part of the library most likely to change shape between releases, since it depends on a Triton kernel rather than pure Pytorch operators.

Where it stops being the right tool

The library operates on tensors you already have. It does not know about your model, your attention mask, or your cache manager, and it will not warn you if you call rotate_queries_or_keys on tensors whose last two dimensions are not (seq_len, head_dim). A silent misapplication produces a model that trains to a worse loss rather than an exception, which is the failure mode to watch for during integration. The XPos path is narrower than the headline suggests: the README restricts it to autoregressive transformers, so it is not available to encoder-only or bidirectional setups. The interpolate_factor path carries the author's own caveat about a negative community report. And the Flash Attention path inherits a Triton dependency for its fast route, with the Pytorch fallback available but not benchmarked in the README. If you need a positional scheme with learned parameters, or you want the framework to own the attention kernel end to end, this package is a partial answer by design.

Compared with using a framework's built-in positional encoding

The obvious alternative is to use whatever positional encoding your training framework already provides. Hugging Face transformer implementations, for instance, construct rotary embeddings inside the attention module, tied to a specific model class and its config. The difference in approach is ownership. In a framework, you configure the model and the rotation is an implementation detail you inherit; you cannot swap the frequency computation without editing the model class, and you cannot apply the same rotation to a custom attention kernel that the framework does not know about. Here you own the call site. That is what makes the axial and fused-attention paths possible at all: they are not model options, they are functions you invoke on your own tensors. The cost is that nothing is wired up for you. You must place the call correctly relative to head splitting and the dot product, and you must handle the KV cache offset yourself or through the provided helper. The trade is configuration convenience for control over where and how the rotation is applied.

Install, versions, and what the MIT licence leaves you

Installation is one command: pip install rotary-embedding-torch. The package is pure Python and depends on Pytorch, which is not pinned in the material supplied here, so check the package metadata for the version range your environment needs rather than assuming the latest Pytorch works. The release cadence visible in the material is active: 0.8.7, 0.8.8, and 0.8.9 shipped within July 2025, and the repository's last push is dated 2026-06-20. Frequent patch releases in a small library usually mean small fixes rather than API churn, but the Flash Attention module is the component most exposed to upstream Triton and Pytorch changes, so pin a version if you depend on that path. The licence is MIT, which permits commercial and closed-source use and modification provided the copyright notice and permission notice are retained. That is a summary of the licence identifier, not legal advice; read the LICENSE file and your own obligations before shipping.

Editorial conclusion

Adopt it if you are writing your own attention loop in Pytorch and want RoPE, XPos, axial frequencies, or a fused Triton attention kernel without vendoring code from a model repo. Do not adopt it if you need a framework-level drop-in: this is a tensor operation, not a model. Before committing, verify that your attention code splits heads before calling rotate_queries_or_keys, that your Pytorch version satisfies the package metadata, and test the interpolate_factor path yourself, since the README states a community report that it does not work well.

Official sources

1. Issues
2. License: MIT
3. lucidrains/rotary-embedding-torch on GitHub
4. README
5. Releases