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

Length Extrapolation

Premium

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

Get code access

The Rotation View: Understanding Extrapolation

Suppose the model is trained with max_seq_len = 4096. What happens if inference length is 8192?

A common intuition is “the position index is out of range, so it is OOD.” But that is imprecise. The position index mmm is unbounded, while RoPE position embeddings are bounded trigonometric functions. The model interacts with embeddings, not indices. To understand OOD, we must analyze the embeddings.

Log in to continue reading

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

Rotating on the Unit Circle

Recall the dot product after RoPE (complex form):

(Rmq)⊤(Rnk)=Re[∑i=0d/2−1q[2i:2i+1] k[2i:2i+1]∗ ei(m−n)θi](R_m q)^\top (R_n k) = \text{Re}\left[\sum_{i=0}^{d/2-1} q_{[2i:2i+1]} \, k^*_{[2i:2i+1]} \, e^{i(m-n)\theta_i}\right](Rm​q)⊤(R

The key is ei(m−n)θie^{i(m-n)\theta_i}ei(m−n)θi​. By Euler’s formula, it is a point on the unit circle. As relative distance m−nm-nm−n grows, this point rotates around the circle. Larger θi\theta_iθ rotates faster; smaller rotates slower.

This is the core of the “rotation view”: whether m−nm-nm−n is OOD is not important; what matters is whether the unit-circle points have been sufficiently covered during training.

High Frequency vs Low Frequency: Coverage Differences

Assume training length LtrainL_{\text{train}}Ltrain​, then m−n∈[0,Ltrain−1]m-n \in [0, L_{\text{train}} - 1]m−n∈[0,Ltrain​−. For each dimension , the number of rotations during training is:

ri=θi⋅Ltrain2πr_i = \frac{\theta_i \cdot L_{\text{train}}}{2\pi}ri​=2πθi​⋅L
  • High-frequency dims (θi\theta_iθi​ large, iii small): fast rotation, many turns during training, covering the circle. At test time, even larger m−nm-nm−n just keeps rotating on already covered points → no OOD
  • Low-frequency dims (θi\theta_iθ small, large): slow rotation, may not complete one full circle during training, covering only a small arc. When test-time exceeds that arc, we enter unseen territory →
Unit Circle Coverage
4096
8192
High freq (i=0)
i = 0, θ = 1.0000
Rotations
651.9
Period
6
Full coverage — safe
Low freq (i=63)
i = 63, θ = 0.0001
Rotations
0.1
Period
54410
Arc only — OOD risk
Trained arcExtrapolation (OOD)

Concrete numbers (d=128d=128d=128, base=10000\text{base}=10000base=10000, Ltrain=4096L_{\text{train}}=4096Ltrain​=4096):

# Highest-frequency dim (i=0): θ₀ = 1.0
# Rotations: 1.0 × 4096 / (2π) ≈ 651 turns → full coverage, safe

# Lowest-frequency dim (i=63): θ₆₃ ≈ 0.00011
# Rotations: 0.00011 × 4096 / (2π) ≈ 0.07 turns → tiny arc, high OOD risk

The core issue is not “the rotation angle is too large,” but insufficient unit-circle coverage in low-frequency dimensions. High-frequency dims are the safest.

From Rotation to Solutions

With this view, the solution becomes clear:

  • Dims with enough rotations (high frequency) → no change, extrapolate directly
  • Dims with insufficient rotations (low frequency) → compress out-of-range angles back into the trained arc (position interpolation)
  • Middle range → smoothly transition between the two

All methods below are variations of how to compress and how much.

RoPE Implementation

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

GPU Programming Basics

Learn CUDA and Triton, and write efficient GPU kernels

Table of Contents

The Rotation View: Understanding Extrapolation
Rotating on the Unit Circle
High Frequency vs Low Frequency: Coverage Differences
From Rotation to Solutions
Position Interpolation (PI)
NTK-aware Scaling
Core Idea
Derivation
Implementation
Advantages of NTK-aware
Dynamic NTK
YaRN
Motivation
YaRN’s Three Components
1. NTK-by-parts (piecewise interpolation)
2. Attention Scaling
YaRN Results
Summary Comparison
Summary
n
​
k
)
=
Re​i=0∑d/2−1​q[2i:2i+1]​k
i​
θi\theta_iθi​
1
]
iii
train
​
​
i
​
iii
m−nm-nm−n
true OOD
[2i:2i+1]
∗
​
ei(m−n)θi​
​