LogoCookLLM Docs
LogoCookLLM Docs
HomeCookLLM

Principles

Tokenization
Tokenization BasicsBPE AlgorithmGPT TokenizersBPE Training Engineering
Model Architecture
Transformer LM
From token ids to logitsEmbedding and LM Head
Attention Mechanisms
From Self-Attention to GQAAttention Sink
Position Encoding
Position Encoding BasicsRoPE Math DerivationRoPE ImplementationLength Extrapolation
GPU Programming Basics
GPU Architecture BasicsTensor LayoutTriton Basics: Vector Add
FlashAttention
Flash Attention PrinciplesFrom Naive Implementation to Auto-TuningBlock Pointers and Multi-Dim SupportCausal Masking OptimizationGrouped Query AttentionBackward Pass Implementation
Distributed Training
Data ParallelismZeRO OptimizerFully Sharded Data ParallelTensor ParallelismPipeline ParallelismMulti-Dimensional Hybrid Parallelism

Hands-on Training

Overview
Pretraining
Pretraining DataTokenizer TrainingModel ArchitectureData PipelineTraining LoopMonitoring and Validation
X (Twitter)
FundamentalsModel ArchitecturePosition Encoding

RoPE Implementation

Premium

Inverse frequency computation, cos/sin caching, and a vectorized apply_rotary_pos_emb

Get code access

Inverse Frequency Precomputation

The previous chapter reduced RoPE to "rotate each two-dimensional subspace." Turning that sentence into code means three concrete steps: work out how fast each subspace rotates, multiply that rate by the position to get an angle, then turn the angle into cos⁡\coscos / sin⁡\sinsin tensors you can multiply directly.

The first step is easy to wave through as a one-line formula, but it is what governs the model's extrapolation behaviour. Recall the formula:

Log in to continue reading

This is premium content. Please log in to access the full article.

θi=base−2i/d,i=0,1,…,d/2−1\theta_i = \text{base}^{-2i/d}, \quad i = 0, 1, \ldots, d/2 - 1θi​=base−2i/d,i=0,1,…,

Note that this is an exponentially decaying sequence: at i=0i = 0i=0, θ0=1\theta_0 = 1θ0​=1 and the pair spins fastest, covering a full radian per position; by i=d/2−1i = d/2 - 1i=d/2−1, is down around and needs tens of thousands of positions to complete one turn. — and that division of labour is exactly what every length-extrapolation method (NTK-aware, YaRN) reaches in and modifies.

In code we usually compute inverse frequency (which is θi\theta_iθi​ itself):

# Step 1: compute inverse frequency
inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
# inv_freq shape: (dim // 2,)
# inv_freq[0] = 1.0, inv_freq[1] ≈ 0.85, ..., inv_freq[-1] ≈ 0.00011

# Step 2: build position-frequency matrix

What does torch.outer do?

# Assume max_seq_len = 4, dim = 6 (3 pairs)
t = [0, 1, 2, 3]
inv_freq = [θ₀, θ₁, θ₂]

# outer product:
# freqs[m, i] = m * θᵢ
freqs = [[

With freqs, take cos/sin for each position mmm:

RoPE Math Derivation

From complex rotations to higher-dimensional generalization, understand the core math of rotary position embeddings

Length Extrapolation

NTK-aware Scaling, YaRN, and other methods to let RoPE handle longer sequences

Table of Contents

Inverse Frequency Precomputation
Two Implementation Styles
Interleaved Style (Original Paper)
Pairing and Rotation Matrix
Pairwise Loop
Complex-Multiply Vectorization
Split-Halves Style (HuggingFace Transformers)
Pairing and Rotation Matrix
rotate_half Vectorization
Vectorized apply_rotary_pos_emb
Equivalence of the Two Styles
A Full RoPE Module
Integrating RoPE Into Attention
Working With KV Cache
Summary
d
/2
−
1
θ\theta
θ
10−410^{-4}10−4
Low dimensions resolve neighbouring tokens, high dimensions encode long-range distance
t
=
torch.arange(max_seq_len).float()
freqs = torch.outer(t, inv_freq) # (max_seq_len, dim // 2)
# freqs[m, i] = m * theta_i
0
·θ₀,
0
·θ₁,
0
·θ₂],
# position 0: no rotation
[1·θ₀, 1·θ₁, 1·θ₂], # position 1
[2·θ₀, 2·θ₁, 2·θ₂], # position 2
[3·θ₀, 3·θ₁, 3·θ₂]] # position 3