Module 10: Long Context

Introduction

Your GPT was trained on sequences up to some length L — say 2048 tokens. Ask it to read 8192 and something strange happens: the output degrades not gradually but sharply, often into gibberish, the moment you pass the length it saw in training. The model did not “run out of memory.” It ran out of positions it recognizes.

Context extension is the set of tricks that let a model trained to length L work at length L' > L without retraining from scratch — usually with only a short fine-tune, sometimes with none at all. Every method in this module is a single idea applied to the Rotary Position Embeddings you built in m04: don’t show the model new position-angles it has never seen; rescale the long sequence so it reuses the angles it already knows.

Why it matters for LLMs:

  • The context-length race — 2k → 4k → 32k → 128k → 1M tokens — was won mostly by these cheap post-hoc tricks, not by training at full length from the start.
  • They cost almost nothing: Position Interpolation and YaRN extend a model 8–16× with a few billion fine-tuning tokens, versus the trillions used to pretrain it.
  • They build directly on RoPE (m04) and pair with the KV-cache and efficient attention (m08, m09) that make a long context affordable to serve.

What You’ll Learn

After this module, you can:

  • Explain why RoPE breaks past the training length — the low-frequency pairs rotate to angles the model never saw.
  • Implement Position Interpolation (PI) from scratch: squeeze positions by the scale factor s = L'/L.
  • Implement NTK-aware scaling: raise the frequency base so fast pairs are preserved and slow pairs are stretched.
  • Implement YaRN: interpolate each frequency pair by how many rotations it completes, plus an attention-softmax temperature.
  • Drive a ScaledRotaryEmbedding that swaps between all four behaviors with one argument.
  • Explain the attention sink — why softmax forces a query to dump spare mass on the initial tokens — and the exact distortion evicting it causes.
  • Build StreamingLLM from scratch: a fixed KV cache that retains a few sink tokens plus a rolling window, with cache-relative RoPE positions.
  • Turn the sink into a learned per-head parameter (gpt-oss) — a trained logit in the softmax denominator — and prove it is an exact, content-aware output gate: \text{out}_{\text{sink}} = (1-a_{\text{sink}})\,\text{out}_{\text{plain}}.
  • See why that learned sink is what lets a model alternate sliding-window and full-attention layers with no positional sink to fall back on.

Prerequisites

This module requires familiarity with:

  • Module 04: Embeddings — Rotary Position Embeddings, the theta_i = base^{-2i/d} frequency schedule, and the rotate-don’t-add idea we are about to rescale.
  • Module 09: Efficient Attention — the KV-cache and GQA that make serving the extended context practical.

Intuition: The Context Wall

Recall RoPE from m04. Each 2-D feature pair i of a query or key is rotated by an angle m\,\theta_i, where m is the token’s position and \theta_i = \text{base}^{-2i/d} is that pair’s fixed rotation speed. Fast pairs (small i) spin almost a full turn every token; slow pairs (large i) creep around over thousands of tokens.

During training on lengths up to L, the model only ever sees position angles in the band [0,\ (L-1)\,\theta_i]. For a fast pair that band already wraps around the circle many times — position L{+}100 looks like some earlier angle it has seen, so extrapolation is harmless. For the slowest pair, though, (L-1)\,\theta_i is a small fraction of a single turn. Push past L and that pair swings into fresh, never-trained angular territory — and that is the pair whose signal the model leans on for long-range position. That is the wall.

There are two ways over it. Either squeeze the long sequence back into the trained band (interpolate), or stretch the slow pairs’ clock so the same band now spans more tokens. Drive the slowest pair and watch plain RoPE walk right off the edge while a scaling method folds it back inside:

NoteKey Insight

Only the slowest pairs hit the wall, and they hit it hard: at position L' they sit at an angle the model was never trained on. The fast pairs are fine — they wrapped the circle many times during training, so a bit more rotation is nothing new. Every method below is a way of leaving the fast pairs alone while pulling the slow pairs back into the trained band.

The Math: Three Ways to Rescale RoPE

All three methods start from RoPE’s per-pair frequencies \theta_i = \text{base}^{-2i/d} and produce a new set. The scale factor is s = L'/L.

1. Position Interpolation (PI). The bluntest fix: divide every position by s. Position L'-1 then lands on the angle the model saw at (L'-1)/s \approx L-1. Because m\,\theta_i = (m/s)\,\theta_i when you fold the 1/s into the frequency, this is identical to dividing every frequency by s:

\theta_i^{\text{PI}} = \frac{\theta_i}{s}.

Simple and effective, but it squeezes the fast pairs too — and those didn’t need help, so you lose some local, high-frequency resolution.

2. NTK-aware scaling. Instead of moving the positions, raise the frequency base so the stretch is spread unevenly across pairs — almost nothing for the fast pairs, a lot for the slow ones:

\text{base}' = \text{base}\cdot s^{\,d/(d-2)}, \qquad \theta_i^{\text{NTK}} = (\text{base}')^{-2i/d}.

The exponent is chosen so the fastest pair (i=0) is left exactly alone while the slowest pair is interpolated by roughly s. No fine-tuning is needed for modest extensions, which is why it spread through the open-source community before it had a paper.

3. YaRN (“NTK-by-parts”). Make the per-pair choice explicit. For pair i, count how many full rotations it completes within the original context:

\lambda_i = \frac{2\pi}{\theta_i}, \qquad r_i = \frac{L}{\lambda_i}.

A pair that spins many times (r_i > \beta) is doing local work — leave it alone (extrapolate). A pair that has not even completed one turn (r_i < \alpha) is doing long-range work — interpolate it fully. Ramp linearly between:

\gamma_i = \operatorname{clamp}\!\left(\frac{r_i - \alpha}{\beta - \alpha},\,0,\,1\right), \qquad \theta_i^{\text{YaRN}} = (1-\gamma_i)\,\frac{\theta_i}{s} + \gamma_i\,\theta_i.

With \alpha=1,\ \beta=32 (the LLaMA values), YaRN keeps the fast pairs, fully interpolates the slow pairs, and blends the handful in between — plus a small attention temperature we cover below. It is the method behind most 128k-token open models.

Here is the whole story in one picture: the ratio of each pair’s rescaled frequency to its original. PI drops every pair by the same factor 1/s; NTK and YaRN leave the fast pairs near 1 and pull only the slow pairs down.

TipTry This
  1. Watch the shapes diverge. PI is a flat line at 1/s — every pair squeezed equally. NTK curves smoothly from 1 down. YaRN sits on top of NTK for the fast pairs (ratio ≈ 1), then drops to the PI line (1/s) for the slow pairs — the best of both.
  2. Crank s to 16. The dashed 1/s floor sinks toward zero; PI drags every pair down with it (blurring local detail), while NTK/YaRN keep the fast pairs pinned near 1.

Code: Scaling from Scratch

long_context.py builds all three schemes as transforms of RoPE’s inverse frequencies. The base schedule is exactly m04’s:

import torch
from long_context import rope_inv_freq

theta = rope_inv_freq(head_dim=64, base=10000.0)
print(f"θ has one entry per pair: {theta.shape}")
print(f"fastest pair θ₀   = {theta[0]:.4f}   (≈ 1 turn / token)")
print(f"slowest pair θ₃₁  = {theta[-1]:.2e}  (≈ 1 turn / {2*3.14159/theta[-1]:.0f} tokens)")
θ has one entry per pair: torch.Size([32])
fastest pair θ₀   = 1.0000   (≈ 1 turn / token)
slowest pair θ₃₁  = 1.33e-04  (≈ 1 turn / 47117 tokens)

Position Interpolation divides every frequency by the scale; NTK raises the base; YaRN blends per pair. Each reduces to plain RoPE when scale=1:

from long_context import (
    linear_interpolation_inv_freq,
    ntk_inv_freq,
    yarn_inv_freq,
)

s = 8.0
pi = linear_interpolation_inv_freq(64, 10000.0, s)
ntk = ntk_inv_freq(64, 10000.0, s)
yarn = yarn_inv_freq(64, 10000.0, s, original_max_position=2048)

print(f"{'pair':>4} {'plain':>10} {'PI':>10} {'NTK':>10} {'YaRN':>10}")
for i in [0, 8, 16, 24, 31]:
    print(f"{i:>4} {theta[i]:>10.2e} {pi[i]:>10.2e} {ntk[i]:>10.2e} {yarn[i]:>10.2e}")
pair      plain         PI        NTK       YaRN
   0   1.00e+00   1.25e-01   1.00e+00   1.00e+00
   8   1.00e-01   1.25e-02   5.85e-02   1.00e-01
  16   1.00e-02   1.25e-03   3.42e-03   1.89e-03
  24   1.00e-03   1.25e-04   2.00e-04   1.25e-04
  31   1.33e-04   1.67e-05   1.67e-05   1.67e-05

Read the top row: pair 0 is untouched by NTK and YaRN (still 1.00) but divided by 8 under PI. Read the bottom row: the slowest pair is interpolated by all three. That is the whole design — preserve fast, interpolate slow — expressed as numbers.

ScaledRotaryEmbedding wraps this into a drop-in replacement for m04’s rotary layer. Pick a method and a scale; it precomputes the cos/sin tables out to the extended length and rotates a query or key tensor by position:

from long_context import ScaledRotaryEmbedding

rope = ScaledRotaryEmbedding(
    head_dim=64, scale=4.0, method="yarn", original_max_position=2048
)

q = torch.randn(1, 8, 8192, 64)   # (batch, heads, seq=8192 > 2048, head_dim)
q_rot = rope(q)
print(f"Rotated a {q.shape[-2]}-token sequence: {tuple(q_rot.shape)}")
print(f"YaRN attention-temperature factor: {rope.attention_scale:.4f}")
Rotated a 8192-token sequence: (1, 8, 8192, 64)
YaRN attention-temperature factor: 1.1386

RoPE’s signature property — attention depends only on the relative offset — survives the rescaling within the extended range. demonstrate_context_extension makes the wall quantitative, tracking the slowest pair’s angle at the target length relative to the trained band:

from long_context import demonstrate_context_extension

ratio = demonstrate_context_extension(
    head_dim=64, train_len=2048, target_len=8192
)
print(f"\nPlain RoPE asks the model to read angles {ratio:.1f}× outside its band.")
==============================================================
CONTEXT EXTENSION - slowest RoPE pair
==============================================================
  head_dim=64, base=10000, train_len=2048, target_len=8192  (scale s=4)

  Trained max angle for the slowest pair: 0.273 rad
  Angle at target_len / trained max  (1.0 = stays in trained band):
    none  (plain RoPE, extrapolates)     4.00x
    linear (Position Interpolation)      1.00x
    ntk   (NTK-aware base)               1.00x
    yarn  (NTK-by-parts)                 1.00x

Plain RoPE asks the model to read angles 4.0× outside its band.

YaRN: Interpolate Per Dimension

The heart of YaRN is that ramp \gamma_i — the decision, pair by pair, of keep versus interpolate. It splits the pairs into three zones by how many rotations they complete within the training length. Step through them:

Beyond the frequencies, YaRN adds one more correction: it multiplies the attention logits by a temperature \sqrt{1/t} = 0.1\ln(s) + 1 before the softmax. Stretching positions slightly flattens the attention distribution; this factor re-sharpens it and recovers a bit of the perplexity lost to interpolation. ScaledRotaryEmbedding folds it into the cos/sin tables, so a YaRN layer is still drop-in:

from long_context import yarn_attention_scale

for scale in (1, 2, 4, 8, 16):
    print(f"s = {scale:>2} →  attention scale √(1/t) = {yarn_attention_scale(scale):.4f}")
s =  1 →  attention scale √(1/t) = 1.0000
s =  2 →  attention scale √(1/t) = 1.0693
s =  4 →  attention scale √(1/t) = 1.1386
s =  8 →  attention scale √(1/t) = 1.2079
s = 16 →  attention scale √(1/t) = 1.2773
NoteKey Insight

NTK-aware scaling and YaRN are the same instinct — don’t touch the fast pairs — at two levels of precision. NTK bakes the taper into a single raised base; YaRN makes the keep/interpolate decision explicit per pair (with a ramp) and adds the softmax temperature. When you see a model card say “128k context via YaRN,” this ramp is what it means.

Measuring It: Needle in a Haystack

Extending the positions is only half the job — you have to check the model can still use the far context. The standard probe is needle-in-a-haystack: hide a single fact (the “needle”) at a controlled depth inside a long filler document (the “haystack”), then ask a question only that sentence answers. Sweep the needle’s depth and the total length, and you get a 2-D grid of pass/fail.

  • Plain RoPE past L: the grid turns red as soon as the needle sits beyond the training length — the model literally cannot address that position.
  • PI / NTK / YaRN: the grid stays green much further out, with YaRN typically holding longest before the far-depth cells start to fade.

Perplexity tells the same story more cheaply: measured on held-out long documents, it stays flat under a good scaling method and explodes the moment plain RoPE crosses L. The interactive at the top of this lesson is exactly that crossing, one frequency pair at a time.

WarningExtension is not free capability

Rescaling positions lets the model attend to far tokens; it does not teach it to reason over them. A short fine-tune at the target length (a few billion tokens) is what turns “can address position 100k” into “can actually use position 100k.” YaRN needs the least of this fine-tuning, which is much of why it won.

Streaming with Attention Sinks

PI, NTK, and YaRN all answer the context wall the same way: make the window bigger by rescaling positions. There is a completely different answer — keep the window small and stream through it. Instead of asking the model to address position 100k, you feed an endless stream through a fixed-size KV cache, dropping old keys and values as new ones arrive. Memory stays flat no matter how long the input; the model never sees a position it wasn’t trained on.

The naive version of this — window attention, evict the oldest KV once the cache is full — is a disaster. Perplexity doesn’t drift up gently as the model forgets old context; it explodes, suddenly, the moment the cache overflows. Xiao et al. (2023) found the cause, and the fix is one line.

Intuition: why naive windowing collapses

Watch a fixed cache of capacity sink + window as tokens stream in. The first few tokens fill a sink region and never leave; every later token pushes the oldest recent token out of a rolling window. The trick is what happens to the positions the surviving tokens report to RoPE.

Two things to notice. First, once the stream is long enough the cache stops growing — it holds exactly sink + window tokens forever. Second, the cache carries a gap: it keeps tokens 0, 1 and then jumps to the recent ones. But the positions handed to RoPE are contiguous — 0, 1, 2, 3, … — because StreamingLLM assigns positions within the cache, not in the original text. This is exactly the RoPE-position lever from the rest of this module: never show the model an angle it hasn’t seen.

NoteKey Insight

StreamingLLM’s positions are cache-relative. If the cache holds original tokens [0,1,2,3,6,7,8] while decoding token 9, RoPE sees positions [0,1,2,3,4,5,6,7] — the model stays inside its trained range no matter how far the stream has run. That is why this method belongs beside PI/NTK/YaRN: it is another way of keeping RoPE honest about position.

The Math: softmax must sum to one

Why do those first tokens matter so much that keeping them fixes everything? The answer is in the softmax. Attention weights are

\text{softmax}(x)_i = \frac{e^{x_i}}{\sum_{j} e^{x_j}},

and this sum is forced to 1. A query that has already gathered what it needs from nearby tokens still has to spend its full unit of attention somewhere. It cannot output “I attend to nothing.” So the model learns a dumping ground for that leftover mass — and the natural choice is the initial tokens, because under causal masking they are visible to every later position and so are the one place every query can always reach. Trained on enough data, those tokens accumulate a large, content-independent share of attention. This is the attention sink: tokens that soak up spare softmax mass, often while contributing almost nothing to the output (their value vectors are near zero).

Now the collapse is obvious. Evict the sink and the softmax has to re-normalize over whatever survives. The mass that would have gone to the sink gets forced onto the recent tokens instead — corrupting the very outputs the model depends on. The distortion is exact. Take a query whose sink holds attention mass a_{\text{sink}} and a zero value (a purely-normalizing sink). Its output is just the content sum \sum_{j\in\text{content}} a_j v_j. Drop the sink and every surviving weight scales by 1/(1-a_{\text{sink}}), so

\text{output}_{\text{no sink}} = \frac{\text{output}_{\text{full}}}{1 - a_{\text{sink}}}.

If the sink held 90% of the mass, dropping it multiplies the output by 10. Drive the sink’s share and watch the surviving weights — and the distortion — blow up:

NoteSoftMax off by one

There is a cleaner fix than working around the sink: give the softmax a legal way to sum to less than 1. SoftMax₁ adds a fixed +1 to the denominator, \text{softmax}_1(x)_i = e^{x_i}/(1 + \sum_j e^{x_j}) — mathematically, a permanent virtual key with logit 0 and value 0. A model trained with it can attend to nothing, so it never needs to grow a real sink. That’s softmax_one in streaming.py, and it equals softmax([x, 0])[:-1] exactly.

Code: the streaming KV cache

Everything above is streaming.py. The mask picks out sinks + a causal window; the cache retains the sinks and rolls the window; positions come back cache-relative. First the mask:

from streaming import streaming_attention_mask

# 1 sink + a window of 2, over a length-6 stream. Row i = what query i attends to.
mask = streaming_attention_mask(seq_len=6, sink_size=1, window_size=2)
for i, row in enumerate(mask.int().tolist()):
    print(f"query {i}: {row}")
query 0: [1, 0, 0, 0, 0, 0]
query 1: [1, 1, 0, 0, 0, 0]
query 2: [1, 1, 1, 0, 0, 0]
query 3: [1, 0, 1, 1, 0, 0]
query 4: [1, 0, 0, 1, 1, 0]
query 5: [1, 0, 0, 0, 1, 1]

Query 5 keeps token 0 (the sink) and tokens 4–5 (the window); tokens 1–3 are gone. Now the cache, streamed past its capacity:

import torch
from streaming import StreamingKVCache

cache = StreamingKVCache(sink_size=4, window_size=3)
for t in range(9):                       # stream tokens 0..8
    cache.append(torch.tensor([float(t)]), torch.tensor([float(t)]))

print("kept text positions:", cache.text_positions())   # the gap is real
print("RoPE positions:     ", cache.positions())         # but RoPE sees 0..len-1
print("cache length:", len(cache), "/", cache.sink_size + cache.window_size)
kept text positions: [0, 1, 2, 3, 6, 7, 8]
RoPE positions:      [0, 1, 2, 3, 4, 5, 6]
cache length: 7 / 7

The one guarantee that makes this trustworthy: streaming a token at a time through the cache gives exactly the same output as computing the streaming mask in one shot. Same anchor as the KV cache in m09 — the incremental path must match the full recompute bit-for-bit.

from streaming import streaming_attention, streaming_attention_incremental

torch.manual_seed(0)
seq_len, d = 30, 8
q, k, v = (torch.randn(seq_len, d) for _ in range(3))
logits = (q @ k.transpose(-1, -2)) / d**0.5

one_shot = streaming_attention(logits, v, sink_size=2, window_size=4)
incremental = streaming_attention_incremental(q, k, v, sink_size=2, window_size=4)
print("incremental == one-shot:", torch.allclose(one_shot, incremental, atol=1e-5))
incremental == one-shot: True

The collapse, measured

Here is the whole point in one plot. Take an attention pattern whose mass sits on the sink plus a couple of recent tokens (the regime StreamingLLM targets), and measure how far each fixed-cache scheme drifts from full attention as the stream grows. streaming_reconstruction_trace in streaming.py computes it.

The two curves lie on top of each other — both perfect — until the stream passes the window length. At that instant the naive window drops the sink and its error jumps and stays high; the streaming cache keeps the sink and its error never leaves zero. Same cache size, same recent window; the only difference is a handful of retained sink tokens. On real models this is the difference between stable perplexity out to millions of tokens and gibberish the moment the cache fills — a 22× speedup over recomputing a growing window, at a fixed memory budget.

TipTry This
  1. Slide the sink logit to 6+ in the normalization widget, then toggle Evict the sink: the content weights and the distortion factor rocket up. The more the model leaned on the sink, the worse dropping it hurts.
  2. Step the cache animator to the end: confirm the cache stops growing and the bottom row (RoPE positions) stays a dense 0,1,2,… even as the top row (text positions) opens a gap.
  3. Shrink the window in streaming_reconstruction_trace(window_size=…): the collapse point in the plot slides left, but streaming stays flat regardless.

The Learned Attention Sink

StreamingLLM keeps the sink by keeping tokens — retain the first few keys and the model’s dumping ground stays reachable. But that spends real cache slots on tokens whose only job is to absorb spare mass, and it only works because the first tokens are visible to everyone. A sliding-window head has no such luxury: its band scrolls forward, so the original sink tokens fall out and there is nothing fixed to dump onto. The softmax_one callout above hinted at the cleaner fix — give the softmax a legal way to sum to less than one — but a permanent +1 is a fixed escape valve: every head, every position, gets the exact same one unit of slack.

The 2025 frontier’s answer, shipped in OpenAI’s gpt-oss, is to make that slack a trained parameter. Give each attention head one learnable scalar s — gpt-oss calls it sinks — and append it as an extra logit in the softmax denominator, with no value attached:

\text{softmax}_{\text{sink}}(x)_i = \frac{e^{x_i}}{\sum_j e^{x_j} + e^{s}}.

The model card describes it exactly this way: “a learned bias in the denominator of the softmax, similar to off-by-one attention and attention sinks.” It is the learnable generalization of Miller’s fixed off-by-one, and it contains both of the schemes you already know as limits:

s \to -\infty \;\Rightarrow\; \text{softmax}_{\text{sink}} = \text{softmax} \qquad\qquad s = 0 \;\Rightarrow\; \text{softmax}_{\text{sink}} = \text{softmax}_1 .

The Math: a content-aware output gate

Here is why this one scalar is worth a section. Because the virtual sink key carries no value, it never appears in the output — it only steals normalization. Let Z = \sum_j e^{x_j} be the usual softmax denominator. The sink absorbs a fraction

a_{\text{sink}} = \frac{e^{s}}{Z + e^{s}} = \sigma\!\big(s - \log Z\big),

and every real weight is the plain softmax weight scaled by Z/(Z+e^{s}) = 1-a_{\text{sink}}. So the whole attention output is the ordinary attention output, gated:

\boxed{\;\text{out}_{\text{sink}} = (1 - a_{\text{sink}})\;\text{out}_{\text{plain}}\;}

Read the gate a_{\text{sink}} = \sigma(s - \log Z) carefully — it depends on the content through Z. When the query finds a strong match, Z is large, a_{\text{sink}} is tiny, and the output passes through untouched. When the query finds nothing, Z collapses toward its floor, a_{\text{sink}} \to 1, and the output is gated toward zero. The head has learned, per token, how much to say — including nothing at all. That is the capability a plain softmax cannot express: it is forced to spend a full unit of attention no matter what, and the escape hatch it invents instead is a massive activation on a sink token (Sun et al., 2024). The learned sink hands the model a clean lever for the same job.

Drive the sink logit

Turn the dial. Left of the marked s = 0 you are in Miller’s off-by-one regime; push far left and the sink vanishes into plain softmax; push right and the head attends to nothing.

NoteKey Insight

A learned sink is a gate, not a memory. It has no value vector, so it can never add information to the output — it can only scale the whole thing down by 1-a_{\text{sink}}. The sink lets a head abstain; it cannot make a head say something new. That is exactly why it fixes the sink pathology without costing the model any expressive power over the real tokens.

Code: the learned sink from scratch

Everything above is sinks.py, built on m10’s streaming.softmax and reusing its softmax_one as the s=0 reference. First the two limits — the new softmax contains both schemes you already know:

import torch
from sinks import softmax_with_sink
from streaming import softmax, softmax_one

x = torch.tensor([1.4, 0.6, 0.1, -0.3])
print("s → −∞  == softmax:    ", torch.allclose(softmax_with_sink(x, -1e30), softmax(x)))
print("s = 0   == softmax₁:   ", torch.allclose(softmax_with_sink(x, 0.0), softmax_one(x)))
s → −∞  == softmax:     True
s = 0   == softmax₁:    True

Now the identity that makes the sink a gate. The output with a sink is exactly the plain attention output times 1-a_{\text{sink}} — no approximation:

from sinks import demonstrate_output_gate

d = demonstrate_output_gate(sink_logit=2.0)
print(f"a_sink = {d['a_sink']:.3f}   gate (1 - a_sink) = {d['gate']:.3f}")
print(f"out_plain = {d['out_plain'].flatten().tolist()}")
print(f"out_sink  = {d['out_sink'].flatten().tolist()}")
print(f"out_sink == (1 - a_sink) · out_plain : {d['gate_identity_holds']}")
a_sink = 0.482   gate (1 - a_sink) = 0.518
out_plain = [2.366149112090229]
out_sink  = [1.2251278906382905]
out_sink == (1 - a_sink) · out_plain : True

The gate is content-aware: hold s fixed and a confident query keeps a small a_{\text{sink}}, while a query that matches nothing gives its mass away.

from sinks import sink_mass

s = 1.0
strong = torch.tensor([5.0, 0.0, 0.0, 0.0])   # one clear match
quiet  = torch.zeros(4)                        # nothing stands out
print(f"strong match → a_sink = {float(sink_mass(strong, s)):.3f}")
print(f"nothing matches → a_sink = {float(sink_mass(quiet, s)):.3f}")
strong match → a_sink = 0.018
nothing matches → a_sink = 0.405

Sliding Windows Need a Sink

Now the payoff, and why gpt-oss ships this. gpt-oss does not enlarge attention; it shrinks most of it. Its layers alternate: a full-context (“dense”) layer, then a 128-token sliding-window (“banded”) layer, and so on — half the layers only ever see the last 128 tokens. That is cheap, but it reintroduces the StreamingLLM problem inside every window layer: the band scrolls, the first tokens are gone, and there is no positional sink to reach. A plain sliding-window head with nothing to attend to is forced, by the sum-to-one softmax, to smear its full unit of mass over whatever junk is in the window. The learned per-head sink is the fix that travels with the head instead of living at a fixed position.

Step a query through a stream that hits a stretch where nothing in its window matches, and compare a plain window head against one with a learned sink:

from sinks import sliding_window_sink_trace, layer_attention_pattern

_trace = sliding_window_sink_trace(seq_len=22, window_size=4, sink_logit=3.0)
_pattern = layer_attention_pattern(n_layers=12, period=2, window=128)
ojs_define(slideTrace = _trace)
ojs_define(layerPattern = _pattern)

At a position with a real match both heads behave the same — the sink stays shut. But drag the query into the quiet stretch: the plain head keeps a full-size output (‖output‖ stays high — it is averaging values it should be ignoring), while the learned-sink head sends most of its mass to the sink (mass → sink fills up) and its output shrinks toward zero. Same window, same values; the sink is the only difference. On the built trace the sink head’s output norm in the quiet stretch drops to about 0.13 against the plain head’s forced 0.81 — roughly a 6× gate — with 83% of the mass parked on the sink. That is how gpt-oss can afford to make half its layers 128-token windows without them going haywire on long, low-signal spans.

from sinks import demonstrate_learned_sink

_ = demonstrate_learned_sink(sink_logit=2.0)
Learned sink s = 2.0:
  gate identity  out_sink == (1 - a_sink) * out_plain: True  (a_sink=0.482, gate=0.518)
  limits         s->-inf == softmax: True   s=0 == softmax_one: True
  in a sliding window with nothing to match:
    plain head  spare mass to sink=0.00  output norm=0.810  (forced to average junk)
    sink head   spare mass to sink=0.83  output norm=0.134  (gated toward zero)
  gpt-oss schedule: 6/12 full layers, rest 128-token sliding.
TipTry This
  1. Slide the sink logit to s = 0 in the first widget: the sink bar and every content weight land exactly on softmax₁. Now push it to −6 and the sink bar vanishes — you are back to plain softmax. One scalar spans both.
  2. Push s past +6: the output gate (1 − a_sink) collapses toward zero — the head is choosing to say nothing, even though the content logits never changed.
  3. Step the stream widget into and out of the quiet stretch: watch the sink head’s mass → sink bar open only when there is nothing to match, then close again the instant a real key returns. The plain head’s bar never moves — it has no sink to open.

Common Pitfalls

When extending context, watch out for:

  1. Forgetting scale = 1 must be identity. A correct implementation leaves a model untouched inside its training length. If your PI/NTK/YaRN frequencies differ from plain RoPE at scale=1, you have a bug — the tests assert this.
  2. Interpolating the fast pairs (plain PI). PI’s uniform squeeze blurs local, high-frequency position information. NTK/YaRN exist precisely to spare those pairs; reach for them past ~4× extension.
  3. Extending without any fine-tuning and expecting magic. NTK-aware buys a modest zero-shot extension; larger jumps (8–16×) need a short fine-tune at the new length or quality still sags.
  4. Dropping YaRN’s attention temperature. The frequency ramp alone recovers most of the quality, but the \sqrt{1/t} softmax scaling is part of the method; omitting it leaves perplexity measurably higher.
  5. Mismatched train/inference scaling. The scale and base used at inference must match what the model was fine-tuned with. A model tuned for YaRN at 8× will misbehave if you serve it with plain RoPE, and vice versa.
  6. Streaming with a naive window (no sinks). Evicting the oldest KV to cap memory looks harmless and destroys the model — perplexity explodes the moment the first tokens fall out. Always retain a few sink tokens (4 is plenty).
  7. Using original text positions in a streaming cache. If a streamed token reports its true position (which grows without bound) instead of its cache-relative one, RoPE sees angles past the training length again — the exact failure this whole module is about. Positions must be assigned within the cache.
  8. Giving the learned sink a value vector. The sink is a virtual key with no value — it belongs only in the softmax denominator. If you accidentally let it contribute a value (e.g. by appending a real zero-vector key and value and then reading it back), you break the exact (1-a_{\text{sink}}) gate and the sink can inject content it was never meant to carry. Add it to the logits, drop it from the weighted sum.

Exercises

Exercise 1: PI folds positions back into the band

import torch
from long_context import linear_interpolation_inv_freq, rope_inv_freq

# Show that Position Interpolation at scale s makes position L' land on roughly the
# trained angle at (L')/s. Compare the slowest pair's angle at L'=8192 under PI(s=4)
# against plain RoPE's angle at L=2048.

# Your implementation here:
# pi = linear_interpolation_inv_freq(64, 10000.0, 4.0)
# ...

Exercise 2: The NTK taper

import torch
from long_context import ntk_inv_freq, rope_inv_freq

# Confirm NTK-aware scaling leaves the fastest pair (index 0) unchanged while
# interpolating the slowest pair. Print the per-pair ratio ntk/plain for a few
# indices and check it decreases from ~1 toward ~1/s.

# Your implementation here:

Exercise 3: Build a YaRN ramp by hand

import torch
from long_context import yarn_ramp

# For head_dim=64, base=10000, L=2048, compute gamma and report: how many pairs are
# kept (gamma≈1), how many are fully interpolated (gamma≈0), and how many are in the
# ramp in between? Change L to 512 and watch the boundaries move.

# Your implementation here:

Exercise 4: The sink is a normalizer, not a memory

from streaming import demonstrate_sink_normalization

# Confirm the exact identity: with a zero-value sink holding mass a_sink, evicting
# it multiplies every content weight (and the output) by 1/(1 - a_sink). Sweep the
# sink logit and print a_sink and the distortion factor; check identity_holds.

# Your implementation here:
# for sink_logit in (1.0, 3.0, 5.0):
#     out = demonstrate_sink_normalization(sink_logit=sink_logit)
#     ...

Exercise 5: The learned sink is an output gate

import torch
from sinks import softmax_with_sink, sink_mass
from streaming import softmax

# Recover the gate identity yourself. For random logits x and values v, and a few
# sink logits s, check that sink_attention equals (1 - a_sink) * plain_attention,
# where a_sink = sigmoid(s - logsumexp(x)). Then find the s that halves the output
# norm and confirm it satisfies s = logsumexp(x) (a_sink = 1/2).

# Your implementation here:
# x, v = torch.randn(6), torch.randn(6, 4)
# for s in (-2.0, 0.0, 2.0):
#     a = float(sink_mass(x, s))
#     out_sink = softmax_with_sink(x, s) @ v
#     out_plain = softmax(x) @ v
#     assert torch.allclose(out_sink, out_plain * (1 - a))
# ...

Summary

Key takeaways:

  1. RoPE breaks past the training length because of the slow pairs — they reach position-angles the model never saw, while the fast pairs (which wrap the circle many times) extrapolate harmlessly.
  2. Every method rescales RoPE’s frequencies \theta_i — the scale factor is s = L'/L, and all methods reduce to plain RoPE at s=1.
  3. Position Interpolation divides every frequency by s — simple and effective, but it squeezes the fast pairs it didn’t need to.
  4. NTK-aware raises the base, \text{base}\cdot s^{d/(d-2)}, tapering the stretch so fast pairs are preserved and slow pairs interpolated — often with no fine-tuning.
  5. YaRN interpolates per pair by rotation count (\gamma_i ramp, \alpha=1, \beta=32) and adds a \sqrt{1/t}=0.1\ln s + 1 attention temperature — the method behind most long-context open models.
  6. Extension enables addressing, not comprehension — a short fine-tune at the target length is what turns reachable far positions into usable ones, and needle-in-a-haystack / long-document perplexity are how you measure it.
  7. Streaming is the other answer to the context wall — instead of enlarging the window, keep it fixed and stream through it. A naive rolling cache collapses; retaining a few attention-sink tokens plus a recent window (StreamingLLM) holds perplexity stable to millions of tokens at flat memory.
  8. The sink exists because softmax must sum to one — spare attention mass has to land somewhere, and the always-visible initial tokens absorb it. Evicting a zero-value sink rescales every other weight by 1/(1-a_{\text{sink}}), and cache-relative positions keep RoPE inside its trained range.
  9. A learned sink turns that pathology into a lever — gpt-oss gives each head a trained logit s in the softmax denominator, making the sink an exact, content-aware output gate \text{out}_{\text{sink}}=(1-a_{\text{sink}})\, \text{out}_{\text{plain}} with a_{\text{sink}}=\sigma(s-\log Z). It contains plain softmax (s\to-\infty) and off-by-one (s=0) as limits, and is what lets a model alternate sliding-window and full-attention layers without a positional sink.

What’s Next

You can now stretch a from-scratch GPT far past its training length. The frontier keeps going: Mixture of Experts grows a model’s knowledge without growing its per-token compute (sparse FFNs and top-k routing), and fast inference (quantization, speculative decoding) makes serving all of it cheap. Both reuse the efficient-attention and KV-cache machinery from m09 that a long context depends on.

Going Deeper

Core Papers:

Practical Resources: