Lesson 10 of 73
The Transformer Block and the Full Model
A language model such as Llama looks complicated from the outside, but its structure is simple. It is one small unit, called a block, repeated many times. At the start sits an embedding table, which turns token ids into vectors. At the end sits one output matrix, which turns vectors back into scores for the next token. "Decoder-only" means the model has only this one kind of stack, and it reads text from left to right.
This lesson builds that structure in Java, piece by piece. First it turns a token id into a vector. Then it adds the parts of a block one at a time: normalization, the MLP (a small two-layer network), and position information. Next it assembles one block and stacks blocks into a full model. It ends by counting the parameters of a model from its config.json and by reading the names that weights carry inside checkpoint files.
Afterwards you can write the forward pass of a Llama-style model in about 100 lines of Java, predict a model's size from its configuration, and match tensor names in a file to the parts of the model. The lesson builds on self-attention, which explains how each position reads the earlier positions, and on the math of transformers, which explains vectors, matrices, and softmax. The table lists the pieces in the order the lesson meets them.
| Idea | What it does |
|---|---|
| Embedding lookup | Turns a token id into a vector by selecting a row |
| Residual stream | The running vector each sublayer adds its result to |
| Pre-norm RMSNorm | Rescales the vector before each sublayer |
| MLP and SwiGLU | Transforms each position on its own |
| RoPE | Encodes position by rotating query and key vectors |
| Block, stack, LM head | One block repeated N times, then a projection to vocabulary logits |
| Config and tensor names | Where the sizes and the weights are written down |
Here is the whole model as a picture, read from top to bottom. Token ids go in and logits come out. A logit is one raw score per vocabulary entry, and the highest score marks the model's favourite next token. In the picture, x is the vector that represents one position, and x = x + ... means a sublayer computes an update and adds it to x (Example 2 shows this). The words in square brackets are the shapes of the data: tokens is the number of input tokens, dim is the length of each vector, and vocab is the number of tokens the model knows. A stage such as Attention(RMSNorm(x)) reads from the inside out: normalize first, then run attention.
token ids
-> embedding lookup [tokens, dim]
-> block 1 ... block N, each one:
x = x + Attention(RMSNorm(x)) (RoPE on queries and keys)
x = x + MLP(RMSNorm(x))
-> final RMSNorm
-> LM head (matrix: vocabulary x dim) [tokens, vocab] logitsThe programs below use tiny sizes: a model dimension dim of 4 to 16, and a vocabulary of 6 to 24 tokens. With sizes this small you can read every number the program prints and check it by hand. Real models run the same code with larger numbers. Weights come from a seeded java.util.Random, so the numbers are the same on every run. They are random, not trained, so the outputs show how the structure works and say nothing about language.
Example 1: token embedding lookup
A neural network multiplies matrices, but a token is only an integer id such as 1 for "the". The model needs a way to turn an id into a list of numbers it can compute with. The solution is a lookup table with one row per token id. Looking up a token copies its row. The row is the token's vector, and its length is the model dimension dim. This table is called the embedding table. In a trained model the rows are learned, so that tokens with related meanings end up with related rows.
The program below uses a vocabulary of 6 tokens and dim = 4. It looks up the sentence "the cat sat on the mat" (with a start token <s> in front) and prints the resulting matrix. A sequence of n tokens becomes an n by dim matrix, with one row per position. This matrix is the first value of the residual stream, which Example 2 explains. Watch rows 1 and 5: both positions hold the same word.
public class EmbeddingLookup {
public static void main(String[] args) {
String[] vocab = {"<s>", "the", "cat", "sat", "on", "mat"};
int dim = 4;
// 6 rows x 4 columns, row-major: row i starts at i * dim
float[] table = {
0.0f, 0.1f, 0.0f, -0.1f, // <s>
0.5f, 0.2f, -0.3f, 0.1f, // the
-0.4f, 0.9f, 0.3f, 0.6f, // cat
0.7f, -0.5f, 0.2f, 0.8f, // sat
0.3f, 0.3f, -0.6f, 0.0f, // on
-0.2f, 0.8f, 0.4f, -0.5f // mat
};
int[] tokens = {0, 1, 2, 3, 4, 1, 5}; // <s> the cat sat on the mat
float[][] x = new float[tokens.length][dim];
for (int t = 0; t < tokens.length; t++) {
System.arraycopy(table, tokens[t] * dim, x[t], 0, dim);
}
System.out.println("shape: [" + x.length + ", " + x[0].length + "]");
for (int t = 0; t < tokens.length; t++) {
System.out.printf("%d %-4s %s%n", t, vocab[tokens[t]], java.util.Arrays.toString(x[t]));
}
System.out.println("rows 1 and 5 equal: " + java.util.Arrays.equals(x[1], x[5]));
}
}shape: [7, 4]
0 <s> [0.0, 0.1, 0.0, -0.1]
1 the [0.5, 0.2, -0.3, 0.1]
2 cat [-0.4, 0.9, 0.3, 0.6]
3 sat [0.7, -0.5, 0.2, 0.8]
4 on [0.3, 0.3, -0.6, 0.0]
5 the [0.5, 0.2, -0.3, 0.1]
6 mat [-0.2, 0.8, 0.4, -0.5]
rows 1 and 5 equal: trueThe table is a flat float[] in row-major order, which means the rows are stored one after another. Row i starts at index i * dim, so the code copies dim numbers starting at tokens[t] * dim. For cat (id 2) that is index 8, and the four numbers there are -0.4, 0.9, 0.3, 0.6, which matches the output.
The shape [7, 4] reads as 7 positions with 4 numbers each. Rows 1 and 5 are equal because both positions hold the token the. This shows a limit of the lookup: at this point the vector for the carries no information about its neighbours or its position. Attention and the later layers add that information. For scale, in Llama 3 8B the table has 128,256 rows and 4096 columns, which is 525,336,576 numbers.
Example 2: the residual stream
A block has two sublayers, attention and the MLP. Each one reads the vector x of a position and computes an update. Where does the update go? It does not replace x. It is added to x: x = x + f(x), where f stands for whatever the sublayer computes. The vector x of each position, which every sublayer reads from and writes to, is called the residual stream. Think of it as a shared notebook for one position. Each sublayer reads the notebook and appends a note, and nobody erases what is already written. The analogy stops being accurate in one way: the "notes" are numbers added into the same slots, so they mix with the existing values instead of sitting apart.
The program applies three made-up sublayers to the embedding of the and keeps the sum of their updates. Each sublayer takes a rotated copy of x and scales it by its strength (0.1, 0.2, 0.0). The third sublayer has strength 0, so it writes nothing. Look at whether the stream changes after each step, and compare the last line with the final stream.
public class ResidualStream {
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
float[] x = {0.5f, 0.2f, -0.3f, 0.1f}; // the embedding
float[] sumOfDeltas = new float[4];
float[] strength = {0.1f, 0.2f, 0.0f}; // three sublayers; the last one writes nothing
for (int k = 0; k < strength.length; k++) {
float[] delta = new float[4];
if (strength[k] != 0) for (int i = 0; i < 4; i++) delta[i] = strength[k] * x[(i + k) % 4]; // f(x)
for (int i = 0; i < 4; i++) { x[i] += delta[i]; sumOfDeltas[i] += delta[i]; } // x = x + f(x)
System.out.printf("sublayer %d: delta %.3f %.3f %.3f %.3f, stream %.3f %.3f %.3f %.3f%n",
k, delta[0], delta[1], delta[2], delta[3], x[0], x[1], x[2], x[3]);
}
System.out.printf("embedding + sum of deltas: %.3f %.3f %.3f %.3f%n",
0.5f + sumOfDeltas[0], 0.2f + sumOfDeltas[1], -0.3f + sumOfDeltas[2], 0.1f + sumOfDeltas[3]);
}
}sublayer 0: delta 0.050 0.020 -0.030 0.010, stream 0.550 0.220 -0.330 0.110
sublayer 1: delta 0.044 -0.066 0.022 0.110, stream 0.594 0.154 -0.308 0.220
sublayer 2: delta 0.000 0.000 0.000 0.000, stream 0.594 0.154 -0.308 0.220
embedding + sum of deltas: 0.594 0.154 -0.308 0.220Each line shows the update (delta) and the stream after adding it. In sublayer 0 the stream starts at 0.5, 0.2, -0.3, 0.1, the update is 0.05, 0.02, -0.03, 0.01, and the new stream is 0.55, 0.22, -0.33, 0.11. The third update is all zeros, so the stream does not change. The last line adds the original embedding and every update by hand and gets exactly the final stream. So the stream is always "the embedding plus every update so far".
The additive form has two consequences that matter for large models.
- A sublayer with nothing useful to add can output zeros, and the information from the embedding still reaches the last layer unchanged.
- During training, the gradient (the signal that tells each weight how to change) also passes through the additions to every earlier layer. This keeps the signal strong enough for stacks of dozens of layers to be trainable.
Example 3: pre-norm with a learned RMSNorm weight
Every sublayer adds to the stream, so the stream grows as layers accumulate. A matrix multiply is sensitive to the size of its input: an input ten times larger gives an output ten times larger. If the size drifts from layer to layer, the layers have to cope with inputs of very different scales. The fix is to rescale the vector to a standard size before each sublayer reads it.
RMSNorm (root mean square normalization) does this in two steps. First it divides the vector by its root mean square, which is the square root of the average of the squared values. This gives the vector a typical size of 1. Then it multiplies each coordinate by a learned weight, so the model can still choose to emphasize some coordinates. In the pre-norm layout used by Llama-style models, the norm is applied only to the input of the sublayer, and the stream itself stays unnormalized: x = x + f(norm(x)). "Pre" means the norm comes before the sublayer.
The program normalizes [3, -4, 12, 0] with a weight of all ones, then with a learned weight, then normalizes a vector ten times larger, and finally shows one residual update. Check whether the ten times larger input gives a different output.
public class PreNorm {
static float[] rmsnorm(float[] x, float[] w) {
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 * w[i];
return y;
}
static String show(float[] v) { return String.format("%.4f %.4f %.4f %.4f", v[0], v[1], v[2], v[3]); }
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
float[] x = {3f, -4f, 12f, 0f}; // root mean square = 6.5
float[] ones = {1f, 1f, 1f, 1f}, learned = {1f, 1f, 0.5f, 2f};
System.out.println("weight 1: " + show(rmsnorm(x, ones)));
System.out.println("learned: " + show(rmsnorm(x, learned)));
float[] big = {30f, -40f, 120f, 0f}; // 10 * x
System.out.println("10 * x: " + show(rmsnorm(big, learned)));
float[] n = rmsnorm(x, learned), next = new float[4];
for (int i = 0; i < 4; i++) next[i] = x[i] + 0.1f * n[(i + 1) % 4]; // x = x + f(norm(x))
System.out.println("stream after: " + show(next));
}
}weight 1: 0.4615 -0.6154 1.8462 0.0000
learned: 0.4615 -0.6154 0.9231 0.0000
10 * x: 0.4615 -0.6154 0.9231 0.0000
stream after: 2.9385 -3.9077 12.0000 0.0462Compute the root mean square by hand. The squares are 9, 16, 144, and 0. Their average is 169 / 4 = 42.25, and the square root of that is 6.5. With weight 1, each value is divided by 6.5, so the first output is 3 / 6.5 = 0.4615. The learned weight [1, 1, 0.5, 2] then scales each coordinate: the third value 12 / 6.5 = 1.8462 is halved to 0.9231, and the fourth stays 0 because 0 * 2 = 0.
The vector 10 * x gives exactly the same output as x. So RMSNorm removes overall scale and keeps direction. The last line shows the stream after the sublayer: it keeps the large raw value (12.0000) and receives a small update, which is why the stream can grow while each sublayer still sees a well-scaled input.
The constant 1e-5 (0.00001) inside the square root keeps the division safe for an all-zero vector, where dividing by zero would otherwise break the result. Real models read it from rms_norm_eps in the config (Example 12).
Example 4: the MLP
Attention lets each position gather information from earlier positions. After that, each position needs to process what it gathered. That is the job of the MLP (multi-layer perceptron, also called the feed-forward network). It works on each position separately. It has three steps. It projects the vector up to a wider hidden size, which gives the model more room to compute. It applies an activation function, a non-linear step that lets the network represent more than straight-line relationships. Then it projects back down to the model dimension. The classic hidden size is four times the model dimension.
The program uses d = 4, so the hidden size is 16. The activation is GELU, a smooth curve that passes large positive values almost unchanged and squashes negative values toward zero. The weights are random. Look at the hidden values, which show GELU at work, and at the parameter counts on the last two lines.
import java.util.Random;
public class Mlp {
static float[] matvec(float[] w, int rows, int cols, float[] x) {
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 gelu(float x) {
return (float) (0.5 * x * (1 + Math.tanh(0.7978845608 * (x + 0.044715 * x * x * x))));
}
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
int d = 4, hidden = 4 * d;
Random r = new Random(1);
float[] up = new float[hidden * d], down = new float[d * hidden];
for (int i = 0; i < up.length; i++) { up[i] = (float) (r.nextGaussian() * 0.5); down[i] = (float) (r.nextGaussian() * 0.5); }
float[] x = {0.5f, -1.0f, 0.25f, 2.0f};
float[] h = matvec(up, hidden, d, x);
for (int i = 0; i < hidden; i++) h[i] = gelu(h[i]);
float[] y = matvec(down, d, hidden, h);
StringBuilder sb = new StringBuilder();
for (float v : h) sb.append(String.format("%.2f ", v));
System.out.println("hidden (" + hidden + " values): " + sb.toString().trim());
System.out.printf("output: %.3f %.3f %.3f %.3f%n", y[0], y[1], y[2], y[3]);
System.out.println("parameters: " + (up.length + down.length) + " = 2 * 4 * 16");
long dm = 768;
System.out.println("d = 768: " + 2 * dm * 4 * dm + " parameters per layer, no biases");
}
}hidden (16 values): -0.15 0.15 -0.02 0.20 -0.10 -0.15 -0.12 -0.17 0.15 1.85 2.42 -0.15 -0.07 0.06 0.44 -0.15
output: 2.883 3.753 2.065 1.021
parameters: 128 = 2 * 4 * 16
d = 768: 4718592 parameters per layer, no biasesThe input has 4 numbers and the hidden vector has 16. GELU keeps large positive values almost unchanged (1.85 and 2.42) and pushes negative values to small numbers near zero (about -0.15). The output has 4 numbers again, so it has the same shape as the stream and can be added to it.
The two matrices hold 2 * d * 4d = 8 * d^2 parameters. For d = 4 that is 8 * 16 = 128. For d = 768 it is 8 * 589,824 = 4,718,592 per layer.
Positions never exchange information in the MLP. Each position goes through the same matrices alone. Only attention moves information between positions.
Example 5: SwiGLU, the MLP of Llama-style models
Most current open models replace the classic MLP with a variant called SwiGLU. The idea is a gate. Two matrices, called gate and up, both project the input to the hidden size. The gate result goes through SiLU, a smooth activation similar to GELU. The SiLU output then multiplies the up result element by element, so each hidden value is scaled by how open its gate is. A third matrix, called down, projects the product back to the model dimension. In one line: down(silu(gate(x)) * up(x)).
The program prints each of the 8 hidden units with its gate value, its SiLU value, its up value, and their product. It then prints the parameter counts of a classic MLP and of SwiGLU for d = 4096. Look for units where the gate is near zero or negative.
import java.util.Random;
public class SwiGlu {
static float[] matvec(float[] w, int rows, int cols, float[] x) {
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 silu(float x) { return (float) (x / (1 + Math.exp(-x))); }
static float[] random(Random r, int n) {
float[] w = new float[n];
for (int i = 0; i < n; i++) w[i] = (float) (r.nextGaussian() * 0.5);
return w;
}
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
int d = 4, hidden = 8;
Random r = new Random(2);
float[] gate = random(r, hidden * d), up = random(r, hidden * d), down = random(r, d * hidden);
float[] x = {0.5f, -1.0f, 0.25f, 2.0f};
float[] g = matvec(gate, hidden, d, x), u = matvec(up, hidden, d, x), act = new float[hidden];
for (int i = 0; i < hidden; i++) {
act[i] = silu(g[i]) * u[i];
System.out.printf("unit %d: gate %6.3f silu %6.3f up %6.3f product %6.3f%n", i, g[i], silu(g[i]), u[i], act[i]);
}
float[] y = matvec(down, d, hidden, act);
System.out.printf("output: %.3f %.3f %.3f %.3f%n", y[0], y[1], y[2], y[3]);
long dm = 4096;
long h = (long) Math.ceil(2.0 * 4 * dm / 3 / 256) * 256; // 2/3 of 4d, rounded up to a multiple of 256
System.out.println("classic MLP, d = 4096: " + 2 * dm * 4 * dm + " parameters");
System.out.println("SwiGLU hidden size: " + h + ", parameters: " + 3 * dm * h);
}
}unit 0: gate 0.168 silu 0.091 up -0.009 product -0.001
unit 1: gate 0.637 silu 0.416 up 0.307 product 0.128
unit 2: gate 1.457 silu 1.181 up -0.126 product -0.149
unit 3: gate 1.825 silu 1.571 up 0.461 product 0.724
unit 4: gate 0.627 silu 0.408 up 0.112 product 0.046
unit 5: gate -1.182 silu -0.277 up -1.228 product 0.341
unit 6: gate -1.824 silu -0.253 up -0.519 product 0.132
unit 7: gate 1.018 silu 0.748 up -0.958 product -0.716
output: -0.104 -0.099 -0.177 -0.478
classic MLP, d = 4096: 134217728 parameters
SwiGLU hidden size: 11008, parameters: 135266304Each line shows one hidden unit. Take unit 1: the gate value 0.637 becomes silu = 0.637 / (1 + e^-0.637) = 0.416, and the product with up = 0.307 is 0.128. A gate near zero (unit 0) silences the unit whatever up holds. A negative gate (units 5 and 6) gives a small negative silu value that flips the sign of up. So the gate decides how much of each hidden value passes through.
Three matrices cost more than two, so SwiGLU shrinks the hidden size to keep the total parameter count close to the classic design. The program applies the rule from the original Llama code: take two thirds of 4 * d, then round up to a multiple of 256. For d = 4096 that is 2/3 * 16,384 = 10,922.67, which rounds up to 11,008. This is the intermediate_size of Llama 2 7B.
Parameters per layer: SwiGLU has 3 * 4096 * 11008 = 135,266,304, and the classic MLP has 134,217,728. The two are within one percent. Other models choose other sizes: Llama 3 8B uses 14,336 at the same d.
Example 6: attention without positions is order blind
Self-attention compares a query with keys using dot products, then mixes values. Nothing in that recipe mentions where a token sits in the sentence. This raises a question: if you shuffle the input tokens, does attention notice? This example tests it. The answer is no, and that is why models need an extra position signal.
The program runs one attention head on three random vectors A, B, C, and then on a permutation (a reordering) of the same vectors. It compares each token's output across the two runs. It does this once without a mask and once with the causal mask, which stops position t from seeing positions after t. The output prints same or different for each comparison.
import java.util.Random;
public class OrderBlind {
static final int D = 4;
static float[][] wq = new float[D][D], wk = new float[D][D], wv = new float[D][D];
static float[] mul(float[][] w, float[] x) {
float[] y = new float[D];
for (int i = 0; i < D; i++) for (int j = 0; j < D; j++) y[i] += w[i][j] * x[j];
return y;
}
// one attention head over the rows of x; nothing adds positions
static float[][] attention(float[][] x, boolean causal) {
int n = x.length;
float[][] out = new float[n][D];
for (int t = 0; t < n; t++) {
int last = causal ? t : n - 1;
double[] w = new double[last + 1];
double sum = 0;
for (int s = 0; s <= last; s++) {
double dot = 0;
float[] q = mul(wq, x[t]), k = mul(wk, x[s]);
for (int i = 0; i < D; i++) dot += q[i] * k[i];
w[s] = Math.exp(dot / Math.sqrt(D));
sum += w[s];
}
for (int s = 0; s <= last; s++) {
float[] v = mul(wv, x[s]);
for (int i = 0; i < D; i++) out[t][i] += (float) (w[s] / sum) * v[i];
}
}
return out;
}
static String same(float[] a, float[] b) {
for (int i = 0; i < D; i++) if (Math.abs(a[i] - b[i]) > 1e-6f) return "different";
return "same";
}
static float[] vec(Random r) {
float[] v = new float[D];
for (int i = 0; i < D; i++) v[i] = (float) r.nextGaussian();
return v;
}
public static void main(String[] args) {
Random r = new Random(3);
for (float[][] m : new float[][][] {wq, wk, wv}) for (int i = 0; i < D; i++) m[i] = vec(r);
float[] a = vec(r), b = vec(r), c = vec(r);
float[][] abc = attention(new float[][] {a, b, c}, false);
float[][] cab = attention(new float[][] {c, a, b}, false);
System.out.println("no mask, [A,B,C] vs [C,A,B]: output of A " + same(abc[0], cab[1])
+ ", B " + same(abc[1], cab[2]) + ", C " + same(abc[2], cab[0]));
float[][] ab = attention(new float[][] {a, b, c}, true);
float[][] ba = attention(new float[][] {b, a, c}, true);
System.out.println("causal, [A,B,C] vs [B,A,C]: position 1 " + same(ab[1], ba[1]) + ", position 2 " + same(ab[2], ba[2]));
}
}no mask, [A,B,C] vs [C,A,B]: output of A same, B same, C same
causal, [A,B,C] vs [B,A,C]: position 1 different, position 2 sameWithout a mask, the order C, A, B produces the same three outputs as A, B, C, only moved to the matching rows. The layer treats its input as a set, not a sequence. "Dog bites man" and "man bites dog" would be indistinguishable.
The causal mask adds a little order information, since position t only sees positions up to t. The second experiment swaps the first two inputs. Position 2 gives the same output, because it attends to the same three vectors with the same query. Position 1 changes, because its own vector changed. So even with the mask, the current position cannot see the order of the tokens before it. Models need an explicit position signal.
The causal mask adds a little order information, since position t only sees positions up to t. The second experiment swaps the first two inputs. Position 2 gives the same output, because it attends to the same three vectors with the same query. Position 1 changes, because its own vector changed. So even with the mask, the order of the tokens before the current one is invisible to the current position. Models need an explicit position signal.
Example 7: learned absolute positions versus rotary positions
Since attention cannot see order, the model must be told where each token sits. There are two common ways.
The first way, used by GPT-2, keeps a second table with one learned row per position (position 0, position 1, and so on) and adds that row to the token embedding. The second way is rotary position embedding (RoPE). It leaves the embeddings alone and instead rotates the query and key vectors by an angle that grows with the position. A rotation can be computed for any position, so no table is needed.
The program first builds the table version for the words "the cat the cat the" and then tries position 4, which the table does not cover. After that it rotates the point (1, 0) by pos * 0.5 radians for four positions, which is the idea behind RoPE. Think of a clock hand that turns a little further for each later token.
public class Positions {
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
float[][] tokenTable = {{0.5f, 0.2f, -0.3f, 0.1f}, {-0.4f, 0.9f, 0.3f, 0.6f}}; // "the", "cat"
float[][] positionTable = {{0f, 0f, 0f, 0f}, {0.1f, -0.1f, 0.1f, -0.1f}, {0.2f, -0.2f, 0.2f, -0.2f}, {0.3f, -0.3f, 0.3f, -0.3f}};
int[] tokens = {0, 1, 0, 1, 0}; // the cat the cat the
for (int pos = 0; pos < tokens.length; pos++) {
if (pos >= positionTable.length) {
System.out.println("position " + pos + " has no row: the table holds " + positionTable.length);
break;
}
float[] x = new float[4];
for (int i = 0; i < 4; i++) x[i] = tokenTable[tokens[pos]][i] + positionTable[pos][i];
System.out.printf("pos %d token %d -> %.2f %.2f %.2f %.2f%n", pos, tokens[pos], x[0], x[1], x[2], x[3]);
}
// the idea behind RoPE: rotate (1, 0) by an angle that grows with the position
for (int pos = 0; pos < 4; pos++) {
double angle = pos * 0.5;
System.out.printf("pos %d: rotated by %.1f rad = (%.3f, %.3f)%n", pos, angle, Math.cos(angle), Math.sin(angle));
}
}
}pos 0 token 0 -> 0.50 0.20 -0.30 0.10
pos 1 token 1 -> -0.30 0.80 0.40 0.50
pos 2 token 0 -> 0.70 0.00 -0.10 -0.10
pos 3 token 1 -> -0.10 0.60 0.60 0.30
position 4 has no row: the table holds 4
pos 0: rotated by 0.0 rad = (1.000, 0.000)
pos 1: rotated by 0.5 rad = (0.878, 0.479)
pos 2: rotated by 1.0 rad = (0.540, 0.841)
pos 3: rotated by 1.5 rad = (0.071, 0.997)With a position table, the at position 0 gives 0.50 0.20 -0.30 0.10 and the at position 2 gives 0.70 0.00 -0.10 -0.10. The vectors differ, so later layers can tell the two positions apart. The table has a fixed number of rows (4 here), so position 4 has nothing to look up. A model with learned absolute positions cannot read inputs longer than its table. GPT-2 has 1,024 positions (n_positions in its config.json).
The last four lines show the idea behind RoPE with one 2-dimensional vector. Position p rotates (1, 0) by p * 0.5 radians, giving (cos, sin) of that angle: (0.878, 0.479) at position 1 and (0.540, 0.841) at position 2. The angle grows with the position, and no table is needed.
The last four lines show the idea behind RoPE with one 2-dimensional vector: position p rotates (1, 0) by p * 0.5 radians. No table is needed, because a rotation can be computed for any position.
Example 8: RoPE and relative distance
Rotation alone is not the goal. The goal is that attention scores depend on how far apart two tokens are. Here is why rotation gives that. A dot product between two vectors depends only on their lengths and the angle between them. If the query at position m is rotated by m * theta and the key at position n is rotated by n * theta, the angle between them changes by (m - n) * theta. So the score depends on the distance m - n, not on m and n separately.
A real head has many coordinates, so RoPE splits the vector into adjacent pairs and rotates each pair as a 2-dimensional point. Each pair uses its own rate: theta_i = base^(-2i/d), where d is the head dimension and base is rope_theta from the config. With d = 8 and base 10,000, the four pairs turn at 1, 0.1, 0.01, and 0.001 radians per position.
The program prints those rates. Then it takes one random query and one random key and computes their score after rotating them for five position pairs. Three of the pairs have the same distance (2) but very different positions.
import java.util.Random;
public class Rope {
// rotates each adjacent pair (v[2i], v[2i+1]) by pos * theta_i, theta_i = base^(-2i/d)
static float[] rope(float[] v, int 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;
}
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;
}
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
int d = 8;
for (int i = 0; i < d / 2; i++) System.out.printf("theta_%d = %.4f%n", i, Math.pow(10000, -2.0 * i / d));
Random r = new Random(5);
float[] q = new float[d], k = new float[d];
for (int i = 0; i < d; i++) { q[i] = (float) r.nextGaussian(); k[i] = (float) r.nextGaussian(); }
System.out.printf("no rotation: score %.4f%n", dot(q, k));
for (int[] p : new int[][] {{3, 1}, {7, 5}, {102, 100}, {10, 4}, {50, 40}})
System.out.printf("query at %3d, key at %3d, distance %2d: score %.4f%n",
p[0], p[1], p[0] - p[1], dot(rope(q, p[0], 10000), rope(k, p[1], 10000)));
}
}theta_0 = 1.0000
theta_1 = 0.1000
theta_2 = 0.0100
theta_3 = 0.0010
no rotation: score -0.2882
query at 3, key at 1, distance 2: score -4.7472
query at 7, key at 5, distance 2: score -4.7472
query at 102, key at 100, distance 2: score -4.7472
query at 10, key at 4, distance 6: score 0.7837
query at 50, key at 40, distance 10: score -2.8823The rates for d = 8 and base 10,000 are 1, 0.1, 0.01, and 0.001. The first pair turns fast and the last pair turns slowly. Fast pairs separate nearby positions, and slow pairs still tell far-apart positions apart, so the model gets both fine and coarse position signals.
The unrotated score is -0.2882. After rotation, the position pairs (3, 1), (7, 5), and (102, 100) all have distance 2 and all give -4.7472. Other distances give other scores. So absolute positions drop out of the score and only the distance remains.
RoPE is applied to queries and keys, never to values, because only queries and keys are compared in the score. Llama 3 sets rope_theta to 500,000 instead of 10,000. A larger base slows every rotation, which helps with longer contexts. A later lesson on attention optimizations and long context covers the scaling methods.
One detail matters when reading engine code. There is more than one way to choose which coordinates form a pair. This program pairs consecutive coordinates (2i, 2i+1), the Llama convention. Other families pair coordinate i with i + d/2. In llama.cpp these are the "normal" and "NEOX" RoPE types, chosen per architecture: the source selects the normal type for LLM_ARCH_LLAMA and the NEOX type for Qwen2 and Qwen3. For Llama checkpoints from Hugging Face, convert_hf_to_gguf.py reorders the rows of the query and key weights once at conversion time, so that the consecutive-pair rotation gives the right result.
Example 9: one full block forward
Now all the parts exist, and this example assembles them into one block. A block has two sublayers. Each has its own norm and its own residual addition. The first is attention, with RoPE on the queries and keys and with grouped-query heads (several query heads share one key/value head, as in the self-attention lesson). The second is the SwiGLU MLP. In code form: x = x + Attention(RMSNorm(x)), then x = x + MLP(RMSNorm(x)).
The program uses a model dimension of 16, 4 query heads sharing 2 key/value heads, a head dimension of 4, and an MLP hidden size of 32. The weights come from new Random(42), each scaled by 1 / sqrt(cols), a common initialization that keeps output sizes stable. The program runs four token positions through the block and prints the root mean square (rms), the typical size, of each position's vector at three moments: at the start, after attention, and after the MLP.
import java.util.Arrays;
import java.util.Locale;
import java.util.Random;
public class OneBlock {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32;
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final Random R = new Random(42);
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;
}
static final float[] EMBED = mat(V, D);
static final Layer LAYER = 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 float[] matvec(float[] w, int rows, int cols, float[] x) {
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 float rms(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
return (float) Math.sqrt(ss / x.length);
}
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;
}
}
// x = x + Attention(RMSNorm(x)) for every position
static void attention(float[][] x, Layer ly) {
int n = x.length;
float[][] q = new float[n][], k = new float[n][], v = new float[n][];
for (int t = 0; t < n; t++) {
float[] h = rmsnorm(x[t]);
q[t] = matvec(ly.wq, D, D, h); k[t] = matvec(ly.wk, KVH * HD, D, h); v[t] = matvec(ly.wv, KVH * HD, D, h);
rope(q[t], H, t);
rope(k[t], KVH, t);
}
for (int t = 0; t < n; t++) {
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH); // the shared key/value head of query head h
double[] w = new double[t + 1];
double sum = 0;
for (int s = 0; s <= t; s++) { // causal: positions 0 to t only
double dot = 0;
for (int i = 0; i < HD; i++) dot += q[t][h * HD + i] * k[s][kv * HD + i];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= t; s++)
for (int i = 0; i < HD; i++) cat[h * HD + i] += (float) (w[s] / sum) * v[s][kv * HD + i];
}
float[] out = matvec(ly.wo, D, D, cat);
for (int i = 0; i < D; i++) x[t][i] += out[i];
}
}
// x = x + MLP(RMSNorm(x)) for every position, SwiGLU
static void mlp(float[][] x, Layer ly) {
for (float[] xt : x) {
float[] h = rmsnorm(xt);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int i = 0; i < F; i++) g[i] = (float) (g[i] / (1 + Math.exp(-g[i]))) * u[i];
float[] out = matvec(ly.down, D, F, g);
for (int i = 0; i < D; i++) xt[i] += out[i];
}
}
public static void main(String[] args) {
Locale.setDefault(Locale.ROOT);
int[] tokens = {3, 17, 5, 9};
float[][] x = new float[tokens.length][];
for (int t = 0; t < tokens.length; t++) x[t] = Arrays.copyOfRange(EMBED, tokens[t] * D, (tokens[t] + 1) * D);
float[] start = new float[x.length], afterAttention = new float[x.length];
for (int t = 0; t < x.length; t++) start[t] = rms(x[t]);
attention(x, LAYER);
for (int t = 0; t < x.length; t++) afterAttention[t] = rms(x[t]);
mlp(x, LAYER);
for (int t = 0; t < x.length; t++)
System.out.printf("position %d (id %2d): rms %.3f, after attention %.3f, after the MLP %.3f%n",
t, tokens[t], start[t], afterAttention[t], rms(x[t]));
System.out.println("output shape: [" + x.length + ", " + x[0].length + "], same as the input");
int weights = LAYER.wq().length + LAYER.wk().length + LAYER.wv().length + LAYER.wo().length
+ LAYER.gate().length + LAYER.up().length + LAYER.down().length;
System.out.println("matrix weights in the block: " + weights);
}
}position 0 (id 3): rms 0.291, after attention 0.877, after the MLP 1.003
position 1 (id 17): rms 0.276, after attention 0.616, after the MLP 0.620
position 2 (id 5): rms 0.213, after attention 0.748, after the MLP 0.738
position 3 (id 9): rms 0.342, after attention 0.507, after the MLP 0.627
output shape: [4, 16], same as the input
matrix weights in the block: 2304Each line follows one position through the block. The rms starts at the size of the token's embedding (0.291 for position 0), changes to 0.877 when attention adds its update, and changes again to 1.003 when the MLP adds its update. The numbers describe random weights only, so do not read meaning into them. The structure is what matters: the block takes four vectors of 16 numbers and returns four vectors of 16 numbers. Input and output have the same shape, so the output can be the input of another block.
Three details of the code carry the structure.
kv = h / (H / KVH)maps query heads 0 and 1 to key/value head 0, and query heads 2 and 3 to key/value head 1. This is grouped-query attention.- The loop
for (int s = 0; s <= t; s++)is the causal mask from the self-attention lesson: positiontscores positions0totonly. - The learned RMSNorm weights from Example 3 are left out. They start at 1.0 in a new model, and multiplying by 1.0 changes nothing.
The block holds 2,304 matrix weights, counted as follows.
- 256 for the query matrix (16 by 16).
- 128 each for the key and value matrices (2 heads of 4 values gives 8 rows, times 16 inputs).
- 256 for the output matrix.
- 512 for each of the three MLP matrices (32 by 16).
The sum is 256 + 128 + 128 + 256 + 3 * 512 = 2,304. Example 13 adds the two norm vectors (16 numbers each) to get 2,336 per layer.
Example 10: stacking blocks, the final norm, and the LM head
A real model is the same block repeated. Layer 0 reads the embeddings. Layer 1 reads the stream written by layer 0, and so on. Every layer runs the same code but has its own weights, so each layer can learn a different job. After the last block, two more steps produce scores. A final RMSNorm rescales the stream. Then the LM head (language model head), a vocab by dim matrix, turns each position's vector into one logit per vocabulary entry.
The program stacks two blocks and uses the same seed as Example 9, so layer 0 has the same weights as the block there. It prints the rms of the last position after attention and after the MLP in each layer. Then it prints the shape of the logits and the three largest logits of the last position.
import java.util.Arrays;
import java.util.Locale;
import java.util.Random;
public class Stack {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2;
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final Random R = new Random(42);
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;
}
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 float[] matvec(float[] w, int rows, int cols, float[] x) {
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 float rms(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
return (float) Math.sqrt(ss / x.length);
}
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;
}
}
static void attention(float[][] x, Layer ly) { // unchanged from Example 9
int n = x.length;
float[][] q = new float[n][], k = new float[n][], v = new float[n][];
for (int t = 0; t < n; t++) {
float[] h = rmsnorm(x[t]);
q[t] = matvec(ly.wq, D, D, h); k[t] = matvec(ly.wk, KVH * HD, D, h); v[t] = matvec(ly.wv, KVH * HD, D, h);
rope(q[t], H, t);
rope(k[t], KVH, t);
}
for (int t = 0; t < n; t++) {
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[t + 1];
double sum = 0;
for (int s = 0; s <= t; s++) {
double dot = 0;
for (int i = 0; i < HD; i++) dot += q[t][h * HD + i] * k[s][kv * HD + i];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= t; s++)
for (int i = 0; i < HD; i++) cat[h * HD + i] += (float) (w[s] / sum) * v[s][kv * HD + i];
}
float[] out = matvec(ly.wo, D, D, cat);
for (int i = 0; i < D; i++) x[t][i] += out[i];
}
}
static void mlp(float[][] x, Layer ly) { // unchanged from Example 9
for (float[] xt : x) {
float[] h = rmsnorm(xt);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int i = 0; i < F; i++) g[i] = (float) (g[i] / (1 + Math.exp(-g[i]))) * u[i];
float[] out = matvec(ly.down, D, F, g);
for (int i = 0; i < D; i++) xt[i] += out[i];
}
}
public static void main(String[] args) {
Locale.setDefault(Locale.ROOT);
int[] tokens = {3, 17, 5, 9};
int n = tokens.length;
float[][] x = new float[n][];
for (int t = 0; t < n; t++) x[t] = Arrays.copyOfRange(EMBED, tokens[t] * D, (tokens[t] + 1) * D);
for (int l = 0; l < L; l++) { // the same code, different weights
attention(x, LAYERS[l]);
System.out.printf("layer %d: last position rms %.3f after attention", l, rms(x[n - 1]));
mlp(x, LAYERS[l]);
System.out.printf(", %.3f after the MLP%n", rms(x[n - 1]));
}
float[][] logits = new float[n][];
for (int t = 0; t < n; t++) logits[t] = matvec(HEAD, V, D, rmsnorm(x[t])); // final norm, then the LM head
System.out.println("logits shape: [" + logits.length + ", " + logits[0].length + "]");
Integer[] ids = new Integer[V];
for (int i = 0; i < V; i++) ids[i] = i;
float[] last = logits[n - 1];
Arrays.sort(ids, (a, b) -> Float.compare(last[b], last[a]));
for (int r = 0; r < 3; r++) System.out.printf("last position, rank %d: id %2d logit %.3f%n", r + 1, ids[r], last[ids[r]]);
}
}layer 0: last position rms 0.507 after attention, 0.627 after the MLP
layer 1: last position rms 0.720 after attention, 0.752 after the MLP
logits shape: [4, 24]
last position, rank 1: id 6 logit 2.305
last position, rank 2: id 21 logit 1.685
last position, rank 3: id 5 logit 1.354The layer 0 line repeats the last row of Example 9 (0.507 and 0.627), because the weights and the input are the same. Layer 1 starts from the stream that layer 0 wrote and adds its own updates. A model with 32 layers runs this loop 32 times.
The shape [4, 24] is 4 positions by 24 vocabulary entries: the forward pass gives one row of logits per input position. Generation needs only the last row, because that row predicts the token after the prompt. The other rows exist because training scores the prediction at every position at once.
The last three lines rank the last row. Id 6 has the largest logit (2.305), so a greedy choice (take the largest logit, from What a Language Model Computes) would append token 6. Example 14 runs that choice in a loop.
The final RMSNorm matters for the same reason as the norms inside the blocks. The stream has grown through every layer, and the LM head expects inputs of a fixed scale.
Example 11: tied embeddings
Look at the shapes. The embedding table maps an id to a vector, and its shape is vocab by dim. The LM head maps a vector to one score per id, and its shape is also vocab by dim. Each logit is the dot product of the hidden vector with one row of the head, so a row acts like a "prototype" for one token: the more the hidden vector points the same way, the higher the score. This suggests reusing one matrix for both jobs. That is called tying the weights. With tying, the logit of token i is the dot product of the hidden vector with the embedding of token i.
The program uses a 5-token vocabulary and a hidden vector close to the embedding of sat. It computes every logit with the embedding table as the head. Then it computes how many parameters tying saves for a Llama 3.2 1B shape.
public class TiedEmbeddings {
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
String[] vocab = {"the", "cat", "sat", "on", "mat"};
float[][] embed = {{0.5f, 0.2f, -0.3f, 0.1f}, {-0.4f, 0.9f, 0.3f, 0.6f}, {0.7f, -0.5f, 0.2f, 0.8f},
{0.3f, 0.3f, -0.6f, 0.0f}, {-0.2f, 0.8f, 0.4f, -0.5f}};
float[] h = {0.6f, -0.4f, 0.3f, 0.7f}; // final hidden state, close to the row of "sat"
for (int i = 0; i < vocab.length; i++) {
float logit = 0; // tied: the output matrix is the embedding matrix
for (int j = 0; j < 4; j++) logit += embed[i][j] * h[j];
System.out.printf("%-4s logit %6.3f%n", vocab[i], logit);
}
long matrix = 128_256L * 2048, total = 1_235_814_400L; // Llama 3.2 1B shape
System.out.println("output matrix saved by tying: " + matrix + " parameters");
System.out.printf("share of the 1.2B model: %.1f%%%n", 100.0 * matrix / total);
}
}the logit 0.200
cat logit -0.090
sat logit 1.240
on logit -0.120
mat logit -0.670
output matrix saved by tying: 262668288 parameters
share of the 1.2B model: 21.3%The hidden vector [0.6, -0.4, 0.3, 0.7] is close to the row for sat, and sat gets the largest logit, 1.240. For example, the logit for sat is 0.6 * 0.7 + (-0.4) * (-0.5) + 0.3 * 0.2 + 0.7 * 0.8 = 1.240. A tied model is trained to place the final hidden vector near the embedding of the token it predicts.
The saving is a whole matrix. For a 128,256-token vocabulary and 2048 dimensions that is 262,668,288 parameters, 21.3% of a 1.2B model. Small models tie their embeddings because the matrix is a large share of their size.
The published config.json of Llama 3.2 1B sets tie_word_embeddings to true, and that of Llama 3 8B sets it to false. A tied GGUF file has no output.weight tensor. The llama.cpp Llama loader marks the output tensor as optional and, when it is missing, reuses token_embd.weight.
Example 12: reading a Llama-style config.json
A downloaded model comes with a file called config.json that lists the shape of the network: how many layers, how wide each vector is, how many heads, and so on. If you can read it, you can compute sizes and memory needs without loading the weights. This example reads such a file in Java.
The program reads an inline copy of a Llama 3 8B style config with Jackson, a Java JSON library, into a record. @JsonIgnoreProperties(ignoreUnknown = true) tells Jackson to skip fields the record does not have, and @JsonProperty maps the snake_case JSON names to the record components. The program then derives the head dimension and the grouped-query ratio from the numbers it read.
import com.fasterxml.jackson.annotation.JsonIgnoreProperties;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
public class ReadConfig {
@JsonIgnoreProperties(ignoreUnknown = true)
record Config(
@JsonProperty("hidden_size") int hiddenSize,
@JsonProperty("intermediate_size") int intermediateSize,
@JsonProperty("num_hidden_layers") int layers,
@JsonProperty("num_attention_heads") int heads,
@JsonProperty("num_key_value_heads") int kvHeads,
@JsonProperty("vocab_size") int vocabSize,
@JsonProperty("max_position_embeddings") int maxPositions,
@JsonProperty("rope_theta") double ropeTheta,
@JsonProperty("rms_norm_eps") double rmsNormEps,
@JsonProperty("tie_word_embeddings") boolean tieWordEmbeddings) {}
public static void main(String[] args) throws Exception {
String json = """
{
"architectures": ["LlamaForCausalLM"],
"hidden_size": 4096,
"intermediate_size": 14336,
"num_hidden_layers": 32,
"num_attention_heads": 32,
"num_key_value_heads": 8,
"vocab_size": 128256,
"max_position_embeddings": 8192,
"rope_theta": 500000.0,
"rms_norm_eps": 1e-05,
"hidden_act": "silu",
"tie_word_embeddings": false,
"torch_dtype": "bfloat16"
}
""";
Config c = new ObjectMapper().readValue(json, Config.class);
System.out.println(c);
int headDim = c.hiddenSize() / c.heads();
System.out.println("head dimension: " + headDim);
System.out.println("query heads per KV head: " + (c.heads() / c.kvHeads()));
System.out.println("query width: " + c.heads() * headDim + ", key/value width: " + c.kvHeads() * headDim);
System.out.println("RoPE pairs per head: " + headDim / 2);
}
}Config[hiddenSize=4096, intermediateSize=14336, layers=32, heads=32, kvHeads=8, vocabSize=128256, maxPositions=8192, ropeTheta=500000.0, rmsNormEps=1.0E-5, tieWordEmbeddings=false]
head dimension: 128
query heads per KV head: 4
query width: 4096, key/value width: 1024
RoPE pairs per head: 64From a handful of numbers the program derives the rest.
- The head dimension is
hidden_size / num_attention_heads = 4096 / 32 = 128. - There are 4 query heads per key/value head (32 / 8). This is the grouped-query ratio.
- The query projection is 4096 wide. The key and value projections are 1024 wide (8 heads of 128), so the key and value caches are a quarter of the size they would be with a full set of heads.
- Each head has 64 RoPE pairs (128 / 2).
The values are those of the Llama 3 8B shape. Other releases of the family differ: the config.json files of Llama 3.1 and 3.2 list max_position_embeddings as 131072 and add a rope_scaling entry. The real file has more fields (bos_token_id, eos_token_id, rope_scaling) that the record ignores. The GGUF format shows where llama.cpp stores the same values.
Example 13: counting parameters from a config
A model's name such as "8B" means about 8 billion parameters, where a parameter is one learned number. You can reproduce that figure by adding up every matrix. Per layer there are the query and output matrices (hidden * heads * head_dim each), the key and value matrices (hidden * kv_heads * head_dim each), three MLP matrices (hidden * intermediate each), and two norm vectors (hidden each). The model adds the embedding table, a final norm vector, and an output matrix unless the embeddings are tied.
The program computes this for three configs: the Llama 3 8B shape, the Llama 3.2 1B shape, and the tiny model of the earlier examples. It prints the parameters per layer, in all layers, and in total.
public class CountParams {
record Config(String name, long hidden, long intermediate, long layers, long heads, long kvHeads, long vocab, boolean tied) {
long headDim() { return hidden / heads; }
long perLayer() {
long attention = hidden * heads * headDim() * 2 + 2 * hidden * kvHeads * headDim(); // Q and O, K and V
return attention + 3 * hidden * intermediate + 2 * hidden; // MLP, two norms
}
long total() {
return vocab * hidden + layers * perLayer() + hidden + (tied ? 0 : vocab * hidden); // embedding, layers, final norm, head
}
}
public static void main(String[] args) {
java.util.Locale.setDefault(java.util.Locale.ROOT);
Config[] configs = {
new Config("Llama 3 8B shape", 4096, 14336, 32, 32, 8, 128256, false),
new Config("Llama 3.2 1B shape", 2048, 8192, 16, 32, 8, 128256, true),
new Config("tiny model", 16, 32, 2, 4, 2, 24, false)};
for (Config c : configs)
System.out.printf("%-20s per layer %,12d all layers %,15d total %,15d%n", c.name(), c.perLayer(), c.layers() * c.perLayer(), c.total());
}
}Llama 3 8B shape per layer 218,112,000 all layers 6,979,584,000 total 8,030,261,248
Llama 3.2 1B shape per layer 60,821,504 all layers 973,144,064 total 1,235,814,400
tiny model per layer 2,336 all layers 4,672 total 5,456Check the 8B shape by hand. Attention holds 2 * 4096 * 4096 + 2 * 4096 * 1024 = 41,943,040. The MLP holds 3 * 4096 * 14336 = 176,160,768. The two norms hold 2 * 4096 = 8,192. The sum is 218,112,000 per layer, and 32 layers hold 6,979,584,000. Adding two matrices of 525,336,576 (embedding and head) and the final norm gives 8,030,261,248, the 8.03B in the model's name.
The Llama 3.2 1B shape has tied embeddings, so it counts the embedding once and gives 1,235,814,400. The tiny model of the earlier examples has 5,456 parameters. The formula returns in Memory Math, where it gives weight memory.
The MLP is the largest part of each layer. In the 8B shape it holds 176,160,768 of the 218,112,000 parameters (81%), and attention holds 41,943,040 (19%). Grouped-query attention makes each key and value matrix a quarter of the query matrix.
Example 14: a tiny random-weight transformer that generates token ids
Now the pieces form a loop. A model that produces logits can write text: pick a token from the last logits, append it to the input, and run the model again. This example does that with the model from Example 10, which has an embedding table of 24 tokens, two blocks, a final norm, and an LM head. Its 5,456 parameters match the count in Example 13.
The forward method runs the whole model and returns the logits of every position. The generation loop takes the last position's logits and picks the largest (the greedy choice). It also computes that token's probability with softmax. Then it appends the id to the sequence and repeats ten times. Look at how the length grows by one per step.
import java.util.Arrays;
import java.util.Locale;
import java.util.Random;
public class TinyModel {
static final int V = 24, D = 16, H = 4, KVH = 2, HD = D / H, F = 32, L = 2;
record Layer(float[] wq, float[] wk, float[] wv, float[] wo, float[] gate, float[] up, float[] down) {}
static final Random R = new Random(42);
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;
}
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 float[] matvec(float[] w, int rows, int cols, float[] x) {
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 float rms(float[] x) {
float ss = 0;
for (float v : x) ss += v * v;
return (float) Math.sqrt(ss / x.length);
}
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;
}
}
static void attention(float[][] x, Layer ly) { // unchanged from Example 9
int n = x.length;
float[][] q = new float[n][], k = new float[n][], v = new float[n][];
for (int t = 0; t < n; t++) {
float[] h = rmsnorm(x[t]);
q[t] = matvec(ly.wq, D, D, h); k[t] = matvec(ly.wk, KVH * HD, D, h); v[t] = matvec(ly.wv, KVH * HD, D, h);
rope(q[t], H, t);
rope(k[t], KVH, t);
}
for (int t = 0; t < n; t++) {
float[] cat = new float[D];
for (int h = 0; h < H; h++) {
int kv = h / (H / KVH);
double[] w = new double[t + 1];
double sum = 0;
for (int s = 0; s <= t; s++) {
double dot = 0;
for (int i = 0; i < HD; i++) dot += q[t][h * HD + i] * k[s][kv * HD + i];
sum += w[s] = Math.exp(dot / Math.sqrt(HD));
}
for (int s = 0; s <= t; s++)
for (int i = 0; i < HD; i++) cat[h * HD + i] += (float) (w[s] / sum) * v[s][kv * HD + i];
}
float[] out = matvec(ly.wo, D, D, cat);
for (int i = 0; i < D; i++) x[t][i] += out[i];
}
}
static void mlp(float[][] x, Layer ly) { // unchanged from Example 9
for (float[] xt : x) {
float[] h = rmsnorm(xt);
float[] g = matvec(ly.gate, F, D, h), u = matvec(ly.up, F, D, h);
for (int i = 0; i < F; i++) g[i] = (float) (g[i] / (1 + Math.exp(-g[i]))) * u[i];
float[] out = matvec(ly.down, D, F, g);
for (int i = 0; i < D; i++) xt[i] += out[i];
}
}
// logits for every position of the token sequence: embedding, blocks, final norm, LM head
static float[][] forward(int[] tokens) {
int n = tokens.length;
float[][] x = new float[n][];
for (int t = 0; t < n; t++) x[t] = Arrays.copyOfRange(EMBED, tokens[t] * D, (tokens[t] + 1) * D);
for (Layer ly : LAYERS) { attention(x, ly); mlp(x, ly); }
float[][] logits = new float[n][];
for (int t = 0; t < n; t++) logits[t] = matvec(HEAD, V, D, rmsnorm(x[t]));
return logits;
}
public static void main(String[] args) {
Locale.setDefault(Locale.ROOT);
int[] seq = {3, 17, 5, 9};
for (int step = 0; step < 10; step++) { // greedy generation, whole sequence each step
float[] last = forward(seq)[seq.length - 1];
int best = 0;
double sum = 0;
for (int i = 0; i < V; i++) { if (last[i] > last[best]) best = i; sum += Math.exp(last[i]); }
System.out.printf("step %d: length %2d -> id %2d (logit %.3f, probability %.3f)%n", step, seq.length, best, last[best], Math.exp(last[best]) / sum);
seq = Arrays.copyOf(seq, seq.length + 1);
seq[seq.length - 1] = best;
}
System.out.println("generated: " + Arrays.toString(Arrays.copyOfRange(seq, 4, seq.length)));
}
}step 0: length 4 -> id 6 (logit 2.305, probability 0.246)
step 1: length 5 -> id 6 (logit 1.846, probability 0.158)
step 2: length 6 -> id 13 (logit 2.023, probability 0.197)
step 3: length 7 -> id 10 (logit 1.536, probability 0.116)
step 4: length 8 -> id 13 (logit 1.563, probability 0.139)
step 5: length 9 -> id 13 (logit 1.851, probability 0.162)
step 6: length 10 -> id 3 (logit 1.950, probability 0.170)
step 7: length 11 -> id 8 (logit 2.658, probability 0.344)
step 8: length 12 -> id 8 (logit 2.307, probability 0.281)
step 9: length 13 -> id 8 (logit 2.271, probability 0.257)
generated: [6, 6, 13, 10, 13, 13, 3, 8, 8, 8]Each step prints the length of the sequence, the chosen id, its logit, and its softmax probability (softmax turns logits into probabilities that sum to 1, from the math lesson). Step 0 picks id 6 with logit 2.305, the rank-1 entry of Example 10. The ids say nothing about language, because the weights are random. They are still the real output of a complete forward pass, and the same prompt gives the same ids in every run.
The probabilities are low (0.116 to 0.344) because random weights spread the probability over many tokens. A trained model is usually far more confident on most steps. Repeats such as 8, 8, 8 appear when one token keeps the highest logit. Sampling covers the methods that vary the choice.
The loop runs the whole sequence through the model at every step, so step 9 processes 13 positions to produce one new token. That is correct but wasteful, because the first 12 positions give the same keys and values as in step 8. Generation, Prefill, Decode, and the KV Cache measures that waste and removes it.
Example 15: tensor names in checkpoints and GGUF files
A model file holds hundreds of matrices, each stored under a name. To load a model, an engine must match every name to a part of the architecture. This example shows the two naming schemes you will meet. A Hugging Face checkpoint (safetensors files) uses path-like names: model.layers.0.self_attn.q_proj.weight is the query matrix of layer 0. GGUF, the file format of llama.cpp, uses shorter names such as blk.0.attn_q.weight.
The program applies the renaming rules for Llama-style models to a list of names and prints each old name next to the new one. Read the output as a map from the parts you built in this lesson to their names in a file.
public class TensorNames {
static final String[][] RULES = {
{"model.embed_tokens", "token_embd"}, {"model.layers.", "blk."}, {"model.norm", "output_norm"}, {"lm_head", "output"},
{".input_layernorm", ".attn_norm"}, {".self_attn.q_proj", ".attn_q"}, {".self_attn.k_proj", ".attn_k"},
{".self_attn.v_proj", ".attn_v"}, {".self_attn.o_proj", ".attn_output"}, {".post_attention_layernorm", ".ffn_norm"},
{".mlp.gate_proj", ".ffn_gate"}, {".mlp.up_proj", ".ffn_up"}, {".mlp.down_proj", ".ffn_down"}};
static String toGguf(String hf) {
for (String[] rule : RULES) hf = hf.replace(rule[0], rule[1]);
return hf;
}
public static void main(String[] args) {
String[] names = {"model.embed_tokens.weight", "model.layers.0.input_layernorm.weight",
"model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.k_proj.weight",
"model.layers.0.self_attn.v_proj.weight", "model.layers.0.self_attn.o_proj.weight",
"model.layers.0.post_attention_layernorm.weight", "model.layers.0.mlp.gate_proj.weight",
"model.layers.0.mlp.up_proj.weight", "model.layers.31.mlp.down_proj.weight", "model.norm.weight", "lm_head.weight"};
for (String n : names) System.out.printf("%-48s -> %s%n", n, toGguf(n));
}
}model.embed_tokens.weight -> token_embd.weight
model.layers.0.input_layernorm.weight -> blk.0.attn_norm.weight
model.layers.0.self_attn.q_proj.weight -> blk.0.attn_q.weight
model.layers.0.self_attn.k_proj.weight -> blk.0.attn_k.weight
model.layers.0.self_attn.v_proj.weight -> blk.0.attn_v.weight
model.layers.0.self_attn.o_proj.weight -> blk.0.attn_output.weight
model.layers.0.post_attention_layernorm.weight -> blk.0.ffn_norm.weight
model.layers.0.mlp.gate_proj.weight -> blk.0.ffn_gate.weight
model.layers.0.mlp.up_proj.weight -> blk.0.ffn_up.weight
model.layers.31.mlp.down_proj.weight -> blk.31.ffn_down.weight
model.norm.weight -> output_norm.weight
lm_head.weight -> output.weightThe rules replace model.layers. with blk., rename the attention matrices (q_proj to attn_q, o_proj to attn_output) and rename the MLP matrices (gate_proj to ffn_gate, up_proj to ffn_up, down_proj to ffn_down). The two norms become attn_norm (before attention) and ffn_norm (before the MLP). The embedding is token_embd, the final norm is output_norm, and the head is output. The names match gguf.get_tensor_name_map in gguf-py 0.19.0 for the Llama architecture.
The GGUF side is visible in the tiny random-weight model used across the course (4 layers, dimension 256, 8 heads, 2 key/value heads, F32). Its text output is meaningless, but its structure is real. The command uses the gguf-dump script that ships with gguf-py, and grep keeps only the tensors of the embedding, layer 0, and the output matrix.
gguf-dump tiny-random-qwen2vocab-f32.gguf | grep -E "token_embd|blk\.0\.| output\.weight"Output (gguf-py 0.19.0; the first column is the tensor index and the second is the element count):
1: 38895616 | 256, 151936, 1, 1 | F32 | token_embd.weight
2: 256 | 256, 1, 1, 1 | F32 | blk.0.attn_norm.weight
3: 65536 | 256, 256, 1, 1 | F32 | blk.0.attn_q.weight
4: 16384 | 256, 64, 1, 1 | F32 | blk.0.attn_k.weight
5: 16384 | 256, 64, 1, 1 | F32 | blk.0.attn_v.weight
6: 65536 | 256, 256, 1, 1 | F32 | blk.0.attn_output.weight
7: 256 | 256, 1, 1, 1 | F32 | blk.0.ffn_norm.weight
8: 180224 | 256, 704, 1, 1 | F32 | blk.0.ffn_gate.weight
9: 180224 | 256, 704, 1, 1 | F32 | blk.0.ffn_up.weight
10: 180224 | 704, 256, 1, 1 | F32 | blk.0.ffn_down.weight
39: 38895616 | 256, 151936, 1, 1 | F32 | output.weightTwo details stand out.
- The key and value matrices have 64 rows: 2 key/value heads times a head dimension of 32. The query matrix has 256 rows because all 8 heads have their own query.
- GGUF lists dimensions innermost first. For
ffn_gate.weightit prints256, 704, so the first number is the input width and the second is the output width. PyTorch checkpoints list the same matrix as[704, 256], output first.
This model has a separate output.weight of the same size as the embedding table (38,895,616 numbers each), so it is untied.
Interview points
- A decoder-only block has two residual sublayers:
x = x + Attention(norm(x)), thenx = x + MLP(norm(x)). Attention is the only place where positions exchange information. - Pre-norm means the norm is applied to the sublayer input while the residual stream stays unnormalized. This leaves an additive path from the embedding to the output.
- RMSNorm divides by the root mean square and multiplies by a learned weight. It has no mean subtraction and no bias, so it costs less than LayerNorm.
- SwiGLU uses three matrices and a SiLU gate. The hidden size is about two thirds of
4 * d, so the parameter count stays near that of a two-matrix MLP. - Attention without positional information is permutation-equivariant, which means shuffling the input only shuffles the output. The causal mask reveals only how many tokens came before, not their order.
- RoPE rotates query and key pairs by
pos * theta_i. The score depends on relative distance only, and no position table limits the length. - Llama-style models rotate consecutive pairs and NEOX-style models rotate pairs
(i, i + d/2). Converters reorder the Q and K rows to match the engine's layout. - The parameter count from a config is: embedding + layers x (Q, K, V, O, three MLP matrices, two norms) + final norm + head. The Llama 3 8B shape gives 8,030,261,248.
- Tying the LM head to the embedding table removes
vocab * dimparameters. A GGUF file withoutoutput.weightis a tied model. - Grouped-query attention shrinks the K and V projections by
heads / kv_heads, and with them the KV cache.
Summary
- A token id selects one row of the embedding table, which is a vector of
hidden_sizenumbers. - The residual stream is one vector per position, and each sublayer adds its output to it.
- RMSNorm computes
x / sqrt(mean(x^2) + eps) * weight. It removes overall scale and keeps direction. - The classic MLP has two matrices and a hidden size of
4 * d. SwiGLU has three matrices:down(silu(gate(x)) * up(x)). - Llama 2 7B has
d = 4096and a SwiGLU hidden size of 11,008. Llama 3 8B has 14,336. - Attention alone cannot see token order. Position information comes from a learned table or from RoPE.
- RoPE rotates each coordinate pair by
pos * base^(-2i/d), so rotated query-key scores depend on the position difference. - A model is an embedding table, N blocks with separate weights, a final RMSNorm, and an LM head. A forward pass gives logits for every position, and generation uses the last one.
config.jsongiveshidden_size,intermediate_size, the layer count, the head counts, the vocabulary size, the RoPE base, and the tying flag. The head dimension ishidden_size / num_attention_heads.- The Llama 3 8B shape has 8,030,261,248 parameters and the Llama 3.2 1B shape has 1,235,814,400.
- GGUF tensor names are shorter than checkpoint names (
blk.0.attn_q.weight), and GGUF lists dimensions innermost first.