# How RoPE Positional Embeddings Work in Needle 2: A Technical Guide

> Discover how Needle 2 uses RoPE positional embeddings for advanced attention. Explore deterministic sinusoidal rotations applied to query and key tensors.

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

---

**Needle 2 implements Rotary Positional Embeddings (RoPE) as deterministic sinusoidal rotations applied directly to query and key tensors within the attention mechanism, using pre-computed cosine and sine tables generated by `precompute_rope_freqs` and applied via `apply_rope` in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py).**

The cactus-compute/needle repository provides a lean implementation of transformer architectures that replaces traditional learned positional encodings with geometric rotation matrices. This article examines how RoPE Positional Embeddings are integrated throughout Needle 2's attention pipeline, from initial frequency pre-computation to real-time application during autoregressive text generation.

## Core RoPE Utilities

The foundation of Needle 2's positional encoding lives in [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py), where two primary functions handle the mathematical operations.

### Pre-computing Frequency Tables

Before the forward pass begins, the system generates rotation frequency tables using the `precompute_rope_freqs` function located at **lines 97-101**. This helper calculates cosine and sine values for every position and head dimension based on the configurable scaling factor `rope_theta` (defaulting to **100,000.0** as defined at lines 69-71).

The function accepts:
- `head_dim`: The dimensionality of each attention head (must be divisible by 2)
- `seq_len`: The maximum sequence length to pre-compute
- `theta`: The RoPE scaling factor controlling rotation frequency

### Applying Rotations to Q/K Tensors

The `apply_rope` function at **lines 104-112** performs the actual rotation. It accepts tensors of shape **`[B, H, T, D]`** (Batch, Heads, Time, Dimension) and mixes the first half of the hidden dimension with the second half using complex number multiplication via the pre-computed cosine and sine tables.

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

# Simulate query tensor (B=2, H=4, T=16, D=32)

q = jnp.ones((2, 4, 16, 32))

# Pre-compute rotation tables

cos, sin = precompute_rope_freqs(head_dim=32, seq_len=16, theta=100_000.0)

# Apply rotary embeddings

q_rotated = apply_rope(q, cos, sin)  # Shape remains (2, 4, 16, 32)

```

## Integrating RoPE into the Attention Pipeline

RoPE is not applied as a separate embedding layer but injected directly into the multi-head attention computation.

### Model-Level RoPE Generation

The top-level `SimpleAttentionNetwork` manages RoPE caching through its private `_rope(seq_len)` method at **lines 504-507**. This method derives the head dimension from the model configuration, invokes `precompute_rope_freqs`, and caches the resulting `(cos, sin)` tuple for reuse across all layers during the forward pass.

```python
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig

cfg = TransformerConfig(vocab_size=8192, d_model=512, num_heads=8)
model = SimpleAttentionNetwork(cfg)

# During __call__, the model computes:

# rope = self._rope(tokens.shape[1])

```

### Attention Block Injection

Inside `MultiHeadAttention.__call__` at **lines 237-241**, the system checks for the optional `rope` argument. When provided, the cached cosine and sine tables are unpacked and passed to `apply_rope` for both the **query** (`q`) and **key** (`k`) tensors before the attention scores are computed.

This ensures that every token receives position-aware representations without introducing additional trainable parameters, as the rotations are purely deterministic functions of the token's index.

## RoPE in Autoregressive Decoding

The cached decoder implementation in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) reuses the same RoPE tables during generation. The JIT-compiled `_attn_cached` function applies `apply_rope` to incoming queries and keys before writing them into the KV cache (referenced at **lines 73, 261, and 387**).

This design ensures that RoPE works identically for both training and fast autoregressive inference, maintaining rotational consistency between cached historical keys and newly computed queries.

## Configuring RoPE Parameters

The `TransformerConfig` dataclass exposes `rope_theta` for customization. Lower values increase rotation frequency, while higher values (such as the **100,000.0** default) extend context length capabilities by slowing the rotation rate.

```python
cfg = TransformerConfig(
    vocab_size=8192,
    d_model=512,
    num_heads=8,
    rope_theta=50_000.0  # Custom scaling for shorter, high-frequency contexts

)
model = SimpleAttentionNetwork(cfg)

```

## Practical Code Examples

### Running Inference with Default RoPE

```python
from needle.model.architecture import SimpleAttentionNetwork, TransformerConfig
from needle.model.run import load_checkpoint, generate
from needle.model.tokenizer import get_tokenizer

cfg = TransformerConfig(vocab_size=8192, d_model=512, num_heads=8)
model = SimpleAttentionNetwork(cfg)

params, _ = load_checkpoint("needle-base.ckpt")
tokenizer = get_tokenizer()

output = generate(model, params, tokenizer, prompt="Explain RoPE")
print(output)

```

### Manual RoPE Application

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

q = jnp.ones((2, 4, 16, 32))
cos, sin = precompute_rope_freqs(head_dim=32, seq_len=16, theta=100_000.0)
q_rotated = apply_rope(q, cos, sin)
print(q_rotated.shape)  # (2, 4, 16, 32)

```

### Customizing the Scaling Factor

```python
cfg = TransformerConfig(
    vocab_size=8192,
    d_model=512,
    num_heads=8,
    rope_theta=10_000.0  # Alternative scaling

)
model = SimpleAttentionNetwork(cfg)

```

## Summary

- **RoPE Positional Embeddings** in Needle 2 are deterministic geometric rotations, not learned parameters.
- `precompute_rope_freqs` at **architecture.py:97-101** generates cosine and sine tables based on `rope_theta`.
- `apply_rope` at **architecture.py:104-112** rotates Q/K tensors of shape **[B, H, T, D]** by mixing dimension halves.
- The `SimpleAttentionNetwork._rope()` method caches frequency tables for the entire forward pass.
- Both the training forward pass and the cached decoder ([`decode.py`](https://github.com/cactus-compute/needle/blob/main/decode.py)) apply consistent rotations before attention scoring.
- The default `rope_theta` is **100,000.0**, configurable via `TransformerConfig`.

## Frequently Asked Questions

### What is the default rope_theta value in Needle 2?

The default `rope_theta` value is **100,000.0**, defined in the `TransformerConfig` dataclass at [`needle/model/architecture.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/architecture.py) lines 69-71. This high value provides slower rotation frequencies, enabling better generalization on longer sequence lengths compared to traditional 10,000.0 defaults.

### Where does RoPE get applied in the Needle 2 architecture?

RoPE is applied inside the `MultiHeadAttention.__call__` method at **architecture.py:237-241**, where the `apply_rope` function rotates both query and key tensors immediately before computing attention scores. This occurs after the initial linear projections but before the dot-product attention calculation.

### How does Needle 2 handle RoPE during text generation?

During autoregressive generation, the cached decoder in [`needle/model/decode.py`](https://github.com/cactus-compute/needle/blob/main/needle/model/decode.py) (specifically within `_attn_cached` at lines 73, 261, and 387) applies `apply_rope` to new queries and keys before updating the KV cache. This ensures positional information remains consistent between cached historical states and newly generated tokens without recomputing past rotations.

### Can I use Needle 2 without RoPE positional embeddings?

No. The `SimpleAttentionNetwork` architecture is designed to always invoke `self._rope()` during the forward pass, generating the `(cos, sin)` tuple that gets passed to attention blocks. While you could theoretically pass `rope=None` to individual `MultiHeadAttention` layers, the standard model pipeline always pre-computes and supplies these tensors.