GPT-2 (124M) to Gemma 3-Style (113M) language models, Part 2: KV caches, RoPE clocks, 4× context extension and attention visualization

This is the second part of my GPT-2 to Gemma 3-style language model series. In Part 1, I trained a GPT-2 124M (Model A) and a Gemma 3-style 113M model (Model B) on the same 9.8 billion tokens and found them tied on quality, with the Gemma recipe slightly slower to train but with a 6.4× smaller KV cache. In this part, I look at the model internals at inference time: how fast they generate text and how they represent position. I then extend Model B’s context length from 1024 to 4096 tokens with YaRN, and finally look at what each model attends to.

Here is the TLDR of what I found:

• On my M1 Pro MacBook, whose memory bandwidth is about 1/17 of an H100’s, Model B generates text up to 1.55× faster than Model A due to its 6.4× smaller KV cache. But on an H100, Model B is 1/3 slower than A. At this size, the H100 spends most of each step on the fixed cost of launching GPU operations, and Model B has more of them. So the smaller cache size pays off at serving scale, with big batches, long contexts and bigger models.
• RoPE gives Model B 32 fixed “clocks” to read position from: pairs of dimensions that rotate at different speeds as the position grows. Model A’s learned position table organizes itself into slow waves instead. The two schemes work differently, but in both models, how similar two positions look depends mainly on how far apart they are.
• YaRN helped stretch Model B from 1024 to 4096 tokens (4×) with no retraining, taking perplexity at 4096 from 37.7 down to 23.2. Continuing the training run for 200 steps (1% of pretraining) brings perplexity to 21.0, close to the original model’s 20.6 at 1024 tokens.
• Model B’s strongest induction head, the circuit that copies text from earlier in the context, sits in a local layer. So when the text to copy is, for example, 768 tokens back, Model A still copies it and Model B can’t. On ordinary text the two models rarely disagree.

The two models

As a reminder, the two models from Part 1 (124M and 113M parameters, 1024 context length) were trained identically on FineWeb-Edu:

• Model A is GPT-2 as nanoGPT ships it: learned position embeddings and full multi-head attention in every layer.
• Model B is Gemma 3-style [9]: RoPE [1], grouped-query attention (12 query heads sharing 3 key/value heads), a GeGLU MLP, sandwich RMSNorm, and local 512-token attention in three layers out of four.

The two models tied on loss and on every benchmark. On paper, Model B needs fewer FLOPs; in practice, it trained 4% slower. Most of the Gemma recipe targets efficiency and longer contexts. Its two attention changes, GQA and local windows, shrink the memory a model needs during inference, and RoPE makes position relative and easier to extend to longer contexts.

Generating text with a KV cache

To measure how the smaller cache helps during text generation, I benchmarked both models with a KV cache on two kinds of hardware. The answer turned out to depend more on the hardware than on the model.

What a KV cache stores

A language model generates one token at a time, and each new token attends to every token before it. Without a cache, every step reprocesses the whole sequence from scratch, which is wasteful because the keys and values of earlier tokens never change. A KV cache stores them once and reuses them, so each step only has to process one new token. Queries aren’t cached because in a causal model, each token’s query is used only once, to compute that token’s own output, which later tokens can never change. Only keys and values are read again by every token that comes after.

Each layer stores one key and one value (the 2 below) per key/value head and per cached position, so the cache for one sequence takes

\[\text{KV cache bytes} = 2 \cdot n_{kv} \cdot h \cdot b \cdot \sum_{\ell=1}^{L} T_\ell\]

where:

• Number of key/value heads, \(n_{kv}\): 12 in Model A, 3 in Model B.
• Head size, \(h\): each key and each value is a vector of 64 numbers.
• Bytes per number, \(b\): 4 in float32, 2 in bf16.
• Number of layers, \(L\): 12 in both models.
• Positions kept by layer \(\ell\), \(T_\ell\): how many past tokens that layer holds keys and values for. A global layer keeps every 1024 token; a local layer keeps at most 511.

Model A keeps all 12 heads for all 1024 positions in all 12 layers: 75.5 MB. Model B keeps only its 3 key/value heads, which its 12 query heads share, a 4x saving. Its local layers never look further back than their window, so 9 of its 12 layers stop at 511 positions. Together that’s 11.8 MB, 6.4x smaller.

What limits generation speed

Horace He’s Making Deep Learning Go Brrrr [12] sorts GPU time by what the chip is waiting on:

• Compute: matrix multiplications and element-wise operations.
• Memory bandwidth: it’s waiting for data to arrive from memory.
• Overhead: it’s waiting for the next instruction. That covers launching kernels, Python, and moving small tensors around.

Which of these dominates decides whether a smaller KV cache makes a model faster. Generating text with a KV cache happens in two phases, prefill and decode, and each runs into a different limit.

Prefill reads the prompt: every prompt token goes through the model in one forward pass, and their keys and values are written to the cache. With enough tokens (many batches and long context length), that’s large matrix multiplications, so prefill leans towards compute. Its speed depends on number of FLOPs more than on cache size. Model B does 9% less FLOPs per token than Model A (Part 1, “Counting FLOPs”).

Decode generates the answer one token per step. So each step does very little math, but it has to read every weight and every sequence’s whole cache. For a large model that makes decode memory-bound, which is where a smaller cache pays off. For a small model there is overhead cost: each step pays a fixed cost to launch a few hundred small kernels, one per operation per layer. If that fixed cost is bigger than the time spent moving data, decode is overhead-bound, and the cache size barely matters.

If you’ve used an LLM API with “cache write” and “cache read”, that is when prefill writes the KV cache and decode reads it.

Mac vs H100

I ran the same benchmark on my M1 Pro MacBook (16-core GPU, 16 GB of unified memory, 200 GB/s memory bandwidth) and on one H100 SXM (80 GB of HBM3 memory, 3.35 TB/s memory bandwidth), both in float32 (the model checkpoints’ precision), taking the median of repeated runs after warm-up. Each generation run starts from a prompt of random token IDs and produces 256 new tokens, always picking the most likely next one. I used random tokens because the content of a prompt doesn’t affect how fast the model generates.

	Model A	Model B	B vs A
KV cache per sequence, 1024 tokens (any machine)	75.5 MB	11.8 MB	6.4× smaller
Mac: decode with cache, batch 8, 512-token prompt	297 tok/s	461 tok/s	1.55×
Mac: prefill, batch 8, 1024 tokens	13.5k tok/s	12.2k tok/s	0.90×
H100: decode with cache, batch 8, 512-token prompt	1,988 tok/s	1,290 tok/s	0.65×
H100: prefill, batch 8, 1024 tokens	212k tok/s	194k tok/s	0.92×

At 1024 tokens, compute-heavy prefill behaves the same on both machines: B is 8–10% slower, in line with Part 1’s training slowdown. On the H100, short batch \(B=1\) prompts are even slower, because they are overhead-bound. The two machines differ substantially in decoding speed. The time per decode step shows why:

Understanding the plots above: I benchmarked four settings per model: a batch of 1 or 8, each generating after a 128-token or a 512-token prompt. That’s why each line has four points. The x-axis is how much data one decode step reads from memory: all the model’s weights, read once and shared by the batch, plus the KV cache of every sequence in the batch. A bigger batch or a longer sequence means more cache, so the point moves right. The y-axis shows how long one step takes.

On the Mac (the left plot), step time rises with the data read. Model A’s step time climbs from about 10 ms to 27 ms as the data read per step grows from about 520 to 875 MB. Across all eight points of both models, the MB read during decoding explains the variation in step time. This is what memory-bound looks like: each step mostly waits for the weights and the cache to arrive. Model B reads less cache per step, so its points sit kind-of-clustered and further left on the x-axis compared to A’s. It’s faster than Model A in every setting except the smallest one (batch 1, 128-token prompt), where both caches are tiny, and its lead grows with the batch.

On the H100 (the right plot), step time is flat. Model A takes 3.8–4.0 ms per step and Model B 6.1–6.3 ms, whatever the batch and prompt length. The data read barely matters here: at the H100’s 3.35 TB/s, reading 460–880 MB takes only 0.14–0.26 ms. The rest of each step is overhead. A decode step runs a few hundred small GPU operations (kernels), and launching each one costs a few microseconds, which adds up to milliseconds. Model B runs more of them per block: two extra normalizations, the gating multiply in its MLP, the rotations and the local-window masks. So on the H100 it’s slower, and its smaller cache does not help it win in speed.

The next figure shows how much the cache itself speeds up generation (throughput), on the Mac M1 Pro. In the two left panels, each model has a light bar (generating without the cache, recomputing every earlier token at each step) and a dark bar (with the cache), and the label above the dark bar is the speed-up. The right panel is prefill throughput (tokens/s) at batch 8, where the cache plays no part.

On the Mac, the cache speeds up generation 2.5–5.3× at batch 1 and 8–23× at batch 8, because recomputing the whole sequence every step is a lot of extra work.

On the H100 (not in the figure), at batch 1 with a short prompt, the cache gives no speed-up at all (0.9–1.0×), because recomputing 128 tokens costs similar launch overhead than processing one. The cache benefit appears only with longer prompts and bigger batches (up to 6.3× for Model A at batch 8).

The memory plot above shows the two savings separately. Model B’s line rises 1/4-th as steeply as Model A’s (GQA), and flattens past 512 tokens, where its local layers stop adding to the cache due to windowed attention. The measured sizes (dots) match the KV cache formula from previous section (lines) exactly, on both machines.

When would the H100 result flip? Model A spends about 4 ms of each step on overhead. In 4 ms, at its peak bandwidth, the H100 can read 3.35 TB/s × 0.004 s ≈ 13 GB. A full 1024-token sequence takes 75.5 MB of Model A’s float32 cache. So for reading the cache to take as long as the overhead, a step would have to read about 13 GB of cache: 13.4 GB / 75.5 MB ≈ 177 full-length sequences of Model A’s float32 cache at once. Real serving systems reach the scale where overhead is no longer the bottleneck, using large batches, long contexts and bigger models, and they also cut the per-step overhead with CUDA graphs and kernel fusion.

The KV cache of a larger model

Arithmetic intensity, the FLOPs done per byte read, tells which part of decoding can ever become compute-bound. A great reference to check out in this regard is CS336’s lectures on resource accounting and inference. For the weights, each parameter read in bf16 (2 bytes) is used for 2 FLOPs per sequence in the batch, so the intensity is about the batch size. An H100 does about 295 bf16 FLOPs per byte it can read, so decode only becomes compute-bound for the weights at batches of a few hundred sequences. Attention over the cache is different. Each cached key and value is read for one sequence only, and is used by every query head that shares it: \(n_h / n_{kv}\) heads, which is 1 in Model A and 4 in Model B. That gives an intensity of 1 FLOP per byte for Model A and 4 for Model B. Reading the cache is memory-bound at any scale, and GQA’s task is to multiply that intensity by the number of query heads per key/value head.

The same formula from ‘What a KV cache stores’ section gives the KV cache per sequence of, let’s say, a 7B-class model (32 layers, \(d_{model}\) = 4096, 32 heads of size 128) in bf16:

Context	Full multi-head attention	GQA (8 KV heads)	GQA + Gemma 3-style local windows
8k tokens	4.3 GB	1.1 GB	0.31 GB
32k tokens	17.2 GB	4.3 GB	0.91 GB
128k tokens	68.7 GB	17.2 GB	3.3 GB

At 32k tokens, one sequence’s full-attention cache would fill 1/5-th of an 80 GB H100. The Gemma 3-style layout fits nearly twenty sequences in the same space. Overhead stops dominating too: reading a 7B model’s bf16 weights alone takes about 4 ms per decode step at 3.35 TB/s (7B parameters × 2 bytes = 14 GB, and 14 GB / 3.35 TB/s ≈ 4.2 ms). That’s the regime the Gemma 3-style architecture targets, where a 6.4x smaller cache means serving 6.4x more sequences in the same GPU memory. At that scale the Mac result, where the smaller cache made decoding faster, is the one that carries over to the H100.

FLOPs at long contexts

Part 1 counted each model’s FLOPs per token, term by term. Attention is the only term that grows with context length, so the same algebra tells us what long contexts do to a model.

Per layer, attention costs \(2d(T + 1)\) FLOPs per token against the MLP’s \(16d^2\), so attention overtakes the MLP when \(T ≈ 8d\), with \(d\) the model dimension. For our \(d\) of 768 that’s about 6000 tokens; for a larger 7B-class model (\(d\) = 4096), it happens at about 33000 tokens. A local layer’s queries see at most \(W\) keys, so its attention cost levels off at \(4d·W\) per token. That only reaches the MLP’s \(16d^2\) if \(W > 4d\) (about 3,000 tokens at our \(d\)). So with Model B’s 512-token window the local layers attention costs never cross over; only the global layers do.

To see this in practice, I scaled each recipe from Part 1 up to a 7B-class and a 14B-class model and calculated the forward FLOPs per token at a 1k and a 128k context. The numbers below are absolute FLOPs per token:

Forward GFLOPs per token	1k context	128k context	Attention scores’ share at 128k
Model A’s recipe, 6.7B	13.6	47.7 (3.5×)	72%
Model A’s recipe, 15.4B	31.2	95.1 (3.0×)	68%
Model B’s recipe, 5.8B	11.9	20.5 (1.7×)	43%
Model B’s recipe, 13.5B	27.3	43.4 (1.6×)	38%

With full attention, stretching the context from 1k to 128k tokens multiplies the work per token by 3–3.5×, almost all of it attention scores. With Model B’s recipe (512-token windows in three layers of four), the same stretch costs 1.6–1.7×, because only the global layers pay for the long context. The window stays at 512 tokens however long the context gets, which is how sliding-window models are built: Gemma 3 keeps its 1024-token window at 128k. A window that grew with the context, say to half of it, would make the local layers quadratic again and remove most of the saving. So at 128k tokens, a 13.5B model built with Gemma 3’s recipe does less work per token than a 6.7B model built with GPT-2’s (43 against 48 GFLOPs). That’s on top of the memory savings from the KV cache section above. Gemma 3’s report [9] gives memory as its main reason for going to five local layers in six at 128k.

The same counting matters for how training efficiency is reported. The PaLM formula (Chowdhery et al., 2022) [17] behind Part 1’s logged MFU charges every token for attending to the whole context; a causal layer computes about half of that, and a local layer only its window. At 1024 tokens, for the 7B-class models above, the formula overstates the work by 2–3%. At 128k tokens it overstates it by 72% for Model A’s recipe and by a factor of 3.9 for Model B’s, because it charges the local layers for the full 128k they never look at. So an MFU number for a long-context or hybrid-attention run is only meaningful if you know which FLOP count it assumes.

A note on attention’s quadratic cost

Per token, attention is linear in context length. And summed over a whole sequence it gives the familiar quadratic, about \(2d·T^2\) FLOPs per layer. FlashAttention [11] makes that fit in GPU memory at long contexts, and enforcing local attention changes the order. The appendix (Attention and FlashAttention) works through the numbers.

How each model represents position

The two models end with the same loss, but they use very different mechanism for the one thing a transformer can’t know by itself: where each token sits in the sequence. Model A learns a vector for every position from scratch. Model B is supplied the position through a fixed rotation. Below, I share how each model handles position information.

RoPE as 32 clocks

Each layer splits the 768 hidden (embedding) dimensions of its queries and keys across 12 heads, so a head gets 768 / 12 = 64 hidden dimensions. RoPE [1] splits our 64-dimensional attention head into 32 pairs and rotates each pair by an angle proportional to the token’s position. Each pair turns at its own speed. When a query and a key score attention, only the difference of their angles survives, so attention depends only on how far apart two tokens are. Here is a mental model that helps me build intuition about RoPE: each attention head reads the distance between two tokens off a set of 32 clocks running at different speeds.

Here is the simple math behind the wavelengths of the fastest and slowest bands. Counting bands from 1, band \(j\) turns by \(\theta_j = 10000^{-2(j-1)/64}\) radians per position, where 10,000 is RoPE’s base and wavelength, the number of positions for one full turn, is \(2\pi/\theta_j\). The fastest band, \(j = 1\), turns 1 radian per position, so a full turn takes \(2\pi \approx 6.3\) positions. The slowest, \(j = 32\), turns \(10000^{-62/64} \approx 1/7500\) radians per position, so a full turn takes \(2\pi \times 7500 \approx 47000\) positions.

As shown in the plot above, only 18 of the 32 clocks complete a full turn within the 1024-token training context. The other 14 never do so during training, the model only ever sees a fraction of their cycle. At 1024 tokens it is harmless but it becomes a problem when the model reads longer inputs. I will show how YaRN fixes this, later in this article.

How an attention head uses the clocks

To see what these rotations do to attention, it helps to treat each pair of dimensions as one complex number. I use \(j\) for the band index below, to keep it apart from the imaginary unit \(\mathrm{i}\).

Take one head, with head size \(h = 64\), and band \(j\): dimensions \(j\) and \(j + h/2\), which RoPE rotates by \(m\theta_j\) at position \(m\), where \(\theta_j\) is the band’s fixed speed. The key at position \(n\) is rotated the same way, by \(n\theta_j\).

Step 1: a rotation is a multiplication. The RoFormer paper [1] uses the same complex form. Write the query’s pair as one complex number

\[z_j^q(m) = q_j(m) + \mathrm{i}\,q_{j+h/2}(m),\]

and the key’s as \(z_j^k(n)\). Rotating a 2-D vector by an angle \(\alpha\) is the same as multiplying its complex number by \(e^{\mathrm{i}\alpha}\). So RoPE turns the query’s pair into \(z_j^q(m)\,e^{\mathrm{i} m \theta_j}\), and the key’s into \(z_j^k(n)\,e^{\mathrm{i} n \theta_j}\).

Step 2: a dot product is the real part of a product. For two 2-D vectors written as complex numbers \(a\) and \(b\), their dot product is \(\mathrm{Re}[a\,\bar{b}]\), where \(\bar{b}\) is the complex conjugate. So band \(j\)’s share of the attention score is

\[\frac{1}{\sqrt{h}}\,\mathrm{Re}\Big[\, z_j^q(m)\,e^{\mathrm{i} m \theta_j} \cdot \overline{z_j^k(n)\,e^{\mathrm{i} n \theta_j}} \,\Big] = \frac{1}{\sqrt{h}}\,\mathrm{Re}\Big[\, z_j^q(m)\,\overline{z_j^k(n)} \; e^{\,\mathrm{i} (m-n)\,\theta_j} \Big]\]

where the two \((q,k)\) rotations combine into one. Only the difference \(\Delta = m - n\) survives. This is what makes RoPE relative.

Step 3: separate content from position. The product \(c_j = z_j^q(m)\,\overline{z_j^k(n)}\) depends only on what the query and key vectors contain before rotation. Writing it as amplitude and an angle, \(c_j = |c_j|\,e^{\mathrm{i}\varphi_j}\) and summing over the 32 bands gives the whole score:

\[\text{score}(\Delta) = \frac{1}{\sqrt{h}} \sum_{j=1}^{32} |c_j| \cos\!\big(\Delta\,\theta_j + \varphi_j\big)\]

This attention score can be interpreted as a sum of 32 waves, each a function of the distance \(\Delta\) between two tokens. The waves’ frequencies \(\theta_j\) are fixed by RoPE. How much of each wave a head uses (amplitude \(|c_j|\)), and where each wave peaks (phase \(\varphi_j\)), are set by content, through the query and key weights the head has learned. So each head builds its own preference for distance by mixing fixed waves, a bit like an equaliser setting the volume of fixed frequency bands.

A head can choose how it treats distance. Fast waves give a sharp preference for nearby tokens. Middle waves give a smooth fade over tens to hundreds of tokens. Slow waves never complete a turn within the context, so they change very little between a key 10 tokens away vs one 900 tokens away. They add roughly the same amount to the score at every distance.

Do Model B’s trained heads actually split like this? To check, I grouped the bands the way the YaRN paper [3] does, by counting how many full turns each band makes within the 1024-token training context:

• Fast bands make 32 turns or more (6 of the 32 bands).
• Middle bands make between 1 and 32 turns (12 bands).
• Slow bands never complete a single turn (the other 14).

The heatmap below answers this from the trained weights:

Each row is one head (12 per layer, 12 layers) and each column one band, fastest on the left. Each cell is one head’s \(**q·k**\) weight in one band. The head’s q·k weight in band \(j\) is the size of the query and key projection rows that feed band \(j\) which caps how much the head can use it on any input. Darker means more of that head’s \(q·k\) weight. The derivation is in the appendix, How the band measures are computed.

The pattern leans toward the slow bands on the right. Its slow-band weight sums that share over the 14 slow bands:

• Overall: the median slow-band weight is 53%, against 44% for an even spread.
• Five heads have a slow-band weight above 75%: one in layer 1, three in layer 6 and one in layer 8. They’re the dark stripes on the right.
• No head concentrates on the fast bands. The lowest slow-band weights, mostly in layer 0, sit at or just below the 44% of an even spread.

This matches what Barbero et al. (2025) [2] found in a much larger Gemma model: a strong preference for the lowest frequencies slow bands, which they interpret as carrying meaning, and the highest frequencies fast bands used for positional patterns.

The weights only show what a head can use. The figure below shows what it does use on real text. I ran held-out FineWeb-Edu through Model B for three heads, averaging the pre-softmax score \(q·k/\sqrt{h}\) over every query–key pair the same distance apart. Using the complex form from the derivation above, that average score splits exactly into one contribution per band (formula in the appendix). The figure’s bottom row is the head’s mean score at each distance, and the top row sums the bands’ contributions within each speed group: fast, middle and slow.

Each column in the above plot is one real head, chosen from the heatmap:

• the highest slow-band weight (layer 6, head 11, 84%);
• the median (layer 9, head 1, 53%);
• the lowest (layer 0, head 6, 42%).

Panel titles mark global or local layers; local heads’ curves stop at \(W = 512\). Each head matches one of the behaviours described earlier:

• Layer 0, head 6 is sharply local. Its score halves from distance 0 to distance 1 and keeps falling. That first drop is spread across all three speed groups.
• Layer 9, head 1 fades smoothly over about 100 tokens, through its middle and slow bands together.
• Layer 6, head 11 barely cares about distance beyond its immediate neighbours: its score is nearly flat from 3 tokens out to the edge of its 512-token window, and its fast bands contribute nothing at all.

One point to note in these plots, they mix two things, RoPE’s rotations and the fact that nearby tokens tend to be similar in content. The steep drop over the first few tokens in the slow bands, in layer 6’s head and layer 0’s, can’t come from rotation, since those bands barely turn in a few tokens. It comes from content, that is a query matching keys from its own neighbourhood. The flat stretch beyond that is where the slow bands’ indifference to position shows.

Model A’s learned position table

Model A’s position table is just 1024 learnt vectors of \(d=768\) dimension each. By its design, there is no encoding about distance, smoothness or waves information. But training gives it structure. Model B has no such table, since RoPE rotates queries and keys instead of adding a vector to each token. But one thing can be compared directly between the two models: how similar two positions look as a function of their distance.

The plot below captures the cosine similarity between the position vectors \(p_m\) and \(p_n\) at positions \(m\) and \(n\) (0 to 1023):

\[\text{sim}(m, n) = \frac{p_m \cdot p_n}{\lVert p_m \rVert \, \lVert p_n \rVert}\]

It is 1 when two positions get vectors pointing the same way, 0 when they are unrelated and negative when they point apart. For RoPE, rotating the same vector to positions \(m\) and \(n\) and comparing gives \(\frac{1}{32}\sum_{j=1}^{32} \cos\big((m-n)\,\theta_j\big)\) when every band carries equal weight, which depends only on the distance \(m - n\).

Left plot above: Model A’s position similarity. The diagonal itself is always 1, since each vector is compared with itself. Training adds the bright band around it where nearby positions have similar vectors, and the similarity fades with distance. An untrained table of random vectors would show only the diagonal, with similarities near 0 everywhere else. To check how much of this pattern comes from distance alone, I first average the similarity over all position pairs the same distance \(\Delta\) apart:

\[\bar{s}(\Delta) = \underset{|m-n| = \Delta}{\text{Avg}} \; (\text{sim}(m, n))\]

If similarity depended on distance and nothing else, every entry would equal \(\bar{s}(|m-n|)\). Coefficient of determination \(R^2\) measures how close the real matrix comes to that:

\[R^2 = 1 - \frac{\sum_{m,n} \big(\text{sim}(m, n) - \bar{s}(|m-n|)\big)^2}{\sum_{m,n} \big(\text{sim}(m, n) - \overline{\text{sim}}\big)^2}\]

The top residual sum of squares is what distance fails to explain, entry by entry. The bottom total sum of squares is the matrix’s total variation around its overall mean \(\overline{\text{sim}}\). \(R^2 = 1\) would mean distance explains everything as \(\bar{s}(|m-n|)\) equals \(\text{sim}(m, n)\), and 0 means it explains nothing. For Model A’s table, \(R^2 = 0.71\) means distance explains 71% of the variation. Model A has learned something close to “position as distance”, a property RoPE builds in algorithmically in Model B. Wang & Chen (2020) [16] found the same kind of structure in the learned position embeddings of pretrained models.

Middle plot above: the top three directions (principal components) along which Model A’s position vectors vary most, plotted against position. Each curve shows how far each position’s vector reaches along one of those directions. The curves are smooth waves of half a cycle, a full cycle and about one and a half. That’s what distance-based similarity predicts: the principal components are the eigenvectors of the matrix of dot products between positions, and when that matrix depends only on distance, its eigenvectors are close to sinusoidal waves.

Right plot above: similarity against distance, for both models. For RoPE, rotate any vector to two positions and compare: band \(j\)’s two copies end up \(\Delta\,\theta_j\) apart in angle, so each band contributes \(\cos(\Delta\,\theta_j)\), with \(\theta_j\) its frequency from RoPE as 32 clocks. Giving every band equal weight, the plotted curve is

\[\text{sim}_{\text{RoPE}}(\Delta) = \frac{1}{32}\sum_{j=1}^{32} \cos(\Delta\,\theta_j),\]

which depends on distance alone by construction. Equal weights show position alone, before a head’s content decides which bands matter. The RoPE curve is jagged because it is an exact sum of the 32 cosines at each whole-number distance. Model A (blue line) is smooth because it averages every pair of positions at that distance. Model A’s shaded blue band, the middle 80% of those pairs, shows that its similarity also depends on where in the window a pair sits. RoPE has no band, because every pair at the same distance gets exactly the same value. Additionally:

• RoPE is sharp for neighboring tokens and blurs out as inter-token distance grows. Its similarity falls from 1.0 to 0.66 within 10 positions, as the fast bands lose alignment, and then drifts slowly down, to about 0.2 by the end of the window, held up by the slow bands that barely turn. Nearby positions are easy to tell apart, while distances of 600 and 900 tokens look much alike.
• Model A’s table is blurry up close and coarse-grained far away. Its similarity is still 0.91 at distance 10. Neighbouring positions get nearly the same vector, but by 500 positions apart the similarity turns negative: the table mostly encodes roughly where in the window a token sits.

RoPE’s sharp view of nearby positions is one plausible reason it was the one change in Part 1’s ladder rungs that improved loss. Its slow bands, which barely turn within the window, are what YaRN rescales later in this article.

Position at longer contexts

At longer contexts, the two schemes behave very differently. Model A’s position table stops at 1024 rows, so it can’t read further at all. For comparison, building a 128k context 7B model would take 537M parameters, each row trained only on the tokens that happen to land at its position. This is part of why long-context models use RoPE or another relative scheme instead.

RoPE needs adjusting too. With the RoPE base of 10,000 used for our model B, every one of a 128-dimensional head’s 64 bands completes at least one full turn within 128k tokens; the slowest takes about 54,000 positions. A band that has turned all the way round gives a distant key the same angle as a near one, so no band changes steadily across the whole context to tell “far” from “very far”. Long-context models therefore raise the base, which slows every band down: Llama 3 [10] uses 500,000, and Gemma 3 uses 1,000,000 on its global layers. That leaves 15 and 17 bands slower than the whole 128k context. Model B’s attention windows save more and more memory as the context grows, but more and more of the context also lies outside them. In Gemma 3’s layout, five layers in six never look beyond 1024 tokens, even at a 128k context, so everything long-range has to travel through the global layers. This is why Gemma 3 keeps one global layer in six.

Extending Model B’s context from 1024 to 4096 with YaRN

Model B has no position embedding table, so in principle RoPE can rotate a query to any position. In practice, a RoPE model run past its training length falls apart quickly. It is becuase the slow bands (seen in RoPE clocks above) reach angles never seen in training, and a head’s mix of waves was only tuned on distances up to 1024.

YaRN is a recipe for fixing this after training. It leaves the fast bands alone, since they already cycled many times and their behaviour at new positions is familiar. It stretches the slow bands so that longer-context positions map into the range of angles the model saw originally, blends smoothly between the two in the middle, and adds a small correction to the attention temperature. The paper shows its zero-shot results hold up well to about 2x the training length and degrade beyond it. A short fine-tune on longer training sequences enable scaling beyond 2x without quality drop. I chose to extend Model B’s context by 4x for our study and here is my two-fold reasoning motivated from the paper’s findings: training data availability and zero-shot vs fine-tuned performance difference. At 2x, zero-shot YaRN results are similar to fine-tuning. At 8x, the target is beyond zero-shot range and pages that long are rare in our training dataset (about 1% of FineWeb-Edu validation pages reach 8k tokens, against about 3% for 4k). At 4x, there’s room for the fine-tune to close the gap between zero-shot and fine-tuning. Below, we will see the results for our Model B.

In terms of the fast, middle and slow bands from earlier, YaRN’s default thresholds leave Model B’s 6 fast bands alone, fully interpolate the 14 slow ones and blend the 12 in between, so 26 of the 32 bands in every head get rescaled.

Each dot is one of Model B’s 32 RoPE bands, fastest on the left. Hollow dots are the wavelengths the model was trained with, filled dots the wavelengths YaRN gives it, and the vertical line joining them is how far YaRN moved that band. The two dashed lines mark the training length (1024) and the new target (4096). Purple (fast) bands don’t move. Gold bands are stretched the full 4x, and every one of them lands above the 4096 line. Teal bands are blended, moving less the faster they are.

Without YaRN, stretching the context 2x, 4x or 8x would make 3, 5 or 7 more bands complete their first full turn, at angles the model never saw in training. With YaRN it’s 18 bands doing at least one full turn again: the same bands that turned within 1024 tokens during training are the only ones that turn within 4096 now.

Design choices the YaRN paper leaves open

• Which layers to stretch. Three of Model B’s four layers use a 512-token local window, so they never look more than 511 tokens back at any context length. Only the global layers ever meet a distance they weren’t trained on. Released models handle this differently. Gemma 3 stretches only its global layers (and with linear interpolation rather than YaRN). Rnj-1.5 and gpt-oss, both mixing local and global layers, apply YaRN to every layer alike. Since there’s no settled convention, I followed the paper, Rnj-1.5 and gpt-oss: one stretched table for every layer. YaRN leaves the fast bands, which do most of a local layer’s work, untouched anyway.

• The attention temperature. YaRN also multiplies every query and key by \(0.1\ln s + 1\), which is 1.139 at \(s = 4\), so attention scores grow by \(1.139^2 ≈ 1.30\) and attention gets slightly sharper. It compensates for the squeezed slow bands and for having 4x as many keys to spread attention over, both of which flatten attention at the new length. Rnj-1.5 blog mentioned they turned this correction off, presumably because a long mid-training stage let the model find its own sharpness. I chose to keep it for Model B because in zero-shot there is no training to absorb the change, and keeping the same setting in both versions means the fine-tune is the only difference between them. Folding the factor into RoPE’s cosine and sine tables applies it to both \(q\) and \(k\) for free.

YaRN without any training

The test set is 500 held-out FineWeb-Edu pages, each at least 4096 tokens long, each read from its start in one pass. Every window is therefore one real page, and position \(p\) always has exactly \(p\) earlier tokens of the same page to draw on. With causal attention, one pass at 4096 tokens also gives the answer for every shorter length, since a token’s loss never depends on what comes after it.

The left panel is the perplexity of a page’s first \(L\) tokens, the way the YaRN paper reports it. The right panel is the average loss at each position. The dashed line is the original context length. Model A’s line stops there because it has no position past 1024 to look up. The teal line is the fine-tuned model B with YaRN, which the next part of this section introduces.

Past its training length, plain RoPE breaks down. Its loss at each position jumps from about 3.0 to 4.2 between 1024 and 4096, and the perplexity of the whole 4096-token page climbs to 37.7. With zero-shot YaRN, it stays at 23.2 which is 14.6 lower, with a 95% interval of [14.1, 15.1] from a paired bootstrap over the pages. So YaRN removes the breakdown. But it has a cost. At 1024 tokens, where the original model needed no help, zero-shot YaRN costs perplexity 20.6 → 22.4. The sharper attention and squeezed slow bands apply at every length, including the short ones. The right panel shows what happens past the training length: zero-shot YaRN’s loss holds level at about where it was before 1024 (3.15 at positions 2048–4096, against 3.13 at 512–1024), but at that slightly worse loss throughout. So the model is stable at 4096 tokens, with no gain yet from the extra context.

Fine-tuning setup

The YaRN paper fine-tuned Llama 2 7B for 400 steps of 64 × 64k-token sequences from long books (about 1.7B tokens, roughly 0.1% of its pretraining). Those numbers don’t translate cleanly to our 113M model B trained on 9.8B tokens. Matching the share of pretraining gives about 10M tokens, about 20 of Model B’s steps. Matching the steps and batch gives about 105M. There is no scaling law for how much a YaRN fine-tune needs, and the paper’s learning rate (2e-5) says little about what a small model needs.

So I chose to continue fine-tuning Model B’s own training run and changed as little as possible:

• I started from Model B’s final checkpoint (step 20,000) with the optimizer’s state restored, so AdamW’s running averages carry over.
• The learning rate stays where the cosine schedule ended, at 6e-5, with no new warmup. A run resumed mid-schedule doesn’t re-warm, and the restored optimizer state already matches the model.
• Each step is still 491,520 tokens as before, with each micro-batch now 15 sequences of 4096 tokens instead of previous 60 of 1024. Weight decay, gradient clipping, precision and seed are unchanged.
• The pretraining stream packed pages back to back, so most 4096-token windows would join several short pages and seldom reward looking far back. The fine-tune reads only pages at least 4096 tokens long, one page per sequence. The model has seen these pages once already. Pretraining made one pass over almost all of the 10B-token sample. Repeating about 1% of the data once is well within what Muennighoff et al. (2023) [4] found to be as good as fresh data. The evaluation data come from the held-out split, which the model never trained on.

The fine-tuning continued for 200 steps, that is about 98M tokens or 1% of pretraining, keeps a snapshot of the weights every 25 steps and took about 23 minutes on one H100. Scoring each snapshot shows how many steps the fine-tune actually needs. The first step’s gradient norm was 1.49, the first gradient the model had ever seen from positions beyond 1024; clipping at 1.0 caught it. Within 40 steps it settled at about 0.28, close to the end of pretraining (about 0.25). Loss on held-out long pages went from 3.13 to 3.03, most of it in the first 50 steps.

Results after 200 fine-tuning steps

After 200 steps, the perplexity of a 4096-token page is 21.0, against 23.2 zero-shot: 2.2 lower, with a 95% interval of [2.1, 2.3] (shown in the plots above). The short-context cost is gone too. At 1024 tokens the fine-tuned model scores 20.5, against 22.4 zero-shot and 20.6 for the original model. The most useful comparison is with original Model B. It reads the first 1024 tokens of these pages at a perplexity of 20.6, and can’t read further. The fine-tuned model reads all 4096 at 21.0, and its loss at positions 2048–4096 (3.05) is within 0.03 of the original model’s at 512–1024. So the model gets 4x the context at nearly the same per-token quality, for 1% of the original training tokens.

One caveat to notice: the fine-tune trained on the kind of text that is the opening 4096 tokens of long pages, so some of its gain is plain adaptation to the page type. It’s better than the original even at the first 256 tokens (21.1 against 21.4), where position plays no part. Also, the per-position loss stays flat rather than falling past 1024, for the fine-tuned model as for the original one within its own range. Later tokens in these pages aren’t easier to predict, so a flat curve is what a model reading the whole page should show, but it can’t prove how much of the distant context the model actually draws on. The YaRN paper [3] checks this with passkey retrieval: it hides a random number deep inside long filler text and asks the model to repeat it, which only works if the model really uses context that far back. That test was beyond my experiment’s budget. A cheaper check, which I didn’t run here, would score each token near the end of a page twice: once with the whole page before it, and once with only the 1024 tokens just before it. If the token is predicted equally well both times, the model isn’t using anything more than 1024 tokens back. If it does better with the whole page, it is.

What my runs do measure is how quickly the fine-tune reaches its final quality, which the plot below tracks step by step.

In the plot above, each point is one snapshot of the fine-tune, scored on the same 500 pages. The teal line is the perplexity of the whole 4096-token page; the gold line counts only its first 1024 tokens, the length the model was originally trained on. Both come from the same pass over the same pages. The dashed line is Model B before YaRN on those same first 1024 tokens. The gold line above it means stretching has cost the model quality inside its old range, below it means that cost is gone.

Most of the gain comes early. The first 25 steps (12M tokens, about 0.1% of pretraining) take whole-page perplexity from 23.2 to 21.5, 77% of the total gain. By step 100 it’s 21.1 (94%), and the last 100 steps does not add much. The gold line is back at the original model’s quality by step 75. Of the two ways of scaling the YaRN paper’s budget down to our small model of 113M, matching the share of pretraining (about 20 steps for our model size) turned out to be the better guide to where the curve above bends.

Visualizing attention

This section looks at what each model attends to while it reads within their original context length of 1024 tokens. Attention weights show where information can flow, which is weaker evidence than showing why a model made a prediction (Jain & Wallace, 2019) [13]. And a weight averaged over a layer’s heads can hide a single head doing something sharp. I treat the views below as a way to look inside the models and not necessarily explain their final output predictions.

A small example first, with emojis 🤗 ! To represent the spirit of my blog series, i picked this sentence: “Learning is Fun 🤗, Keep exploring 🔭”

GPT-2’s BPE tokenizer works on bytes, and an emoji is 4 bytes of UTF-8. For our sentence, each emoji becomes 3 tokens: one they share (a space and the first two bytes, which every emoji begins with) and two of their own. Reading the shared token, a model can’t know which emoji is coming as only the later tokens decide. In greedy generation mode, when asked to continue the line always picking the most likely next token, both models copy it. Each starts a new line and repeats the sentence. Below and in the plots, ↵ marks a line break, which would otherwise be invisible.

	Greedy continuation of “Learning is Fun 🤗, Keep exploring 🔭”
Model A	↵ Learning is Fun 🤗, Keep exploring 🤗 ↵ Learning is Fun 🤗, Keep exploring 🤗 ↵ …
Model B	↵ Learning is Fun 🤗, Keep exploring 🔭 ↵ Learning is Fun 🤗, Keep exploring 🔭 ↵ …

Model B copies the line exactly. Model A copies it with the wrong emoji at the end, every time. Small models picking greedily fall into loops like this; that’s a known property of greedy decoding (Holtzman et al., 2019) [8]. With top-k sampling instead (temperature 0.8, top-k 200), both write ordinary on-topic text. Example, Model A continues “for more learning! 🔨 This is the first part of our series of lessons!”, Model B “for more fun learning activities. 1. The Elements of Play”.

The copying is visible in the attention itself. For each model, the figure below shows every layer’s attention, averaged over its heads, on the prompt (read after the <|endoftext|> token that starts every page in training) and the first line of the model’s own greedy continuation.

Each panel above is one layer averaged over all its heads. Rows (query) are the token doing the looking, columns (keys) the tokens it looks at; darker means more attention, and the grey triangle is the future, which a causal model never sees. The top half of each panel is the prompt; the bottom half is the copied line. A copying layer shows a stripe running parallel to the diagonal in the bottom half: each token of the copy looks back at the token right after its twin in the first line which is the next token to generate.

The copy stripe is strongest in one layer per model (with a weaker echo in layer 8): layer 7 in Model A and layer 6 in Model B, the same layers that hold each model’s strongest induction head (see below sub-sections). The other layers do simpler things. Model B’s layer 0 mostly attends to the token itself and its layer 2 to the previous token, the kind of local bookkeeping a copying head needs as an input. Both models also park a lot of attention on the opening <|endoftext|> token, Model A especially in its later layers.

A single short sentence never reaches Model B’s 512-token attention windows, so for the rest of this section I measure every head on full-length text and test one mechanism: copying, directly. None of the plots from here on generate text, so there is no prompt or decoding mode. Each model reads a fixed input in one forward pass, and I record its attention weights: held-out FineWeb-Edu windows for attention distance and sinks, and sequences built from random tokens for induction heads.

Attention distance and attention sinks

Each dot in the plots above is one head and the line is the layer’s average. Blue is Model A, orange Model B, the shaded columns Model B’s three global layers, and the dashed line Model B’s window. The left panel is attention distance [7]: how many tokens back a head looks on average, weighting each earlier token by its attention (a measure from the Vision Transformer paper). The equations for the same are written below. It uses only queries in the second half of 32 held-out 1024-token windows, so every head could look at least 512 tokens back. The other panels measure attention sinks [6], a fixed place models park attention when nothing else needs it: (middle panel) the share on the window’s first position, whatever token sits there, and on document separators (right panel) the token that starts every page.

In symbols, with \(p_{q,k}\) the attention a head puts from the query at position \(q\) on the key at position \(k \le q\) (each query’s weights add up to 1), each panel averages one of these over the queries:

\[\text{distance}_q = \sum_{k \le q} p_{q,k}\,(q - k) \qquad \text{first-position share}_q = p_{q,0} \qquad \text{separator share}_q = \sum_{k \le q,\ x_k = \texttt{EOT}} p_{q,k}\]

The separator share is averaged only over the windows that contain a separator.

Model A’s heads look further back on average: 132 tokens against Model B’s 81. A has three layers averaging more than 200 tokens (0, 7 and 10) and single heads reaching almost 480. B’s global layers don’t compensate: its most far-sighted layer, global layer 8, averages 148 tokens, and its farthest-looking head 235. On ordinary text, B looks less far, and none of its heads comes close to the 512-token window that would stop it.

The two sink panels show the clearest difference between the recipes. Model A rests attention on position 0 (up to 13% of a layer’s attention, in its later layers) and even more on document separators (up to 29%). Model B’s global layers never rest on position 0 (its local layers can’t reach it from these queries), but B does rest on separators, though less than A (at most 11% of a layer, in layers 6, 8 and 10). BERT’s heads behave similarly with its [SEP] separator token, which Clark et al. (2019) [15] interpreted as a way for a head to do nothing. The likely reason for the difference between A and B is how each model knows where it is. A’s learned position table gives position 0 its own vector, so the first position is recognisable whatever token sits there. RoPE gives position 0 no vector of its own, so B has nothing built in that marks it, and it rests on a token it can recognise instead.

Induction heads and the local window

One of the best-studied circuits in language models is the induction head, first described by Elhage et al. (2021) [14] and studied in depth by Olsson et al. (2022) [5]: at a token it has seen before, it attends to the token that followed the earlier occurrence, so the model can copy what came next. The standard test, from those papers, removes every other clue: a block of 64 random tokens, random filler, then the same 64 tokens again. Nothing can predict a random token except the earlier copy. Varying the filler puts the repeat either 256 tokens back, inside Model B’s local window, or 768 tokens back, beyond it. Model A can attend 768 tokens back in any layer. Model B can only do it in its three global layers.

Each head gets an induction score: the attention it puts on exactly the token it should copy, averaged over the repeated block. Near 1 is a clean induction head, near 0 is none. In symbols: if the block’s first copy starts at position \(s\) and its repeat sits \(D\) tokens later (256 or 768), then the token at position \(t\) in the second copy has its twin at \(t - D\), and the token to copy is the one right after it, at \(t - D + 1\). A head’s induction score averages its attention there over the \(n - 1\) tokens of the second copy that have a next token to copy (\(n = 64\)), and over the test sequences:

\[\text{induction score} = \frac{1}{n-1}\sum_{t=s+D}^{s+D+n-2} p_{t,\;t-D+1}\]

The above plot is the same format as the emoji figure, for each model’s strongest induction head reading 12 random tokens and then the same 12 again. In the first copy the head rests on the first token; in the second, each token attends to the token right after its earlier twin, the stripe one step right of the earlier copy.

In the two left panels, each dot is one head, and its height is its induction score. The shaded columns are Model B’s global layers again. The right panel asks whether the copying works: the loss on the second copy of the block, against the dashed line for the first copy, which neither model has any way to predict. A loss near zero means the model copies.

The loss on each copy is Part 1’s cross-entropy, \(-\ln p(x_{t+1} \mid x_{\le t})\), averaged over the copy’s tokens.

With the repeat 256 tokens back, both models have clean induction heads: Model A’s strongest is layer 7, head 11 (score 0.93), with 5 heads above 0.3; Model B’s is layer 6, head 11 (0.90), with 4. Both copy, and Model A copies much better. The second copy of the block costs Model A 1.6 nats per token and Model B 6.1, against about 13.3 for the unseen first copy.

With the repeat 768 tokens back, Model A is barely affected: the same head still scores 0.90, and the second copy costs 3.0 nats. Model B’s induction heads are gone. No head scores above 0.3, its best manages 0.05, and the second copy costs 12.3 nats, almost as much as tokens it has never seen. Its strongest induction head lives in layer 6, a local layer, and the earlier copy is now out of that layer’s reach. None of its three global layers took over the job.

Average loss doesn’t show this cost of the local/global recipe: Model B’s attention windows save memory, and the model loses the ability to copy from beyond them.

Where the two models disagree

Average loss can hide where two models differ, so I compared their predictions token by token on held-out text, using the KL divergence between the two next-token distributions, KL(A ‖ B): the extra surprise, in nats, of predicting with Model B’s distribution when Model A’s is the reference. Zero means identical predictions. For one position, summing over every token \(v\) in the vocabulary:

\[\mathrm{KL}(A \,\|\, B) = \sum_{v} p_A(v)\,\big(\ln p_A(v) - \ln p_B(v)\big)\]

The left panel above counts tokens by how much the models disagree, on a log scale: most tokens sit in the middle, a few in the long tail. The right two heatmaps take the single strongest disagreement and show, for each model and layer, where its attention went from that token, averaged over heads. The first column collects the attention that went to tokens before the passage shown, so each row adds up to 1.

Over 20,480 held-out tokens, the typical disagreement is small: a mean KL of 0.29 nats and a median of 0.21. Only 3% of tokens exceed 1 nat and 0.1% exceed 3. On these 20,480 tokens, average loss is same as expected from Part 1 training runs: 2.98 for A, 2.97 for B.

In the tail, B does slightly better: among the 1,000 strongest disagreements, B has the lower loss on 59%. The strongest few are mostly completions where one model predicts the phrase and the other doesn’t:

Context (end)	True next token	Model A’s guess	Model B’s guess
“…capable of supporting up to 127 daisy”	”-“ (daisy-chained)	“ph” 48%	”-“ 59%
“…whether individuals use their dominant versus their”	” nond(ominant)”	” dominant” 36%	” nond” 99%
“…the measurement of these inflammatory ‘medi”	“ators”	“as” 63%	“ators” 87%
“…if you have 40 CALs … The 41st”	” user”	” mode” 12%	” CAL” 96%

The last row goes the other way: Model B is very sure of a word that’s wrong. So the two models look at similar places and draw different conclusions from them, which attention weights alone can’t show.

Caveats, and what to expect at larger scale

The inference numbers in the KV cache analysis section come from one model size, two machines and float32 without kernel fusion or CUDA graphs, all of which push a 124M model towards being overhead-bound. The RoPE position analysis covers one context length. So here i discuss what should carry over to larger models (my go-to pick being a 7B model size) and longer contexts, and what is the basis of my understanding:

Finding at 124M	With bigger models	With longer contexts	Basis
KV cache 6.4× smaller for Model B compared to Model A	Same ratio per layer; the absolute size grows with \(d·L\)	The ratio grows (≈19× at 32k for a 7B-class model with Gemma 3’s layout)	Mathematical formula above
Decode speed depends on the machine; overhead-bound on H100	Becomes memory-bound (a 7B model’s weights alone take ~4 ms to read per step), so the smaller cache pays off	Cache traffic dominates	Arithmetic + both measurements
Reading the cache: 1 FLOP per byte for A, 4 for B	Unchanged: it depends only on \(n_h/n_{kv}\)	Unchanged, but the cache grows	Arithmetic
RoPE with base 10,000; 18 of 32 bands turn within the context	—	The base must rise (Llama 3: 500k; Gemma 3 global: 1M), or bands must be rescaled (YaRN)	Arithmetic + papers [3], [9], [10]
YaRN at 4×: zero-shot removes the breakdown; 25 fine-tuning steps (~0.1% of pretraining) give most of the gain	The paper’s own results are at 7B and above; the 0.1% budget matched here too	Larger stretches need more fine-tuning, or the base raised in pretraining	Experimental findings above and the YaRN paper
Model B’s induction heads sit in a local layer, so it can’t copy from 768 tokens back	Unknown: larger hybrid models may place copying heads in global layers	The window caps copying unless global layers learn it	Experimental findings above
Model A sinks attention on position 0, Model B (RoPE) on separator tokens	Attention sinks are widespread at scale (Xiao et al., 2023) [6]	Sinks on a token need that token within reach of local layers	Experimental findings above

A second seed of Model B could, for instance, place its induction head in a global layer, which would change the copying result. The attention analysis is also descriptive: it shows where attention goes. Showing that the model needs it there requires a causal check, such as switching a head off and measuring what breaks.

Appendix

Attention and FlashAttention

This section works through three claims behind the main text’s note on attention’s quadratic cost: attention’s compute grows with the square of the context length. FlashAttention removes that square from memory but not from compute. Besides, by limiting which keys a query sees, as Model B’s local layers do, lowers the compute itself.

• Cost per token. A query at position \(p\) (counting from 0) is compared with the \(p + 1\) keys up to and including its own. Each attention score is a dot product of length \(h\) in each of the \(n_h\) heads, \(2h·n_h = 2d\) FLOPs, and multiplying with value vector costs another \(2d\). So one token costs \(4d·(p + 1)\) FLOPs per layer, linear in its position (Part 1, Counting the FLOPs). Across a sequence of \(T\) tokens, the queries at positions 0, 1, …, \(T - 1\) see 1, 2, …, \(T\) keys, for a total of \(T(T + 1)/2\). Dividing by the \(T\) queries, the average query sees \(\bar{k} = (T + 1)/2\) keys, so the average token costs \(4d·(T + 1)/2 = 2d(T + 1)\) FLOPs, the per-token figure used in FLOPs at long contexts.

• Cost per sequence. Adding up all \(T\) tokens gives

  \[\sum_{p=0}^{T-1} 4d·(p + 1) = 4d·\frac{T(T + 1)}{2} \approx 2d·T^2\]

  FLOPs per layer, and gives back the famous quadratic. So \(x\) times the context costs \(x^2\) times the work.

• Memory without FlashAttention: A standard implementation writes each head’s \(T × T\) matrix of scores to GPU memory, reads it back for the softmax, writes the probabilities and reads them again for the weighted sum of values, then keeps them for the backward pass. In bf16 (2 bytes per number), one layer’s matrices for a batch of \(B\) sequences take

  \[B · n_h · T^2 · 2 \text{ bytes.}\]

  In our runs (\(B = 60\) sequences per GPU, \(n_h = 12\), \(T = 1024\)) that’s about 1.5 GB per layer, or about 18 GB for all 12 layers. At 128k tokens (128 × 1024 = 131,072), a single head’s \(131{,}072^2\) matrix for a single sequence takes about 34 GB in bf16, more than a third of an H100’s 80 GB.

• What FlashAttention changes: FlashAttention [11] computes the same exact attention without ever storing the \(T × T\) matrix. It splits the queries and keys into tiles small enough for the GPU’s on-chip memory and, in one kernel, computes scores, softmax and weighted sum for one tile at a time, keeping a running maximum and sum so the softmax comes out exact. Memory then grows with \(T\) instead of \(T^2\), which is what makes long-context training fit at all. The FLOPs don’t go down: every query-key pair is still scored, and the backward pass recomputes the scores instead of reading them back, which adds a little arithmetic. The time it saves is time the GPU would have spent waiting on memory.

• What lowers compute: Only limiting which keys a query sees. A query’s cost is \(4d\) per key it sees, so what matters is \(\bar{k}\), the average number of keys per query. A global layer, as seen above, has \(\bar{k} = (T + 1)/2\). In a local layer, the first \(W\) positions see 1, 2, …, \(W\) keys, which adds up to \(W(W + 1)/2\), and the other \(T - W\) positions see \(W\) keys each, adding \((T - W)W\). Dividing the total by the \(T\) queries gives the average (Part 1’s appendix):

  \[\bar{k}_W = \frac{W(W + 1)/2 + (T - W)\,W}{T}.\]

  Since \(\bar{k}_W\) can never exceed \(W\), a local layer’s cost per token stops growing once \(T\) passes \(W\), and its cost per sequence, \(4d·T·\bar{k}_W\), grows linearly in \(T\) instead of quadratically. With Model B’s \(W = 512\):

  Context \(T\)	Global \(\bar{k}\)	Local \(\bar{k}_W\)	Local layer’s saving
  1,024	512.5	384.25	1.33×
  131,072 (128k)	65,536.5	511.0	128×

  That saving only shows up in time if the kernel skips the scores outside the window, as FlexAttention’s block masks [18] and FlashAttention’s window option do. With an ordinary dense mask, all \(T^2\) scores are computed and most are then thrown away, which is why the dense-mask Model B in Part 1’s ladder trained slower. And switching to FlexAttention won back about 5% in time.

• During generation. With a KV cache, each new token is a single query against the cached keys: \(4d·(p + 1)\) FLOPs at position \(p\) in a global layer, at most \(4d·W\) in a local one. Generating \(T\) tokens therefore still adds up to about \(2d·T^2\) per global layer, spread across the \(T\) steps.

How the band measures are computed

The position section uses two measures of how a head uses RoPE’s 32 bands. The first comes from the weights alone and says how much a head can use each band: it is behind the band-usage heatmap and the slow-band weight. The second comes from running real text and says how much each band does add to the score at each distance: it is behind the three real heads. The notation follows How an attention head uses the clocks.

From the weights. The RMSNorm in front of attention scales every token’s vector to the same length, \(\sqrt{d}\), and then multiplies it by a learned gain. Let \(x\) be the vector after the scaling and before the gain, so \(\lVert x \rVert = \sqrt{d}\), and fold the gain into the projection weights. Band \(j\) of a query uses two dimensions, \(j\) and \(j + h/2\), so it comes from two rows of the head’s query projection. Call those two rows \(W_q^{(j)}\), a \(2 × d\) matrix, so that band \(j\)’s pair is \(W_q^{(j)} x\). Define \(W_k^{(j)}\) the same way from the key projection; under grouped-query attention, that’s the key head the query head shares with three others.

A pair’s length can’t exceed its matrix’s Frobenius norm (the square root of the sum of its squared entries) times the input’s length, because each of the two rows’ dot product with \(x\) is at most that row’s length times \(\sqrt{d}\). So \(\lvert z_j^q \rvert \le \lVert W_q^{(j)} \rVert \sqrt{d}\), and likewise for the key. In the score formula, band \(j\)’s wave has amplitude \(\lvert c_j \rvert / \sqrt{h} = \lvert z_j^q \rvert\, \lvert z_j^k \rvert / \sqrt{h}\), since the length of a product of complex numbers is the product of their lengths. So for any two tokens

\[\frac{\lvert c_j \rvert}{\sqrt{h}} \;\le\; \frac{d}{\sqrt{h}}\, w_j, \qquad w_j = \big\lVert W_q^{(j)} \big\rVert \cdot \big\lVert W_k^{(j)} \big\rVert.\]

So \(w_j\) caps how loud band \(j\)’s wave can be in this head. Dividing by the head’s total gives each band’s share, and adding the shares of the 14 slow bands gives the slow-band weight:

\[s_j = \frac{w_j}{\sum_{l=1}^{32} w_l}, \qquad \text{slow-band weight} = \sum_{j \,\in\, \text{slow}} s_j.\]

Each cell of the heatmap is one head’s \(s_j\). A head that weighted all 32 bands equally would have a slow-band weight of 14/32 ≈ 44%, the “even spread” in the main text. The cap is only an upper bound: real inputs may not line up with the rows that give a band its weight, which is why the second measure looks at real text.

On real text. I ran 64 held-out windows of \(T = 1024\) tokens through Model B and recorded each head’s queries and keys before RoPE’s rotation. As in the main text, \(z_j^q(m)\) and \(z_j^k(n)\) are band \(j\)’s pairs at positions \(m\) and \(n\), written as complex numbers. By Step 2 of How an attention head uses the clocks, band \(j\)’s share of the score between a query and a key \(\Delta = m - n\) positions apart is \(\mathrm{Re}\big[ z_j^q(m)\, \overline{z_j^k(n)}\, e^{\mathrm{i}\Delta\theta_j} \big] / \sqrt{h}\). A window has \(T - \Delta\) query-key pairs at distance \(\Delta\), and averaging over them gives

\[S_j(\Delta) = \frac{1}{T - \Delta} \sum_{n=0}^{T-1-\Delta} \frac{1}{\sqrt{h}}\, \mathrm{Re}\Big[\, z_j^q(n+\Delta)\; \overline{z_j^k(n)}\; e^{\,\mathrm{i}\,\Delta\,\theta_j} \Big],\]

which is then averaged over the 64 windows. The factor \(e^{\mathrm{i}\Delta\theta_j}\) is RoPE’s rotation at that distance; everything else is content. Since a head’s score is the sum of its 32 bands’ shares, and an average of a sum is the sum of the averages, the bands add up exactly to the head’s mean pre-softmax score at each distance:

\[S(\Delta) = \sum_{j=1}^{32} S_j(\Delta).\]

The figure’s bottom row plots \(S(\Delta)\). The top row splits the same sum into the fast, middle and slow bands, so the three top curves add up to the bottom one.

Code and models

The models are public on Hugging Face: Model A, Model B and the YaRN-extended Model B (4096-token context, after the 200-step fine-tune). The KV cache, the attention-weights path, the YaRN fine-tune and the analysis behind every figure are in the CuriosLM repository; its README explains how to run them.

Acknowledgement

Everyone I acknowledged in Part 1 shaped this part too, especially Percy Liang’s Stanford course CS336: Language Modeling from Scratch. Horace He’s Making Deep Learning Go Brrrr From First Principles gave me the way to think about generation speed here: compute, memory bandwidth or overhead. Aleksa Gordić’s Inside the Transformer: The Life of a Token is the best place to read about RoPE, YaRN and hybrid attention in depth.

Platforms and tools I used for this part:

• I rented H100s on RunPod for the inference benchmark and the YaRN fine-tune.
• Weights & Biases tracked the fine-tuning run.
• Hugging Face Hub hosts all three models, including the YaRN-extended Model B.
• Claude Code, Anthropic’s coding assistant, worked alongside me: it helped build and test the KV cache and the analysis scripts, debug the fine-tune setup, and draft and fact-check this article.

If you spot a mistake, or have comments or critiques, I’d be glad to hear them on LinkedIn or X.

References

1. Su et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding.
2. Barbero et al. (2025). Round and Round We Go! What Makes Rotary Positional Encodings Useful?.
3. Peng et al. (2023). YaRN: Efficient Context Window Extension of Large Language Models.
4. Muennighoff et al. (2023). Scaling Data-Constrained Language Models.
5. Olsson et al. (2022). In-context Learning and Induction Heads.
6. Xiao et al. (2023). Efficient Streaming Language Models with Attention Sinks.
7. Dosovitskiy et al. (2020). An Image is Worth 16x16 Words.
8. Holtzman et al. (2019). The Curious Case of Neural Text Degeneration.
9. Gemma Team (2025). Gemma 3 Technical Report.
10. Llama Team (2024). The Llama 3 Herd of Models.
11. Dao et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.
12. Horace He (2022). Making Deep Learning Go Brrrr From First Principles.
13. Jain & Wallace (2019). Attention is not Explanation.
14. Elhage et al. (2021). A Mathematical Framework for Transformer Circuits.
15. Clark et al. (2019). What Does BERT Look At? An Analysis of BERT’s Attention.
16. Wang & Chen (2020). What Do Position Embeddings Learn? An Empirical Study of Pre-Trained Language Model Positional Encoding.
17. Chowdhery et al. (2022). PaLM: Scaling Language Modeling with Pathways.
18. Dong et al. (2024). Flex Attention: A Programming Model for Generating Optimized Attention Kernels.

Citation

@misc{bhattacharyya2026gpt_to_gemma_2,
  author       = {Bhattacharyya, Shinibali},
  title        = {GPT-2 (124M) to Gemma 3-Style (113M) language models, Part 2: KV caches, RoPE clocks, 4× context extension and attention visualization},
  year         = {2026},
  url          = {https://shinibali.github.io/gpt2-to-gemma-part-2/},
}