How to push transformer context windows past 128K tokens without retraining: RoPE interpolation, NTK-aware scaling, YaRN's non-uniform approach, ALiBi, sliding window with sink tokens, and ring attention for million-token sequences.
Standard scaled dot-product attention computes Q KT / √dk, a T×T matrix where T is sequence length. Memory and FLOPs both scale as O(T2 · dk · H) where H is number of heads. At T=128K with H=32 heads and dk=128:
| Context length T | Attention matrix size (bytes, BF16) | Single-layer FLOPs |
|---|---|---|
| 4K | 4,096 × 4,096 × 32 × 2 = 1.07 GB | ~8.6 GFLOP |
| 32K | 67 GB | ~550 GFLOP |
| 128K | 1.07 TB | ~8.8 TFLOP |
| 1M | 64 TB | ~536 TFLOP |
Flash Attention (Dao et al. 2022) reduces the memory to O(T) by tiling the computation, but the FLOPs remain O(T2). For 1M tokens, the FLOPs alone make full attention impractical — a different approach is required.
Problem 1 — Positional generalisation: Models trained on 4K-token sequences use positional encodings calibrated for 0–4095. At position 100000 the encodings are out-of-distribution → perplexity spikes. Problem 2 — Compute: Even with perfect positional encodings, O(T2) attention is unaffordable at 1M tokens. Slides 02–05 address Problem 1; slides 06–07 address Problem 2.
Rotary Position Embedding (Su et al. 2021, arXiv 2104.09864) is the positional encoding used by LLaMA, Mistral, Qwen, Gemma, and most modern open-weights models. It encodes absolute position as a rotation applied to Q and K, such that dot products Qm · Kn depend only on the relative position m−n.
# For dimension pair (2i, 2i+1), base θ = 10000:
θi = 10000-2i/d # i = 0..d/2-1
Rm,i = [ cos(m θi) -sin(m θi) ]
[ sin(m θi) cos(m θi) ]
# Applied to query/key pairs before dot-product attention:
q̃m = Rm qm, k̃n = Rn kn
# Relative encoding property:
q̃m · k̃n = (Rm qm)T (Rn kn) = qmT Rm-n knAt position m ≫ Ltrain, the rotation angles mθi land in regions never seen during training. For high-frequency dimensions (large i, small θi), the angle wraps many times — positional aliasing. For low-frequency dimensions (small i, large θi), the rotation is nearly trivial and generalises fine. This asymmetry is the key to YaRN’s non-uniform treatment.
RoPE adds no parameters, keeps the relative-encoding property without needing a separate relative-attention mechanism, is trivially compatible with KV caching, and has clean mathematical structure that makes extension tractable. ALiBi (slide 05) avoids RoPE entirely — it’s simpler but harder to extend post-hoc.
Position Interpolation (PI) (Chen et al. 2023, arXiv 2306.15595, from Meta) takes the simplest possible approach: if the model was trained up to position Ltrain and you want to extend to Ltarget, just compress position indices by the factor s = Ltarget / Ltrain:
# Original RoPE: angle at position m, dimension i:
φm,i = m / 100002i/d
# PI: divide m by the extension factor s = L_target / L_train
φm,iPI = (m / s) / 100002i/d
# Equivalently: multiply all θ_i by s (slow them down)
θiPI = 100002i/d / sPI keeps all angles within the training distribution at the cost of packing more positions into the same angular range — nearby tokens become harder to distinguish. After just 1000 fine-tuning steps on long documents, LLaMA extended cleanly from 4K to 32K tokens.
The Neural Tangent Kernel analogy observes that position interpolation degrades high-frequency dimensions most (they rotate fastest). Instead of scaling all dimensions uniformly, NTK-aware scaling scales the base θ of RoPE rather than the positions:
# NTK-Aware: change the base from 10000 to a larger value
θiNTK = (base · sd/(d-2))-2i/d
# s = context extension factor, d = head dimension
# For s=8 extension, LLaMA-2 d=128: base ≈ 10000 · 8^(128/126) ≈ 82,000A further refinement partitions RoPE dimensions into three regions: low-frequency dims (extrapolated without change), high-frequency dims (interpolated like PI), and mid-frequency dims (linearly interpolated). This is the basis of YaRN. Dynamic NTK sets s equal to the current sequence length / Ltrain at inference time, allowing the model to adapt without retraining.
YaRN (Yet another RoPE extensioN method, Peng et al. 2023, arXiv 2309.00071) is the technique used by Mistral, Qwen, and others for cheap post-hoc context extension. It combines three elements:
RoPE dimensions are split into three groups by their wavelength λi = 2π / θi:
• High-freq (λ < Ltrain): no interpolation (extrapolate)
• Low-freq (λ > Ltarget): standard PI
• Mid-freq: linear ramp between the two
Longer sequences increase the entropy of attention distributions. YaRN corrects this by multiplying attention logits by a temperature factor t = 0.1 · ln(s) + 1, where s is the extension factor. This prevents attention from sharpening artificially on nearby tokens.
| Step | Action | Cost |
|---|---|---|
| 1 | Set YaRN hyperparameters: s = Ltarget/Ltrain, α = 1, β = 32 (LLaMA default) | Zero |
| 2 | Fine-tune on 400–1000 steps with long-context data (e.g. SlimPajama, Books3) at the target length | ~0.1% of original training cost |
| 3 | Optionally fine-tune on instruction data at target length | Additional 100–200 steps |
Mistral-7B (base: 8K) extended to 128K using YaRN + 400-step fine-tune achieved perplexity within 0.5 nats of models trained from scratch at 128K. Qwen-1.5 and Qwen-2 use YaRN for their 128K and 1M context variants. The method has largely superseded PI for practical extension tasks.
ALiBi (Attention with Linear Biases, Press et al. 2022, arXiv 2108.12409) takes a completely different approach: rather than encoding position in Q/K, it adds a scalar bias to attention logits that grows linearly with the distance between tokens.
# Standard attention:
Aij = softmax( (qi kjT) / √d )
# ALiBi attention:
Aij = softmax( (qi kjT) / √d - mh · (i - j) )
# m_h: per-head slope, geometric series: 2^(-8/H), 2^(-16/H), ...
# (i - j) >= 0 (causal: only attend to past)• No positional parameters to learn or interpolate.
• Generalises smoothly past training length without any fine-tuning (biases just grow linearly).
• Heads with smaller mh attend over longer ranges — implicit multi-scale coverage.
• Strictly worse than RoPE on standard benchmarks at equal context lengths.
• Cannot be retrofitted onto RoPE-pretrained models.
• The linear decay forces all heads toward a recency bias; hard to suppress for tasks needing distant recall.
• Used by BLOOM, MPT — but not by any major model post-2023.
ALiBi’s key contribution was demonstrating that position can be expressed as a relative bias on attention scores without any positional vectors — a conceptually clean design. Its decline in favour of RoPE-with-YaRN reflects the ecosystem’s choice of retrofittability over theoretical elegance.
Sliding-window attention restricts each token to attending only to the W most recent tokens. This reduces the attention FLOPs from O(T2) to O(T · W) — linear in T for a fixed window W. Mistral-7B (2023) uses W=4096 on alternating layers to cap KV cache growth.
Xiao et al. (StreamingLLM, 2023) observed that transformer models trained with positional encodings assign anomalously high attention weight to the first few tokens regardless of their semantic content. These “attention sinks” serve as rest positions for attention heads that have nothing relevant to attend to. Without sinks in the KV cache, evicting early tokens causes attention distributions to go unstable. With 4 sink tokens retained, streaming generation can proceed indefinitely with constant memory.
Sliding-window + sinks solves the memory problem but not the recall problem: a token evicted from the sliding window is genuinely lost. For tasks requiring facts from 50K tokens ago, a different approach (ring attention, full KV cache, or Mamba hybrid) is needed.
Ring Attention (Liu et al. 2023, arXiv 2310.01889) achieves true full-context attention at 1M+ tokens by distributing the KV cache and attention computation across multiple devices in a ring topology. Each device holds a slice of the sequence and passes its KV block to the next device while simultaneously processing the current block.
Megatron-LM added Context Parallelism (CP) as a first-class parallelism dimension in 2024. CP splits the sequence dimension across GPUs, sending Q/K/V slices via all_gather and reduce_scatter. NVIDIA Tensor Parallel (TP) handles head parallelism; CP handles sequence parallelism. The two compose orthogonally: a model can be TP=4, CP=8, DP=16 simultaneously.
Gemini 1.5 Pro (1M context, 2024) and Claude 3 (200K context) both require sequence parallelism at inference. The ring attention idea underlies most production implementations under the name “context parallelism”. NVLink bandwidth (600 GB/s bidirectional per A100 pair) makes intra-node CP essentially free; cross-node CP via InfiniBand (400 Gb/s) dominates latency above 32 GPUs.
A pragmatic synthesis: use Mamba (or another SSM) for the majority of layers where O(T) cost matters, but retain a small number of full-attention layers at regular intervals to preserve in-context recall. This hybrid strategy is used by Jamba, Zamba, and several internal models.
| Model | Attn:Mamba ratio | Context window | KV cache cost |
|---|---|---|---|
| Pure transformer (LLaMA-3 8B) | 32:0 | 128K (with YaRN) | High — grows with T |
| Jamba (AI21, 52B total) | 1:7 | 256K | Very low — 4 attn layers only |
| Zamba (Zyphra, 7B) | 1:6 | 32K | Low — shared single attn layer |
| Mamba-pure (3B) | 0:48 | Theoretically ∞ | Zero (constant SSM state) |
Even a single attention layer per 7 Mamba layers has unlimited receptive field within its layer. This is sufficient to anchor the model to distant facts — think of the attention layer as a “lookup” escape hatch for anything not captured by Mamba’s state.
86% of layers have O(1) inference cost. The KV cache for 1M tokens consists only of the 4–6 attention layers’ KVs — a 7× reduction vs a pure transformer at the same total layer count.
Deck 04 leaves the autoregressive paradigm entirely and asks: what if instead of generating tokens left-to-right, we generate all tokens in parallel via a diffusion process? SEDD and LLaDA represent the current frontier of diffusion language models.