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

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

The RoPE implementation centers on two utilities located in 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.

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/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.

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/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.

Frequency Table Initialization

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

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/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:

# 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

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

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

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	Houses RoPE utilities precompute_rope_freqs and apply_rope
needle/model/decode.py	Invokes RoPE operations during self-attention forward pass
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, 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) 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.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}