Lesson 42 of 73
Attention Optimizations and Long Context
Attention is the step where every token looks at the tokens before it. To do that, the model compares each token with every earlier token and writes the results into a table of scores. The table has one row and one column per token, so it grows with the square of the text length. A 4,000-token prompt needs a table of about 16 million scores for each attention head. A 128,000-token prompt needs about 17 billion.
This growth causes three separate problems, and this lesson covers one family of fixes for each:
- The score table is too large to store and too slow to move. Online softmax and tiled attention (the idea behind FlashAttention) compute the same result without ever storing the table.
- The KV cache is too large. The KV cache (key-value cache) is the memory where the model keeps the keys and values of all earlier tokens so it does not recompute them. Fewer KV heads (GQA and MQA), a compressed latent (MLA), sliding windows, and attention sinks all shrink or manage it.
- The model has only seen short texts. A model trained on 4,096-token texts has never seen a token at position 16,000. RoPE scaling (position interpolation, NTK-aware scaling, YaRN) changes how positions are encoded so long positions look familiar.
After this lesson you can derive online softmax and tiled attention by hand, compute the KV cache bytes per token for MHA, GQA, MQA, MLA, and interleaved sliding-window models, read the related llama.cpp and vLLM settings, and explain what linear interpolation, NTK-aware scaling, and YaRN each change in RoPE. These skills matter when you pick a context size such as -c 131072 for a coding agent: the choice decides how much memory the cache needs and how well the model still finds information in the prompt.
| Idea | Plain definition | Examples |
|---|---|---|
| Score matrix | The table of all query-key scores for one head: n x n values for n tokens |
1, 5 |
| Online softmax | Softmax computed in one pass, with a running maximum and a running sum | 2, 3 |
| Tiled attention (FlashAttention) | Attention computed block by block in fast memory, so the score table is never stored | 4, 5 |
| GQA, MQA | Several query heads share one key-value head | 6 |
| MLA | Keys and values are stored as one small compressed vector per token | 7, 8 |
| Sliding-window attention | A layer looks only at the last W tokens |
9, 10 |
| Attention sinks | The first tokens of a text, which receive a large share of attention and must stay in the cache | 11 |
| RoPE scaling | Changing the rotation speeds so longer positions look like trained ones | 12 to 15 |
| Long-context cost | Prefill compute grows with the square of the length, cache memory grows linearly, and retrieval tests measure quality | 16, 17 |
A short reminder of the words used throughout. Each token produces three vectors: a query (what this token looks for), a key (what this token offers for matching), and a value (the information it hands over). A score is the dot product of one query and one key, divided by sqrt(d), where d is the head dimension (the length of the vectors). Softmax turns a row of scores into probabilities that add up to 1. The attention output for a token is the sum of all value vectors, each multiplied by its probability. A head is one independent copy of this process, and a layer has many heads. Prefill is the step that processes the whole prompt at once. Decode is the step that produces one new token at a time.
The lesson builds on attention from Self-Attention from Scratch, the KV cache from Generation, Prefill, Decode, and the KV Cache, KV bytes from Memory Math, RoPE (rotary position embedding, the way the model encodes where a token sits) from The Transformer Block and the Full Model, and the -fa, --swa-full, and RoPE flags introduced in Tuning llama.cpp. Engine facts come from llama.cpp commit 0c1e57098 and vLLM 0.30.0. Real runs use tiny models with random weights: their text is meaningless, but their logs and buffer sizes are real.
Example 1: memory traffic of naive attention
The straightforward way to compute attention for one head has three steps. First, compute the score matrix S = Q K^T / sqrt(d), which holds one score for every query-key pair, and store it. Second, apply softmax to each row to get the probability matrix P, and store that too. Third, multiply P by V to get the output. This is called naive attention.
The sizes follow from counting. With 4 tokens, S has 4 x 4 = 16 entries. With 1,000 tokens it has one million. For n tokens, S and P each hold n * n values. The program computes that size for the 32 query heads of a Llama 3 8B-shaped model, from 4K to 128K tokens. It also computes the bytes that pass through memory when each matrix is written once and read once.
import java.util.Locale;
public class ScoreMatrixSize {
static double gib(double bytes) { return bytes / (1024.0 * 1024 * 1024); }
public static void main(String[] args) {
int heads = 32, layers = 32, bytesPerScore = 4; // Llama 3 8B shape, FP32 scores
System.out.println("tokens scores/head one layer traffic/layer 32 layers");
for (long n : new long[] {4096, 8192, 32768, 131072}) {
long perHead = n * n;
double layer = (double) perHead * heads * bytesPerScore;
double traffic = 4 * layer; // write S, read S, write P, read P
System.out.printf(Locale.ROOT, "%6d %12d %7.0f GiB %10.0f GiB %7.0f GiB%n",
n, perHead, gib(layer), gib(traffic), gib(traffic * layers));
}
long n = 131072, ubatch = 512;
System.out.printf(Locale.ROOT, "prefill micro-batch of %d queries at %d keys: %.0f GiB per layer%n",
ubatch, n, gib((double) ubatch * n * heads * bytesPerScore));
System.out.printf(Locale.ROOT, "decode, 1 query at %d keys: %.0f MiB per layer%n",
n, n * heads * bytesPerScore / (1024.0 * 1024));
}
}tokens scores/head one layer traffic/layer 32 layers
4096 16777216 2 GiB 8 GiB 256 GiB
8192 67108864 8 GiB 32 GiB 1024 GiB
32768 1073741824 128 GiB 512 GiB 16384 GiB
131072 17179869184 2048 GiB 8192 GiB 262144 GiB
prefill micro-batch of 512 queries at 131072 keys: 8 GiB per layer
decode, 1 query at 131072 keys: 16 MiB per layerRead the table from left to right. The second column is n * n: 16.8 million scores per head at 4,096 tokens and 17.2 billion at 131,072 tokens. Multiply by 32 heads and 4 bytes per score, and one layer needs 2 GiB at 4K and 2,048 GiB at 128K. The size grows with the square of the length: 32 times more tokens give 1,024 times more scores (32 x 32).
The traffic column counts four passes over each layer's scores: write S, read S for the softmax, write P, and read P for the product with V. At 4K tokens that is 8 GiB of traffic per layer and 256 GiB for the 32 layers of one prefill. This matters because a GPU spends its time waiting for memory, not only doing arithmetic (see the roofline lesson). At 128K tokens, the full matrix would need 2 TiB for one layer, which is more memory than any single GPU has.
Engines therefore never build the full prefill matrix at once. llama.cpp processes the prompt in micro-batches (-ub, default 512 tokens). A micro-batch is a slice of the prompt processed together, so without flash attention the engine stores a 512 x n slice instead of n x n. At 128K that slice is still 8 GiB for one layer, which is more than the weights of the 8B model in 4-bit form. During decode there is one query per step, so the score row is only 16 MiB per layer at 128K. That is small, but it is still written and read again for every generated token.
Each score is used exactly twice, once in the softmax and once in the weighted sum, and then thrown away. Storing it is therefore wasted work. The next three examples remove the need to store it.
Example 2: online softmax in one pass
Softmax turns scores x_1 ... x_n into probabilities with the formula p_i = exp(x_i) / sum_j exp(x_j). Computers cannot use this form directly, because exp of a large number overflows. The fix is to subtract the largest score m first: p_i = exp(x_i - m) / sum_j exp(x_j - m). The result is the same, because the factor exp(-m) appears in the top and the bottom and cancels. After the subtraction, every exponent is zero or negative, so every exp is between 0 and 1.
A small case with two scores shows the steps:
x = {1, 3} largest value m = 3
exp(1 - 3) = 0.1353 exp(3 - 3) = 1.0000
sum = 1.1353
p = {0.1353 / 1.1353, 1.0000 / 1.1353} = {0.1192, 0.8808}The standard method needs two passes over the row. The first pass finds the maximum. The second pass computes the sum. This is a problem for tiled attention, which sees the scores in pieces and cannot wait until it has seen all of them to know the maximum.
Online softmax solves this by computing both numbers in one pass. It keeps a running maximum m and a running sum s. When a new value raises the maximum, the old sum was computed relative to the old maximum, so it is multiplied by exp(m_old - m_new) to move it to the new reference point. The program compares both methods on six values and also shows what happens without the subtraction.
import java.util.Locale;
public class OnlineSoftmax {
public static void main(String[] args) {
double[] x = {1.0, 3.0, 2.0, 5.0, 0.5, 4.0};
double max = Double.NEGATIVE_INFINITY; // pass 1: the maximum
for (double v : x) max = Math.max(max, v);
double sum = 0; // pass 2: sum of exp(x - max)
for (double v : x) sum += Math.exp(v - max);
System.out.printf(Locale.ROOT, "two-pass: max %.1f, sum %.6f%n", max, sum);
double m = Double.NEGATIVE_INFINITY, s = 0; // one pass: running max and sum
for (int i = 0; i < x.length; i++) {
double mNew = Math.max(m, x[i]);
s = s * Math.exp(m - mNew) + Math.exp(x[i] - mNew); // rescale the old sum to the new max
m = mNew;
System.out.printf(Locale.ROOT, "step %d: x=%.1f m=%.1f s=%.6f%n", i, x[i], m, s);
}
for (int i = 0; i < x.length; i++)
System.out.printf(Locale.ROOT, "p[%d] two-pass %.6f one-pass %.6f%n",
i, Math.exp(x[i] - max) / sum, Math.exp(x[i] - m) / s);
System.out.println("exp(1000) without subtracting the max: " + Math.exp(1000));
}
}two-pass: max 5.0, sum 1.582426
step 0: x=1.0 m=1.0 s=1.000000
step 1: x=3.0 m=3.0 s=1.135335
step 2: x=2.0 m=3.0 s=1.503215
step 3: x=5.0 m=5.0 s=1.203438
step 4: x=0.5 m=5.0 s=1.214547
step 5: x=4.0 m=5.0 s=1.582426
p[0] two-pass 0.011574 one-pass 0.011574
p[1] two-pass 0.085524 one-pass 0.085524
p[2] two-pass 0.031462 one-pass 0.031462
p[3] two-pass 0.631941 one-pass 0.631941
p[4] two-pass 0.007020 one-pass 0.007020
p[5] two-pass 0.232478 one-pass 0.232478
exp(1000) without subtracting the max: InfinityThe two-pass method finds the maximum 5.0, then the sum 1.582426 of exp(x - 5). The one-pass steps show how the sum is rescaled each time the maximum grows:
- Step 0:
m = 1,s = exp(1 - 1) = 1. - Step 1:
x = 3raises the maximum. The old sum, computed relative to 1, is multiplied byexp(1 - 3) = 0.135335, and thenexp(3 - 3) = 1is added:s = 1.135335. - Step 2:
x = 2does not change the maximum, soexp(2 - 3) = 0.367879is added:s = 1.503215. - Step 3:
x = 5raises the maximum again. The sum is multiplied byexp(3 - 5) = 0.135335and 1 is added:s = 1.203438. - Steps 4 and 5 add
exp(0.5 - 5)andexp(4 - 5)and reach 1.582426, the same sum as the two-pass method.
The six probabilities are identical to six decimals. The reason is that s always equals the sum of exp(x_j - m) over the values seen so far, for the current m. Rescaling only changes the reference point, and the identity exp(a - b) * exp(b - c) = exp(a - c) makes the change exact, up to floating-point rounding.
The last line shows why the maximum is subtracted at all. exp(1000) is too large for a double, so Java returns Infinity. Attention scores stored in low-precision formats overflow at much smaller values.
The paper "Self-attention Does Not Need O(n^2) Memory" (Rabe and Staats, 2021) and the FlashAttention paper (Dao et al., 2022) both use this rescaling to compute attention in chunks. llama.cpp's CPU flash-attention kernel (ggml_compute_forward_flash_attn_ext_f16 in ggml/src/ggml-cpu/ops.cpp) cites the first paper in a comment that reads // online softmax / attention.
Example 3: the online attention row with a value accumulator
Attention needs more than the probabilities. The output is the weighted sum of the value vectors, so the program must also add up probability x value for each key. The same rescaling idea works for this sum. The program keeps an accumulator acc, a running weighted sum of value vectors, next to m and s. Whenever the maximum grows, it multiplies acc by the same factor as s. It divides by s only once, at the very end. The scores are the six values of Example 2, and each key has a 2-dimensional value vector. A plain reference computation (full softmax first, then the weighted sum) runs alongside to check the answer.
import java.util.Locale;
public class OnlineAttentionRow {
public static void main(String[] args) {
double[] scores = {1.0, 3.0, 2.0, 5.0, 0.5, 4.0}; // q . k_j / sqrt(d) for 6 keys
double[][] values = {{1, 0}, {0, 1}, {1, 1}, {2, 0}, {0, 2}, {1, -1}};
double max = Double.NEGATIVE_INFINITY, sum = 0; // reference: full softmax, then weighted sum
for (double s : scores) max = Math.max(max, s);
for (double s : scores) sum += Math.exp(s - max);
double[] ref = new double[2];
for (int j = 0; j < scores.length; j++)
for (int c = 0; c < 2; c++) ref[c] += Math.exp(scores[j] - max) / sum * values[j][c];
double m = Double.NEGATIVE_INFINITY, s = 0; // streaming: one key at a time
double[] acc = new double[2];
for (int j = 0; j < scores.length; j++) {
double mNew = Math.max(m, scores[j]);
double scaleOld = Math.exp(m - mNew), w = Math.exp(scores[j] - mNew);
for (int c = 0; c < 2; c++) acc[c] = acc[c] * scaleOld + w * values[j][c];
s = s * scaleOld + w;
m = mNew;
System.out.printf(Locale.ROOT, "key %d: m=%.1f s=%.4f acc=(%.4f, %.4f)%n", j, m, s, acc[0], acc[1]);
}
System.out.printf(Locale.ROOT, "streaming output (%.6f, %.6f)%n", acc[0] / s, acc[1] / s);
System.out.printf(Locale.ROOT, "reference output (%.6f, %.6f)%n", ref[0], ref[1]);
}
}key 0: m=1.0 s=1.0000 acc=(1.0000, 0.0000)
key 1: m=3.0 s=1.1353 acc=(0.1353, 1.0000)
key 2: m=3.0 s=1.5032 acc=(0.5032, 1.3679)
key 3: m=5.0 s=1.2034 acc=(2.0681, 0.1851)
key 4: m=5.0 s=1.2145 acc=(2.0681, 0.2073)
key 5: m=5.0 s=1.5824 acc=(2.4360, -0.1605)
streaming output (1.539397, -0.101451)
reference output (1.539397, -0.101451)Each line of the output shows the state after one key. Key 0 gives weight 1 to its value (1, 0). Key 1 raises the maximum from 1 to 3, so the accumulator (1, 0) and the sum 1 are both multiplied by exp(-2) = 0.1353. The new value (0, 1) is then added with weight 1, which gives acc = (0.1353, 1.0000) and s = 1.1353. Key 3 raises the maximum to 5, and the whole accumulator shrinks by exp(-2) again before (2, 0) is added with weight 1. This is why the second coordinate drops from 1.3679 to 0.1851.
After the last key, acc / s = (1.539397, -0.101451), which equals the reference that computed all six probabilities first. The streaming version never held more than one score at a time. Its state per query is only m, s, and one vector of the value dimension. This loop is the core of every flash-attention kernel. In llama.cpp's CPU kernel the three variables are named M, S, and VKQ32, and the two branches are commented // V = V*expf(Mold - M) and // V += v*expf(s - M).
The final division by s happens once per query, after all keys. Dividing earlier would force the code to undo the division at every rescale.
Example 4: tiled attention: the FlashAttention algorithm in Java
A GPU has two kinds of memory. The large one is HBM (high bandwidth memory), the GPU's main memory, which holds many gigabytes but is slow compared with the speed of the arithmetic units. The small one is on the chip itself: shared memory, also called SRAM, which is very fast but tiny. NVIDIA documents up to 164 KB per streaming multiprocessor (one compute unit of the GPU) on the A100 and 228 KB on the H100.
A workbench and a warehouse describe the difference. The workbench (SRAM) is small and everything on it is within arm's reach. The warehouse (HBM) holds everything but each trip takes time. The analogy fails in one respect: a GPU runs thousands of workers at once, so the cost is the total number of trips, not the time of one trip. The goal is to carry each item to the workbench as few times as possible.
Tiled attention does this. It loads a block of queries onto the workbench, then streams blocks of keys and values past it. For each query it updates m, s, and acc with the rule from Example 3. When all key blocks have passed, it writes only the finished output block back to the warehouse. The score matrix exists only as small pieces inside fast memory and is never written out. This is the FlashAttention algorithm (Dao et al., 2022).
for each block of queries Qi (loaded once into fast memory)
m, s, acc = start values
for each block of keys/values Kj, Vj (streamed past Qi)
scores = Qi * Kj^T (a small tile, stays in fast memory)
update m, s, acc with the online rule
write Oi = acc / s (the only write to slow memory)The program runs naive and tiled attention on 64 tokens with head dimension 8 and blocks of 16. It counts the floats each method moves to and from slow memory. It then applies the same counting formulas to real sizes.
import java.util.Locale;
import java.util.Random;
public class TiledAttention {
static long reads, writes; // floats moved to and from slow memory
static float[][] rand(Random r, int n, int d) {
float[][] a = new float[n][d];
for (float[] row : a) for (int i = 0; i < d; i++) row[i] = (float) r.nextGaussian();
return a;
}
static float dot(float[] a, float[] b) { float s = 0; for (int i = 0; i < a.length; i++) s += a[i] * b[i]; return s; }
static float[][] naive(float[][] q, float[][] k, float[][] v) {
int n = q.length, d = q[0].length;
float scale = (float) (1 / Math.sqrt(d));
float[][] s = new float[n][n]; // the full score matrix
for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) s[i][j] = dot(q[i], k[j]) * scale;
reads += 2L * n * d; writes += (long) n * n; // read Q and K, write S
for (float[] row : s) { // softmax in place: read S, write P
float mx = Float.NEGATIVE_INFINITY, sum = 0;
for (float x : row) mx = Math.max(mx, x);
for (int j = 0; j < n; j++) { row[j] = (float) Math.exp(row[j] - mx); sum += row[j]; }
for (int j = 0; j < n; j++) row[j] /= sum;
}
reads += (long) n * n; writes += (long) n * n;
float[][] o = new float[n][d];
for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) for (int c = 0; c < d; c++) o[i][c] += s[i][j] * v[j][c];
reads += (long) n * n + (long) n * d; writes += (long) n * d; // read P and V, write O
return o;
}
static float[][] tiled(float[][] q, float[][] k, float[][] v, int tile) {
int n = q.length, d = q[0].length;
float scale = (float) (1 / Math.sqrt(d));
float[][] o = new float[n][d];
for (int i0 = 0; i0 < n; i0 += tile) {
reads += (long) tile * d; // load a block of queries
float[] m = new float[tile], l = new float[tile];
java.util.Arrays.fill(m, Float.NEGATIVE_INFINITY);
float[][] acc = new float[tile][d];
for (int j0 = 0; j0 < n; j0 += tile) {
reads += 2L * tile * d; // load a block of keys and values
for (int i = 0; i < tile; i++) for (int j = 0; j < tile; j++) {
float x = dot(q[i0 + i], k[j0 + j]) * scale;
float mNew = Math.max(m[i], x);
float old = (float) Math.exp(m[i] - mNew), w = (float) Math.exp(x - mNew);
for (int c = 0; c < d; c++) acc[i][c] = acc[i][c] * old + w * v[j0 + j][c];
l[i] = l[i] * old + w;
m[i] = mNew;
}
}
for (int i = 0; i < tile; i++) for (int c = 0; c < d; c++) o[i0 + i][c] = acc[i][c] / l[i];
writes += (long) tile * d; // store the finished output block
}
return o;
}
public static void main(String[] args) {
int n = 64, d = 8, tile = 16;
Random r = new Random(42);
float[][] q = rand(r, n, d), k = rand(r, n, d), v = rand(r, n, d);
float[][] a = naive(q, k, v);
System.out.println("naive: reads " + reads + ", writes " + writes + ", total " + (reads + writes));
reads = writes = 0;
float[][] b = tiled(q, k, v, tile);
System.out.println("tiled: reads " + reads + ", writes " + writes + ", total " + (reads + writes));
float maxDiff = 0;
for (int i = 0; i < n; i++) for (int c = 0; c < d; c++) maxDiff = Math.max(maxDiff, Math.abs(a[i][c] - b[i][c]));
System.out.printf(Locale.ROOT, "max |naive - tiled| = %.1e%n", maxDiff);
System.out.println("per head, d = 128 (floats moved, millions):");
long dd = 128;
for (long[] c : new long[][] {{4096, 128}, {131072, 128}, {131072, 64}, {131072, 256}}) {
long len = c[0], br = c[1];
long naiveFloats = 3 * len * dd + 2 * len * len + 2 * len * len + len * dd;
long tiledFloats = len * dd + (len / br) * 2 * len * dd + len * dd;
System.out.printf(Locale.ROOT, "n=%6d block %3d: naive %8.0f tiled %8.0f ratio %.1f score buffer %,d floats, tile %,d%n",
len, br, naiveFloats / 1e6, tiledFloats / 1e6, (double) naiveFloats / tiledFloats, len * len, br * br);
}
}
}naive: reads 9728, writes 8704, total 18432
tiled: reads 4608, writes 512, total 5120
max |naive - tiled| = 3.0e-07
per head, d = 128 (floats moved, millions):
n= 4096 block 128: naive 69 tiled 35 ratio 2.0 score buffer 16,777,216 floats, tile 16,384
n=131072 block 128: naive 68787 tiled 34393 ratio 2.0 score buffer 17,179,869,184 floats, tile 16,384
n=131072 block 64: naive 68787 tiled 68753 ratio 1.0 score buffer 17,179,869,184 floats, tile 4,096
n=131072 block 256: naive 68787 tiled 17213 ratio 4.0 score buffer 17,179,869,184 floats, tile 65,536Both methods produce the same output. The largest difference over all 64 x 8 outputs is 3.0e-07, which is the size of float rounding error. The order of additions differs between the methods, so the last bits differ.
The counts for 64 tokens follow from the code:
- Naive reads Q and K (2 x 512 floats), reads
Sfor the softmax (4,096), and readsPand V (4,096 + 512). That is 9,728 reads. It writesS,P, and O: 4,096 + 4,096 + 512 = 8,704 writes. - Tiled reads each query block once (512 floats in total) and, for each of the 4 query blocks, reads all keys and values (4 x 1,024). That is 4,608 reads. It writes only O: 512 floats.
The score and probability matrices never leave fast memory in the tiled version, so writes drop from 8,704 to 512. The price is that keys and values are read once per query block instead of once overall. Total traffic still falls from 18,432 to 5,120 floats.
The bottom lines apply the formulas to one head with d = 128. Naive traffic is about 4 * n^2, which is two writes and two reads of an n x n matrix. Tiled traffic is about 2 * n^2 * d / B, where B is the query block size. The reason is that each of the n / B query blocks reads all n keys and values, and each key-value pair has 2d floats. Dividing the two gives the ratio 4 n^2 / (2 n^2 d / B) = 2B / d. With d = 128, blocks of 128 give a ratio of 2.0, blocks of 64 give 1.0, and blocks of 256 give 4.0. For n = 4096 and B = 128, the naive count is about 69 million floats and the tiled count about 35 million, as the table shows.
Larger query blocks need more fast memory, which is why FlashAttention's tile sizes depend on the GPU's shared memory. The FlashAttention paper states the number of HBM accesses as O(n^2 d^2 / M) for fast memory of size M, against O(n d + n^2) for standard attention.
The column that changes the most is memory size. Naive attention needs a buffer of n^2 floats per head (17 billion at 128K). Tiled attention needs a tile of B^2 floats (16,384 for B = 128) regardless of n. The tiled version is also one fused kernel instead of three separate ones, so it avoids the launch cost and the intermediate writes between them. With a causal mask (the rule that a token cannot look at later tokens), key blocks that lie entirely after a query block are skipped, which removes almost half of the blocks in prefill. The program leaves the mask out for brevity.
Example 5: flash attention in llama.cpp and vLLM
llama.cpp implements flash attention as one fused graph node, FLASH_ATTN_EXT. A graph node is one operation in the list of computations the engine runs for each step. The switch is -fa on|off|auto (Tuning llama.cpp, Example 1). The question here is how much memory the score matrix costs when attention is not fused. The commands below run the tiny model four times at a 65,536-token context, with micro-batches of 128 and 512 tokens, and with flash attention on and off. The grep keeps only the lines that show the settings and the compute buffer size.
for fa in on off; do for ub in 128 512; do
llama-completion -m tiny-random-qwen2vocab-f32.gguf -p hello -n 2 -no-cnv -lv 4 -c 65536 -ub "$ub" -fa "$fa" 2>&1 \
| grep -E 'n_ubatch|flash_attn |compute buffer size'
done; doneThe output below comes from the pinned build, with timestamps removed. Blank lines separate the four runs. In each run, look at the flash_attn setting and at the CPU compute buffer size line.
I llama_context: n_ubatch = 128
I llama_context: flash_attn = enabled
I sched_reserve: CPU compute buffer size = 74.44 MiB
I llama_context: n_ubatch = 512
I llama_context: flash_attn = enabled
I sched_reserve: CPU compute buffer size = 297.75 MiB
I llama_context: n_ubatch = 128
I llama_context: flash_attn = disabled
I sched_reserve: CPU compute buffer size = 288.69 MiB
I llama_context: n_ubatch = 512
I llama_context: flash_attn = disabled
I sched_reserve: CPU compute buffer size = 1154.76 MiBThe compute buffer is the memory that holds the temporary tensors of one graph evaluation, such as the score matrix. The next program computes the two largest temporaries for each micro-batch: the F32 (32-bit float) score matrix, ub x 65536 x 8 heads x 4 bytes, and the F32 logits, ub x 151936 x 4 bytes. Logits are the final scores over the vocabulary, one number for each of the 151,936 possible next tokens. The program compares both sizes with the logged buffer sizes.
import java.util.Locale;
public class ScoreBuffer {
public static void main(String[] args) {
long ctx = 65536, heads = 8, vocab = 151936; // the tiny model at -c 65536
double mib = 1024 * 1024;
double[][] logged = {{128, 74.44, 288.69}, {512, 297.75, 1154.76}}; // ubatch, fa on, fa off
for (double[] row : logged) {
long ub = (long) row[0];
double scores = ub * ctx * heads * 4 / mib; // F32 score matrix of one micro-batch
double logits = ub * vocab * 4 / mib; // F32 logits of one micro-batch
System.out.printf(Locale.ROOT, "-ub %3d: scores %7.2f MiB, logits %6.2f MiB | log: fa on %7.2f, fa off %7.2f | off - scores %6.2f%n",
ub, scores, logits, row[1], row[2], row[2] - scores);
}
}
}-ub 128: scores 256.00 MiB, logits 74.19 MiB | log: fa on 74.44, fa off 288.69 | off - scores 32.69
-ub 512: scores 1024.00 MiB, logits 296.75 MiB | log: fa on 297.75, fa off 1154.76 | off - scores 130.76With flash attention on, the compute buffer is the logits buffer plus less than 1.1 MiB: 74.44 MiB against 74.19 MiB of logits for -ub 128, and 297.75 MiB against 296.75 MiB for -ub 512. The fused node needs only a small workspace for each thread, so the largest tensor in the graph is the output projection to the 151,936-token vocabulary.
With flash attention off, the score matrix of one micro-batch must exist in memory all at once. That is 256 MiB for 128 queries against 65,536 keys, and 1,024 MiB for 512 queries. The logged buffers are larger than the score matrix by 32.69 MiB and 130.76 MiB. The extra memory belongs to other tensors that are alive at the same time, for example the permuted copy of V that the non-fused path makes. The effect is clear: at -ub 512, turning flash attention off takes the buffer from 297.75 MiB to 1,154.76 MiB, nearly four times more.
vLLM computes attention in an attention backend, a class that wraps one family of kernels. The backends are listed in AttentionBackendEnum (vllm/v1/attention/backends/registry.py). The flag --attention-backend NAME forces one. Without it, the platform picks the first valid backend from a priority list. For NVIDIA GPUs, _get_backend_priorities in vllm/platforms/cuda.py gives these orders for models without MLA (MLA is covered in Example 7). The compute capability is a version number that tells which features an NVIDIA GPU supports.
| Compute capability | Priority order |
|---|---|
| 10.x (Blackwell data-center GPUs) | FLASHINFER, FLASH_ATTN, TRITON_ATTN, FLEX_ATTENTION, TURBOQUANT |
| Others | FLASH_ATTN, FLASHINFER, TRITON_ATTN, FLEX_ATTENTION, TURBOQUANT |
Each backend reports what it supports, and the selector skips backends that cannot run the model. FlashAttentionBackend (flash_attn.py) requires compute capability 8.0 or higher, and a head size that is a multiple of 8 and at most 256 (512 with FlashAttention 4). Models with learned sinks (Example 11) need compute capability 9.0. MLA models (Example 7) have their own list (FLASH_ATTN_MLA, FLASHMLA, FLASHINFER_MLA, TRITON_MLA, and others), and the CPU platform uses CPU_ATTN. The selection is logged with a format string from cuda.py. The line below fills it in for an illustrative case and was not captured from a run.
Using FLASH_ATTN attention backend out of potential backends: ['FLASH_ATTN', 'FLASHINFER', 'TRITON_ATTN', 'FLEX_ATTENTION'].Example 6: MQA and GQA: KV bytes and the head mapping
The KV cache stores one key vector and one value vector per token, per layer, per KV head. In the original design, called multi-head attention (MHA), every query head has its own KV head. That is a lot of memory. Self-Attention from Scratch introduced grouped-query attention (GQA): several query heads share one key head and one value head, so the model stores fewer KV heads. Multi-query attention (MQA) is the extreme case, with one KV head shared by all query heads.
The cache stores only the KV heads, so it shrinks by the group size (the number of query heads per KV head). The program does two things. First, it prints which KV head each query head of the tiny model reads, for 8 query heads and 2 KV heads. Second, it prints the KV bytes per token of a Llama 3 8B-shaped model (32 layers, head dimension 128, F16 cache, which means 2 bytes per number) with 32, 8, and 1 KV heads.
import java.util.Locale;
public class GqaBytes {
static long kvBytesPerToken(int layers, int kvHeads, int headDim, int bytes) {
return 2L * layers * kvHeads * headDim * bytes; // K and V
}
public static void main(String[] args) {
int heads = 8, kvHeads = 2; // the tiny model
int group = heads / kvHeads;
StringBuilder map = new StringBuilder("tiny model, query head -> KV head:");
for (int h = 0; h < heads; h++) map.append(" ").append(h).append("->").append(h / group);
System.out.println(map);
int layers = 32, headDim = 128, ctx = 32768; // Llama 3 8B shape, F16 cache
for (int kv : new int[] {32, 8, 1}) {
long perToken = kvBytesPerToken(layers, kv, headDim, 2);
String name = kv == 32 ? "MHA" : kv == 1 ? "MQA" : "GQA";
System.out.printf(Locale.ROOT, "%s (%2d KV heads): %6d bytes/token, %5.2f GiB at %d tokens, %2d query heads per KV head%n",
name, kv, perToken, perToken * (double) ctx / (1 << 30), ctx, 32 / kv);
}
}
}tiny model, query head -> KV head: 0->0 1->0 2->0 3->0 4->1 5->1 6->1 7->1
MHA (32 KV heads): 524288 bytes/token, 16.00 GiB at 32768 tokens, 1 query heads per KV head
GQA ( 8 KV heads): 131072 bytes/token, 4.00 GiB at 32768 tokens, 4 query heads per KV head
MQA ( 1 KV heads): 16384 bytes/token, 0.50 GiB at 32768 tokens, 32 query heads per KV headThe first line is the mapping kvHead = queryHead / (heads / kvHeads). Here the group size is 8 / 2 = 4, so query heads 0 to 3 read KV head 0, and query heads 4 to 7 read KV head 1. ggml's flash-attention kernel computes the same index (ik2 = iq2 / rk2, where rk2 is the ratio of query heads to KV heads). A GQA model therefore needs no copy of the keys for each query head.
The bytes per token follow the formula 2 x layers x kvHeads x headDim x 2 bytes, where the first 2 counts K and V and the last 2 is the bytes per F16 number. For MHA with 32 KV heads: 2 x 32 x 32 x 128 x 2 = 524,288 bytes, which is 512 KiB per token and 16 GiB for 32,768 tokens. Llama 3 8B's actual layout has 8 KV heads and needs a quarter of that: 128 KiB per token, 4 GiB at 32K. MQA would need 16 KiB per token, 0.5 GiB at 32K.
Fewer KV heads also mean fewer bytes read per decode step, because every step reads the whole cache (FLOPs, Bandwidth, and the Roofline Model, Example 8). The cost is capacity for information: with MQA, all 32 query heads must find what they need in the same keys and values. The GQA paper (Ainslie et al., 2023) reports quality close to multi-head attention with speed comparable to MQA. Llama 3 8B uses 4 query heads per KV head and Llama 3 70B uses 8.
Example 7: MLA: a compressed latent instead of keys and values
GQA shrinks the cache by sharing heads. Multi-head latent attention (MLA) takes a different route: it compresses. It was introduced in the DeepSeek-V2 paper (2024) and is used in DeepSeek-V3. Instead of per-head keys and values, MLA stores one small latent vector c per token and layer. A latent vector is a short list of numbers from which the keys and values can be rebuilt. A down-projection matrix W_DKV maps the hidden state (the token's vector inside the model) to c. Per-head up-projection matrices W_UK (for keys) and W_UV (for values) recreate keys and values from c when needed.
A compression analogy: a zip file is small and the original can be rebuilt from it. The mapping is that c is the zip file and W_UK and W_UV are the unzip step. Unlike a zip file, the compression here is lossy in general (the 2-number latent holds less information than the 12 original numbers), and the model is trained to work with it.
Rebuilding every key from c at every step would cost extra work. Implementations avoid this by absorbing W_UK into the query. The identity is q . (W_UK c) = (W_UK^T q) . c, where . is the dot product and ^T is the transpose of a matrix. The left side decompresses the key and then scores it. The right side transforms the query once and scores it directly against the cached latent, so keys never need to be rebuilt. The program checks this identity with a hidden size of 8, a 2-dimensional latent, and two heads of dimension 3.
MHA caches per token: K head 0 | V head 0 | K head 1 | V head 1 | ... (many numbers)
MLA caches per token: c (one short vector)
score for head h: q_h . (W_UK_h c) = (W_UK_h^T q_h) . cimport java.util.Locale;
import java.util.Random;
public class LatentAttention {
static double[] mul(double[][] w, double[] x) { // w is rows x cols, x has cols entries
double[] y = new double[w.length];
for (int r = 0; r < w.length; r++) for (int c = 0; c < x.length; c++) y[r] += w[r][c] * x[c];
return y;
}
static double[] mulT(double[][] w, double[] x) { // w transposed times x
double[] y = new double[w[0].length];
for (int r = 0; r < w.length; r++) for (int c = 0; c < y.length; c++) y[c] += w[r][c] * x[r];
return y;
}
static double dot(double[] a, double[] b) { double s = 0; for (int i = 0; i < a.length; i++) s += a[i] * b[i]; return s; }
static double[][] rand(Random r, int rows, int cols) {
double[][] w = new double[rows][cols];
for (double[] row : w) for (int c = 0; c < cols; c++) row[c] = Math.round(r.nextGaussian() * 10) / 10.0;
return w;
}
public static void main(String[] args) {
Random r = new Random(7);
int dModel = 8, latent = 2, headDim = 3, heads = 2;
double[][] down = rand(r, latent, dModel); // W_DKV: hidden state -> latent c
double[][][] upK = {rand(r, headDim, latent), rand(r, headDim, latent)}; // W_UK per head
double[] hidden = rand(r, 1, dModel)[0]; // the token being cached
double[][] q = {rand(r, 1, headDim)[0], rand(r, 1, headDim)[0]}; // a later query, per head
double[] c = mul(down, hidden); // the only thing MLA caches
System.out.printf(Locale.ROOT, "cached latent c = (%.2f, %.2f)%n", c[0], c[1]);
for (int h = 0; h < heads; h++) {
double[] k = mul(upK[h], c); // decompress the key, then score
double direct = dot(q[h], k);
double[] qAbsorbed = mulT(upK[h], q[h]); // fold W_UK into the query instead
double absorbed = dot(qAbsorbed, c);
System.out.printf(Locale.ROOT, "head %d: q.(W_UK c) = %.4f (W_UK^T q).c = %.4f%n", h, direct, absorbed);
}
System.out.println("floats cached per token: MHA K and V = " + heads * headDim * 2 + ", MLA latent = " + latent);
}
}cached latent c = (4.62, 3.03)
head 0: q.(W_UK c) = -7.2519 (W_UK^T q).c = -7.2519
head 1: q.(W_UK c) = 9.8355 (W_UK^T q).c = 9.8355
floats cached per token: MHA K and V = 12, MLA latent = 2The cached latent for the token is (4.62, 3.03): two numbers instead of the 12 that two heads of keys and values (3 + 3 per head, times 2 heads) would need. For each head, the score computed by first decompressing the key (q . (W_UK c)) equals the score computed with the absorbed query ((W_UK^T q) . c): -7.2519 for head 0 and 9.8355 for head 1. The values work the same way: W_UV can be applied after the attention-weighted sum of latents, or folded into the output projection.
With absorption, every head scores against the same cached vector, so MLA behaves like MQA with one wide KV head. The llama.cpp source says so in src/models/deepseek2.cpp: "MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)". vLLM's MLAAttentionSpec (vllm/v1/kv_cache_interface.py) states the storage rule: "MLA stores a single latent vector per state; there is no separate V."
One part does not fit this trick. RoPE rotates keys by an angle that depends on the position, and a position-dependent rotation sitting between W_UK and c cannot be moved into the query once. DeepSeek's design solves this by adding a small separate key part that carries RoPE (the decoupled RoPE key). All heads share it, and it is cached next to the latent. The next example adds its size.
Example 8: MLA bytes per token for a DeepSeek-V3 config
The next program uses the real numbers of a large model. The DeepSeek-V3 config.json published with the model lists num_hidden_layers 61, num_attention_heads 128, qk_nope_head_dim 128, qk_rope_head_dim 64, v_head_dim 128, and kv_lora_rank 512 (the latent size). The program computes the cache per token for MLA, which is the latent plus the decoupled RoPE key, 512 + 64 values per layer. It also computes the cache for hypothetical MHA and GQA versions of the same model, with keys of 128 + 64 dimensions and values of 128 dimensions, so the saving can be compared.
import java.util.Locale;
public class MlaBytes {
public static void main(String[] args) {
// DeepSeek-V3 published config.json values
int layers = 61, heads = 128, qkNope = 128, qkRope = 64, vDim = 128, kvLoraRank = 512;
int bytes = 2; // BF16 or F16 cache
long mla = (long) layers * (kvLoraRank + qkRope) * bytes;
long mha = (long) layers * heads * ((qkNope + qkRope) + vDim) * bytes;
long gqa8 = (long) layers * 8 * ((qkNope + qkRope) + vDim) * bytes;
long llama8b = 2L * 32 * 8 * 128 * bytes;
long ctx = 131072;
System.out.printf(Locale.ROOT, "MLA latent + RoPE key: %9d bytes/token %8.2f GiB at 128K%n", mla, mla * (double) ctx / (1 << 30));
System.out.printf(Locale.ROOT, "same model as MHA: %9d bytes/token %8.2f GiB at 128K%n", mha, mha * (double) ctx / (1 << 30));
System.out.printf(Locale.ROOT, "same model as GQA-8: %9d bytes/token %8.2f GiB at 128K%n", gqa8, gqa8 * (double) ctx / (1 << 30));
System.out.printf(Locale.ROOT, "Llama 3 8B (GQA-8): %9d bytes/token %8.2f GiB at 128K%n", llama8b, llama8b * (double) ctx / (1 << 30));
System.out.printf(Locale.ROOT, "MHA / MLA = %.1f%n", (double) mha / mla);
}
}MLA latent + RoPE key: 70272 bytes/token 8.58 GiB at 128K
same model as MHA: 4997120 bytes/token 610.00 GiB at 128K
same model as GQA-8: 312320 bytes/token 38.13 GiB at 128K
Llama 3 8B (GQA-8): 131072 bytes/token 16.00 GiB at 128K
MHA / MLA = 71.1MLA stores 576 values per layer, 61 layers, 2 bytes each: 576 x 61 x 2 = 70,272 bytes per token, and 8.58 GiB for a 131,072-token context. The same model with full multi-head attention would store 128 heads of 192-dimensional keys and 128-dimensional values. That is about 4.8 MiB per token and 610 GiB at 128K, which is 71.1 times more. Even a GQA version with 8 KV heads would need 38.13 GiB at 128K.
For comparison, DeepSeek-V3 has about 671 billion parameters (its published model card), yet its KV cache per token is about half that of the much smaller Llama 3 8B (70,272 against 131,072 bytes).
llama.cpp follows the same layout when it runs the absorbed path. src/llama-kv-cache.cpp allocates only a K tensor for MLA models (the source sets has_v to false for MLA), and its row is the latent plus the RoPE part. The V input to attention is a view of the latent, not a separate copy. vLLM's MLA layer (mla_attention.py) sets its cache head size to kv_lora_rank + qk_rope_head_dim, which is one 576-wide entry per token per layer. A K cache in a quantized type (Example 2 of Tuning llama.cpp) shrinks this further.
Example 9: sliding-window attention: the mask and a ring buffer
The previous techniques make each token's cache entry smaller. A different idea is to keep fewer tokens. In sliding-window attention (SWA), a layer lets each query attend only to the last W positions, itself included. W is the window size. The causal mask from Self-Attention from Scratch gains a second condition: a key at position k is visible to a query at position q only if k <= q and q - k < W.
This has a direct consequence for memory. No later query looks further back than W positions, so the cache for such a layer needs to hold only W entries. It can be a ring buffer, a fixed array used in a circle: position p goes to slot p mod W (the remainder after dividing by W), which overwrites the oldest entry. For example, with W = 4, position 4 goes to slot 0 and replaces position 0. The program prints the mask for 10 tokens with W = 4, and the ring buffer after each token from position 3 on.
public class SlidingWindow {
public static void main(String[] args) {
int n = 10, window = 4;
System.out.println("mask (q = query position, x = may attend, . = masked):");
for (int q = 0; q < n; q++) {
StringBuilder row = new StringBuilder(String.format("q%-2d ", q));
for (int k = 0; k < n; k++) row.append(k <= q && q - k < window ? 'x' : '.');
System.out.println(row);
}
int[] slots = new int[window]; // ring buffer of positions
java.util.Arrays.fill(slots, -1);
for (int pos = 0; pos < n; pos++) {
int slot = pos % window;
int evicted = slots[slot];
slots[slot] = pos;
if (pos >= 3) System.out.println("after pos " + pos + ": slot " + slot + " gets " + pos
+ (evicted >= 0 ? " (evicts " + evicted + ")" : "") + ", cache = " + java.util.Arrays.toString(slots));
}
}
}mask (q = query position, x = may attend, . = masked):
q0 x.........
q1 xx........
q2 xxx.......
q3 xxxx......
q4 .xxxx.....
q5 ..xxxx....
q6 ...xxxx...
q7 ....xxxx..
q8 .....xxxx.
q9 ......xxxx
after pos 3: slot 3 gets 3, cache = [0, 1, 2, 3]
after pos 4: slot 0 gets 4 (evicts 0), cache = [4, 1, 2, 3]
after pos 5: slot 1 gets 5 (evicts 1), cache = [4, 5, 2, 3]
after pos 6: slot 2 gets 6 (evicts 2), cache = [4, 5, 6, 3]
after pos 7: slot 3 gets 7 (evicts 3), cache = [4, 5, 6, 7]
after pos 8: slot 0 gets 8 (evicts 4), cache = [8, 5, 6, 7]
after pos 9: slot 1 gets 9 (evicts 5), cache = [8, 9, 6, 7]Read the mask one row at a time. Each row is a query, and the x marks show which positions it may attend to. The mask has at most 4 x marks per row. Rows q0 to q3 look like the plain causal mask, because fewer than 4 earlier positions exist. From q4 on, the band moves to the right: q4 sees positions 1 to 4, and q9 sees positions 6 to 9.
The ring buffer holds exactly the positions the next query needs. After position 4, slot 0 holds 4 instead of 0, and the cache is [4, 1, 2, 3]. That is the window of query 4 in a rotated order. The order of keys in memory does not matter to attention, because softmax and the weighted sum are sums over a set. RoPE, however, needs each key's original position, so the cache stores a position with each entry. llama.cpp keeps a position per KV cell (pos) for this reason.
Memory per SWA layer is constant: W entries, however long the conversation grows. The cost is information. A token that has left the window is gone from this layer and cannot be attended to again, even if a later question is about it.
Stacked layers widen the reach. Information from position 0 can move into position 3's hidden state in layer 1, then into position 6's in layer 2, and so on. So L stacked layers of window W can carry information about L x (W - 1) positions back. The next example shows the design that most current SWA models use instead of relying on this indirect reach.
Example 10: interleaved local and global layers, and --swa-full
Models such as Gemma 2, Gemma 3, and gpt-oss mix local layers (sliding-window) with global layers (full attention over the whole context). The local layers keep memory small. The global layers let information travel the whole distance directly. In llama.cpp, the function set_swa_pattern(n) in src/llama-hparams.cpp marks layer il as local when il % n < n - 1. With n = 6, layers 0 to 4 are local, layer 5 is global, layers 6 to 10 are local, layer 11 is global, and so on. src/models/gemma3.cpp loads the pattern with a default of 6, and the Gemma 2 and gpt-oss files use 2. The local layers get a small cache whose size llama-kv-cache-iswa.cpp computes as GGML_PAD(min(n_ctx, n_swa * n_seq + n_ubatch), 256) cells (Tuning llama.cpp, Example 9). In that formula, n_ctx is the context size, n_swa is the window, n_seq is the number of parallel sequences, and GGML_PAD rounds up to a multiple of 256.
The course's tiny model has no window. For this example, a second tiny random-weight model was written with gguf-py in the same way: Gemma 3 architecture, 6 layers, 2 KV heads of dimension 32, and sliding_window 128 in its metadata. The commands below run it three times at a 32,768-token context: with the defaults, with --swa-full, and with -ub 64.
llama-completion -m tiny-gemma3-swa.gguf -p hello -n 2 -no-cnv -lv 4 -c 32768
llama-completion -m tiny-gemma3-swa.gguf -p hello -n 2 -no-cnv -lv 4 -c 32768 --swa-full
llama-completion -m tiny-gemma3-swa.gguf -p hello -n 2 -no-cnv -lv 4 -c 32768 -ub 64Selected lines from the pinned build follow, with timestamps removed. Blank lines separate the three runs. In each run, compare the number of cells and layers of the non-SWA (global) cache with those of the SWA (local) cache.
I print_info: n_swa = 128
I print_info: is_swa_any = 1
I llama_kv_cache_iswa: creating non-SWA KV cache, size = 32768 cells
I llama_kv_cache: size = 8.00 MiB ( 32768 cells, 1 layers, 1/1 seqs), K (f16): 4.00 MiB, V (f16): 4.00 MiB
I llama_kv_cache_iswa: creating SWA KV cache, size = 768 cells
I llama_kv_cache: size = 0.94 MiB ( 768 cells, 5 layers, 1/1 seqs), K (f16): 0.47 MiB, V (f16): 0.47 MiB
W llama_kv_cache_iswa: using full-size SWA cache (ref: https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)
I llama_kv_cache_iswa: creating non-SWA KV cache, size = 32768 cells
I llama_kv_cache: size = 8.00 MiB ( 32768 cells, 1 layers, 1/1 seqs), K (f16): 4.00 MiB, V (f16): 4.00 MiB
I llama_kv_cache_iswa: creating SWA KV cache, size = 32768 cells
I llama_kv_cache: size = 40.00 MiB ( 32768 cells, 5 layers, 1/1 seqs), K (f16): 20.00 MiB, V (f16): 20.00 MiB
I llama_kv_cache_iswa: creating non-SWA KV cache, size = 32768 cells
I llama_kv_cache: size = 8.00 MiB ( 32768 cells, 1 layers, 1/1 seqs), K (f16): 4.00 MiB, V (f16): 4.00 MiB
I llama_kv_cache_iswa: creating SWA KV cache, size = 256 cells
I llama_kv_cache: size = 0.31 MiB ( 256 cells, 5 layers, 1/1 seqs), K (f16): 0.16 MiB, V (f16): 0.16 MiBThe program reproduces these sizes and then applies the same rules to a larger configuration. That configuration has 62 layers (the Gemma 3 27B layer count in gemma3.cpp), 16 KV heads of dimension 128, and a window of 1,024 (values from Gemma 3 27B's published config.json), at 131,072 tokens.
import java.util.Locale;
public class InterleavedSwa {
static long pad(long x, long to) { return (x + to - 1) / to * to; }
// llama-kv-cache-iswa.cpp: GGML_PAD(min(n_ctx, n_swa * n_seq + n_ubatch), 256) cells for windowed layers
static long swaCells(long nCtx, long nSwa, long nSeq, long nUbatch) { return pad(Math.min(nCtx, nSwa * nSeq + nUbatch), 256); }
public static void main(String[] args) {
double mib = 1 << 20;
StringBuilder pattern = new StringBuilder();
for (int il = 0; il < 12; il++) pattern.append(il % 6 < 5 ? 'L' : 'G'); // set_swa_pattern(6)
System.out.println("layers 0-11, pattern 6: " + pattern);
// the tiny Gemma 3 style model: 6 layers, 2 KV heads of 32, F16 cache, n_swa = 128
long perCellLayer = 2L * 2 * 32 * 2;
for (long[] run : new long[][] {{32768, 512}, {32768, 64}}) {
long cells = swaCells(run[0], 128, 1, run[1]);
System.out.printf(Locale.ROOT, "tiny, -c %d -ub %d: global 1 x %d cells = %.2f MiB, local 5 x %d cells = %.2f MiB%n",
run[0], run[1], run[0], run[0] * perCellLayer / mib, cells, 5 * cells * perCellLayer / mib);
}
System.out.printf(Locale.ROOT, "tiny, --swa-full: local 5 x 32768 cells = %.2f MiB%n", 5 * 32768 * perCellLayer / mib);
// 62 layers, 16 KV heads of 128, n_swa = 1024 (Gemma 3 27B values), F16 cache, 131072 tokens
long perCell = 2L * 16 * 128 * 2, ctx = 131072;
int global = 0, local = 0;
for (int il = 0; il < 62; il++) if (il % 6 < 5) local++; else global++;
long cells = swaCells(ctx, 1024, 1, 512);
double g = global * ctx * perCell / (double) (1L << 30), l = local * cells * perCell / (double) (1L << 30);
System.out.printf(Locale.ROOT, "62 layers: %d global, %d local; global %.2f GiB + local %.2f GiB (%d cells) = %.2f GiB%n",
global, local, g, l, cells, g + l);
System.out.printf(Locale.ROOT, "all 62 layers global (or --swa-full): %.2f GiB%n", 62 * ctx * perCell / (double) (1L << 30));
}
}layers 0-11, pattern 6: LLLLLGLLLLLG
tiny, -c 32768 -ub 512: global 1 x 32768 cells = 8.00 MiB, local 5 x 768 cells = 0.94 MiB
tiny, -c 32768 -ub 64: global 1 x 32768 cells = 8.00 MiB, local 5 x 256 cells = 0.31 MiB
tiny, --swa-full: local 5 x 32768 cells = 40.00 MiB
62 layers: 10 global, 52 local; global 10.00 GiB + local 0.61 GiB (1536 cells) = 10.61 GiB
all 62 layers global (or --swa-full): 62.00 GiBThe pattern line shows set_swa_pattern(6): five local layers (L), then one global layer (G), repeated. The tiny model's 6 layers are therefore 5 local and 1 global, which matches 1 layers and 5 layers in the log.
The tiny model stores 2 KV heads x 32 dimensions x 2 (K and V) x 2 bytes = 256 bytes per cell per layer. The global layer has all 32,768 cells: 32,768 x 256 bytes = 8.00 MiB. The local layers have pad(min(32768, 128 + 512), 256) = 768 cells, so 5 layers x 768 cells x 256 bytes = 0.94 MiB. The window is 128, but the cache must also hold the micro-batch being processed, which is why n_ubatch appears in the formula. With -ub 64, 128 + 64 = 192 pads up to 256 cells and 0.31 MiB. The option --swa-full gives the local layers all 32,768 cells, and their cache grows from 0.94 MiB to 40.00 MiB, which the program also reproduces.
For the 62-layer configuration, 10 layers are global and 52 are local. The global layers need 10 GiB at 128K tokens. The local layers need only 0.61 GiB, for 1,536 cells each. The total, 10.61 GiB, is about a sixth of the 62 GiB that the same model would need with full attention in every layer, or with --swa-full.
--swa-full exists because of prefix reuse. A small local cache keeps only the last window of each sequence. When a new request shares a long prefix with an earlier one, the local layers no longer hold the keys of the early tokens, so llama-server cannot resume from an arbitrary point of the prefix. The default behavior instead saves context checkpoints, covered in Paged KV Cache and Prefix Caching. vLLM handles mixed layers with its hybrid KV cache manager. It gives sliding-window layers their own block groups and frees the blocks that fall out of the window. The flag --disable-hybrid-kv-cache-manager makes it allocate the same size for all layers (SchedulerConfig in vllm/config/scheduler.py).
Example 11: attention sinks: keeping the first tokens while evicting
A sliding window seems safe, but it hides a problem. The StreamingLLM paper (Xiao et al., "Efficient Streaming Language Models with Attention Sinks", 2023) observed that trained models give a large share of attention to the first few tokens of a sequence, whatever those tokens are. It calls them attention sinks. The reason follows from softmax: the weights must add up to 1, so when no token is relevant to a head, the head still has to put its weight somewhere. Trained heads put the leftover weight on the first tokens, which act as a drain.
A cache that keeps only the last W tokens evicts the sinks. The paper reports that model quality then collapses, because the leftover weight is forced onto tokens that matter. Keeping the first 4 tokens plus a recent window restores quality.
The program builds the scores of one query over 20 positions, with a sink at position 0 (high score, value near 0). It compares three caches: all 20 keys, the last 8 keys, and 4 sinks plus the last 8. It prints how much weight the recent positions 12 to 19 receive and what the output is. It then runs a 12-entry streaming cache over 20 tokens to show which positions stay.
import java.util.ArrayList;
import java.util.List;
import java.util.Locale;
import java.util.Random;
public class AttentionSinks {
static double[] score = new double[20], value = new double[20];
static void attend(String name, List<Integer> keep) {
double max = Double.NEGATIVE_INFINITY, sum = 0, out = 0;
for (int p : keep) max = Math.max(max, score[p]);
for (int p : keep) sum += Math.exp(score[p] - max);
for (int p : keep) out += Math.exp(score[p] - max) / sum * value[p];
double recent = 0;
for (int p : keep) if (p >= 12) recent += Math.exp(score[p] - max) / sum;
System.out.printf(Locale.ROOT, "%-18s keys %2d weight on positions 12-19 %.3f output %+.4f%n", name, keep.size(), recent, out);
}
static List<Integer> range(int from, int to) { List<Integer> l = new ArrayList<>(); for (int i = from; i < to; i++) l.add(i); return l; }
public static void main(String[] args) {
Random r = new Random(3);
score[0] = 6.0; // position 0 draws most of the attention
value[0] = 0.0; // and contributes almost nothing
for (int i = 1; i < 20; i++) {
score[i] = r.nextDouble() * 2;
value[i] = Math.round((r.nextDouble() * 2 - 1) * 100) / 100.0;
}
attend("full cache", range(0, 20));
attend("last 8 only", range(12, 20));
List<Integer> sinks = range(0, 4);
sinks.addAll(range(12, 20));
attend("4 sinks + last 8", sinks);
int nSink = 4, nRecent = 8; // a streaming cache of 12 entries
List<Integer> cache = new ArrayList<>();
for (int pos = 0; pos < 20; pos++) {
cache.add(pos);
if (cache.size() > nSink + nRecent) cache.remove(nSink); // evict the oldest non-sink entry
if (pos == 11 || pos == 15 || pos == 19)
System.out.println("after token " + pos + ": cache " + cache + ", positions used for RoPE 0.." + (cache.size() - 1));
}
}
}full cache keys 20 weight on positions 12-19 0.055 output -0.0172
last 8 only keys 8 weight on positions 12-19 1.000 output +0.0705
4 sinks + last 8 keys 12 weight on positions 12-19 0.059 output -0.0018
after token 11: cache [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], positions used for RoPE 0..11
after token 15: cache [0, 1, 2, 3, 8, 9, 10, 11, 12, 13, 14, 15], positions used for RoPE 0..11
after token 19: cache [0, 1, 2, 3, 12, 13, 14, 15, 16, 17, 18, 19], positions used for RoPE 0..11With the full cache, positions 12 to 19 receive only 5.5% of the attention. Most of the rest goes to the sink at position 0, whose value is 0, so it adds nothing to the output. The output is -0.0172.
With only the last 8 keys, softmax must distribute 100% of the weight over positions 12 to 19. Each recent key's weight grows by a factor of about 18 (1 / 0.055), and the output becomes +0.0705. That is different in sign and four times larger in magnitude.
With 4 sinks plus the last 8, the recent tokens again receive 5.9% of the weight, close to the full cache's 5.5%, and the output is -0.0018. The remaining difference comes from positions 4 to 11, which both partial caches have evicted.
The streaming cache keeps positions 0 to 3 forever, and the most recent 8 positions after them. After token 15, positions 4 to 7 have been evicted. After token 19, positions 4 to 11 have been evicted. The paper assigns RoPE positions by place in the cache rather than by place in the original text. The model therefore always sees positions 0 to 11, which are inside its trained range. The printed positions used for RoPE 0..11 reflects that rule.
Two related mechanisms exist in engines:
- Learned sinks. Some models, gpt-oss among them, include a learned per-head sink value. This is an extra logit that joins the softmax denominator but has no value vector, so a head can put its weight "nowhere". llama.cpp's flash-attention kernel accepts it as an optional input (
sinksinggml_compute_forward_flash_attn_ext_f16) and adds it to the running sum with the same rescaling rule as Example 2. vLLM'sFlashAttentionBackendreports sink support only on compute capability 9.0 or higher. --keepwith context shift. llama.cpp's--context-shiftkeeps the first--keep Ntokens and discards half of the rest (Tuning llama.cpp, Example 9). This preserves the start of the prompt in the same way as attention sinks.
Example 12: RoPE frequencies beyond the training length
The remaining examples address the third problem: a model trained on short texts used on long ones. First, a reminder of how RoPE works. RoPE splits each query and key vector into pairs of numbers. It rotates pair i by the angle position x theta_i, where theta_i = base^(-2i/d) (The Transformer Block and the Full Model, Example 8). Pair 0 has the largest theta and rotates fastest. The last pair rotates slowest. The base is a constant, 10,000 in the original design.
A clock gives the picture. The second hand, minute hand, and hour hand turn at different speeds. The fast hand tells nearby moments apart, and the slow hand tells distant moments apart. Mapping: each pair is one hand, and the angle is the position times the hand's speed. The comparison breaks in one place: a real clock repeats after 12 hours, so every angle is familiar. In RoPE, a slow pair that never completes a turn during training shows angles in training only up to its largest trained value. Beyond that, the angles are new.
The wavelength of pair i is 2 pi / theta_i, the number of positions needed for one full turn. A model trained on 4,096 positions has seen every angle of the fast pairs many times. Slow pairs never complete a turn within 4,096 positions, so the model has seen only part of their angles. The program prints, for selected pairs of a head with d = 128 and base 10,000, the wavelength, the number of turns inside 4,096 positions, and the largest angle seen in training against the angle at position 16,383.
import java.util.Locale;
public class RopeBeyondTraining {
public static void main(String[] args) {
int d = 128, train = 4096, target = 16384;
double base = 10000;
System.out.println("pair theta wavelength turns in 4096 max angle trained angle at 16383");
for (int i : new int[] {0, 8, 16, 32, 45, 46, 52, 63}) {
double theta = Math.pow(base, -2.0 * i / d);
double wavelength = 2 * Math.PI / theta;
System.out.printf(Locale.ROOT, "%4d %.3e %10.1f %13.2f %17.2f %14.2f%n",
i, theta, wavelength, train / wavelength, (train - 1) * theta, (target - 1) * theta);
}
int unseen = 0;
for (int i = 0; i < d / 2; i++) if (2 * Math.PI / Math.pow(base, -2.0 * i / d) > train) unseen++;
System.out.println("pairs that never complete one turn within 4096 positions: " + unseen + " of " + d / 2);
}
}pair theta wavelength turns in 4096 max angle trained angle at 16383
0 1.000e+00 6.3 651.90 4095.00 16383.00
8 3.162e-01 19.9 206.15 1294.95 5180.76
16 1.000e-01 62.8 65.19 409.50 1638.30
32 1.000e-02 628.3 6.52 40.95 163.83
45 1.540e-03 4080.2 1.00 6.31 25.23
46 1.334e-03 4711.7 0.87 5.46 21.85
52 5.623e-04 11173.3 0.37 2.30 9.21
63 1.155e-04 54410.1 0.08 0.47 1.89
pairs that never complete one turn within 4096 positions: 18 of 64Pair 0 turns 1 radian per position, which is a wavelength of 6.3 positions, and completes 651.9 turns within the training length. Every angle between 0 and 2 pi appeared thousands of times. Position 16,383 therefore produces an angle (16,383 radians, or 2.74 radians after removing whole turns) that the model has seen.
The slow pairs are different. Pair 46 has a wavelength of 4,711.7 positions and completes 0.87 turns in training, so its largest trained angle is 5.46 radians. At position 16,383 its angle is 21.85 radians, which includes angles between 5.46 and 2 pi that never occurred in training. Pair 63 rotates at most 0.47 radians in training and 1.89 at position 16,383, four times further than anything it saw.
18 of the 64 pairs never complete one turn within 4,096 positions. These pairs encode coarse, long-range position, and the model's weights learned to read them only inside the trained range. Past that range, attention scores that depend on these pairs take values the model never produced. This is the measured drop in quality when a context exceeds n_ctx_train (the context length the model was trained on), and llama.cpp warns possible training context overflow. The next three examples are three ways to keep all angles inside the trained range.
Pair 45 (wavelength 4,080.2, 1.00 turn) sits at the boundary. The value of the base sets where that boundary falls. Llama 3 uses base 500,000, which slows every pair and makes more of them long-wavelength.
Example 13: linear position interpolation
The simplest fix is to shrink the positions. Position interpolation (PI; Chen et al., "Extending Context Window of Large Language Models via Positional Interpolation", 2023) divides every position by the scale factor s = target length / trained length before RoPE. To extend 4,096 positions to 16,384, s = 4, so position 16,383 is rotated as if it were position 4,095.75. Every angle then stays inside the range the model saw. llama.cpp applies this with --rope-scaling linear --rope-scale N, which sets freq_scale = 1/N (Tuning llama.cpp, Example 8). The program extends 4,096 positions to 16,384 and prints the angle per token and the angle at the last position for four pairs.
import java.util.Locale;
public class PositionInterpolation {
public static void main(String[] args) {
int d = 128, train = 4096, target = 16384;
double base = 10000, s = (double) target / train; // scale factor 4
System.out.printf(Locale.ROOT, "s = %.0f: position 16383 is rotated as position %.2f; trained positions 0 to %d%n",
s, (target - 1) / s, train - 1);
for (int i : new int[] {0, 16, 46, 63}) {
double theta = Math.pow(base, -2.0 * i / d);
System.out.printf(Locale.ROOT, "pair %2d: angle per token %.6f -> %.6f, angle at 16383 %8.2f (at 4095 in training %8.2f)%n",
i, theta, theta / s, (target - 1) * theta / s, (train - 1) * theta);
}
System.out.printf(Locale.ROOT, "pair 0: neighbours now differ by %.2f rad; in training the smallest difference was 1.00 rad%n", 1 / s);
}
}s = 4: position 16383 is rotated as position 4095.75; trained positions 0 to 4095
pair 0: angle per token 1.000000 -> 0.250000, angle at 16383 4095.75 (at 4095 in training 4095.00)
pair 16: angle per token 0.100000 -> 0.025000, angle at 16383 409.58 (at 4095 in training 409.50)
pair 46: angle per token 0.001334 -> 0.000333, angle at 16383 5.46 (at 4095 in training 5.46)
pair 63: angle per token 0.000115 -> 0.000029, angle at 16383 0.47 (at 4095 in training 0.47)
pair 0: neighbours now differ by 0.25 rad; in training the smallest difference was 1.00 radAfter scaling, position 16,383 is rotated like position 4,095.75, inside the trained range of 0 to 4,095. Every pair's largest angle is now within a quarter radian of its trained maximum: pair 46 reaches 5.46 radians, as in training, and pair 63 reaches 0.47.
The cost shows in the angle per token. Pair 0 now turns 0.25 radians between neighbouring tokens instead of 1.00. In training, the smallest difference between two positions on this pair was 1 radian. Now the model must tell apart tokens whose angles differ by a quarter of that. All pairs are compressed by the same factor 4, including the fast pairs that did not need it, because they had already seen every angle. The result is lower resolution for nearby tokens, which matters for local word order.
The PI paper reports that a short fine-tune on the longer length (within 1,000 steps in its experiments) lets the model adapt to the compressed positions. It also reports that the extended models keep quality on tasks within the original context window relatively well. The next two methods compress the slow pairs and leave the fast ones alone.
Example 14: NTK-aware scaling
PI slows every pair by the same factor. NTK-aware scaling slows the pairs by different amounts. It changes the base instead of the positions: base' = base x s^(d/(d-2)). A larger base makes theta_i = base^(-2i/d) smaller for every pair except pair 0, and the slow pairs are slowed the most. It was first published in 2023 as a community post. The YaRN paper (Peng et al., 2023) describes it as a step toward its own method, and vLLM's implementation credits Reddit users /u/bloc97 and /u/emozilla. ("NTK" refers to the neural tangent kernel, the theory the post used as motivation.) With the new base, pair i slows down by s^(2i/(d-2)). Pair 0 is not slowed at all, and the last pair (i = d/2 - 1) is slowed by exactly s. The program computes the new base for s = 4 and d = 128 and the slowdown of selected pairs.
import java.util.Locale;
public class NtkAwareScaling {
public static void main(String[] args) {
int d = 128, train = 4096, target = 16384;
double base = 10000, s = 4;
double newBase = base * Math.pow(s, (double) d / (d - 2));
System.out.printf(Locale.ROOT, "base %.0f -> %.1f%n", base, newBase);
System.out.println("pair theta theta' slowdown angle at 16383 trained max");
for (int i : new int[] {0, 16, 32, 46, 56, 63}) {
double theta = Math.pow(base, -2.0 * i / d), thetaNew = Math.pow(newBase, -2.0 * i / d);
System.out.printf(Locale.ROOT, "%4d %.4e %.4e %8.3f %14.2f %11.2f%n",
i, theta, thetaNew, theta / thetaNew, (target - 1) * thetaNew, (train - 1) * theta);
}
}
}base 10000 -> 40889.9
pair theta theta' slowdown angle at 16383 trained max
0 1.0000e+00 1.0000e+00 1.000 16383.00 4095.00
16 1.0000e-01 7.0323e-02 1.422 1152.10 409.50
32 1.0000e-02 4.9453e-03 2.022 81.02 40.95
46 1.3335e-03 4.8462e-04 2.752 7.94 5.46
56 3.1623e-04 9.2222e-05 3.429 1.51 1.29
63 1.1548e-04 2.8870e-05 4.000 0.47 0.47The base grows from 10,000 to 40,889.9. Pair 0 keeps theta = 1, so the fast pairs keep their full resolution. This is the difference from linear interpolation. Pair 63 slows down by exactly 4.000, so its angle at position 16,383 is 0.47 radians, equal to its trained maximum. In between, the slowdown grows smoothly: 1.422 for pair 16, 2.022 for pair 32, and 2.752 for pair 46.
The middle pairs show the method's limit. Pair 46 reaches 7.94 radians at position 16,383 and pair 56 reaches 1.51, both beyond their trained maxima (5.46 and 1.29). These pairs are slowed by less than 4 but still have wavelengths near or above the training length, so they partly extrapolate, which means they still see angles the model never saw. The YaRN paper reports that NTK-aware scaling extends the context without fine-tuning better than linear interpolation, and that some dimensions are still extrapolated to angles outside the trained range. The next example is the paper's answer to that.
vLLM implements an ntk type (ntk_scaling_rope.py) and a dynamic type (dynamic_ntk_scaling_rope.py), selected by the rope_scaling entry of the model's config. The dynamic class computes its base once, as base x ((s x L / L_train) - (s - 1))^(d/(d-2)), with L and L_train taken from the configuration. llama.cpp's --rope-scaling accepts none, linear, and yarn at this commit (the library also knows longrope from model metadata), so it has no NTK mode. The flag --rope-freq-base sets the base directly, and setting it to the computed base' gives the fixed NTK-aware variant.
Example 15: YaRN: a per-frequency ramp and an attention temperature
The two earlier methods treat all pairs by one rule. YaRN (Peng et al., "YaRN: Efficient Context Window Extension of Large Language Models", 2023) treats each pair by its wavelength, which is the property that decides whether the pair needs help. It has three zones:
- Pairs that complete many turns in the trained length keep their frequency. This is called extrapolation here, and it means unchanged.
- Pairs whose wavelength is longer than the trained length are divided by
s, as in PI. This is interpolation. - Pairs in between get a linear mix of the two.
The boundaries are the pairs that complete beta_fast turns and beta_slow turns in the original length. llama.cpp's defaults are 32 and 1 (src/llama-hparams.h). The position of a pair inside the ramp is called ramp: 1 means fully unchanged, 0 means fully divided by s.
YaRN also scales the attention logits through a factor mscale = 1 + 0.1 ln(s), where ln is the natural logarithm. The program reproduces ggml's code: ggml_rope_yarn_corr_dim in ggml/src/ggml.c, and rope_yarn_ramp and rope_yarn in ggml/src/ggml-cpu/ops.cpp.
import java.util.Locale;
public class Yarn {
// ggml.c: the pair index at which a frequency completes nRot turns over nCtxOrig positions
static double corrDim(int nDims, int nCtxOrig, double nRot, double base) {
return nDims * Math.log(nCtxOrig / (nRot * 2 * Math.PI)) / (2 * Math.log(base));
}
public static void main(String[] args) {
int d = 128, orig = 4096;
double base = 10000, s = 4, betaFast = 32, betaSlow = 1, freqScale = 1 / s;
double low = Math.max(0, Math.floor(corrDim(d, orig, betaFast, base)));
double high = Math.min(d - 1, Math.ceil(corrDim(d, orig, betaSlow, base)));
System.out.printf(Locale.ROOT, "corr dims: low %.0f (32 turns), high %.0f (1 turn)%n", low, high);
System.out.println("pair ramp theta used / theta wavelength");
for (int i : new int[] {0, 10, 20, 25, 30, 35, 40, 46, 55, 63}) {
double y = (i - low) / Math.max(0.001, high - low); // ggml-cpu ops.cpp rope_yarn_ramp, i = i0 / 2
double ramp = 1 - Math.min(1, Math.max(0, y));
double extrap = Math.pow(base, -2.0 * i / d), interp = freqScale * extrap;
double theta = interp * (1 - ramp) + extrap * ramp;
System.out.printf(Locale.ROOT, "%4d %.3f %18.4f %10.1f%n", i, ramp, theta / extrap, 2 * Math.PI / extrap);
}
double mscale = 1 + 0.1 * Math.log(s);
System.out.printf(Locale.ROOT, "mscale = 1 + 0.1 ln(%.0f) = %.4f; scores scale by mscale^2 = %.4f%n", s, mscale, mscale * mscale);
}
}corr dims: low 20 (32 turns), high 46 (1 turn)
pair ramp theta used / theta wavelength
0 1.000 1.0000 6.3
10 1.000 1.0000 26.5
20 1.000 1.0000 111.7
25 0.808 0.8558 229.4
30 0.615 0.7115 471.2
35 0.423 0.5673 967.6
40 0.231 0.4231 1986.9
46 0.000 0.2500 4711.7
55 0.000 0.2500 17206.0
63 0.000 0.2500 54410.1
mscale = 1 + 0.1 ln(4) = 1.1386; scores scale by mscale^2 = 1.2965The correction dimensions are the pair indices where a frequency completes 32 turns and 1 turn within 4,096 positions: floor(20.94) = 20 and ceil(45.03) = 46. The table follows from them:
- Pairs 0 to 20 have
ramp = 1and keep their frequency (ratio 1.0000). Their wavelengths, 111.7 positions or less, fit 32 or more times into the training length, so every angle was seen. - Pairs 46 to 63 have
ramp = 0and are divided by 4 (ratio 0.2500), exactly as in position interpolation. Their wavelengths, 4,711.7 positions and more, exceed the training length. - Pairs 25 to 40 lie on the ramp. Pair 30, for example, has
ramp = 0.615, so its frequency is0.615 x 1 + 0.385 x 0.25 = 0.7115of the original.
Unlike linear interpolation, fast pairs keep full resolution. Unlike NTK-aware scaling, every pair with a wavelength above the training length is interpolated by the full factor, so pair 46 no longer extrapolates.
The last line is the attention temperature. Interpolation makes attention distributions flatter at long lengths, which means the probabilities spread out more evenly, and YaRN compensates by scaling the scores up. ggml multiplies both cosine and sine by mscale. Queries and keys are both rotated, so a score is multiplied by mscale^2 = 1.2965 for s = 4. The paper writes the same correction as a temperature t with sqrt(1/t) = 0.1 ln(s) + 1.
The YaRN paper reports that the method extends Llama 2 to 64K and 128K contexts with a fine-tune of about 0.1% of the original pretraining data. Many current models are trained or fine-tuned with YaRN and ship its settings in their config. Tuning llama.cpp, Example 8, shows llama.cpp reading them from GGUF metadata or from --rope-scaling yarn --rope-scale N --yarn-orig-ctx N. The DeepSeek-V3 config, for example, lists a rope_scaling entry of type yarn with factor 40 and an original length of 4,096, and vLLM implements that variant in deepseek_scaling_rope.py.
Example 16: long-context cost: prefill FLOPs and KV memory
A model that can accept a long context still pays for it, and the two costs grow at different rates. FLOPs (floating-point operations) measure compute. Prefill FLOPs have two parts. The first part grows linearly with the length: 2 FLOPs per parameter per token (FLOPs, Bandwidth, and the Roofline Model). The second part is attention, which grows with the square of the length: Q K^T and P V each cost 2 x n x n x qDim FLOPs per layer without a mask, and the causal mask halves both. The KV cache grows linearly with the length, and a decode step reads all of it. The program computes both costs for a Llama 3 8B shape at 4K, 32K, and 128K tokens.
import java.util.Locale;
public class LongContextCost {
public static void main(String[] args) {
double params = 8.03e9; // Llama 3 8B shape
int layers = 32, qDim = 4096; // 32 heads x 128
long kvPerToken = 131072; // F16 cache, from Example 6
System.out.println("tokens weight TFLOP attention TFLOP attention share KV cache (read per decoded token)");
for (long n : new long[] {4096, 32768, 131072}) {
double weights = 2 * params * n; // 2 FLOPs per parameter per token
double attention = 2.0 * layers * qDim * n * n; // QK^T and PV, causal half of 4 n^2 d per layer
double kv = kvPerToken * (double) n / (1 << 30);
System.out.printf(Locale.ROOT, "%6d %12.1f %15.1f %14.0f%% %5.1f GiB%n",
n, weights / 1e12, attention / 1e12, 100 * attention / (weights + attention), kv);
}
}
}tokens weight TFLOP attention TFLOP attention share KV cache (read per decoded token)
4096 65.8 4.4 6% 0.5 GiB
32768 526.3 281.5 35% 4.0 GiB
131072 2105.0 4503.6 68% 16.0 GiBAt 4,096 tokens, attention is 4.4 of 70.2 TFLOP (6%), so prompt processing time is dominated by the weights, as earlier lessons assumed. At 32,768 tokens, attention is 281.5 TFLOP, which is 35% of the total. At 131,072 tokens, attention (4,503.6 TFLOP) is more than twice the weight FLOPs (2,105.0), which makes it 68% of the total. A 32-fold longer prompt costs about 94 times more compute (6,608.6 against 70.2 TFLOP).
FlashAttention does not change these FLOPs. Tiled attention computes every score that naive attention computes. It removes the memory traffic and the n^2 buffer from Example 1. Sliding windows and sparse attention (attending to a chosen subset of tokens) are the methods that reduce the FLOPs themselves.
The KV cache column is the other cost. At 128K tokens it is 16 GiB, which is twice the size of the model's 8B weights in 8-bit form. A decode step reads it in full, so generation at the end of a long context reads 16 GiB of cache in addition to the weights. Examples 6 to 10 are the design answers to this column: fewer KV heads, latents, and windows. A quantized KV cache (-ctk, -ctv in llama.cpp, --kv-cache-dtype in vLLM) halves it again.
Example 17: retrieval tests: a needle-in-a-haystack grid
A model that accepts 128K tokens may not use all of them well. The common first check is needle in a haystack. The test hides one fact (the needle) at a chosen depth in long filler text and asks the model for it. A depth of 0% is the start of the text, and 100% is the end. A grid of context lengths and depths shows where retrieval fails.
The program builds the prompts and scores the answers of a fake model. The fake model finds the needle only if it is within the last 3,000 words or among the first 64 words, which mimics the edge-heavy pattern of real models. A real test sends the same prompts to llama-server with the Project 11 client from The Chat Completions Protocol and Streaming in Java.
import java.util.ArrayList;
import java.util.List;
import java.util.function.BiFunction;
public class NeedleGrid {
// builds a prompt of `length` filler words with the needle at `depth` percent
static List<String> prompt(int length, int depth, String needle) {
List<String> words = new ArrayList<>();
for (int i = 0; i < length; i++) words.add("w" + (i % 50));
words.add(Math.min(length, length * depth / 100), needle);
return words;
}
public static void main(String[] args) {
String needle = "CODE-4721";
// a fake model: it finds the needle only if it is within 3000 words of the end or in the first 64 words
BiFunction<List<String>, String, String> fakeModel = (words, question) -> {
int at = words.indexOf(needle);
return (words.size() - at <= 3000 || at < 64) ? needle : "I could not find it";
};
int[] depths = {0, 25, 50, 75, 100};
System.out.println("length depth: 0 25 50 75 100");
for (int length : new int[] {1000, 2000, 4000, 8000, 16000}) {
StringBuilder row = new StringBuilder(String.format("%6d ", length));
int correct = 0;
for (int depth : depths) {
String answer = fakeModel.apply(prompt(length, depth, needle), "What is the code?");
boolean ok = answer.equals(needle);
if (ok) correct++;
row.append(ok ? " ok" : " miss");
}
System.out.println(row + " " + correct + "/" + depths.length);
}
}
}length depth: 0 25 50 75 100
1000 ok ok ok ok ok 5/5
2000 ok ok ok ok ok 5/5
4000 ok miss ok ok ok 4/5
8000 ok miss miss ok ok 3/5
16000 ok miss miss miss ok 2/5Each row is one context length and each column one depth, from the start (0%) to the end (100%) of the filler. Up to 2,000 words, the fake model answers every cell, because the whole prompt lies within its 3,000-word range. At 4,000 words, the needle at 25% sits 3,000 words from the end, which is one word outside the range, so it is missed. The depths 50% and 75% are close enough to the end. At 16,000 words, only the start (inside the first 64 words) and the end are found.
The grid makes the fake model's pattern visible: good at the edges, weak in the middle, and worse as the length grows. Published evaluations of real models report a similar shape. The "Lost in the Middle" paper (Liu et al., 2023) reports that accuracy on multi-document question answering is highest when the relevant document is at the start or end of the context. The RULER benchmark paper (Hsieh et al., 2024) reports that many models that pass the simple needle test lose accuracy on its harder tasks (several needles, tracing variables, aggregation) as the length grows. The length at which a model stays accurate can therefore be shorter than the length it accepts.
For a coding agent, the context length setting is an upper bound, not a guarantee. Measuring retrieval at the lengths the agent uses, with the same quantization and RoPE settings, gives the usable length.
Interview points
- Naive attention stores an
n x nscore matrix per head. At 128K tokens that is 64 GiB per head in FP32, so long contexts are impossible without tiling. - Online softmax keeps a running maximum
mand a running sums, and rescales byexp(m_old - m_new)whenever the maximum grows. With a value accumulator rescaled the same way, attention for one query needs onlyO(d)state. - FlashAttention tiles queries, keys, and values into fast memory, never writes the score matrix to HBM, and computes exactly the same result (up to rounding). It saves memory and traffic, not FLOPs.
- vLLM chooses an attention backend per platform from a priority list (
FLASH_ATTNfirst on most NVIDIA GPUs,FLASHINFERfirst on compute capability 10.x), and--attention-backendoverrides the choice. - GQA divides the KV bytes by the number of query heads per KV head. Llama 3 8B uses 8 KV heads for 32 query heads, which is 128 KiB per token in F16.
- MLA caches one latent (
kv_lora_rank512) plus a decoupled RoPE key (64) per layer. With weight absorption it behaves like MQA, and DeepSeek-V3 needs 70,272 bytes per token. - Sliding-window layers need a cache of only
Wpositions. Interleaved local and global layers cut a 62-layer model's 128K cache from 62 GiB to about 10.6 GiB, and--swa-fullgives that saving back in exchange for prefix reuse. - Attention sinks: evicting the first tokens changes the softmax normalization and breaks streaming, and keeping 4 initial tokens plus a recent window restores it (StreamingLLM).
- RoPE pairs with wavelengths longer than the trained length see new angles past
n_ctx_train. Linear interpolation compresses all pairs, NTK-aware scaling changes the base, and YaRN ramps between extrapolation and interpolation by wavelength and addsmscale = 1 + 0.1 ln s. - Long contexts raise two separate costs: attention FLOPs grow with the square of the length, and the KV cache grows linearly. Flash attention reduces neither the FLOPs nor the cache size, and a model's accepted length can exceed the length at which it retrieves reliably.
Summary
- The score matrix has
n^2entries per head. At 4,096 tokens and 32 heads it is 2 GiB per layer in FP32, and at 131,072 tokens it is 2 TiB. - Online softmax gives the same probabilities as two-pass softmax, because rescaling by
exp(m_old - m_new)changes only the reference point. - Tiled attention reads keys and values once per query block and writes only the output. For one head, traffic falls from about
4 n^2to2 n^2 d / Bfloats, and the buffer falls fromn^2toB^2. - On the tiny model at a 65,536-token context, the compute buffer is 297.75 MiB with flash attention and 1,154.76 MiB without it at
-ub 512. - DeepSeek-V3's published config gives 61 layers and a 512 + 64 cache entry per layer. That is 8.58 GiB at 128K tokens, 71 times less than an MHA layout of the same model.
- llama.cpp marks layer
ilas sliding-window whenil % n < n - 1(pattern 6 for Gemma 3, 2 for Gemma 2 and gpt-oss), and sizes those layers' caches aspad(min(n_ctx, n_swa x n_seq + n_ubatch), 256)cells. - The tiny Gemma 3-style model at
-c 32768uses 8.00 MiB for its global layer and 0.94 MiB for five local layers.--swa-fullraises the local part to 40.00 MiB. - With base 10,000 and
d = 128, 18 of 64 RoPE pairs do not complete one turn within 4,096 positions. - Position interpolation divides positions by
s. NTK-aware scaling setsbase' = base x s^(d/(d-2)). YaRN keeps pairs belowbeta_fast(32 turns), interpolates pairs beyondbeta_slow(1 turn), and ramps between them. - For an 8B model, attention is 6% of prefill FLOPs at 4K tokens and 68% at 128K tokens.
- Needle-in-a-haystack grids over length and depth measure where retrieval fails. Reported results show weaker retrieval in the middle of long contexts.