# How Needle 2 Implements Rotary Position Encoding (RoPE)

> Discover how Needle 2 implements Rotary Position Encoding RoPE by injecting relative positional info into query key tensors with pre-computed frequencies and rotary transformations.

- Repository: [Cactus Compute, Inc./needle](https://github.com/cactus-compute/needle)
- Tags: internals
- Published: 2026-08-24

---

**Needle 2 injects relative positional information directly into query and key tensors by pre-computing sinusoidal frequency tables and applying a rotary transformation during the attention forward pass.**

Needle 2, the JAX-native transformer library from Cactus Compute, implements Rotary Position Embedding (RoPE) as a parameter-free mechanism to encode token positions within its multi-head attention blocks. Unlike additive positional encodings, RoPE rotates the query and key vectors by angles derived from their position indices, allowing the model to learn relative positional relationships naturally. The implementation is contained primarily in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) and integrates seamlessly with the library's functional JAX design.

## Core RoPE Architecture

The RoPE system in Needle 2 consists of three coordinated components: frequency pre-computation, rotary application, and attention integration. Each component is implemented as a pure function compatible with JAX's compilation and automatic differentiation.

### Pre-computing Sinusoidal Frequency Tables

The `precompute_rope_freqs` function generates cosine and sine lookup tables based on the head dimension, sequence length, and a configurable base frequency θ (theta). According to the source code in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 97-101):

```python
def precompute_rope_freqs(head_dim, seq_len, theta=10000.0):
    freqs = 1.0 / (theta ** (jnp.arange(0, head_dim, 2).astype(jnp.float32) / head_dim))
    t = jnp.arange(seq_len).astype(jnp.float32)
    angles = jnp.outer(t, freqs)
    return jnp.cos(angles), jnp.sin(angles)

```

This function calculates rotation frequencies using the standard RoPE formula where θ defaults to `10000.0` but can be overridden via `TransformerConfig.rope_theta` (line 69). The output tuples `(cos, sin)` represent pre-computed angles for every position and dimension pair, enabling efficient vectorized operations during inference.

### Applying the Rotary Transformation

The `apply_rope` function executes the actual rotation by interleaving cosine and sine values with the query or key tensor. As implemented in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 104-112):

```python
def apply_rope(x, cos, sin):
    T = x.shape[2]
    half = x.shape[-1] // 2
    cos = cos[:T][None, None, :, :]
    sin = sin[:T][None, None, :, :]
    x1 = x[..., :half]
    x2 = x[..., half:]
    return jnp.concatenate([x1 * cos - x2 * sin,
                            x2 * cos + x1 * sin], axis=-1).astype(x.dtype)

```

This operation splits the input tensor into two halves, treats them as complex number components, and performs a 2D rotation. The result maintains the original `dtype` (typically `bfloat16` in Needle 2 configurations), ensuring numerical precision is preserved throughout the attention computation.

### Integration with Multi-Head Attention

Each transformer block receives an optional `rope` argument containing the pre-computed cosine and sine tables. When provided, the attention layer applies the rotary transformation to both query (Q) and key (K) tensors before computing attention scores. From [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) (lines 237-240):

```python
if rope is not None:
    cos, sin = rope
    q = apply_rope(q, cos, sin)
    k = apply_rope(k, cos, sin)

```

This integration ensures that every attention head processes position-aware representations without introducing additional trainable parameters or inference-time latency beyond the initial rotation computation.

## Using RoPE in Practice

While Needle 2 models can automatically generate RoPE tables via the `_rope` helper method (lines 504-507), you can manually pre-compute and pass them for fine-grained control:

```python
import jax.numpy as jnp
from needle.model.architecture import TransformerConfig, precompute_rope_freqs

# Configure model with custom theta if desired

cfg = TransformerConfig(d_model=768, num_heads=12, max_seq_len=1024, rope_theta=10000.0)

# Prepare input tokens (batch_size=1, seq_len=16)

tokens = jnp.arange(16)[None, :]  # Shape: (1, 16)

# Pre-compute RoPE tables for current sequence length

head_dim = cfg.d_model // cfg.num_heads
cos, sin = precompute_rope_freqs(head_dim, tokens.shape[1], theta=cfg.rope_theta)

# Pass RoPE tables to model forward pass

# outputs, hidden = model(tokens, rope=(cos, sin))

```

The model's `_rope` method automatically handles this computation when the `rope` argument is omitted, generating tables using the current sequence length and the configured `rope_theta` value.

## Configuration and Customization

**Theta Scaling:** Modify the base frequency θ through `TransformerConfig.rope_theta` to adjust how aggressively the model distinguishes between distant token positions. Higher values (e.g., `100000.0`) may improve performance on very long sequences by slowing the rotation frequency decay.

**Sequence Length Flexibility:** Because RoPE tables are computed on-the-fly via `_rope(self, seq_len)`, Needle 2 supports variable sequence lengths during inference without requiring fixed positional embedding matrices. This enables efficient processing of sequences shorter or longer than the training maximum.

## Summary

- **Needle 2** implements RoPE through pure JAX functions in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), keeping the implementation compatible with JIT compilation and gradient checkpointing.
- **`precompute_rope_freqs`** generates sinusoidal tables using configurable θ values, defaulting to `10000.0`.
- **`apply_rope`** performs the actual 2D rotation on query and key tensors, preserving the model's `bfloat16` dtype.
- **Attention blocks** accept optional `rope` tuples and apply rotations automatically when present, requiring no architectural changes to existing transformer stacks.

## Frequently Asked Questions

### What is rotary position encoding (RoPE)?

Rotary Position Embedding (RoPE) is a method for encoding relative token positions by rotating query and key vectors in attention mechanisms using sinusoidal functions of their position indices. Unlike absolute positional embeddings, RoPE allows the model to naturally express relative positional relationships through the dot-product operation, improving generalization to sequence lengths unseen during training.

### How do I configure the RoPE theta parameter in Needle 2?

Set the `rope_theta` field when instantiating `TransformerConfig`. The default value is `10000.0`, but you can increase it (e.g., to `100000.0`) for models processing very long contexts. This value is passed to `precompute_rope_freqs` during the forward pass via the `_rope` method.

### Can I use RoPE with different sequence lengths than the training max?

Yes. Needle 2 computes RoPE tables dynamically based on the actual input sequence length rather than relying on fixed embedding matrices. The `_rope` method generates fresh cosine and sine tables for each forward pass, allowing efficient extrapolation to longer sequences or compression for shorter ones.

### Does RoPE add trainable parameters to the model?

No. RoPE is a parameter-free encoding scheme. The rotation angles are derived entirely from deterministic sinusoidal functions of position indices and dimension indices. This reduces model size and eliminates the risk of overfitting to specific absolute positions in the training data.