Learn: The Long-Context Race
What you're learning
“Long context” is really two problems wearing one name. (1) Positional generalization — a model trained at 4k produces garbage at 100k because RoPE meets rotation angles it never saw. (2) Compute & memory — attention is and the KV cache is , so even if the positions worked, you couldn’t afford to run it. Shipping real long context means solving both. By the end you’ll be able to derive Position Interpolation and YaRN’s per-frequency scaling from the RoPE wavelength argument, explain why Ring Attention makes context scale with device count, and reason about why a model’s advertised context is usually much longer than its usable one.
This lesson closes Tier 3. It assumes Tiers 1–2 — especially RoPE (03-positional-information), the quadratic bottleneck and FlashAttention (04-quadratic-bottleneck-and-io-aware), and the KV-cache / GQA / MLA story (05-mqa-gqa-mla) — plus lesson 08 (08-attention-in-vits-and-vlms). Vocabulary follows the attention-glossary.
1. Learning map
graph TD A["Want: train at 4k,<br/>serve at 128k+"] --> B["Problem 1:<br/>positional generalization"] A --> C["Problem 2:<br/>compute & memory"] B --> B1["RoPE angles go OOD:<br/>slow pairs never<br/>completed a rotation"] B1 --> D["Position Interpolation<br/>(PI): m → m/s"] B1 --> E["NTK-aware / YaRN:<br/>per-frequency scaling<br/>(wavelength argument)"] C --> C1["attention O(s²) FLOPs<br/>→ FlashAttention (L04)"] C --> C2["KV cache O(s) bytes<br/>→ GQA / MLA (L05)"] C1 --> F["Ring Attention:<br/>shard sequence,<br/>rotate KV blocks"] F --> G["context ∝ #devices<br/>(context parallelism, GPU L10)"] D --> H["Both solved →<br/>usable long context"] E --> H C2 --> H G --> H H --> I["But: measure it!<br/>NIAH / RULER,<br/>advertised ≠ usable"] style B fill:#a44,color:#fff style C fill:#a44,color:#fff style E fill:#44a,color:#fff style F fill:#483,color:#fff style I fill:#a83,color:#fff
2. Intuition: two problems, one name
Suppose you train a RoPE model at context length and feed it 100k tokens at inference. Two independent things go wrong.
Problem 1 — positional generalization. Recall from lesson 03 that RoPE spins each 2D coordinate pair of by an angle , with frequencies (). The slowest pairs turn so gradually that within 4k tokens they never complete even a fraction of a full rotation — the model has only ever seen those pairs at small angles. At position 100k, those same pairs are at angles larger, out of the distribution the network was trained on. The attention scores become nonsense; perplexity explodes. Nothing is “broken” in the FLOPs sense — the math runs fine — the model simply has never learned to read those rotations.1
Problem 2 — compute & memory. Even with positions fixed, the bill comes due. Attention is in FLOPs (lesson 04) and the KV cache is in bytes, read every decode step (lesson 05). At 128k the score computation is enormous and the KV cache for a single sequence can rival the model weights in HBM. You cannot serve what you cannot fit or feed.
Solving one is not solving the other
These are orthogonal. YaRN can make a model coherent at 128k while it remains unaffordable to serve without FlashAttention + GQA/MLA + sequence sharding. Conversely, Ring Attention lets you run attention over a million tokens, but if the RoPE angles are OOD the output is still garbage. A real long-context system stacks a positional-extension recipe on top of a systems recipe. This lesson does them in that order.
Quiz (retrieval): A 4k-trained RoPE model is run at 64k with an exact FlashAttention kernel that fits in memory. FLOPs and bytes are all fine. Why can the output still be incoherent, and which RoPE frequencies are to blame?
Answer
Because positional generalization is a separate failure from compute/memory. RoPE rotates each pair by ; the low-frequency (slow) pairs — large , tiny — have wavelengths comparable to or longer than the 4k training window, so within training they only ever swept a small arc of angles. At those pairs sit at angles beyond anything seen in training: out-of-distribution rotations the model never learned to interpret. The high-frequency pairs already wrapped many times within 4k, so they’re fine — it’s the slow, global-position-carrying pairs that break. An exact kernel changes none of this; it faithfully computes attention over garbage positional encodings.
3. Math — positional extension
The fix is to keep the RoPE angles inside the trained range at inference. Two philosophies: rescale position uniformly (PI), or rescale frequencies non-uniformly (NTK-aware / YaRN).
3.1 Position Interpolation (PI)
Chen et al. observe the simplest possible fix: if the model was trained on positions and you want to serve with , linearly compress the position index back into the trained range.2 Define the scale factor and replace every position by . Concretely, the rotation angle becomes
Now position maps to angle — exactly the largest angle seen in training. No angle is ever OOD. The paper shows this needs only light fine-tuning (~1000 steps) to adapt, versus naive extrapolation which degrades catastrophically and never recovers.2
PI's flaw: it crushes all frequencies, including the ones that were fine
Dividing every by slows down the high-frequency pairs too — but those were never the problem; they already resolved fine local offsets within the training window. Compressing them means adjacent tokens (offset 1) now differ by a rotation smaller, blurring the model’s ability to distinguish nearby positions. You buy long-range coherence by degrading short-range resolution. This is exactly what NTK-aware scaling and YaRN fix.
3.2 NTK-aware scaling and YaRN — the wavelength argument
The key object is the wavelength of pair : the number of tokens for one full rotation,
High-frequency pairs (small ) have short wavelengths : within the training window they complete many full rotations, so the network has seen every phase — safe to extrapolate, do not touch. Low-frequency pairs (large ) have : they never complete even one rotation in training, so any longer position is OOD — must be interpolated. PI’s mistake was treating both the same.
NTK-aware scaling implements this by changing the base instead of the position.3 Since , raising slows all frequencies but non-uniformly: the fastest pair () has regardless of (untouched), while the slowest pair is slowed the most (interpolated). Choosing
makes the lowest frequency get interpolated by roughly the factor while the highest is essentially preserved — spreading the “interpolation pressure” across dimensions by frequency. This works with little or no fine-tuning where PI needs a full pass.3
YaRN (“NTK-by-parts”) makes the per-frequency decision explicit rather than implicit in a base tweak.3 Define the number of full rotations a pair completes in the training window, , and pick thresholds (the paper uses ). Build a ramp
and set the extended frequency to a blend of “extrapolate” () and “interpolate” ():
- Pairs with many rotations (, high freq) → → : preserved (local resolution intact).
- Pairs with less than one rotation (, low freq) → → : fully interpolated (never OOD).
- In between → smooth ramp.
YaRN also multiplies the attention logits by a length-dependent temperature with , which empirically recovers the entropy of the softmax as context grows (attention over more tokens otherwise flattens).3 The net result: YaRN extends context further than PI with an order of magnitude less fine-tuning, and is the recipe behind many production long-context RoPE models.3
Why per-frequency scaling wins
PI applies one uniform ruler to a bank of frequencies that are in wildly different regimes. The wavelength argument says the right action depends on whether a pair has “seen its whole period.” High frequencies have — leave them to resolve local structure; low frequencies haven’t — interpolate them so they stay in range. NTK-aware/YaRN act on each pair according to its own , preserving the local detail PI would blur while still taming the global pairs that actually go OOD.
Quiz (the one to nail): Why do NTK-aware scaling and YaRN beat naive Position Interpolation, mechanistically? Reference the wavelength.
Answer
PI compresses every frequency by the same factor (), including the high-frequency pairs whose wavelength . Those pairs already completed many rotations within the training window, so they were never OOD — compressing them just blurs local positional resolution (offset-1 tokens now differ by a much smaller angle). The low-frequency pairs () are the only ones that go OOD at long range, because they never completed a full rotation in training. NTK-aware/YaRN use the per-pair rotation count to interpolate selectively: preserve high-frequency pairs (, keep ) so local detail survives, and interpolate only the low-frequency pairs (, use ) that actually need it. Same OOD problem solved, without sacrificing the short-range acuity PI throws away — which is why YaRN extends further with far less (or no) fine-tuning.
4. Systems for long context
Positions fixed, now make it runnable. Three levers, all previously introduced, now aimed at length.
FlashAttention (recap, lesson 04) — the memory wall. Naive attention materializes the score matrix in HBM: traffic and storage. At that matrix alone is tens of GiB — it cannot exist. FlashAttention’s online-softmax tiling never materializes it: same FLOPs, traffic, running state per query row.4 This is the precondition for everything below — without it there is no long-context prefill at all.
Ring Attention — near-unbounded context by sharding the sequence. FlashAttention fits attention on one device up to that device’s memory. To go further, shard the sequence across devices: device holds query block and its local .5 Arrange the devices in a ring. Each step, every device computes the (blockwise, online-softmax) attention of its local against the K/V block it currently holds, then passes that K/V block to its neighbor and receives the next one — after steps each query has attended to all keys. The trick: the P2P send/recv of the next K/V block overlaps with the compute on the current block. If per-block compute time per-block communication time (true once blocks are large enough), communication is fully hidden and the ring adds ~zero wall-clock overhead. Per-device activation memory is , so maximum context scales linearly with device count — “near-infinite” context is a matter of adding GPUs. It is exact attention, not an approximation.5
This is precisely context parallelism (CP) in the parallelism taxonomy — see 10-model-parallelism, where Ring Attention is the mechanism behind splitting the sequence dimension for attention (distinct from sequence parallelism, which splits the LayerNorm/residual/dropout regions). CP composes with DP × TP × PP: the sequence-sharding lives on a fast interconnect ring because its P2P traffic is on the critical path.
KV cache — the affordability floor (recap, lesson 05). The cache is bytes and grows linearly with , read in full every decode step.6 At long context it dominates HBM and decode bandwidth. GQA (few KV heads) and MLA (a low-rank latent cache) are what make long context economically servable — they shrink the per-token bytes so a 128k sequence fits and each decode step reads less. Positional extension makes long context coherent; KV-cache reduction makes it affordable.
flowchart LR subgraph Ring["Ring of N devices — each holds Q_j and rotates K/V blocks"] D0["Dev 0<br/>Q_0, local K/V"] D1["Dev 1<br/>Q_1, local K/V"] D2["Dev 2<br/>Q_2, local K/V"] D3["Dev 3<br/>Q_3, local K/V"] end D0 -->|"send K/V,<br/>overlap w/ compute"| D1 D1 -->|K/V| D2 D2 -->|K/V| D3 D3 -->|K/V| D0
Figure 1 — Ring Attention. Each device keeps its query block resident and passes K/V blocks around the ring; after hops every query has seen every key. Communication of the next block hides behind computation on the current one, so context grows with at near-zero overhead. Exact attention, sequence sharded — this is context parallelism.
Quiz (retrieval): Ring Attention runs attention over sequences far larger than one GPU's memory. How does it dodge the memory wall, and what condition makes its communication "free"?
Answer
It never holds the full sequence — or the full score matrix — on any one device. The sequence is sharded across devices ( activations each), and within each device attention is computed blockwise with online softmax (the FlashAttention recurrence), so the scores are never materialized anywhere — each device only ever forms the small tile on-chip. K/V blocks rotate around the ring; after steps every query has attended to every key, giving exact attention. Communication is “free” when the P2P transfer of the next K/V block fully overlaps the compute on the current block — i.e. per-block compute time per-block comm time, which holds once block size is large enough. Then adding devices raises the max context linearly with negligible wall-clock cost (context parallelism, GPU lesson 10).
5. Evaluation — advertised vs. usable context
A model can be trained/extended to a huge context and still not use it. Measuring this correctly is its own problem.
Why perplexity is insufficient. Long-context perplexity is dominated by local prediction — most next-token loss is decided by the last few hundred tokens. A model can post excellent perplexity at 128k while completely ignoring information 100k tokens back, because averaging loss over the sequence hides the rare tokens that actually require long-range retrieval. Low long-context PPL is necessary but nowhere near sufficient.7
Needle-in-a-haystack (NIAH). Plant a random fact (the “needle”) at a controlled depth inside a long filler “haystack,” then ask for it. Sweeping needle depth × context length yields the familiar heatmap of where retrieval holds and where it collapses (often a weak middle — “lost in the middle”). NIAH is a huge step up from perplexity but is still just single-fact verbatim retrieval — easy relative to real use.
RULER. RULER generalizes NIAH into a battery of synthetic tasks with controllable length: multi-key and multi-value needles, multi-hop variable tracking, aggregation/frequent-word extraction, and long-context QA — probing aggregation and reasoning, not just retrieval.7 Its headline finding: effective context length is usually far below the advertised one. Many models that claim 128k+ maintain performance only to a fraction of that before degrading — the gap between the number on the box and the length you can actually trust.7 [verify] the exact arXiv ID and the specific per-model effective-length numbers below before quoting them.
Advertised context ≠ usable context
A “1M-token” model is a claim about the maximum sequence the code will accept, not a guarantee it reasons over that span. Positional extension (YaRN etc.) can keep the model coherent far out, and Ring Attention can run it, but whether it can retrieve and combine information across the whole window is an empirical question that only RULER-style aggregation/reasoning tasks answer — and the honest answer is often a small fraction of the advertised length.
[verify]current frontier claims (Gemini/GPT/Claude/Llama long-context tiers move fast); treat vendor context numbers as upper bounds, not usable budgets.
Quiz (retrieval): A model reports strong perplexity at 128k and passes a single-needle NIAH probe. Why might it still fail real 128k tasks, and what would RULER catch that neither test does?
Answer
Perplexity is dominated by local next-token prediction, so it can look great while the model ignores distant context — the few tokens that truly need long-range info are drowned out in the average. Single-needle NIAH only tests verbatim retrieval of one planted fact; a model can do that via surface matching without any real long-range integration. RULER adds tasks that require aggregation and multi-step reasoning over the whole window — multi-key/multi-hop tracking, counting/aggregation, variable resolution — which expose models that can locate one string but cannot combine several facts spread across 128k tokens. That’s why RULER’s effective context is routinely far below the advertised length even for models that ace PPL and single-needle NIAH.
6. FLOPs & memory — feeling the cost of 128k
Restating the two scalings that define the whole race. Let =batch, =sequence length, =head dim, =layers, =KV heads, =bytes/element.
| Quantity | Scaling | At long |
|---|---|---|
| Attention-core FLOPs (prefill) | per layer | quadratic — dominates compute |
| KV cache | bytes | linear, but read every decode step |
| Score-matrix HBM (naive) | the wall FlashAttention removes |
The compute is why prefilling a long prompt is expensive; the KV cache is why decoding with a long prompt is memory-bandwidth-bound. FlashAttention removes the storage/traffic; GQA/MLA shrink the cache; Ring Attention distributes the activations. Positional extension is orthogonal and nearly free in FLOPs (it just retunes RoPE frequencies).
Try it: watch the costs explode at 128k
Open the Attention FLOPs & KV-cache calculator. Fix a 70B-class config (, query heads, , ) and sweep : 4k → 32k → 128k → 1M. Watch (1) prefill attention FLOPs grow , and (2) the KV cache grow — flip between MHA (64) and GQA () to see the cache shrink . The gap between “coherent at 128k” (a RoPE-scaling problem) and “affordable at 128k” (a KV-cache/FlashAttention/Ring problem) is exactly the two axes of this lesson.
Practice problem: KV cache + attention FLOPs at 128k
A Llama-3-70B-style model: layers, query heads, GQA with KV heads, (so ), bf16 (). Take , . (a) KV-cache size for one sequence. (b) Attention-core prefill FLOPs (score + matmuls, all layers). (c) What does each number tell you about the binding constraint?
Solution.
(a) bytes — for a single sequence, comparable to a big chunk of the model weights. With plain MHA () it would be 320 GiB, physically impossible on one GPU; GQA is what makes 128k servable. (This is why long context needs GQA/MLA — lesson 05.)
(b) Attention core FLOPs PFLOP — just for the and matmuls of one 128k prefill. For scale, the non-attention work is FLOPs — the quadratic attention term has overtaken the entire rest of the model at this length. That is the wall made concrete.
(c) Prefill is compute-bound and quadratic (the 45 PFLOP), so FlashAttention/Ring matter for feasibility and the term eventually dominates everything. Decode is memory-bound: every step re-reads the 40 GiB cache, so KV-cache bytes (GQA/MLA) set throughput. Long context stresses both axes at once.
7. Primary sources & going deeper
- Position Interpolation (primary): Chen, Wong, Chen, Tian — “Extending Context Window of Large Language Models via Positional Interpolation” — arXiv:2306.15595.2 Read §2: the rescaling and why extrapolation fails but interpolation + light fine-tuning succeeds.
- YaRN / NTK-aware (primary): Peng, Quesnelle, Fan, Shippole — “YaRN: Efficient Context Window Extension of Large Language Models” — arXiv:2309.00071.3 Work the wavelength/“NTK-by-parts” ramp and the attention-temperature term.
- Ring Attention (primary): Liu, Zaharia, Abbeel — “Ring Attention with Blockwise Transformers for Near-Infinite Context” — arXiv:2310.01889.5 The ring rotation + compute/comm overlap; note it builds directly on blockwise/online softmax (lesson 04).
- RULER (evaluation): Hsieh et al. — “RULER: What’s the Real Context Size of Your Long-Context Language Models?” — arXiv:2404.06654
[verify].7 - Systems cross-link: 10-model-parallelism (context/sequence parallelism); KV-cache serving in 12-llm-inference-systems.
Ask the teacher
If any step didn’t land — the NTK base-scaling derivation , YaRN’s ramp thresholds, the compute/comm overlap condition in Ring Attention, or how PI/YaRN interact with GQA/MLA caches — ask me for a follow-up. Good targets: “derive the NTK base formula from the wavelength requirement,” “walk through a RULER task family,” or “how does YaRN change if the base model uses ALiBi instead of RoPE?“
8. What’s next
- 10-sota-synthesis (next, closes the curriculum) — pull every thread together: how frontier models combine RoPE+YaRN, GQA/MLA, FlashAttention, and sparse/sliding-window/linear variants into a coherent long-context stack.
- 08-attention-in-vits-and-vlms (back) — long context isn’t only text: high-resolution images and video produce enormous token counts, so the same /KV pressures and positional-extension tricks reappear in the vision/VLM setting.
Further threads: sparse & sliding-window attention (06-sparse-and-sliding-window) and linear/subquadratic models (07-linear-and-subquadratic) as the other bet — change the algorithm to kill outright, trading exactness for asymptotics.
Reference topic: attention-mechanisms | Glossary: attention-glossary | Concepts: rope-scaling-yarn, ring-attention, long-context | Filed: 2026-09-10
References
Footnotes
-
[established] Jianlin Su, Yu Lu, Shengfeng Pan, Ahmed Murtadha, Bo Wen, Yunfeng Liu, “RoFormer: Enhanced Transformer with Rotary Position Embedding” — arXiv:2104.09864 (2021; Neurocomputing 2024). arxiv.org/abs/2104.09864. Source of RoPE, the per-pair frequencies , and the relative-position identity that the extension methods retune. See lesson 03 for the full derivation. ↩
-
[established] Shouyuan Chen, Sherman Wong, Liangjian Chen, Yuandong Tian, “Extending Context Window of Large Language Models via Positional Interpolation” — arXiv:2306.15595 (2023). arxiv.org/abs/2306.15595. Source of linear Position Interpolation (, scale ) and the finding that interpolation + ~1000 fine-tuning steps far outperforms naive extrapolation. ↩ ↩2 ↩3
-
[recent] Bowen Peng, Jeffrey Quesnelle, Honglu Fan, Enrico Shippole, “YaRN: Efficient Context Window Extension of Large Language Models” — arXiv:2309.00071 (2023). arxiv.org/abs/2309.00071. Source of the wavelength argument, NTK-aware base scaling (), the “NTK-by-parts” per-frequency ramp with and , and the attention-temperature term . Reports stronger extension than PI with far less fine-tuning. ↩ ↩2 ↩3 ↩4 ↩5 ↩6
-
[established] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, Christopher Ré, “FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness” — arXiv:2205.14135 (NeurIPS 2022). arxiv.org/abs/2205.14135. The IO-aware online-softmax tiling that removes the score-matrix storage/traffic. Full treatment in lesson 04. ↩
-
[recent] Hao Liu, Matei Zaharia, Pieter Abbeel, “Ring Attention with Blockwise Transformers for Near-Infinite Context” — arXiv:2310.01889 (2023). arxiv.org/abs/2310.01889. Source of sequence sharding across a device ring with K/V-block rotation and compute/communication overlap, yielding exact attention with per-device memory and context that scales with device count. This is context parallelism (GPU lesson 10). ↩ ↩2 ↩3
-
[established] KV-cache formula and the GQA/MLA reductions are developed in 05-mqa-gqa-mla; GQA from Ainslie et al. (arXiv:2305.13245), MLA from DeepSeek-V2 (arXiv:2405.04434). ↩
-
[recent]
↩ ↩2 ↩3 ↩4[verify]Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, et al. (NVIDIA), “RULER: What’s the Real Context Size of Your Long-Context Language Models?” — arXiv:2404.06654 (2024). arxiv.org/abs/2404.06654. Extends needle-in-a-haystack to multi-key/multi-hop/aggregation/QA tasks; reports that effective context length is typically well below the advertised length.[verify]the exact arXiv ID, author list, and specific per-model effective-length figures before quoting; frontier context claims move fast.