Lesson 11 of 73
Generation, Prefill, Decode, and the KV Cache
A language model produces one token per forward pass (one run of the whole network). A chat answer of a few hundred tokens therefore takes a few hundred passes. If every pass redid all the work of the previous one, answers would become slower with every token. This lesson shows where that waste comes from and how engines remove it with the KV cache (key-value cache), a store of earlier results.
The lesson also explains why the first pass over a prompt behaves differently from the passes after it, what limits the number of tokens per second, and how an engine decides when to stop. It uses the tiny transformer from The Transformer Block and the Full Model, so you can count the work by hand.
Afterwards you can read the timing lines of llama-completion and llama-server and explain each number. The table lists the terms the lesson defines, in the order they appear.
| Term | Meaning |
|---|---|
| Autoregressive loop | Run the model, choose one token, append it, repeat |
| Prefill | Processing the prompt tokens in one batch before the first output token |
| Decode | Producing one new token per pass, using the cache |
| KV cache | Stored keys and values of every earlier position, per layer |
| TTFT | Time to first token: how long the reader waits before output starts |
| Tokens per second | Decode speed after the first token |
The programs in the first examples use the transformer of lesson 10: dimension 16, 4 query heads, 2 key/value heads, 2 layers, a 24-token vocabulary, and seed 42. A counter adds one for every multiply-add (one multiplication followed by an addition) in the weight matrices and in attention. Counting operations measures work independently of how fast the machine is, so the numbers are the same on every computer.
Example 1: the autoregressive loop
A language model predicts one token at a time. To write a sentence, a program feeds in the prompt, takes the token with the highest logit as the next token (the greedy choice), appends it, and runs the model again on the longer sequence. A loop that feeds each output back in as input is called autoregressive. "Auto" means self: the model's own output becomes its next input.
This program runs that loop for 10 tokens with the tiny transformer of lesson 10 and prints the sequence the model sees at each step. Watch how the input grows by one token on every line.
The method process(tokens, start, cache) is the core of every example in this lesson. It runs the given tokens at positions start, start + 1, and so on. It stores their keys and values in a Cache object. Then it returns the logits of the last token. A key and a value are the two vectors each position produces for attention (see the self-attention lesson). In Examples 1 to 3 a fresh cache is created for every call, so nothing is reused. The model code is the same in Examples 1 to 5, and only main changes. The counter ops is used from Example 2 on.
import java.util.Arrays;
import java.util.Random;
public class Autoregressive {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2, MAX = 64;
static final Random R = new Random(42);
static long ops; // multiply-adds counted so far
static float[] mat(int rows, int cols) {
float[] w = new float[rows * cols];
for (int i = 0; i < w.length; i++) w[i] = (float) (R.nextGaussian() / Math.sqrt(cols));
return w;
}
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final float[] EMBED = mat(V, D);
static final Layer[] LAYERS = new Layer[L];
static {
for (int l = 0; l < L; l++)
LAYERS[l] = new Layer(mat(D, D), mat(KVH * HD, D), mat(KVH * HD, D), mat(D, D), mat(F, D), mat(F, D), mat(D, F));
}
static final float[] HEAD = mat(V, D);
static class Cache { // keys and values of every layer and position
final float[][][] k = new float[L][MAX][], v = new float[L][MAX][];
int len;
}
static float[] matvec(float[] w, int rows, int cols, float[] x) {
ops += (long) rows * cols;
float[] y = new float[rows];
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++) y[i] += w[i * cols + j] * x[j];
return y;
}
static float[] rmsnorm(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
float scale = (float) (1.0 / Math.sqrt(ss / x.length + 1e-5));
float[] y = new float[x.length];
for (int i = 0; i < x.length; i++) y[i] = x[i] * scale;
return y;
}
static void rope(float[] v, int heads, int pos) {
for (int h = 0; h < heads; h++)
for (int i = 0; i < HD / 2; i++) {
double a = pos * Math.pow(10000, -2.0 * i / HD);
float c = (float) Math.cos(a), s = (float) Math.sin(a), x0 = v[h * HD + 2 * i], x1 = v[h * HD + 2 * i + 1];
v[h * HD + 2 * i] = x0 * c - x1 * s;
v[h * HD + 2 * i + 1] = x0 * s + x1 * c;
}
}
// Runs `tokens` at positions start, start+1, ... ; keys and values go into the cache; returns the last logits.
static float[] process(int[] tokens, int start, Cache c) {
int n = tokens.length;
float[][] x = new float[n][];
for (int i = 0; i < n; i++) x[i] = Arrays.copyOfRange(EMBED, tokens[i] * D, (tokens[i] + 1) * D);
for (int l = 0; l < L; l++) {
Layer ly = LAYERS[l];
float[][] q = new float[n][];
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
q[i] = matvec(ly.wq, D, D, h);
float[] k = matvec(ly.wk, KVH * HD, D, h), v = matvec(ly.wv, KVH * HD, D, h);
rope(q[i], H, start + i);
rope(k, KVH, start + i);
c.k[l][start + i] = k;
c.v[l][start + i] = v;
}
for (int i = 0; i < n; i++) {
int pos = start + i;
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[pos + 1];
double sum = 0;
for (int s = 0; s <= pos; s++) {
double dot = 0;
for (int j = 0; j < HD; j++) dot += q[i][h * HD + j] * c.k[l][s][kv * HD + j];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= pos; s++)
for (int j = 0; j < HD; j++) cat[h * HD + j] += (float) (w[s] / sum) * c.v[l][s][kv * HD + j];
ops += 2L * HD * (pos + 1);
}
float[] out = matvec(ly.wo, D, D, cat);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int j = 0; j < F; j++) g[j] = (float) (g[j] / (1 + Math.exp(-g[j]))) * u[j];
float[] out = matvec(ly.down, D, F, g);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
}
c.len = start + n;
return matvec(HEAD, V, D, rmsnorm(x[n - 1])); // only the last position's logits are needed
}
static int argmax(float[] v) {
int best = 0;
for (int i = 1; i < v.length; i++) if (v[i] > v[best]) best = i;
return best;
}
public static void main(String[] args) {
int[] seq = {3, 17, 5, 9}; // the prompt
for (int step = 0; step < 10; step++) {
int next = argmax(process(seq, 0, new Cache())); // run the model on the whole sequence
System.out.println("step " + step + ": " + Arrays.toString(seq) + " -> " + next);
seq = Arrays.copyOf(seq, seq.length + 1); // append the chosen id and repeat
seq[seq.length - 1] = next;
}
System.out.println("generated ids: " + Arrays.toString(Arrays.copyOfRange(seq, 4, seq.length)));
}
}step 0: [3, 17, 5, 9] -> 6
step 1: [3, 17, 5, 9, 6] -> 6
step 2: [3, 17, 5, 9, 6, 6] -> 13
step 3: [3, 17, 5, 9, 6, 6, 13] -> 10
step 4: [3, 17, 5, 9, 6, 6, 13, 10] -> 13
step 5: [3, 17, 5, 9, 6, 6, 13, 10, 13] -> 13
step 6: [3, 17, 5, 9, 6, 6, 13, 10, 13, 13] -> 3
step 7: [3, 17, 5, 9, 6, 6, 13, 10, 13, 13, 3] -> 8
step 8: [3, 17, 5, 9, 6, 6, 13, 10, 13, 13, 3, 8] -> 8
step 9: [3, 17, 5, 9, 6, 6, 13, 10, 13, 13, 3, 8, 8] -> 8
generated ids: [6, 6, 13, 10, 13, 13, 3, 8, 8, 8]Each line shows the input of one forward pass and the id the model chose. The input grows by one token per step: step 0 sees the 4 prompt tokens, and step 9 sees 13. The chosen id of each step becomes the last input of the next step, so the steps depend on each other and cannot run in parallel. The ten generated ids match the greedy run in lesson 10, because the weights, the prompt, and the choice rule are the same.
Example 2: counting the work without a cache
The loop of Example 1 feeds the whole sequence into the model at every step. This example measures what that costs. The counter ops adds one for every multiply-add in the weight matrices and in attention. The program prints the count for each step and the total. Look at how the count changes as the sequence grows.
import java.util.Arrays;
import java.util.Random;
public class NoCache {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2, MAX = 64;
static final Random R = new Random(42);
static long ops; // multiply-adds counted so far
static float[] mat(int rows, int cols) {
float[] w = new float[rows * cols];
for (int i = 0; i < w.length; i++) w[i] = (float) (R.nextGaussian() / Math.sqrt(cols));
return w;
}
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final float[] EMBED = mat(V, D);
static final Layer[] LAYERS = new Layer[L];
static {
for (int l = 0; l < L; l++)
LAYERS[l] = new Layer(mat(D, D), mat(KVH * HD, D), mat(KVH * HD, D), mat(D, D), mat(F, D), mat(F, D), mat(D, F));
}
static final float[] HEAD = mat(V, D);
static class Cache { // keys and values of every layer and position
final float[][][] k = new float[L][MAX][], v = new float[L][MAX][];
int len;
}
static float[] matvec(float[] w, int rows, int cols, float[] x) {
ops += (long) rows * cols;
float[] y = new float[rows];
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++) y[i] += w[i * cols + j] * x[j];
return y;
}
static float[] rmsnorm(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
float scale = (float) (1.0 / Math.sqrt(ss / x.length + 1e-5));
float[] y = new float[x.length];
for (int i = 0; i < x.length; i++) y[i] = x[i] * scale;
return y;
}
static void rope(float[] v, int heads, int pos) {
for (int h = 0; h < heads; h++)
for (int i = 0; i < HD / 2; i++) {
double a = pos * Math.pow(10000, -2.0 * i / HD);
float c = (float) Math.cos(a), s = (float) Math.sin(a), x0 = v[h * HD + 2 * i], x1 = v[h * HD + 2 * i + 1];
v[h * HD + 2 * i] = x0 * c - x1 * s;
v[h * HD + 2 * i + 1] = x0 * s + x1 * c;
}
}
// Runs `tokens` at positions start, start+1, ... ; keys and values go into the cache; returns the last logits.
static float[] process(int[] tokens, int start, Cache c) {
int n = tokens.length;
float[][] x = new float[n][];
for (int i = 0; i < n; i++) x[i] = Arrays.copyOfRange(EMBED, tokens[i] * D, (tokens[i] + 1) * D);
for (int l = 0; l < L; l++) {
Layer ly = LAYERS[l];
float[][] q = new float[n][];
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
q[i] = matvec(ly.wq, D, D, h);
float[] k = matvec(ly.wk, KVH * HD, D, h), v = matvec(ly.wv, KVH * HD, D, h);
rope(q[i], H, start + i);
rope(k, KVH, start + i);
c.k[l][start + i] = k;
c.v[l][start + i] = v;
}
for (int i = 0; i < n; i++) {
int pos = start + i;
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[pos + 1];
double sum = 0;
for (int s = 0; s <= pos; s++) {
double dot = 0;
for (int j = 0; j < HD; j++) dot += q[i][h * HD + j] * c.k[l][s][kv * HD + j];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= pos; s++)
for (int j = 0; j < HD; j++) cat[h * HD + j] += (float) (w[s] / sum) * c.v[l][s][kv * HD + j];
ops += 2L * HD * (pos + 1);
}
float[] out = matvec(ly.wo, D, D, cat);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int j = 0; j < F; j++) g[j] = (float) (g[j] / (1 + Math.exp(-g[j]))) * u[j];
float[] out = matvec(ly.down, D, F, g);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
}
c.len = start + n;
return matvec(HEAD, V, D, rmsnorm(x[n - 1])); // only the last position's logits are needed
}
static int argmax(float[] v) {
int best = 0;
for (int i = 1; i < v.length; i++) if (v[i] > v[best]) best = i;
return best;
}
public static void main(String[] args) {
int[] seq = {3, 17, 5, 9};
long total = 0;
for (int step = 0; step < 10; step++) {
ops = 0;
int next = argmax(process(seq, 0, new Cache())); // every position is recomputed
System.out.printf("step %d: %2d positions computed, %,7d multiply-adds%n", step, seq.length, ops);
total += ops;
seq = Arrays.copyOf(seq, seq.length + 1);
seq[seq.length - 1] = next;
}
System.out.printf("total: %,d multiply-adds for %d generated tokens%n", total, 10);
}
}step 0: 4 positions computed, 19,456 multiply-adds
step 1: 5 positions computed, 24,384 multiply-adds
step 2: 6 positions computed, 29,376 multiply-adds
step 3: 7 positions computed, 34,432 multiply-adds
step 4: 8 positions computed, 39,552 multiply-adds
step 5: 9 positions computed, 44,736 multiply-adds
step 6: 10 positions computed, 49,984 multiply-adds
step 7: 11 positions computed, 55,296 multiply-adds
step 8: 12 positions computed, 60,672 multiply-adds
step 9: 13 positions computed, 66,112 multiply-adds
total: 424,000 multiply-adds for 10 generated tokensStep 0 processes 4 positions and costs 19,456 multiply-adds. Step 9 processes 13 positions and costs 66,112. The cost per step grows with the sequence length because every position is computed again at every step. The total of 424,000 grows roughly with the square of the length: doubling the number of steps about quadruples the work. A 4,000-token conversation answered this way would recompute thousands of positions for every new token.
Most of each step's work is the weight matrices: every position multiplies its vector by every matrix of every layer. Attention adds a part that grows with the number of earlier positions, because each position scores all positions before it.
Example 6: the keys and values of earlier positions do not change
The recomputation in Example 2 is only worth removing if the old positions produce the same results every time. This example tests that. It runs the model on three inputs: the prompt, the prompt plus one token, and the prompt plus two different tokens. Then it compares the keys and values stored for positions 0 to 3 in both layers across the runs. The last line compares position 4 between two different tokens as a control.
import java.util.Arrays;
import java.util.Random;
public class StableKeys {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2, MAX = 64;
static final Random R = new Random(42);
static long ops; // multiply-adds counted so far
static float[] mat(int rows, int cols) {
float[] w = new float[rows * cols];
for (int i = 0; i < w.length; i++) w[i] = (float) (R.nextGaussian() / Math.sqrt(cols));
return w;
}
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final float[] EMBED = mat(V, D);
static final Layer[] LAYERS = new Layer[L];
static {
for (int l = 0; l < L; l++)
LAYERS[l] = new Layer(mat(D, D), mat(KVH * HD, D), mat(KVH * HD, D), mat(D, D), mat(F, D), mat(F, D), mat(D, F));
}
static final float[] HEAD = mat(V, D);
static class Cache { // keys and values of every layer and position
final float[][][] k = new float[L][MAX][], v = new float[L][MAX][];
int len;
}
static float[] matvec(float[] w, int rows, int cols, float[] x) {
ops += (long) rows * cols;
float[] y = new float[rows];
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++) y[i] += w[i * cols + j] * x[j];
return y;
}
static float[] rmsnorm(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
float scale = (float) (1.0 / Math.sqrt(ss / x.length + 1e-5));
float[] y = new float[x.length];
for (int i = 0; i < x.length; i++) y[i] = x[i] * scale;
return y;
}
static void rope(float[] v, int heads, int pos) {
for (int h = 0; h < heads; h++)
for (int i = 0; i < HD / 2; i++) {
double a = pos * Math.pow(10000, -2.0 * i / HD);
float c = (float) Math.cos(a), s = (float) Math.sin(a), x0 = v[h * HD + 2 * i], x1 = v[h * HD + 2 * i + 1];
v[h * HD + 2 * i] = x0 * c - x1 * s;
v[h * HD + 2 * i + 1] = x0 * s + x1 * c;
}
}
// Runs `tokens` at positions start, start+1, ... ; keys and values go into the cache; returns the last logits.
static float[] process(int[] tokens, int start, Cache c) {
int n = tokens.length;
float[][] x = new float[n][];
for (int i = 0; i < n; i++) x[i] = Arrays.copyOfRange(EMBED, tokens[i] * D, (tokens[i] + 1) * D);
for (int l = 0; l < L; l++) {
Layer ly = LAYERS[l];
float[][] q = new float[n][];
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
q[i] = matvec(ly.wq, D, D, h);
float[] k = matvec(ly.wk, KVH * HD, D, h), v = matvec(ly.wv, KVH * HD, D, h);
rope(q[i], H, start + i);
rope(k, KVH, start + i);
c.k[l][start + i] = k;
c.v[l][start + i] = v;
}
for (int i = 0; i < n; i++) {
int pos = start + i;
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[pos + 1];
double sum = 0;
for (int s = 0; s <= pos; s++) {
double dot = 0;
for (int j = 0; j < HD; j++) dot += q[i][h * HD + j] * c.k[l][s][kv * HD + j];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= pos; s++)
for (int j = 0; j < HD; j++) cat[h * HD + j] += (float) (w[s] / sum) * c.v[l][s][kv * HD + j];
ops += 2L * HD * (pos + 1);
}
float[] out = matvec(ly.wo, D, D, cat);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int j = 0; j < F; j++) g[j] = (float) (g[j] / (1 + Math.exp(-g[j]))) * u[j];
float[] out = matvec(ly.down, D, F, g);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
}
c.len = start + n;
return matvec(HEAD, V, D, rmsnorm(x[n - 1])); // only the last position's logits are needed
}
static int argmax(float[] v) {
int best = 0;
for (int i = 1; i < v.length; i++) if (v[i] > v[best]) best = i;
return best;
}
static boolean samePrefix(Cache a, Cache b, int positions) {
boolean same = true;
for (int l = 0; l < L; l++)
for (int p = 0; p < positions; p++) same &= Arrays.equals(a.k[l][p], b.k[l][p]) && Arrays.equals(a.v[l][p], b.v[l][p]);
return same;
}
public static void main(String[] args) {
Cache a = new Cache(), b = new Cache(), c = new Cache();
process(new int[] {3, 17, 5, 9}, 0, a); // the prompt
process(new int[] {3, 17, 5, 9, 6}, 0, b); // the prompt plus token 6
process(new int[] {3, 17, 5, 9, 21, 2}, 0, c); // the prompt plus two other tokens
System.out.println("positions 0-3, prompt vs prompt + [6]: identical " + samePrefix(a, b, 4));
System.out.println("positions 0-3, prompt vs prompt + [21, 2]: identical " + samePrefix(a, c, 4));
System.out.println("position 4, token 6 vs token 21: identical keys " + Arrays.equals(b.k[0][4], c.k[0][4]));
}
}positions 0-3, prompt vs prompt + [6]: identical true
positions 0-3, prompt vs prompt + [21, 2]: identical true
position 4, token 6 vs token 21: identical keys falseThe keys and values of positions 0 to 3 are identical in all three runs, whatever tokens follow. The reason is the causal mask from the self-attention lesson: position p sees only positions 0 to p, so nothing computed at position p, in any layer, depends on later tokens. Layer 1 reads layer 0's output, but layer 0's output at position p also depends only on positions up to p, so the argument holds layer after layer.
The last line is the control. Token 6 and token 21 at position 4 give different keys, so the new position does depend on its own token. This means each step has exactly one new position to compute, and the keys and values of all earlier positions can be kept from the step before.
The last line shows that the new position itself does depend on its token: token 6 and token 21 at position 4 give different keys. So each step has exactly one new position to compute, and the keys and values of all earlier positions can be kept from the step before.
Example 7: adding a KV cache
The KV cache is the idea that follows from Example 6. Instead of recomputing the keys and values of old positions, store them. It keeps the keys and values of every layer for every position already processed. A new token then needs only its own query, key, and value, and its attention reads the stored keys and values of the earlier positions. The cost per step no longer depends on the sequence length, apart from the attention reads.
This program generates the same 10 tokens as Example 1, but with one cache that lives across all steps. The first call processes the whole prompt. Each later call processes one token at the next position (cache.len). The program compares the ids and the operation counts with the uncached run.
import java.util.Arrays;
import java.util.Random;
public class WithCache {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2, MAX = 64;
static final Random R = new Random(42);
static long ops; // multiply-adds counted so far
static float[] mat(int rows, int cols) {
float[] w = new float[rows * cols];
for (int i = 0; i < w.length; i++) w[i] = (float) (R.nextGaussian() / Math.sqrt(cols));
return w;
}
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final float[] EMBED = mat(V, D);
static final Layer[] LAYERS = new Layer[L];
static {
for (int l = 0; l < L; l++)
LAYERS[l] = new Layer(mat(D, D), mat(KVH * HD, D), mat(KVH * HD, D), mat(D, D), mat(F, D), mat(F, D), mat(D, F));
}
static final float[] HEAD = mat(V, D);
static class Cache { // keys and values of every layer and position
final float[][][] k = new float[L][MAX][], v = new float[L][MAX][];
int len;
}
static float[] matvec(float[] w, int rows, int cols, float[] x) {
ops += (long) rows * cols;
float[] y = new float[rows];
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++) y[i] += w[i * cols + j] * x[j];
return y;
}
static float[] rmsnorm(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
float scale = (float) (1.0 / Math.sqrt(ss / x.length + 1e-5));
float[] y = new float[x.length];
for (int i = 0; i < x.length; i++) y[i] = x[i] * scale;
return y;
}
static void rope(float[] v, int heads, int pos) {
for (int h = 0; h < heads; h++)
for (int i = 0; i < HD / 2; i++) {
double a = pos * Math.pow(10000, -2.0 * i / HD);
float c = (float) Math.cos(a), s = (float) Math.sin(a), x0 = v[h * HD + 2 * i], x1 = v[h * HD + 2 * i + 1];
v[h * HD + 2 * i] = x0 * c - x1 * s;
v[h * HD + 2 * i + 1] = x0 * s + x1 * c;
}
}
// Runs `tokens` at positions start, start+1, ... ; keys and values go into the cache; returns the last logits.
static float[] process(int[] tokens, int start, Cache c) {
int n = tokens.length;
float[][] x = new float[n][];
for (int i = 0; i < n; i++) x[i] = Arrays.copyOfRange(EMBED, tokens[i] * D, (tokens[i] + 1) * D);
for (int l = 0; l < L; l++) {
Layer ly = LAYERS[l];
float[][] q = new float[n][];
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
q[i] = matvec(ly.wq, D, D, h);
float[] k = matvec(ly.wk, KVH * HD, D, h), v = matvec(ly.wv, KVH * HD, D, h);
rope(q[i], H, start + i);
rope(k, KVH, start + i);
c.k[l][start + i] = k;
c.v[l][start + i] = v;
}
for (int i = 0; i < n; i++) {
int pos = start + i;
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[pos + 1];
double sum = 0;
for (int s = 0; s <= pos; s++) {
double dot = 0;
for (int j = 0; j < HD; j++) dot += q[i][h * HD + j] * c.k[l][s][kv * HD + j];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= pos; s++)
for (int j = 0; j < HD; j++) cat[h * HD + j] += (float) (w[s] / sum) * c.v[l][s][kv * HD + j];
ops += 2L * HD * (pos + 1);
}
float[] out = matvec(ly.wo, D, D, cat);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int j = 0; j < F; j++) g[j] = (float) (g[j] / (1 + Math.exp(-g[j]))) * u[j];
float[] out = matvec(ly.down, D, F, g);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
}
c.len = start + n;
return matvec(HEAD, V, D, rmsnorm(x[n - 1])); // only the last position's logits are needed
}
static int argmax(float[] v) {
int best = 0;
for (int i = 1; i < v.length; i++) if (v[i] > v[best]) best = i;
return best;
}
public static void main(String[] args) {
int[] prompt = {3, 17, 5, 9};
int[] example1 = {6, 6, 13, 10, 13, 13, 3, 8, 8, 8}; // ids generated without a cache in Example 1
long uncachedOps = 424_000; // multiply-adds counted in Example 2
ops = 0;
Cache cache = new Cache();
int[] ids = new int[10];
int next = argmax(process(prompt, 0, cache)); // the prompt in one pass
for (int step = 0; step < 10; step++) {
ids[step] = next;
long before = ops;
if (step < 9) {
next = argmax(process(new int[] {next}, cache.len, cache)); // one new position
System.out.printf("step %d: 1 position computed, %,6d multiply-adds%n", step + 1, ops - before);
}
}
long cachedOps = ops;
System.out.println("cached ids: " + Arrays.toString(ids) + ", same as Example 1: " + Arrays.equals(ids, example1));
System.out.printf("multiply-adds: %,d uncached, %,d cached (%.1f times fewer)%n", uncachedOps, cachedOps, (double) uncachedOps / cachedOps);
System.out.println("cache holds " + cache.len + " positions: " + 2 * L * cache.len * KVH * HD * 4 + " bytes in F32");
}
}step 1: 1 position computed, 5,312 multiply-adds
step 2: 1 position computed, 5,376 multiply-adds
step 3: 1 position computed, 5,440 multiply-adds
step 4: 1 position computed, 5,504 multiply-adds
step 5: 1 position computed, 5,568 multiply-adds
step 6: 1 position computed, 5,632 multiply-adds
step 7: 1 position computed, 5,696 multiply-adds
step 8: 1 position computed, 5,760 multiply-adds
step 9: 1 position computed, 5,824 multiply-adds
cached ids: [6, 6, 13, 10, 13, 13, 3, 8, 8, 8], same as Example 1: true
multiply-adds: 424,000 uncached, 69,568 cached (6.1 times fewer)
cache holds 13 positions: 1664 bytes in F32Each step after the prompt now computes one position and costs between 5,312 and 5,824 multiply-adds. The weight matrices cost the same at every step. The growth of 64 per step is the attention part, which reads one more cached position each time. The total is the prompt pass (19,456, the same as step 0 of Example 2) plus the nine one-position steps.
The generated ids are identical to those of Example 1, so the cache changes the cost and not the result. The total drops from 424,000 (Example 2) to 69,568 multiply-adds, which is 6.1 times fewer. The gap widens as the sequence grows, because the cached cost per step is nearly constant.
The cache holds 13 positions: 4 from the prompt and 9 from later steps, since the tenth generated token has not been fed back yet. Its 1,664 bytes are 2 (K and V) * 2 layers * 13 positions * 8 values * 4 bytes in F32 (32-bit floating point numbers). The cache is the price paid for the saved work: it uses memory that grows with the sequence.
Example 8: prefill in one batch versus token by token
The two kinds of pass in a generation have names. The first pass over the prompt is called prefill: all prompt tokens go through the model, and their keys and values are stored. Each later pass is a decode step: one token, one new position. Prefill has many tokens available at once, which opens a chance to process them together as one batch.
This example compares two ways to prefill the same prompt. The first passes all four tokens in one call. The second passes them one at a time, each at its own position, as if they were decode steps. The program prints the first new token, the number of calls, the operation count, and the largest difference between the two caches.
import java.util.Arrays;
import java.util.Random;
public class PrefillBatch {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2, MAX = 64;
static final Random R = new Random(42);
static long ops; // multiply-adds counted so far
static float[] mat(int rows, int cols) {
float[] w = new float[rows * cols];
for (int i = 0; i < w.length; i++) w[i] = (float) (R.nextGaussian() / Math.sqrt(cols));
return w;
}
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final float[] EMBED = mat(V, D);
static final Layer[] LAYERS = new Layer[L];
static {
for (int l = 0; l < L; l++)
LAYERS[l] = new Layer(mat(D, D), mat(KVH * HD, D), mat(KVH * HD, D), mat(D, D), mat(F, D), mat(F, D), mat(D, F));
}
static final float[] HEAD = mat(V, D);
static class Cache { // keys and values of every layer and position
final float[][][] k = new float[L][MAX][], v = new float[L][MAX][];
int len;
}
static float[] matvec(float[] w, int rows, int cols, float[] x) {
ops += (long) rows * cols;
float[] y = new float[rows];
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++) y[i] += w[i * cols + j] * x[j];
return y;
}
static float[] rmsnorm(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
float scale = (float) (1.0 / Math.sqrt(ss / x.length + 1e-5));
float[] y = new float[x.length];
for (int i = 0; i < x.length; i++) y[i] = x[i] * scale;
return y;
}
static void rope(float[] v, int heads, int pos) {
for (int h = 0; h < heads; h++)
for (int i = 0; i < HD / 2; i++) {
double a = pos * Math.pow(10000, -2.0 * i / HD);
float c = (float) Math.cos(a), s = (float) Math.sin(a), x0 = v[h * HD + 2 * i], x1 = v[h * HD + 2 * i + 1];
v[h * HD + 2 * i] = x0 * c - x1 * s;
v[h * HD + 2 * i + 1] = x0 * s + x1 * c;
}
}
// Runs `tokens` at positions start, start+1, ... ; keys and values go into the cache; returns the last logits.
static float[] process(int[] tokens, int start, Cache c) {
int n = tokens.length;
float[][] x = new float[n][];
for (int i = 0; i < n; i++) x[i] = Arrays.copyOfRange(EMBED, tokens[i] * D, (tokens[i] + 1) * D);
for (int l = 0; l < L; l++) {
Layer ly = LAYERS[l];
float[][] q = new float[n][];
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
q[i] = matvec(ly.wq, D, D, h);
float[] k = matvec(ly.wk, KVH * HD, D, h), v = matvec(ly.wv, KVH * HD, D, h);
rope(q[i], H, start + i);
rope(k, KVH, start + i);
c.k[l][start + i] = k;
c.v[l][start + i] = v;
}
for (int i = 0; i < n; i++) {
int pos = start + i;
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[pos + 1];
double sum = 0;
for (int s = 0; s <= pos; s++) {
double dot = 0;
for (int j = 0; j < HD; j++) dot += q[i][h * HD + j] * c.k[l][s][kv * HD + j];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= pos; s++)
for (int j = 0; j < HD; j++) cat[h * HD + j] += (float) (w[s] / sum) * c.v[l][s][kv * HD + j];
ops += 2L * HD * (pos + 1);
}
float[] out = matvec(ly.wo, D, D, cat);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
for (int i = 0; i < n; i++) {
float[] h = rmsnorm(x[i]);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int j = 0; j < F; j++) g[j] = (float) (g[j] / (1 + Math.exp(-g[j]))) * u[j];
float[] out = matvec(ly.down, D, F, g);
for (int j = 0; j < D; j++) x[i][j] += out[j];
}
}
c.len = start + n;
return matvec(HEAD, V, D, rmsnorm(x[n - 1])); // only the last position's logits are needed
}
static int argmax(float[] v) {
int best = 0;
for (int i = 1; i < v.length; i++) if (v[i] > v[best]) best = i;
return best;
}
public static void main(String[] args) {
int[] prompt = {3, 17, 5, 9};
Cache batch = new Cache(), single = new Cache();
ops = 0;
int fromBatch = argmax(process(prompt, 0, batch)); // all four tokens in one call
long batchOps = ops;
ops = 0;
int fromSingle = -1;
for (int p = 0; p < prompt.length; p++)
fromSingle = argmax(process(new int[] {prompt[p]}, p, single)); // one token per call
long singleOps = ops;
float maxDiff = 0;
for (int l = 0; l < L; l++)
for (int p = 0; p < prompt.length; p++)
for (int j = 0; j < KVH * HD; j++) {
maxDiff = Math.max(maxDiff, Math.abs(batch.k[l][p][j] - single.k[l][p][j]));
maxDiff = Math.max(maxDiff, Math.abs(batch.v[l][p][j] - single.v[l][p][j]));
}
System.out.println("one call: first new token " + fromBatch + ", calls 1, multiply-adds " + batchOps);
System.out.println("token by token: first new token " + fromSingle + ", calls " + prompt.length + ", multiply-adds " + singleOps);
System.out.println("largest difference between the two caches: " + maxDiff);
}
}one call: first new token 6, calls 1, multiply-adds 19456
token by token: first new token 6, calls 4, multiply-adds 20608
largest difference between the two caches: 0.0Both ways give the same first new token and caches with a largest difference of 0.0, so the results are the same. The token-by-token way costs 1,152 more multiply-adds (20,608 against 19,456). The reason is the last line of process: every call computes the LM head for its last token, and the three extra calls each add a 24 x 16 product (384 multiply-adds) whose logits are thrown away. Apart from that, the two ways do the same arithmetic in the same order, because this Java code multiplies each position's vector by each matrix separately even inside one call.
The difference appears in a real engine. A batched prefill multiplies each weight matrix by all prompt vectors at once (a matrix-matrix product). Every weight is then read from memory once for the whole prompt, instead of once per token. Example 12 counts those reads. That is why engines prefill in batches, and why the timing lines of Example 13 show prefill running many times faster per token than decode.
Example 6: a decode step reads every weight once
Why is generation slow even on fast hardware? The answer is memory. Each decode step multiplies one vector by every weight matrix. A matrix-vector product uses each weight exactly once, so the processor must read the whole model from memory for each new token. Reading a few gigabytes takes time, and that time is often longer than the arithmetic. The program computes how many weights one step reads.
The embedding table is the exception, since a lookup reads one row only. The LM head, when it is a separate matrix, is read in full. The table converts the weight count to megabytes in four number formats. Real files mix formats, so the table assumes every weight has the same number of bits. F32 is 32 bits and F16 is 16 bits. Q8_0 is 8.5 bits (34 bytes for 32 weights), and Q4_0 is 4.5 bits (18 bytes for 32 weights). The last two are quantized formats, which store weights with fewer bits. They are built in Building Q8_0 and Q4_0 Like ggml.
public class BytesPerToken {
record Config(String name, long hidden, long inter, long layers, long heads, long kvHeads, long vocab, boolean tied) {
long perLayer() {
long headDim = hidden / heads;
return 2 * hidden * heads * headDim + 2 * hidden * kvHeads * headDim + 3 * hidden * inter + 2 * hidden;
}
long total() { return vocab * hidden + layers * perLayer() + hidden + (tied ? 0 : vocab * hidden); }
// a decode step reads every weight except the embedding table, of which it reads one row
long readPerToken() { return tied ? total() : total() - vocab * hidden + hidden; }
}
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
Config[] configs = {
new Config("tiny model", 16, 32, 2, 4, 2, 24, false),
new Config("Llama 3.2 1B shape", 2048, 8192, 16, 32, 8, 128256, true),
new Config("Llama 3 8B shape", 4096, 14336, 32, 32, 8, 128256, false)};
double[] bitsPerWeight = {32, 16, 8.5, 4.5}; // F32, F16, Q8_0, Q4_0
String[] formats = {"F32", "F16", "Q8_0", "Q4_0"};
System.out.printf("%-20s %16s", "model", "weights read");
for (String f : formats) System.out.printf("%10s", f + " MB");
System.out.println();
for (Config c : configs) {
System.out.printf("%-20s %,16d", c.name(), c.readPerToken());
for (double bits : bitsPerWeight) System.out.printf("%10.3f", c.readPerToken() * bits / 8 / 1e6);
System.out.println();
}
}
}model weights read F32 MB F16 MB Q8_0 MB Q4_0 MB
tiny model 5,088 0.020 0.010 0.005 0.003
Llama 3.2 1B shape 1,235,814,400 4943.258 2471.629 1313.053 695.146
Llama 3 8B shape 7,504,928,768 30019.715 15009.858 7973.987 4221.522The tiny model reads 5,088 weights per token. The Llama 3 8B shape reads 7,504,928,768. That is the 8,030,261,248 total, minus the 525,336,576 embedding table, plus one row of 4096. At F16 this is about 15.0 GB per token, at Q8_0 about 8.0 GB, and at Q4_0 about 4.2 GB. Fewer bits per weight means fewer bytes to read, which is why quantization speeds up decoding. The Llama 3.2 1B shape is tied, so the table doubles as the LM head and the whole 1.24B parameters are read.
A decode step does about 2 FLOPs (floating-point operations: one multiply and one add) for every weight it reads. That is a small amount of arithmetic per byte, so memory bandwidth, the number of bytes per second the memory can deliver, sets the pace instead of arithmetic speed. FLOPs, Bandwidth, and the Roofline Model develops this.
Example 7: time to first token and tokens per second from counts
Now the byte counts can predict speed. Decode speed is memory bandwidth divided by the bytes read per token. If a step reads 4.22 GB and the memory delivers 60 GB per second, one step takes about 70 ms. Prefill is different. The weights are read once for the whole prompt, but the arithmetic grows with the prompt length, at about 2 * weights * prompt tokens FLOPs. So prefill takes the larger of two times: the compute time and the byte-reading time. Time to first token (TTFT), how long the user waits before output starts, is about the prefill time.
The machines below are assumptions with round numbers, not measurements or published specifications. The model is the 8B shape in Q4_0, and attention FLOPs are ignored. The program prints the decode speed of each machine, and then the TTFT and the total time of a 200-token answer for prompts of 100, 1,000, and 8,000 tokens.
The machines below are assumptions with round numbers, not measurements or published specifications. The model is the 8B shape in Q4_0, and attention FLOPs are ignored.
public class Speeds {
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
double params = 7_504_928_768.0; // weights one decode step reads, Llama 3 8B shape
double bytes = params * 4.5 / 8; // Q4_0: 4.5 bits per weight
// assumed machines, not measurements
String[] names = {"laptop CPU (60 GB/s, 1 TFLOP/s)", "GPU (1000 GB/s, 100 TFLOP/s)"};
double[] bandwidth = {60e9, 1000e9}, flops = {1e12, 100e12};
int[] promptLengths = {100, 1000, 8000};
System.out.printf("bytes per decode step: %.2f GB%n", bytes / 1e9);
for (int m = 0; m < names.length; m++) {
double decodeSeconds = bytes / bandwidth[m];
System.out.printf("%s%n decode: %.1f ms per token = %.1f tokens per second%n", names[m], decodeSeconds * 1000, 1 / decodeSeconds);
for (int n : promptLengths) {
double prefillFlops = 2 * params * n; // about 2 FLOPs per weight per prompt token
double prefillSeconds = Math.max(prefillFlops / flops[m], bytes / bandwidth[m]);
double total200 = prefillSeconds + 199 * decodeSeconds;
System.out.printf(" prompt %5d tokens: TTFT %7.2f s, 200-token answer done after %7.2f s%n", n, prefillSeconds, total200);
}
}
}
}bytes per decode step: 4.22 GB
laptop CPU (60 GB/s, 1 TFLOP/s)
decode: 70.4 ms per token = 14.2 tokens per second
prompt 100 tokens: TTFT 1.50 s, 200-token answer done after 15.50 s
prompt 1000 tokens: TTFT 15.01 s, 200-token answer done after 29.01 s
prompt 8000 tokens: TTFT 120.08 s, 200-token answer done after 134.08 s
GPU (1000 GB/s, 100 TFLOP/s)
decode: 4.2 ms per token = 236.9 tokens per second
prompt 100 tokens: TTFT 0.02 s, 200-token answer done after 0.86 s
prompt 1000 tokens: TTFT 0.15 s, 200-token answer done after 0.99 s
prompt 8000 tokens: TTFT 1.20 s, 200-token answer done after 2.04 sOn the assumed CPU, decode runs at 14.2 tokens per second (4.22 GB / 60 GB/s = 70.4 ms per token). A 1,000-token prompt needs 2 * 7.5e9 * 1000 = 1.5e13 FLOPs, and at 1 TFLOP/s that takes 15 seconds before the first token. On the assumed GPU, decode runs at 236.9 tokens per second and the same prompt takes 0.15 seconds.
Prefill time grows in proportion to the prompt length, and decode time grows in proportion to the answer length. This gives two cases. A long prompt with a short answer is dominated by TTFT. A short prompt with a long answer is dominated by the decode speed. Real engines add attention work, overheads, and kernels that fall short of the peak numbers, so measured values are lower. A later lesson on benchmarking covers measurement.
Example 8: KV cache bytes per token
The cache speeds up generation but costs memory, and the amount depends on the model shape. Each token stores one key vector and one value vector per layer per key/value head: 2 * layers * kv_heads * head_dim * bytes_per_value. The 2 is for key and value. For example, the Llama 3 8B shape has 32 layers, 8 key/value heads, a head dimension of 128, and 2 bytes per value in F16, so one token needs 2 * 32 * 8 * 128 * 2 = 131,072 bytes.
This program evaluates the formula for five shapes at 4,096 and 32,768 tokens. Memory Math gives the full treatment, including quantized caches.
public class KvSize {
record Shape(String name, int layers, int kvHeads, int headDim, int bytesPerValue) {
long perToken() { return 2L * layers * kvHeads * headDim * bytesPerValue; } // K and V
}
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
Shape[] shapes = {
new Shape("tiny model of lesson 7 (F32)", 2, 2, 4, 4),
new Shape("tiny GGUF of the course (F16)", 4, 2, 32, 2),
new Shape("Llama 3.2 1B shape (F16)", 16, 8, 64, 2),
new Shape("Llama 3 8B shape (F16)", 32, 8, 128, 2),
new Shape("Llama 2 7B shape (F16)", 32, 32, 128, 2)};
System.out.printf("%-30s %12s %12s %12s%n", "model", "per token", "4096 tokens", "32768 tokens");
for (Shape s : shapes)
System.out.printf("%-30s %9d B %9.2f MiB %9.2f MiB%n", s.name(), s.perToken(),
s.perToken() * 4096 / 1048576.0, s.perToken() * 32768 / 1048576.0);
}
}model per token 4096 tokens 32768 tokens
tiny model of lesson 7 (F32) 128 B 0.50 MiB 4.00 MiB
tiny GGUF of the course (F16) 1024 B 4.00 MiB 32.00 MiB
Llama 3.2 1B shape (F16) 32768 B 128.00 MiB 1024.00 MiB
Llama 3 8B shape (F16) 131072 B 512.00 MiB 4096.00 MiB
Llama 2 7B shape (F16) 524288 B 2048.00 MiB 16384.00 MiBThe tiny GGUF row is a check against a real engine. The course's tiny model (4 layers, 2 key/value heads, head dimension 32, F16 cache) needs 1,024 bytes per token, which is 4.00 MiB at 4,096 tokens. The llama-completion log in Example 13 prints K (f16): 2.00 MiB, V (f16): 2.00 MiB for the same model and context, so the formula matches the real program.
For the Llama 3 8B shape the cache costs 131,072 bytes (128 KiB) per token, so 4,096 tokens take 512 MiB and 32,768 tokens take 4 GiB. The Llama 2 7B shape has 32 key/value heads instead of 8, and costs four times as much per token: 0.5 MiB. Grouped-query attention, where several query heads share one key/value head, is the difference.
Example 9: what happens when the context window is full
A model has a maximum number of positions it can hold in the cache, called the context window and written n_ctx. What should an engine do when a conversation fills it? This example shows the choices. The default in llama.cpp is to stop. In the pinned commit, context shift is off by default in both llama-completion and llama-server (the flag is --context-shift). llama-completion then logs "context full and context shift is disabled => stopping", and the server ends the request with stop type limit.
With context shift enabled, the engine makes room by dropping old tokens. It keeps the first n_keep tokens, discards half of the tokens after them (n_left / 2), and moves the remaining tokens down by the discarded count. The server accepts n_keep and n_discard in the request.
There is a catch. The surviving tokens get new positions, but their cached keys were rotated by RoPE for their old positions. Recomputing them would defeat the purpose of the cache. Because a RoPE rotation by an angle can be undone by the opposite angle, the engine instead rotates each cached key by -n_discard positions. The program below applies the shift rule to a list of positions with n_ctx = 12 and n_keep = 2. Then it checks the rotation trick on one key vector.
With context shift enabled, the engine keeps the first n_keep tokens, discards half of the tokens after them (n_left / 2), and moves the remaining tokens down by the discarded count. The server accepts n_keep and n_discard in the request. The program below applies the same rule to a list of positions. The positions of the surviving tokens change, and their cached keys carry rotations for their old positions. Because a RoPE rotation by an angle can be undone by the opposite angle, the engine rotates each cached key by -n_discard positions instead of recomputing it. The program checks that on one key vector.
import java.util.ArrayList;
import java.util.List;
public class ContextShift {
// rotates adjacent pairs by pos * theta_i (the RoPE of lesson 7); a negative pos rotates backwards
static float[] rope(float[] v, double pos, double base) {
float[] out = new float[v.length];
for (int i = 0; i < v.length / 2; i++) {
double a = pos * Math.pow(base, -2.0 * i / v.length);
float c = (float) Math.cos(a), s = (float) Math.sin(a);
out[2 * i] = v[2 * i] * c - v[2 * i + 1] * s;
out[2 * i + 1] = v[2 * i] * s + v[2 * i + 1] * c;
}
return out;
}
public static void main(String[] args) {
int nCtx = 12, nKeep = 2;
List<Integer> positions = new ArrayList<>(); // position of each cached token
for (int p = 0; p < nCtx; p++) positions.add(p);
System.out.println("full context: " + positions);
int nLeft = positions.size() - nKeep;
int nDiscard = nLeft / 2; // llama.cpp discards half of what follows n_keep
List<Integer> kept = new ArrayList<>(positions.subList(0, nKeep));
for (int p : positions.subList(nKeep + nDiscard, positions.size())) kept.add(p - nDiscard);
System.out.println("discard " + nDiscard + " tokens after the first " + nKeep + ": positions now " + kept);
// a cached key was rotated for its old position; rotating it by -nDiscard positions gives the new one
float[] key = {0.3f, -1.2f, 0.8f, 0.5f, -0.4f, 0.9f, 1.1f, -0.7f};
float[] cachedAt9 = rope(key, 9, 10000);
float[] shifted = rope(cachedAt9, -nDiscard, 10000);
float[] freshAt4 = rope(key, 9 - nDiscard, 10000);
float maxDiff = 0;
for (int i = 0; i < key.length; i++) maxDiff = Math.max(maxDiff, Math.abs(shifted[i] - freshAt4[i]));
System.out.println("key cached at position 9, rotated by -" + nDiscard + ", versus a key rotated for position " + (9 - nDiscard)
+ ": largest difference " + (maxDiff < 1e-6f ? "below 1e-6" : maxDiff));
}
}full context: [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
discard 5 tokens after the first 2: positions now [0, 1, 2, 3, 4, 5, 6]
key cached at position 9, rotated by -5, versus a key rotated for position 4: largest difference below 1e-6Five tokens are dropped after the first two (n_left = 12 - 2 = 10, and half of 10 is 5), and the seven survivors take positions 0 to 6. A key cached at position 9 and rotated back by 5 positions matches a key rotated for position 4, with a largest difference below 1e-6. So the engine can shift the cache without recomputing keys. The check uses RoPE from lesson 10.
A shifted cache is not the same as a fresh prefill of the surviving text. The cached keys and values were computed while the discarded tokens were visible, and the discarded content is gone for later tokens. Agents that need long histories compact or summarize instead of relying on shifts. Context engineering, later in this course, covers that.
Example 10: stop conditions
A generation loop needs a rule for when to stop, or it would run forever. This example shows the three conditions that apply.
- The model can emit an end-of-sequence token (id 2 in the program), which means "I am done".
- The text can contain a stop string supplied by the caller, such as
\n\norUser:. - The count can reach the maximum number of new tokens, set by the caller.
The program checks them in this order: the end token first, then stop strings after each token is appended, then the limit. It generates from a fixed list of token pieces and runs four settings, so you can see which condition fires in each. The output names the reason for each stop.
public class StopConditions {
static final int EOS = 2;
static String esc(String s) { return "'" + s.replace("\n", "\\n") + "'"; }
// pieces[i] is the text of token ids[i]; returns the generated text and the reason it ended
static String generate(int[] ids, String[] pieces, String[] stops, int maxTokens) {
StringBuilder text = new StringBuilder();
for (int n = 0; n < ids.length; n++) {
if (ids[n] == EOS) return esc(text.toString()) + " after " + n + " tokens, stop: eos";
text.append(pieces[n]);
for (String stop : stops) {
int at = text.indexOf(stop);
if (at >= 0) return esc(text.substring(0, at)) + " after " + (n + 1) + " tokens, stop: word (" + stop.replace("\n", "\\n") + ")";
}
if (n + 1 >= maxTokens) return esc(text.toString()) + " after " + (n + 1) + " tokens, stop: limit";
}
return esc(text.toString()) + " stream ended";
}
public static void main(String[] args) {
int[] ids = {10, 11, 12, 13, 14, EOS, 15};
String[] pieces = {"The", " answer", " is", " 42", "\n\n", "<eos>", " never shown"};
System.out.println("no stop string: " + generate(ids, pieces, new String[0], 100));
System.out.println("stop string \\n\\n: " + generate(ids, pieces, new String[] {"\n\n"}, 100));
System.out.println("max tokens 3: " + generate(ids, pieces, new String[0], 3));
System.out.println("stop string ' 42': " + generate(ids, pieces, new String[] {" 42"}, 100));
}
}no stop string: 'The answer is 42\n\n' after 5 tokens, stop: eos
stop string \n\n: 'The answer is 42' after 5 tokens, stop: word (\n\n)
max tokens 3: 'The answer is' after 3 tokens, stop: limit
stop string ' 42': 'The answer is' after 4 tokens, stop: word ( 42)The first run ends at the end token, with \n\n still in the text. The second run ends when the text contains \n\n, and the stop string is cut from the output. The third run ends at the token limit of 3. The fourth run ends at the stop string 42, after 4 tokens, and the text before it is kept.
Engines report the reason. OpenAI-style APIs use finish_reason: stop for an end token or stop string, and length for the limit. The llama-server native /completion response uses stop_type with the values eos, word, and limit, plus stopping_word. A request to the tiny model with "stop": ["clang"] returned stop_type: "word" and stopping_word: "clang" after 2 generated tokens, and the content ended before the stop string.
Chat models often end a turn with a special token other than the vocabulary's end-of-sequence id (such as <|im_end|>). llama.cpp treats a set of such tokens as end of generation (llama_vocab_is_eog), which is why a wrong chat template can break stopping (Chat Templates, Roles, and Special Tokens).
Example 11: matching a stop string split across streamed pieces
Streaming sends text to the user as tokens arrive, and a stop string can span several tokens. Suppose the stop string is \nUser: and the model produces the pieces \nUs, er, and :. If the program shows each piece at once, the reader sees \nUs before the program knows a stop string is coming. The text has already been shown and cannot be taken back.
The solution is to hold back the longest tail of the text that could still become the stop string. When the next piece shows that the tail cannot become the stop string, the tail is released. When the full stop string appears, the program stops and shows only the text before it. The class below does this for one stop string. It also has flush, which releases the held text when generation ends without a match. The program runs three cases: a full match, a near match that fails, and an ending while text is held.
public class StopMatcher {
private final String stop;
private final StringBuilder held = new StringBuilder();
boolean stopped;
StopMatcher(String stop) { this.stop = stop; }
// returns the text that is safe to show now
String push(String piece) {
held.append(piece);
int at = held.indexOf(stop);
if (at >= 0) { // full match: show what precedes it and stop
stopped = true;
String out = held.substring(0, at);
held.setLength(0);
return out;
}
int keep = 0; // longest suffix of the buffer that starts the stop string
for (int k = Math.min(stop.length() - 1, held.length()); k > 0; k--)
if (held.substring(held.length() - k).equals(stop.substring(0, k))) { keep = k; break; }
String out = held.substring(0, held.length() - keep);
held.delete(0, held.length() - keep);
return out;
}
String flush() { String out = held.toString(); held.setLength(0); return out; }
static String esc(String s) { return "'" + s.replace("\n", "\\n") + "'"; }
static void run(String stop, String... pieces) {
StopMatcher m = new StopMatcher(stop);
StringBuilder shown = new StringBuilder();
System.out.println("stop string " + esc(stop));
for (String piece : pieces) {
if (m.stopped) break;
String out = m.push(piece);
shown.append(out);
System.out.printf(" piece %-8s shows %-8s holds %s%n", esc(piece), esc(out), esc(m.held.toString()));
}
if (!m.stopped) shown.append(m.flush()); // generation ended without a match: release the held text
System.out.println(" shown " + esc(shown.toString()) + ", stopped: " + m.stopped);
}
public static void main(String[] args) {
run("\nUser:", "Hello", " there", "\nUs", "er", ":", " ignored");
run("\nUser:", "Hi", "\nUs", "ing", " it", "\n");
run("\nUser:", "Done", "\nUs");
}
}stop string '\nUser:'
piece 'Hello' shows 'Hello' holds ''
piece ' there' shows ' there' holds ''
piece '\nUs' shows '' holds '\nUs'
piece 'er' shows '' holds '\nUser'
piece ':' shows '' holds ''
shown 'Hello there', stopped: true
stop string '\nUser:'
piece 'Hi' shows 'Hi' holds ''
piece '\nUs' shows '' holds '\nUs'
piece 'ing' shows '\nUsing' holds ''
piece ' it' shows ' it' holds ''
piece '\n' shows '' holds '\n'
shown 'Hi\nUsing it\n', stopped: false
stop string '\nUser:'
piece 'Done' shows 'Done' holds ''
piece '\nUs' shows '' holds '\nUs'
shown 'Done\nUs', stopped: falseIn the first run the pieces \nUs, er, and : arrive separately. The matcher holds \nUs, then \nUser, and then sees the full stop string and stops. It shows only Hello there, so the reader never sees any part of the stop string.
In the second run the held \nUs is released as soon as the next piece (ing) breaks the match, and the text is shown in full. In the third run the generation ends while \nUs is held, and flush releases it, because no stop string followed. The server in llama.cpp uses a helper named string_find_partial_stop for the same purpose while streaming.
Example 12: batching, one weight read for several sequences
Example 6 showed that a decode step is limited by reading the weights. This leads to a trick: if the weights are being read anyway, use them for several requests at once. That is batching. The program multiplies one weight matrix by a batch of vectors in two ways. The first runs each sequence separately and reads the matrix once per sequence. The second reads each weight once and applies it to every sequence's vector. The results are identical, and the weight reads are divided by the batch size.
The first table counts weight loads and multiply-adds for batch sizes 1 to 16. The second table uses real-model numbers and shows the catch: the weights are shared, but each sequence has its own KV cache.
import java.util.Random;
public class BatchedDecode {
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
int rows = 32, cols = 16;
Random r = new Random(9);
float[] w = new float[rows * cols];
for (int i = 0; i < w.length; i++) w[i] = (float) r.nextGaussian();
System.out.println("batch weight loads (one by one / batched) multiply-adds per weight load same results");
for (int batch : new int[] {1, 2, 4, 8, 16}) {
float[][] x = new float[batch][cols];
for (float[] v : x) for (int j = 0; j < cols; j++) v[j] = (float) r.nextGaussian();
float[][] one = new float[batch][rows], together = new float[batch][rows];
long loadsOne = 0, loadsTogether = 0, macs = 0;
for (int b = 0; b < batch; b++) // one sequence at a time: reads W again each time
for (int i = 0; i < rows; i++)
for (int j = 0; j < cols; j++) { one[b][i] += w[i * cols + j] * x[b][j]; loadsOne++; }
for (int i = 0; i < rows; i++) // batched: each weight is read once for all sequences
for (int j = 0; j < cols; j++) {
float weight = w[i * cols + j];
loadsTogether++;
for (int b = 0; b < batch; b++) { together[b][i] += weight * x[b][j]; macs++; }
}
boolean same = java.util.Arrays.deepEquals(one, together);
System.out.printf("%5d %12d / %-12d %14d %16.0f %s%n", batch, loadsOne, loadsTogether, macs, (double) macs / loadsTogether, same);
}
// bytes read per decode step: weights are shared by the batch, each sequence has its own KV cache
double weights = 4.22e9, kvPerSequence = 2000 * 131072.0; // 8B shape, Q4_0 weights, 2000 tokens of F16 KV
System.out.println("batch bytes per step (GB) bytes per generated token (GB)");
for (int batch : new int[] {1, 4, 16, 64}) {
double step = weights + batch * kvPerSequence;
System.out.printf("%5d %18.2f %28.2f%n", batch, step / 1e9, step / batch / 1e9);
}
}
}batch weight loads (one by one / batched) multiply-adds per weight load same results
1 512 / 512 512 1 true
2 1024 / 512 1024 2 true
4 2048 / 512 2048 4 true
8 4096 / 512 4096 8 true
16 8192 / 512 8192 16 true
batch bytes per step (GB) bytes per generated token (GB)
1 4.48 4.48
4 5.27 1.32
16 8.41 0.53
64 21.00 0.33For a batch of 16, the matrix is loaded 512 times instead of 8,192, and each load feeds 16 multiply-adds. The results are the same (true in the last column). This is the reason a server handles many sequences at once.
The second table counts bytes for the 8B shape with Q4_0 weights and 2,000 tokens of F16 cache per sequence. The weights (4.22 GB) are shared by the batch, but each sequence reads its own cache (262 MB). Bytes per generated token fall from 4.48 GB at batch 1 to 0.33 GB at batch 64. Then they flatten near the per-sequence cache size, which no batching removes. A later lesson on batching and scheduling covers how servers join and leave batches.
The second table counts bytes for the 8B shape with Q4_0 weights and 2,000 tokens of F16 cache per sequence. The weights (4.22 GB) are shared by the batch, but each sequence reads its own cache (262 MB). Bytes per generated token fall from 4.48 GB at batch 1 to 0.33 GB at batch 64, and then flatten near the per-sequence cache size, which no batching removes. A later lesson on batching and scheduling covers how servers join and leave batches.
Example 13: reading llama.cpp prompt-eval and eval timings (and the server timings)
Every idea above appears in the output of the real tools. This example reads that output. The command runs the tiny random-weight model (Llama architecture, 4 layers, F32) with llama-completion and keeps only the cache-size and timing lines. The model has random weights, so the generated text is meaningless, but the log lines are real.
llama-completion -m tiny-random-qwen2vocab-f32.gguf -p "Write a function that reverses a string" \
-n 32 --seed 1 -no-cnv --temp 0 -v 2>&1 | grep -E "llama_kv_cache: size|common_perf_print"The -v flag (verbose logging) is needed for the cache line. The output below is representative, from the pinned llama.cpp build. Timestamps are removed, the verbose log prints the cache line twice and one copy is shown, and timings differ by machine and by load.
llama_kv_cache: size = 4.00 MiB ( 4096 cells, 4 layers, 1/1 seqs), K (f16): 2.00 MiB, V (f16): 2.00 MiB
common_perf_print: sampling time = 11.14 ms
common_perf_print: samplers time = 4.57 ms / 40 tokens
common_perf_print: load time = 208.75 ms
common_perf_print: prompt eval time = 4.02 ms / 8 tokens ( 0.50 ms per token, 1990.05 tokens per second)
common_perf_print: eval time = 66.52 ms / 31 runs ( 2.15 ms per token, 466.01 tokens per second)
common_perf_print: total time = 84.04 ms / 39 tokensThe prompt eval line is prefill: 8 prompt tokens processed in one batch. The eval line is decode: 31 runs of one token each, 2.15 ms per token. The command asked for 32 new tokens, but the log shows 31 decode runs, because the first new token comes from the logits of the prefill pass. That is why total time counts 39 tokens (8 + 31).
The ratio of the two speeds (1,990 against 466 tokens per second) shows that prefill is faster per token than decode, as Examples 5 and 12 explain: a batch reads each weight once for all prompt tokens.
The cache line confirms Example 8: 4096 cells (one per position), 4 layers, 2.00 MiB for keys and 2.00 MiB for values. The load time is the time to load the model from disk and is not part of generation.
llama-server returns the same data in a timings object inside every /completion response. A slot is one independent sequence the server can work on, with its own KV cache. Two requests to the tiny model whose prompts share a prefix (15 tokens in the second) gave this object for the second request:
{"cache_n": 12, "prompt_n": 3, "prompt_ms": 5.393, "predicted_n": 16, "predicted_ms": 65.112,
"predicted_per_token_ms": 4.3408, "predicted_per_second": 230.37228160707707}cache_n is the number of prompt tokens the server found already in the slot's KV cache from the previous request. prompt_n is the number it had to process. Here 12 of 15 tokens were reused, so prefill covered only 3 tokens. This is prefix reuse: a request that starts with the same tokens as an earlier request in the same slot skips their prefill. Later lessons on paged KV caches, prefix caching, and llama-server slots cover it in depth. Chat Templates, Roles, and Special Tokens shows how to write prompts that keep a stable prefix.
Two more details help when reading the numbers. In the pinned source the per-second figure divides by predicted_n - 1 decode steps, since the first token is free: 15 steps in 65.112 ms is 230.37 steps per second. And prompt_ms is the prefill time, so a client sees a time to first token close to prompt_ms plus network time.
Interview points
- Autoregressive generation runs one forward pass per new token, and each pass depends on the token chosen in the previous one.
- Without a cache the total work grows with the square of the length, because each step recomputes every position. With a cache each step computes one position.
- The cache is valid because of the causal mask. The keys and values of a position depend only on earlier tokens, so appending a token never changes them.
- Prefill processes the prompt as one batch and is limited by arithmetic. Decode processes one token per pass and is limited by memory bandwidth.
- A decode step reads every weight once, so tokens per second is about bandwidth divided by model bytes. Quantization raises decode speed by shrinking the bytes.
- TTFT (time to first token) is about the prefill time. It grows with prompt length and falls with compute speed.
- KV cache bytes per token are
2 * layers * kv_heads * head_dim * bytes_per_value. Grouped-query attention reduces this byheads / kv_heads. - When the context is full, an engine stops or discards tokens. At the pinned llama.cpp commit it stops by default and shifts when
--context-shiftis on, rotating the cached keys with RoPE. - Stop conditions are the end token, stop strings, and the token limit. Streaming with stop strings needs a holdback buffer for partial matches.
- Batching shares weight reads across sequences but not the KV cache, so large batches become limited by cache reads.
Summary
- The autoregressive loop appends one chosen token per forward pass.
- Keys and values of earlier positions do not change when tokens are appended, so they can be cached. A cache turns repeated full-sequence passes into one new position per step with identical output.
- Prefill batches the prompt and decode handles one token at a time. Batch and token-by-token prefill give the same cache.
- A decode step reads all weights except the embedding table (one row only). For the Llama 3 8B shape that is 7,504,928,768 weights, about 4.2 GB at 4.5 bits and 15.0 GB at 16 bits.
- Decode speed is about memory bandwidth divided by bytes per token. Prefill time is about
2 * weights * prompt tokensdivided by compute speed. - KV cache per token is
2 * layers * kv_heads * head_dim * bytes. It is 131,072 bytes for the Llama 3 8B shape in F16. - At the pinned llama.cpp commit context shift is off by default. When on, it keeps
n_keeptokens, discards half of the rest, and rotates cached keys by-n_discard. - Stop conditions are the end token, stop strings, and the token limit.
llama-serverreturns astop_typeofeos,word, orlimit. - Batching divides weight reads by the batch size and leaves per-sequence cache reads unchanged.
llama-completionreports prompt eval (prefill) and eval (decode) separately, and the first new token is not counted as a decode run.llama-servertimings includecache_n(tokens reused from the slot cache) andprompt_n(tokens processed).