# How Needle 2 Implements Rotary Position Embedding (RoPE): A Complete Code Walkthrough

> Explore Needle 2's Rotary Position Embedding RoPE implementation with a code walkthrough. Understand how precomputed frequencies rotate tensors for relative position encoding.

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

---

**Needle 2 implements Rotary Position Embedding (RoPE) through two core functions—`precompute_rope_freqs` and `apply_rope`—that rotate query and key tensors by pre-computed angle frequencies to encode relative position information.**

RoPE has become the standard positional encoding method in modern transformer architectures. In Needle 2, the implementation follows the original paper's formulation while optimized for JAX/Flax execution. This article examines the complete RoPE pipeline, from frequency pre-computation to tensor rotation in the attention forward pass.

## Core RoPE Functions in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py)

The RoPE implementation centers on two utilities located in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py). These functions handle the mathematical operations that embed positional information into attention projections.

### `precompute_rope_freqs`: Building the Frequency Tables

The `precompute_rope_freqs` function generates cosine and sine lookup tables for all positions up to a maximum sequence length. This pre-computation amortizes the angle calculation cost across the forward pass.

```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)

```

Source: [[`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py#L97-L101)](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L97-L101)

The function operates in four steps:

- **Frequency computation**: Divides the head dimension into pairs and computes `θ^(-2i/d)` for indices `i`
- **Position vector**: Creates a range `[0, seq_len)` representing token positions
- **Outer product**: Combines positions with frequencies to produce all angle values
- **Trigonometric outputs**: Returns `(cos(angles), sin(angles))` for later use

The default **theta value of 10000.0** matches the original RoPE paper, with common variants using 100000.0 for longer context windows.

### `apply_rope`: Rotating Query and Key Tensors

The `apply_rope` function performs the actual rotary transformation on attention tensors. It accepts pre-computed cosine/sine tables and applies the 2D rotation to each position-head pair.

```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)

```

Source: [[`architecture.py`](https://github.com/cactus-compute/needle/blob/main/architecture.py#L104-L111)](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py#L104-L111)

The rotation follows standard complex multiplication: treating `(x1, x2)` as a complex number, it multiplies by `(cos, sin)` to rotate by angle `θ_t`. The implementation broadcasts efficiently across batch and head dimensions.

## Integration in the Attention Forward Pass

RoPE is applied to queries and keys immediately after linear projection in the decoder. The relevant logic resides in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py).

### Frequency Table Initialization

Cosine and sine tables are generated once per forward pass based on configuration:

```python
cos, sin = precompute_rope_freqs(head_dim, max_len, config.rope_theta)  # line 261

```

Source: [[`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py#L261-L263)](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py#L261-L263)

### Application to Q and K Tensors

After projecting input to query and key spaces, both tensors undergo rotary transformation:

```python

# In needle/model/decode.py

q = apply_rope(q, cos_s, sin_s)   # line 73

k = apply_rope(k, cos_s, sin_s)   # line 73

```

This placement ensures positional information enters **before** the attention score computation, enabling the relative position inductive bias that makes RoPE effective.

## Practical Usage Examples

### Generating RoPE Tables

```python
from needle.model.architecture import precompute_rope_freqs

head_dim = 64                      # per-head dimension (must be even)

seq_len  = 1024                    # maximum sequence length for the batch

cos, sin = precompute_rope_freqs(head_dim, seq_len, theta=100000.0)

```

### Applying RoPE to Attention Tensors

```python
from needle.model.architecture import apply_rope
import jax.numpy as jnp

# Dummy query tensor: shape (batch, heads, time, dim)

q = jnp.ones((2, 8, 16, head_dim))

# Apply rotary embedding (cos & sin must cover the time dimension)

q_rot = apply_rope(q, cos, sin)

```

### Full Inference Context

```python
from needle.model.decode import decode_cfg, init_kv_cache
from needle.model.architecture import precompute_rope_freqs, apply_rope

cfg = TransformerConfig(rope_theta=100000.0, num_heads=8, d_model=512, max_seq_len=2048)
head_dim = cfg.d_model // cfg.num_heads
cos, sin = precompute_rope_freqs(head_dim, cfg.max_seq_len, cfg.rope_theta)

# Inside the attention block

q = (x @ lp["q"]).reshape(B, S, H, head_dim).transpose(0, 2, 1, 3)
k = (x @ lp["k"]).reshape(B, S, cfg.num_kv_heads, head_dim).transpose(0, 2, 1, 3)

q = apply_rope(q, cos, sin)
k = apply_rope(k, cos, sin)

```

## Key Implementation Files

| File | Purpose |
|------|---------|
| [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) | Houses **RoPE** utilities `precompute_rope_freqs` and `apply_rope` |
| [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) | Invokes RoPE operations during self-attention forward pass |
| [`needle/model/run.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/run.py) | Orchestrates inference pipeline utilizing RoPE-enabled attention |

## Summary

- **RoPE in Needle 2** uses a clean two-function design: frequency pre-computation and tensor rotation
- **`precompute_rope_freqs`** generates position-dependent cosine/sine tables with configurable base frequency
- **`apply_rope`** implements exact 2D rotation on query and key tensors of shape `[B, H, T, D]`
- The implementation integrates at **lines 73 and 261-263** of [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py), applied post-projection and pre-attention
- Both **flash-attention and fallback kernels** are supported through JAX-compatible operations

## Frequently Asked Questions

### What is the default theta value in Needle 2's RoPE implementation?

The default **theta value is 10000.0**, matching the original RoPE paper. This can be overridden via `config.rope_theta` to support longer contexts—common alternatives include 100000.0 for extended sequence lengths.

### Why does `apply_rope` split the tensor in half?

The split enables **pairwise 2D rotation**. RoPE treats adjacent dimension pairs as complex numbers; splitting into `x1` and `x2` allows applying the rotation matrix `[cos, -sin; sin, cos]` to each pair independently.

### Does Needle 2 cache RoPE frequencies across layers?

Frequencies are **pre-computed once per forward pass** (line 261 in [`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py)) and reused for all layers in that pass. This design avoids redundant computation while maintaining flexibility for varying sequence lengths.

### Is RoPE applied to value tensors in Needle 2?

No—**only queries and keys receive RoPE**. Values retain absolute position encoding implicitly through their dependence on positionally-encoded queries and keys in the attention operation. This follows standard practice in RoPE-based architectures.