The Forward Pass in Java

A language model is, at its core, a function. You give it one token (a piece of text, such as a word or part of a word, represented by a number) and it gives back one score for every token in its vocabulary. The highest score marks the token the model thinks should come next. The forward pass is the name for running that function once: the input flows forward through the model's layers, from the token at the start to the scores at the end. The scores are called logits (raw, unnormalized scores; a larger logit means a more likely token).

Until now this course described the pieces (the transformer block, attention, the KV cache, the GGUF file format). This lesson turns them into working Java. It writes the forward pass for Llama-style models in plain Java, one operation per example, starting from the smallest operations (copying a row of numbers) and building up to a full model. Then it checks the Java code against llama.cpp, the C++ program that most local-LLM tools use. It compares tensor by tensor (a tensor is just a block of numbers, here an array) with a tool called llama-eval-callback, and token by token with greedy generation (always pick the highest-scoring token).

After this lesson, you can do four things. You can map every line of llama.cpp's Llama graph to a line of Java code. You can explain why two correct implementations still differ in the last digits. You can run a GGUF model from start to finish with the classes in this section. And you can use the Llama class of Project 18 as the base for the next lessons.

The table below lists the steps in the order the model performs them. Each row is one example in this lesson.

The theory comes from earlier lessons, and this lesson gives a short reminder where each idea is used:

• The block structure, RMSNorm (rescaling a vector to a standard size), SwiGLU (the feed-forward block), and RoPE (rotary position embedding, which encodes where a token is in the text) come from The Transformer Block.
• Queries, keys, values, and grouped-query attention (several query heads sharing one key/value head) come from Self-Attention.
• The KV cache (key-value cache, which stores earlier tokens' keys and values so they are not recomputed) comes from Generation and the KV Cache.
• The model file is read with GgufModel from Loading GGUF with the FFM API.
• Prompts are turned into token ids with GgufTokenizer from A Real Tokenizer from GGUF Metadata.

The reference implementation is the Llama graph in src/models/llama.cpp at the pinned llama.cpp commit (the exact version of llama.cpp this course uses, so results can be repeated).

The real-model examples use a tiny random-weight model: Llama architecture, 4 layers, width 256, 8 query heads, 2 key/value heads, the Qwen2 vocabulary, and all weights stored as F32 (32-bit floating-point numbers). Because the weights are random, the text it produces is meaningless. That does not matter for this lesson. Every number it computes comes from a real forward pass, so if the Java engine is correct, it must produce the same numbers as llama.cpp.

A picture of the whole pass helps before the details. Each line below is one step from the table, and the layer steps repeat once per layer (4 times in the tiny model).

token id
  |
  v
embedding lookup            -> vector x (256 numbers)
  |
  |   repeat for every layer:
  |   +--------------------------------------------------+
  |   | RMSNorm -> Q, K, V projections -> RoPE on Q, K   |
  |   | -> store K, V in the cache -> attention          |
  |   | -> output projection -> add to x   (residual)    |
  |   | RMSNorm -> SwiGLU feed-forward     -> add to x   |
  |   +--------------------------------------------------+
  v
final RMSNorm -> output projection -> logits (one score per vocabulary token)
  |
  v
greedy choice: the token with the largest logit

The word residual means the block's result is added to its own input (x = x + result), so each layer adjusts the vector instead of replacing it.

Example 1: the model configuration from metadata keys

Before the model can compute anything, the program must know the model's shape: how many layers it has, how wide its vectors are, how many attention heads it uses, and so on. These numbers are called hyperparameters (settings chosen when the model was designed, as opposed to weights, which are learned in training). A GGUF file stores them as metadata keys (named values in the file header). Each key is named {general.architecture}.{name}, so a Llama model has llama.embedding_length and llama.attention.head_count, and a Qwen2 model has qwen2.embedding_length. The architecture name comes first so that one reader can handle many model families.

Some keys are optional. When one is missing, llama.cpp falls back to a default, and the Java code must do the same:

This program builds a Config record (a small immutable Java class that holds the settings) from two metadata maps. The first holds the tiny model's values (printed in lesson 49, Example 12). The second holds the qwen2.* keys stored in llama.cpp's ggml-vocab-qwen2.gguf. For each one it also derives a few sizes and counts the weights. Look for the head dimension, the group size, and the weight totals in the output.

import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.Function;

public class ModelConfig {
    record Config(String arch, int dim, int layers, int heads, int kvHeads, int headDim, int ffn, int vocab,
                  int contextLength, float eps, float ropeBase, int ropeDims, boolean neox) {
        static Config from(Map<String, Object> md) {
            String a = (String) md.get("general.architecture");
            Function<String, Number> num = k -> (Number) md.get(a + "." + k);
            int dim = num.apply("embedding_length").intValue(), heads = num.apply("attention.head_count").intValue();
            int kvHeads = md.containsKey(a + ".attention.head_count_kv") ? num.apply("attention.head_count_kv").intValue() : heads;
            int vocab = md.containsKey(a + ".vocab_size") ? num.apply("vocab_size").intValue() : ((List<?>) md.get("tokenizer.ggml.tokens")).size();
            float base = md.containsKey(a + ".rope.freq_base") ? num.apply("rope.freq_base").floatValue() : 10000f;
            int ropeDims = md.containsKey(a + ".rope.dimension_count") ? num.apply("rope.dimension_count").intValue() : dim / heads;
            boolean neox = switch (a) {                    // llama_model_rope_type in src/llama-model.cpp
                case "llama" -> false;
                case "qwen2" -> true;
                default -> throw new IllegalArgumentException("architecture not supported here: " + a);
            };
            return new Config(a, dim, num.apply("block_count").intValue(), heads, kvHeads, dim / heads,
                    num.apply("feed_forward_length").intValue(), vocab, num.apply("context_length").intValue(),
                    num.apply("attention.layer_norm_rms_epsilon").floatValue(), base, ropeDims, neox);
        }
    }

    public static void main(String[] args) {
        Map<String, Object> tiny = Map.ofEntries(Map.entry("general.architecture", "llama"),   // the tiny random model (lesson 46)
                Map.entry("llama.context_length", 4096L), Map.entry("llama.embedding_length", 256L), Map.entry("llama.block_count", 4L),
                Map.entry("llama.feed_forward_length", 704L), Map.entry("llama.attention.head_count", 8L),
                Map.entry("llama.attention.head_count_kv", 2L), Map.entry("llama.rope.dimension_count", 32L),
                Map.entry("llama.attention.layer_norm_rms_epsilon", 1e-5f), Map.entry("llama.rope.freq_base", 500000f),
                Map.entry("tokenizer.ggml.tokens", Collections.nCopies(151936, "")));
        Map<String, Object> qwen = Map.of("general.architecture", "qwen2",                   // keys of ggml-vocab-qwen2.gguf
                "qwen2.block_count", 32L, "qwen2.context_length", 32768L, "qwen2.embedding_length", 4096L,
                "qwen2.feed_forward_length", 11008L, "qwen2.attention.head_count", 32L, "qwen2.attention.head_count_kv", 32L,
                "qwen2.rope.freq_base", 1000000f, "qwen2.attention.layer_norm_rms_epsilon", 1e-6f,
                "tokenizer.ggml.tokens", Collections.nCopies(151936, ""));
        for (Map<String, Object> md : List.of(tiny, qwen)) {
            Config c = Config.from(md);
            System.out.println(c);
            long kvDim = (long) c.kvHeads() * c.headDim();
            long perLayer = 2L * c.dim() * c.dim() + 2 * c.dim() * kvDim + 3L * c.dim() * c.ffn() + 2 * c.dim();
            long total = c.layers() * perLayer + 2L * c.vocab() * c.dim() + c.dim();
            System.out.printf("  head dim %d, KV dim %d, %d query heads per KV head, %,d weights per layer, %,d in total%n",
                    c.headDim(), kvDim, c.heads() / c.kvHeads(), perLayer, total);
        }
    }
}

Config[arch=llama, dim=256, layers=4, heads=8, kvHeads=2, headDim=32, ffn=704, vocab=151936, contextLength=4096, eps=1.0E-5, ropeBase=500000.0, ropeDims=32, neox=false]
  head dim 32, KV dim 64, 4 query heads per KV head, 705,024 weights per layer, 80,611,584 in total
Config[arch=qwen2, dim=4096, layers=32, heads=32, kvHeads=32, headDim=128, ffn=11008, vocab=151936, contextLength=32768, eps=1.0E-6, ropeBase=1000000.0, ropeDims=128, neox=true]
  head dim 128, KV dim 4096, 1 query heads per KV head, 202,383,360 weights per layer, 7,720,931,328 in total

The first record is the tiny model. The values come straight from the keys: width (dim) 256, 4 layers, 8 heads, 2 key/value heads, feed-forward width 704, vocabulary 151,936, context 4,096 tokens, epsilon 1e-5 (a tiny constant used in RMSNorm, Example 4), RoPE base 500,000, and 32 rotated dimensions. Three values are derived from them:

• Head dimension dim / heads = 256 / 8 = 32. The 256 numbers of a vector are split into 8 slices, one per query head, so each head works on 32 numbers.
• KV dimension kvHeads * headDim = 2 * 32 = 64. This is the size of one position's key vector (and of its value vector). The key and value projections therefore produce 64 numbers instead of 256.
• Group size heads / kvHeads = 4. Four query heads share each key/value head. This is grouped-query attention.

neox=false is the RoPE layout for llama (Example 6).

The next line of output counts weights. The per-layer count adds four parts:

Four layers, two embedding-sized matrices (token_embd and output, 151,936 by 256 each), and the final norm vector give 80,611,584. This is the same count that GgufInfo reported for this model in The GGUF Format. The two vocabulary matrices make up 96% of this model. That share is typical only for tiny models with large vocabularies; in a real 7B model the layers dominate.

The second record shows the defaults at work. The qwen2 keys of that vocabulary file have no rope.dimension_count and no vocab_size, so the record uses the head dimension (128) and the token count (151,936). kvHeads=32 equals heads, so this configuration has one query head per key/value head (no grouping), and neox=true selects the other RoPE layout. Its 7,720,931,328 weights are what a 7B-class model with these shapes holds ("7B" means about 7 billion weights).

Example 2: dequantizing weights to float[] (F32, F16, Q8_0)

A model file stores weights in different number formats to save space. The engine in this lesson does its arithmetic in Java float values, so before computing it converts every weight tensor to a float[]. This conversion is called dequantizing (turning compressed numbers back into ordinary floats). The cost depends on the format:

• F32: each weight is already a 4-byte float, so the bytes are copied in bulk.
• F16: each weight is a 2-byte half-precision float, so each value is converted one at a time.
• Q8_0: weights are stored in blocks of 32. A block holds one scale d (a 2-byte float) and 32 small integers qs[j] (1 byte each), 34 bytes in total. Each weight is rebuilt as d * qs[j].

These are the three cases of GgufModel.floats from lesson 49. This program stores the same 64 values in all three formats in memory segments, converts each back, and prints the largest error against the original values. It then computes how much Java heap the float[] approach needs for three model sizes. Look at how the error grows as the format gets smaller, and at the GiB (gibibyte, 1024^3 bytes) figures in the last three lines.

import java.lang.foreign.Arena;
import java.lang.foreign.MemorySegment;
import java.lang.foreign.ValueLayout;
import java.nio.ByteOrder;
import java.util.Locale;

public class Dequantize {
    static final ValueLayout.OfFloat F32 = ValueLayout.JAVA_FLOAT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);
    static final ValueLayout.OfShort U16 = ValueLayout.JAVA_SHORT_UNALIGNED.withOrder(ByteOrder.LITTLE_ENDIAN);

    // The three cases of GgufModel.floats (lesson 46, Project 16).
    static float[] dequantize(int type, MemorySegment s, int n) {
        float[] out = new float[n];
        switch (type) {
            case 0 -> MemorySegment.copy(s, F32, 0, out, 0, n);                                  // F32: a bulk copy
            case 1 -> { for (int i = 0; i < n; i++) out[i] = Float.float16ToFloat(s.get(U16, 2L * i)); }   // F16
            case 8 -> {                                                                           // Q8_0: d * qs[j]
                for (int b = 0; b < n / 32; b++) {
                    float d = Float.float16ToFloat(s.get(U16, 34L * b));
                    for (int j = 0; j < 32; j++) out[b * 32 + j] = d * s.get(ValueLayout.JAVA_BYTE, 34L * b + 2 + j);
                }
            }
            default -> throw new IllegalArgumentException("type " + type);
        }
        return out;
    }

    static int roundf(float v) { return (int) Math.copySign(Math.floor(Math.abs((double) v) + 0.5), v); }

    public static void main(String[] args) {
        int n = 64;
        float[] w = new float[n];
        for (int i = 0; i < n; i++) w[i] = (float) Math.cos(i * 0.7) * 0.08f;
        try (Arena arena = Arena.ofConfined()) {
            MemorySegment f32 = arena.allocate(4L * n), f16 = arena.allocate(2L * n), q8 = arena.allocate(34L * n / 32);
            for (int i = 0; i < n; i++) { f32.setAtIndex(F32, i, w[i]); f16.setAtIndex(U16, i, Float.floatToFloat16(w[i])); }
            for (int b = 0; b < n / 32; b++) {                                                    // quantize_row_q8_0_ref
                float amax = 0;
                for (int j = 0; j < 32; j++) amax = Math.max(amax, Math.abs(w[b * 32 + j]));
                float d = amax / 127, id = d != 0 ? 1 / d : 0;
                q8.set(U16, 34L * b, Float.floatToFloat16(d));
                for (int j = 0; j < 32; j++) q8.set(ValueLayout.JAVA_BYTE, 34L * b + 2 + j, (byte) roundf(w[b * 32 + j] * id));
            }
            Object[][] cases = {{"F32", 0, f32}, {"F16", 1, f16}, {"Q8_0", 8, q8}};
            for (Object[] c : cases) {
                float[] x = dequantize((Integer) c[1], (MemorySegment) c[2], n);
                double maxErr = 0;
                for (int i = 0; i < n; i++) maxErr = Math.max(maxErr, Math.abs(x[i] - w[i]));
                System.out.printf(Locale.ROOT, "%-4s %3d bytes -> float[%d] (%d bytes): x[1..3] = %.7f %.7f %.7f, max error %.7f%n",
                        c[0], ((MemorySegment) c[2]).byteSize(), n, 4 * n, x[1], x[2], x[3], maxErr);
            }
        }
        for (long params : new long[] {80_611_584L, 7_720_931_328L, 8_030_261_248L})
            System.out.printf(Locale.ROOT, "%,15d weights: %8.2f GiB as float[], %7.2f GiB as Q8_0%n",
                    params, params * 4 / Math.pow(1024, 3), params * 34.0 / 32 / Math.pow(1024, 3));
    }
}

F32  256 bytes -> float[64] (256 bytes): x[1..3] = 0.0611874 0.0135974 -0.0403877, max error 0.0000000
F16  128 bytes -> float[64] (256 bytes): x[1..3] = 0.0611877 0.0135956 -0.0403748, max error 0.0000304
Q8_0  68 bytes -> float[64] (256 bytes): x[1..3] = 0.0611005 0.0138578 -0.0403137, max error 0.0003149
     80,611,584 weights:     0.30 GiB as float[],    0.08 GiB as Q8_0
  7,720,931,328 weights:    28.76 GiB as float[],    7.64 GiB as Q8_0
  8,030,261,248 weights:    29.92 GiB as float[],    7.95 GiB as Q8_0

The three conversion lines start from 256, 128, and 68 bytes and all produce a float[64] of 256 bytes. The Q8_0 size follows from the block layout: 64 values are 2 blocks, and 2 blocks of 34 bytes are 68 bytes. So Q8_0 uses about a quarter of the F32 space, and F16 uses half.

• F32 is exact: the values are copied bit for bit, with error 0.
• F16 changes the values in the fourth or fifth significant digit, for example 0.0611874 becomes 0.0611877. The maximum error is 0.0000304.
• Q8_0 rounds each value to one of 255 steps of its block's scale. 0.0135974 becomes 0.0138578, and the maximum error is 0.0003149, ten times larger than F16 but still small relative to values up to 0.08.

The smaller the format, the larger the rounding error. That is the trade: less memory for slightly less accurate weights.

The last three lines show what the float[] approach costs in memory. They are the reason Quantized Kernels, the Vector API, and Threads computes on quantized blocks directly. Dequantized to float[], the tiny model (80.6 million weights) needs 0.30 GiB of heap. The 7.7-billion-weight configuration of Example 1 needs 28.76 GiB as float[] against 7.64 GiB in Q8_0. An 8-billion-weight model, such as Llama 3 8B, needs 29.92 GiB as float[] instead of 7.95 GiB as Q8_0, and a Java array also cannot hold more than about 2^31 elements, so a large token_embd matrix would not even fit in one array.

Dequantizing everything at load time is the simplest correct design, and it is what Project 18 does. Its cost is four bytes of heap per weight and a slower start. Quantized Kernels, the Vector API, and Threads keeps the weights as Q8_0 blocks in the mapped file and multiplies them without expanding them.

Example 3: embedding lookup

A token id is only a number, such as 9707 for Hello. The model cannot do arithmetic with a bare id in any meaningful way, so the first operation replaces the id with a vector (a list of numbers) that represents the token. The model learned one such vector per vocabulary entry during training. They are stored together in the tensor token_embd.weight, which is a table with one row per token. Looking up a token's vector is called an embedding lookup.

GGUF stores that tensor with dimensions [dim, vocab]. GGUF lists the fastest-changing dimension first, so this means vocab rows of dim values each, stored one row after another. Row t therefore occupies floats t * dim to t * dim + dim - 1. This program builds a 6-token, 4-dimensional table from a seeded random generator (so the numbers repeat on every run) and looks up three ids. Check that the output floats' positions match the formula.

import java.util.Arrays;
import java.util.Locale;
import java.util.Random;

public class EmbeddingLookup {
    public static void main(String[] args) {
        int dim = 4, vocab = 6;
        float[] tokenEmbd = new float[vocab * dim];                // GGUF dims [4, 6]: 6 rows of 4
        Random random = new Random(7);
        for (int i = 0; i < tokenEmbd.length; i++) tokenEmbd[i] = Math.round(random.nextGaussian() * 100) / 100f;
        for (int token : new int[] {3, 1, 3}) {
            float[] x = new float[dim];
            System.arraycopy(tokenEmbd, token * dim, x, 0, dim);   // row `token` starts at token * dim
            System.out.println("token " + token + " -> floats " + token * dim + ".." + (token * dim + dim - 1) + " = " + Arrays.toString(x));
        }
        long realRow = 9707L * 256;
        System.out.printf(Locale.ROOT, "tiny model: token 9707 starts at float %,d, byte %,d of token_embd.weight%n", realRow, realRow * 4);
    }
}

token 3 -> floats 12..15 = [0.69, -0.72, -0.4, 0.79]
token 1 -> floats 4..7 = [1.34, -0.95, 0.6, 1.2]
token 3 -> floats 12..15 = [0.69, -0.72, -0.4, 0.79]
tiny model: token 9707 starts at float 2,484,992, byte 9,939,968 of token_embd.weight

Token 3 maps to floats 12 to 15 and token 1 to floats 4 to 7, because each row is 4 floats long (3 * 4 = 12, 1 * 4 = 4). The lookup is a copy, so the second occurrence of token 3 gets exactly the same vector as the first. Before attention runs, the model has no way to tell the two occurrences apart. Position information is added later, by RoPE (Example 6).

System.arraycopy is the whole operation. llama.cpp names it GET_ROWS in its graph (embd = GET_ROWS(token_embd.weight, inp_tokens) in the llama-eval-callback output of Example 11), which describes the same thing: select rows of a matrix by index.

The last line applies the formula to the tiny model, whose width is 256: token 9707 (Hello) starts at float 9707 * 256 = 2,484,992. Each float is 4 bytes, so that is byte 9,939,968 of the tensor. Lesson 49 read the first values of this row through a memory-segment view and found the same values that llama.cpp prints for embd.

Example 4: RMSNorm

As numbers pass through many layers, they can grow or shrink until later operations behave badly. A normalization step fixes this by rescaling the vector to a standard size before each major operation. RMSNorm (root mean square normalization) is the version Llama-style models use. The root mean square (RMS) of a vector is a measure of its typical size: square every element, take the average, then take the square root.

RMSNorm divides the vector by its RMS so the result has an RMS of 1, then multiplies each element by a learned weight:

y[i] = x[i] / sqrt(mean(x^2) + eps) * weight[i]

Here eps (epsilon) is a tiny constant such as 1e-5 that prevents division by zero. For x = [1, 2, 3, 4] the steps are: square each element to get [1, 4, 9, 16], average them to get 7.5, take sqrt(7.5 + 0.00001) = 2.7386, and divide each element by that. Every block applies RMSNorm twice, before attention and before the feed-forward block, and once more after the last layer.

ggml's CPU kernel accumulates the sum of squares in double precision (ggml_float, a 64-bit number, to lose less accuracy while adding many values), and this program does the same. The program runs five cases: a plain vector, the same vector with weights, the same vector times 10, a random 256-value vector, and a vector so small that eps matters.

import java.util.Arrays;
import java.util.Locale;
import java.util.Random;

public class RmsNormDemo {
    // y[i] = x[i] / sqrt(mean(x^2) + eps) * weight[i]; the sum is accumulated in double as ggml does.
    static void rmsnorm(float[] out, float[] x, float[] weight, float eps) {
        double ss = 0;
        for (float v : x) ss += (double) v * v;
        float scale = (float) (1.0 / Math.sqrt(ss / x.length + eps));
        for (int i = 0; i < x.length; i++) out[i] = x[i] * scale * weight[i];
    }

    static double rms(float[] v) { double s = 0; for (float x : v) s += (double) x * x; return Math.sqrt(s / v.length); }

    public static void main(String[] args) {
        float[] x = {1, 2, 3, 4}, w = {1, 1, 0.5f, 2}, out = new float[4], ones = {1, 1, 1, 1};
        rmsnorm(out, x, ones, 1e-5f);
        System.out.println("weight 1   : " + Arrays.toString(out));
        rmsnorm(out, x, w, 1e-5f);
        System.out.println("weight w   : " + Arrays.toString(out));
        float[] big = {10, 20, 30, 40};
        rmsnorm(out, big, ones, 1e-5f);
        System.out.println("x times 10 : " + Arrays.toString(out));
        Random random = new Random(1);
        float[] h = new float[256], y = new float[256], unit = new float[256];
        Arrays.fill(unit, 1f);
        for (int i = 0; i < 256; i++) h[i] = (float) (random.nextGaussian() * 0.02);
        rmsnorm(y, h, unit, 1e-5f);
        System.out.printf(Locale.ROOT, "256 values: rms before %.5f, after %.5f%n", rms(h), rms(y));
        float[] tiny = new float[256];
        Arrays.fill(tiny, 1e-4f);
        rmsnorm(y, tiny, unit, 1e-5f);
        System.out.printf(Locale.ROOT, "all 1e-4  : rms after %.5f (eps matters when mean(x^2) is near %.0e)%n", rms(y), 1e-5);
    }
}

weight 1   : [0.36514813, 0.73029625, 1.0954444, 1.4605925]
weight w   : [0.36514813, 0.73029625, 0.5477222, 2.921185]
x times 10 : [0.36514837, 0.73029673, 1.0954452, 1.4605935]
256 values: rms before 0.02072, after 0.98856
all 1e-4  : rms after 0.03161 (eps matters when mean(x^2) is near 1e-05)

For x = [1, 2, 3, 4], the mean of the squares is (1 + 4 + 9 + 16) / 4 = 7.5, and 1 / sqrt(7.5 + 0.00001) = 0.3651481. With a weight of all ones, the output is x times that scale: 0.365, 0.730, 1.095, 1.461. The result has a root mean square of 1, so every layer receives inputs of a predictable size. The relative sizes of the elements are unchanged (4 is still four times 1); only the overall scale moves.

The weight [1, 1, 0.5, 2] scales the elements afterwards: the third element halves to 0.548 and the fourth doubles to 2.921. These weights are learned, and in the tiny random model they are all 1, so norm-0 and attn_norm-0 in Example 11 have equal sums.

x times 10 gives almost the same output as x: RMSNorm removes the overall scale of the input, which keeps the residual stream from growing without bound through many layers. The tiny differences in the last digits come from eps, which matters less as the input grows.

The fourth line uses 256 random values with a root mean square of 0.02072, typical of an embedding row. After the norm, it is 0.98856 instead of exactly 1, because eps = 1e-5 is not negligible next to mean(x^2) = 0.00043 (the mean of squares is the RMS squared: 0.02072^2 = 0.00043). The last line makes eps dominate: 256 values of 1e-4 have mean(x^2) = 1e-8, a thousand times smaller than eps, so the output stays small (0.03161) instead of being amplified to 1. eps exists to avoid dividing by zero and to keep near-zero vectors from being blown up into noise.

Example 5: matrix-vector product in GGUF's dimension order

Most of the work in a forward pass is one operation: turning an input vector into a new vector by multiplying it with a table of learned numbers (a weight matrix). This is called a matrix-vector product. Each output number is the dot product of one row of the matrix with the input vector. A dot product multiplies two lists element by element and adds the results. For example, [1, 2, 3] dot [1, 2, 3] is 1*1 + 2*2 + 3*3 = 14.

The matrix is stored in the file as one flat list of floats, so the program must know where each row starts. GGUF lists a tensor's dimensions innermost first: [ne0, ne1] means ne1 rows of ne0 values each. For a linear layer (a layer that applies one matrix product), ne0 is the number of inputs and ne1 is the number of outputs. Output row r is stored starting at float r * ne0. PyTorch (a common training library) lists the same weight as [out, in] = [ne1, ne0], the reverse order, and the bytes are identical (The GGUF Format, Example 13).

This program multiplies a 2-by-3 weight by a 3-value input, shows what a wrong reading of the bytes gives, and lists the shapes of the tiny model's matrices with their multiply-add counts.

import java.util.Arrays;

public class MatVec {
    // out[r] = sum over i of w[r * cols + i] * x[i]. A GGUF tensor with dims [cols, rows] stores row r at r * cols.
    static void matvec(float[] out, float[] w, float[] x, int rows, int cols) {
        for (int r = 0; r < rows; r++) {
            float s = 0;
            for (int i = 0; i < cols; i++) s += w[r * cols + i] * x[i];
            out[r] = s;
        }
    }

    public static void main(String[] args) {
        long[] dims = {3, 2};                                      // GGUF: ne0 = 3 inputs, ne1 = 2 outputs
        float[] w = {1, 2, 3, 4, 5, 6};                            // as stored in the file
        float[] x = {1, 2, 3}, out = new float[2];
        matvec(out, w, x, (int) dims[1], (int) dims[0]);
        System.out.println("dims " + Arrays.toString(dims) + ": out = " + Arrays.toString(out));
        System.out.println("PyTorch shape of the same weight: [" + dims[1] + ", " + dims[0] + "], same bytes, same result");
        float[] wrong = new float[2];                              // reading the bytes as columns instead of rows
        for (int r = 0; r < 2; r++) for (int i = 0; i < 3; i++) wrong[r] += w[i * 2 + r] * x[i];
        System.out.println("bytes read as columns: " + Arrays.toString(wrong));
        long[][] tiny = {{256, 256}, {256, 64}, {256, 704}, {704, 256}, {256, 151936}};
        String[] names = {"attn_q", "attn_k", "ffn_up", "ffn_down", "output"};
        for (int i = 0; i < tiny.length; i++)
            System.out.printf("%-8s dims %-13s -> %6d outputs from %3d inputs, %,10d multiply-adds%n",
                    names[i], Arrays.toString(tiny[i]), tiny[i][1], tiny[i][0], tiny[i][0] * tiny[i][1]);
    }
}

dims [3, 2]: out = [14.0, 32.0]
PyTorch shape of the same weight: [2, 3], same bytes, same result
bytes read as columns: [22.0, 28.0]
attn_q   dims [256, 256]    ->    256 outputs from 256 inputs,     65,536 multiply-adds
attn_k   dims [256, 64]     ->     64 outputs from 256 inputs,     16,384 multiply-adds
ffn_up   dims [256, 704]    ->    704 outputs from 256 inputs,    180,224 multiply-adds
ffn_down dims [704, 256]    ->    256 outputs from 704 inputs,    180,224 multiply-adds
output   dims [256, 151936] -> 151936 outputs from 256 inputs, 38,895,616 multiply-adds

The stored floats 1 2 3 4 5 6 with dims [3, 2] are two rows, [1, 2, 3] and [4, 5, 6]. With x = [1, 2, 3], the outputs are 1 + 4 + 9 = 14 and 4 + 10 + 18 = 32. The rule in code is rows = dims[1], cols = dims[0], and the inner loop runs over cols, the contiguous dimension, so the memory access is sequential.

Reading the same bytes as columns (element (r, i) at i * 2 + r) gives [22, 28], a different result with no error message. A model with transposed weights still produces numbers and tokens, only wrong ones, which is why Example 11 compares intermediate tensors and not only the final text.

llama.cpp writes the same operation as MUL_MAT(weight, input), and its shapes in Example 11 follow the same rule: MUL_MAT(blk.0.attn_k.weight{256, 64, 1, 1}, attn_norm-0{256, 2, 1, 1}) = {64, 2, 1, 1} turns 2 positions of 256 values into 2 positions of 64 values.

The last five lines of output list the shapes in the tiny model. Each line also gives the number of multiply-adds (one multiplication plus one addition, the basic unit of work), which is inputs * outputs. For attn_q, that is 256 * 256 = 65,536. These counts are per token. The query projection is 65,536 multiply-adds, the key projection 16,384, each feed-forward matrix 180,224, and the output projection 38,895,616. In this model, the output projection over the 151,936-token vocabulary costs more than all four layers together (4 times 705,024 = 2.8 million), so the tiny model spends almost all of its time computing logits. A decode step reads every weight once, and the matrix-vector product is where the time goes; the next lesson speeds it up with the Vector API and threads.

Example 6: RoPE, normal versus NEOX, chosen per architecture

Attention compares tokens with each other, but the embedding of a token says nothing about where in the text it sits. The model needs a way to know that one token is 3 positions before another. RoPE (rotary position embedding) does this by rotating the query and key vectors by an angle that grows with the token's position. After rotation, the attention score between two tokens depends on the distance between them (The Transformer Block, Example 8).

The rotation works on pairs of values. Think of each pair (a, b) as a point on a flat plane. Rotating the point by an angle theta gives (a*cos(theta) - b*sin(theta), a*sin(theta) + b*cos(theta)). The angle of pair i is pos * base^(-2i / n_rot), where pos is the token position, base is rope.freq_base, and n_rot is rope.dimension_count. The first pairs rotate fast and the last pairs rotate slowly, so together they encode position at many scales (like the hands of a clock).

The remaining question is which two values form a pair. Two layouts exist:

• The normal layout pairs adjacent values (2i, 2i+1).
• The NEOX layout (named after GPT-NeoX, a model that used it) pairs value i with value i + n_rot/2, the first half with the second half.

A model must run with the layout it was trained with. llama.cpp picks the layout per architecture in llama_model_rope_type (src/llama-model.cpp): LLM_ARCH_LLAMA is in the LLAMA_ROPE_TYPE_NORM group and LLM_ARCH_QWEN2 in the LLAMA_ROPE_TYPE_NEOX group. This program rotates the vector [1, 2, ..., 8] at position 1 both ways, then checks that the two layouts are the same rotation under a reordering of values.

import java.util.Arrays;
import java.util.Locale;

public class RopeLayouts {
    // Rotates each head of vec in place. Normal: pairs (2i, 2i+1). NEOX: pairs (i, i + ropeDims/2).
    // theta_i = pos * base^(-2i/ropeDims), computed by repeated multiplication in float as ggml does.
    static void rope(float[] vec, int nHeads, int headDim, int ropeDims, int pos, float base, boolean neox) {
        float thetaScale = (float) Math.pow(base, -2.0 / ropeDims);
        for (int h = 0; h < nHeads; h++) {
            float theta = pos;
            for (int i = 0; i < ropeDims / 2; i++) {
                float cos = (float) Math.cos(theta), sin = (float) Math.sin(theta);
                int a = h * headDim + (neox ? i : 2 * i), b = neox ? a + ropeDims / 2 : a + 1;
                float x0 = vec[a], x1 = vec[b];
                vec[a] = x0 * cos - x1 * sin;
                vec[b] = x0 * sin + x1 * cos;
                theta *= thetaScale;
            }
        }
    }

    static String fmt(float[] v) {
        StringBuilder sb = new StringBuilder();
        for (float f : v) sb.append(String.format(Locale.ROOT, "%7.3f", f));
        return sb.toString();
    }

    public static void main(String[] args) {
        float[] x = {1, 2, 3, 4, 5, 6, 7, 8};
        float[] normal = x.clone(), neox = x.clone();
        rope(normal, 1, 8, 8, 1, 10000f, false);
        rope(neox, 1, 8, 8, 1, 10000f, true);
        System.out.println("input       " + fmt(x));
        System.out.println("normal, p=1 " + fmt(normal) + "   pairs (0,1) (2,3) (4,5) (6,7)");
        System.out.println("NEOX,   p=1 " + fmt(neox) + "   pairs (0,4) (1,5) (2,6) (3,7)");
        int[] perm = {0, 4, 1, 5, 2, 6, 3, 7};                     // interleave the two halves
        float[] px = new float[8], pNeox = new float[8];
        for (int i = 0; i < 8; i++) { px[i] = x[perm[i]]; pNeox[i] = neox[perm[i]]; }
        rope(px, 1, 8, 8, 1, 10000f, false);
        System.out.println("normal(permuted x) == permuted NEOX(x): " + Arrays.equals(px, pNeox));
        float[] p0 = x.clone();
        rope(p0, 1, 8, 8, 0, 10000f, false);
        System.out.println("position 0 leaves the vector unchanged: " + Arrays.equals(p0, x));
        for (String arch : new String[] {"llama", "qwen2"})
            System.out.println(arch + " -> " + (arch.equals("llama") ? "LLAMA_ROPE_TYPE_NORM" : "LLAMA_ROPE_TYPE_NEOX"));
    }
}

input         1.000  2.000  3.000  4.000  5.000  6.000  7.000  8.000
normal, p=1  -1.143  1.922  2.586  4.280  4.940  6.050  6.992  8.007   pairs (0,1) (2,3) (4,5) (6,7)
NEOX,   p=1  -3.667  1.391  2.930  3.992  3.543  6.170  7.030  8.004   pairs (0,4) (1,5) (2,6) (3,7)
normal(permuted x) == permuted NEOX(x): true
position 0 leaves the vector unchanged: true
llama -> LLAMA_ROPE_TYPE_NORM
qwen2 -> LLAMA_ROPE_TYPE_NEOX

At position 1 with base 10,000 and 8 dimensions, the four angles are 1, 0.1, 0.01, and 0.001 radians. The first pair turns by a full radian and the last pair barely moves, which is visible in both rows.

• Normal. The pairs are (1, 2), (3, 4), (5, 6), (7, 8). The first pair turns by 1 radian: 1 * cos 1 - 2 * sin 1 = -1.143 and 1 * sin 1 + 2 * cos 1 = 1.922. The last pair, (7, 8), turns by 0.001 radians and becomes (6.992, 8.007).
• NEOX. The pairs are (1, 5), (2, 6), (3, 7), (4, 8). The first pair turns by 1 radian: 1 * cos 1 - 5 * sin 1 = -3.667 in position 0 and 1 * sin 1 + 5 * cos 1 = 3.543 in position 4.

The outputs differ, so a model must be run with the layout it was trained with. The line normal(permuted x) == permuted NEOX(x): true shows that the two layouts are the same rotation on a different ordering of the values: interleaving the two halves ([x0, x4, x1, x5, ...]), applying the normal layout, and comparing with the interleaved NEOX result gives true, bit for bit. This is why Llama models can use the normal layout in llama.cpp although the Hugging Face implementation uses the half-split layout: convert_hf_to_gguf.py permutes the rows of the query and key weights of Llama models once at conversion time (LlamaModel.permute in conversion/llama.py at the pinned commit), and the permuted weights produce interleaved pairs. Qwen2 weights are not permuted, so llama.cpp runs them with NEOX. Choosing the layout from the architecture name, as the Config record does, reproduces that decision.

Position 0 leaves the vector unchanged, because every angle is 0. The program computes the angles as ggml does in ggml_rope_cache_init: theta_scale = base^(-2/n_rot) once, then theta *= theta_scale per pair, in float. Computing each angle with Math.pow in double gives slightly different values in the last digits, which would show up as tiny differences in Example 11.

Example 7: attention with a KV cache and grouped-query heads

Attention is how a token gathers information from the tokens before it. Each token produces three vectors: a query (what this token is looking for), a key (what this token offers), and a value (the information it passes on). For the newest token, the model:

1. scores its query against the key of every earlier position (a dot product, divided by sqrt(head dimension) to keep scores in a sensible range),
2. turns the scores into weights that add up to 1 with softmax (exp of each score, divided by the sum of those),
3. returns the weighted average of the value vectors.

A high weight on a position means "this token matters for me". Each query head repeats this independently on its own slice of the vector.

Computing keys and values for all earlier positions at every step would waste work, because they never change. The KV cache stores them the first time, so each is computed only once (Generation and the KV Cache). It is just two flat arrays, one for keys and one for values, filled one position at a time.

With grouped-query attention, several query heads share one key/value head: query head h reads key/value head h / group, which saves cache memory. This program runs three positions through a 4-query-head, 2-KV-head attention with head dimension 2, so the group size is 2. Each step stores the new key and value in the flat cache arrays, scores the query against all cached keys, applies softmax, and mixes the cached values. The values are set to pos + 1 + i / 10 so the output shows which positions contributed. For example, position 0's value is around 1 and position 1's value is around 2, so an output near 1.5 means the head listened to both equally.

import java.util.Locale;
import java.util.Random;

public class CachedAttention {
    public static void main(String[] args) {
        int heads = 4, kvHeads = 2, hd = 2, group = heads / kvHeads, kvDim = kvHeads * hd, ctx = 8;
        float[] keyCache = new float[ctx * kvDim], valueCache = new float[ctx * kvDim];
        Random random = new Random(3);
        for (int pos = 0; pos < 3; pos++) {
            float[] q = new float[heads * hd], k = new float[kvDim], v = new float[kvDim];
            for (int i = 0; i < q.length; i++) q[i] = Math.round(random.nextGaussian() * 10) / 10f;
            for (int i = 0; i < kvDim; i++) { k[i] = Math.round(random.nextGaussian() * 10) / 10f; v[i] = pos + 1 + i / 10f; }
            System.arraycopy(k, 0, keyCache, pos * kvDim, kvDim);   // this position's K and V go into the cache
            System.arraycopy(v, 0, valueCache, pos * kvDim, kvDim);
            float[] out = new float[heads * hd], att = new float[pos + 1];
            StringBuilder weights = new StringBuilder();
            for (int h = 0; h < heads; h++) {
                int kvOff = (h / group) * hd;                         // query head h reads KV head h / group
                float max = Float.NEGATIVE_INFINITY, sum = 0;
                for (int t = 0; t <= pos; t++) {                      // causal: only positions 0..pos exist in the cache
                    float s = 0;
                    for (int i = 0; i < hd; i++) s += q[h * hd + i] * keyCache[t * kvDim + kvOff + i];
                    att[t] = s / (float) Math.sqrt(hd);
                    max = Math.max(max, att[t]);
                }
                for (int t = 0; t <= pos; t++) { att[t] = (float) Math.exp(att[t] - max); sum += att[t]; }
                weights.append(" h").append(h).append("->kv").append(h / group).append(" [");
                for (int t = 0; t <= pos; t++) {
                    float a = att[t] / sum;
                    weights.append(String.format(Locale.ROOT, t == 0 ? "%.2f" : " %.2f", a));
                    for (int i = 0; i < hd; i++) out[h * hd + i] += a * valueCache[t * kvDim + kvOff + i];
                }
                weights.append(']');
            }
            System.out.println("pos " + pos + ":" + weights);
            System.out.printf(Locale.ROOT, "       head 0 output %.3f %.3f, head 2 output %.3f %.3f%n", out[0], out[1], out[4], out[5]);
        }
        System.out.println("cache floats per position: " + 2 * kvDim + " (K and V), without GQA: " + 2 * heads * hd);
    }
}

pos 0: h0->kv0 [1.00] h1->kv0 [1.00] h2->kv1 [1.00] h3->kv1 [1.00]
       head 0 output 1.000 1.100, head 2 output 1.200 1.300
pos 1: h0->kv0 [0.45 0.55] h1->kv0 [0.42 0.58] h2->kv1 [0.52 0.48] h3->kv1 [0.55 0.45]
       head 0 output 1.549 1.649, head 2 output 1.684 1.784
pos 2: h0->kv0 [0.47 0.37 0.16] h1->kv0 [0.51 0.35 0.14] h2->kv1 [0.01 0.96 0.04] h3->kv1 [0.15 0.09 0.77]
       head 0 output 1.690 1.790, head 2 output 2.229 2.329
cache floats per position: 8 (K and V), without GQA: 16

Each pos line lists, for every query head, the KV head it reads and its attention weights over the cached positions. The weights in each bracket add up to 1. The second line of each pair shows the mixed output of heads 0 and 2:

• Position 0. Every head has one cached position, so every weight is 1.00, and each head's output equals that position's value: head 0 reads KV head 0 (1.0, 1.1), head 2 reads KV head 1 (1.2, 1.3).
• Position 1. Two positions are cached. Head 0 gives 0.45 to position 0 and 0.55 to position 1, so its output is between the two values: 0.45 * 1.0 + 0.55 * 2.0 = 1.549 for the first element.
• Position 2. Heads 0 and 1 share KV head 0 but have different queries, so their weights differ ([0.47 0.37 0.16] and [0.51 0.35 0.14]). Heads 2 and 3 share KV head 1 and attend to different positions: head 2 puts 0.96 on position 1, head 3 puts 0.77 on position 2. Head 2's output, 2.229, is close to position 1's value 2.2 for that reason.

The cache holds keys and values in the order positions arrived, pos * kvDim + kvHead * hd + i. Attention reads only positions 0..pos, which also implements the causal mask (the rule that a token may not look at later tokens): future positions are not in the cache yet, so there is nothing to look at.

The last line is the memory effect of grouping: each position stores 8 floats (2 KV heads times 2 values, for K and for V) instead of 16 without grouping. In the tiny model, the group of 4 makes the cache four times smaller than it would be with 8 KV heads.

The code divides each score by sqrt(hd) and subtracts the maximum score before exp. llama.cpp passes the same scale as kq_scale = 1.0f/sqrtf(float(n_embd_head)) in src/models/llama.cpp. Subtracting the maximum keeps exp from overflowing (exp(100) is too large for a float) without changing the softmax result, because the same factor cancels between the numerator and the denominator.

Example 8: the SwiGLU feed-forward block

Attention mixes information between tokens. The feed-forward block then processes each token's vector on its own, which is where much of the model's stored knowledge is thought to live. It first expands the vector to a wider size (704 in the tiny model, 256 before), applies a non-linear step, and shrinks it back. The non-linear step matters: without it, stacked matrix products collapse into a single matrix product and the model could not learn complex patterns.

Llama-style models use SwiGLU. Two projections of the same input, called gate and up, are combined as silu(gate) * up, and a third projection, down, maps the result back to the model width. The gate decides how much of each up value passes through. The function silu(x) = x / (1 + e^-x) is the non-linear step: it is close to x for large positive inputs and close to 0 for large negative ones. This program prints silu at seven points and runs one tiny SwiGLU block step by step, with a width of 2 and a hidden width of 3 so every number can be checked by hand.

import java.util.Locale;

public class SwiGluBlock {
    static float silu(float x) { return x / (1f + (float) Math.exp(-x)); }

    static void matvec(float[] out, float[] w, float[] x, int rows, int cols) {
        for (int r = 0; r < rows; r++) { float s = 0; for (int i = 0; i < cols; i++) s += w[r * cols + i] * x[i]; out[r] = s; }
    }

    static String fmt(float[] v) {
        StringBuilder sb = new StringBuilder();
        for (float f : v) sb.append(String.format(Locale.ROOT, "%7.3f", f));
        return sb.toString();
    }

    public static void main(String[] args) {
        for (float z : new float[] {-4, -2, -0.5f, 0, 0.5f, 2, 4})
            System.out.printf(Locale.ROOT, "silu(%4.1f) = %7.4f%n", z, silu(z));
        int dim = 2, ffn = 3;
        float[] x = {1, -1};
        float[] gate = {2, 0, 0, 2, 1, 1};                         // dims [2, 3]: 3 rows of 2
        float[] up = {1, 1, 1, -1, 0.5f, 0};
        float[] down = {1, 0, 1, 0, 1, -1};                        // dims [3, 2]: 2 rows of 3
        float[] g = new float[ffn], u = new float[ffn], h = new float[ffn], out = new float[dim];
        matvec(g, gate, x, ffn, dim);
        matvec(u, up, x, ffn, dim);
        for (int i = 0; i < ffn; i++) h[i] = silu(g[i]) * u[i];
        matvec(out, down, h, dim, ffn);
        System.out.println("gate       " + fmt(g));
        System.out.println("silu(gate) " + fmt(new float[] {silu(g[0]), silu(g[1]), silu(g[2])}));
        System.out.println("up         " + fmt(u));
        System.out.println("h          " + fmt(h));
        System.out.println("down(h)    " + fmt(out));
    }
}

silu(-4.0) = -0.0719
silu(-2.0) = -0.2384
silu(-0.5) = -0.1888
silu( 0.0) =  0.0000
silu( 0.5) =  0.3112
silu( 2.0) =  1.7616
silu( 4.0) =  3.9281
gate         2.000 -2.000  0.000
silu(gate)   1.762 -0.238  0.000
up           0.000  2.000  0.500
h            0.000 -0.477  0.000
down(h)      0.000 -0.477

One value by hand: silu(2) = 2 / (1 + e^-2) = 2 / (1 + 0.1353) = 1.7616. The silu table shows its shape. For large positive inputs it approaches the input (silu(4) = 3.928); for large negative inputs it approaches 0 from below (silu(-4) = -0.072); and it is slightly negative between, with its minimum near x = -1.28 (silu(-2) = -0.238, silu(-0.5) = -0.189). Unlike ReLU, it has no sharp corner at 0, and small negative values still pass a little signal.

The block has width 2 and feed-forward width 3, with input x = [1, -1] and weights chosen to show the three cases. Multiplying x by the gate matrix gives gate = [2, -2, 0], and by the up matrix gives up = [0, 2, 0.5]. Then each unit is processed separately:

• Unit 0: gate = 2, silu(2) = 1.762, but up = 0, so h = 0. The up projection can switch a unit off even when the gate is open.
• Unit 1: gate = -2, silu(-2) = -0.238, up = 2, so h = -0.477. A negative gate lets a small, sign-flipped amount through.
• Unit 2: gate = 0, silu(0) = 0, so h = 0 whatever up is.

down maps the three hidden values back to 2: [0, -0.477]. In the tiny model, the hidden width is 704 (feed_forward_length), 2.75 times the model width, and the block is three matrix-vector products and one element-wise step. llama.cpp's graph contains the same steps as ffn_gate, ffn_up, ffn_swiglu (one fused operation), and ffn_out.

Example 9: logits and the greedy choice

After the last layer, the vector (the hidden state, the model's internal representation of the text so far) must become a prediction. It goes through the final RMSNorm (output_norm.weight) and then the output projection (output.weight, dims [dim, vocab]). This is a matrix-vector product with one row per vocabulary token, so it produces one logit per token: the dot product of the hidden state with that token's row.

Some models save space by reusing the embedding table as the output matrix. This is called tied embeddings: when output.weight is absent, the model uses token_embd.weight instead, and llama.cpp does the same (TENSOR_DUPLICATED in load_arch_tensors).

Greedy decoding picks the token with the largest logit. Other strategies sample randomly from the probabilities; greedy always gives the same answer for the same input, which makes it ideal for testing. This program computes logits for a 6-token vocabulary, converts them to probabilities with softmax, and checks the tie rule (which token wins when two logits are equal).

import java.util.Arrays;
import java.util.Locale;

public class GreedyChoice {
    // The first index with the largest value, as llama.cpp's temperature sampler does for temp <= 0.
    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 dim = 4, vocab = 6;
        float[] output = {0.5f, 0, 0, 0,  0, 1, 0, 0,  0, 0, 1, 0,  1, 1, 0, 0,  0, 0, 0, 1,  -1, 0, 0, 0};   // dims [4, 6]
        float[] h = {0.8f, 0.6f, -0.2f, 0.1f};                       // the normalized final hidden state
        float[] logits = new float[vocab];
        for (int t = 0; t < vocab; t++) for (int i = 0; i < dim; i++) logits[t] += output[t * dim + i] * h[i];
        double max = Arrays.stream(new double[] {logits[0], logits[1], logits[2], logits[3], logits[4], logits[5]}).max().getAsDouble();
        double[] p = new double[vocab];
        double sum = 0;
        for (int t = 0; t < vocab; t++) sum += p[t] = Math.exp(logits[t] - max);
        for (int t = 0; t < vocab; t++)
            System.out.printf(Locale.ROOT, "token %d: logit %6.3f, probability %.3f%n", t, logits[t], p[t] / sum);
        System.out.println("greedy choice: token " + argmax(logits));
        System.out.println("tie [0.5, 2.0, 2.0, 1.0] -> token " + argmax(new float[] {0.5f, 2f, 2f, 1f}));
    }
}

token 0: logit  0.400, probability 0.153
token 1: logit  0.600, probability 0.187
token 2: logit -0.200, probability 0.084
token 3: logit  1.400, probability 0.416
token 4: logit  0.100, probability 0.113
token 5: logit -0.800, probability 0.046
greedy choice: token 3
tie [0.5, 2.0, 2.0, 1.0] -> token 1

Each logit is a dot product of a vocabulary row with the hidden state [0.8, 0.6, -0.2, 0.1]. Token 3's row is [1, 1, 0, 0], so its logit is 0.8 + 0.6 = 1.4; token 5's row is [-1, 0, 0, 0], so its logit is -0.8. The softmax turns the logits into probabilities that add to 1: token 3 gets 0.416, the largest share.

Greedy choice does not need the softmax: softmax keeps the order of the values, so the largest logit is also the largest probability, and argmax(logits) (the index of the largest value) picks token 3 directly. That saves an exp over 151,936 values per token in a real model.

The tie line follows llama.cpp. With --temp 0, the temperature sampler (llama_sampler_temp_impl in src/llama-sampler.cpp) keeps the first index whose logit is strictly greater than all before it and sets the rest to minus infinity. For [0.5, 2.0, 2.0, 1.0] that is index 1, and argmax with a strict > gives the same. With random weights and a 151,936-token vocabulary, exact ties are rare, but an implementation that picked the last maximum would diverge from llama.cpp at the first tie.

Example 10: an F16 KV cache

The KV cache grows by one entry per token, so for long texts it becomes large. One way to shrink it is to store each number with fewer bits. F16 uses 16 bits instead of the 32 of F32, which halves the cache, at the price of rounding each stored number a little.

llama.cpp stores the KV cache in F16 by default (-ctk, --cache-type-k and -ctv, --cache-type-v, default f16 in common/common.h), so every cached key and value is rounded to 16 bits. The Java engine has the same option: Llama's f16Cache flag rounds each value through Float.floatToFloat16 before storing it. This program measures what the rounding changes in one attention head of dimension 32 over 4, 64, and 1,024 random positions, and prints the cache size per token of the tiny model. It runs attention twice per case, once with exact values and once with rounded values, and prints the largest difference between the two outputs.

import java.util.Locale;
import java.util.Random;

public class F16Cache {
    static float f16(float v) { return Float.float16ToFloat(Float.floatToFloat16(v)); }

    // One attention head over n cached positions; round = true stores K and V as F16 first.
    static float[] attend(float[] q, float[][] k, float[][] v, boolean round) {
        int n = k.length, hd = q.length;
        float[] att = new float[n], out = new float[hd];
        float max = Float.NEGATIVE_INFINITY, sum = 0;
        for (int t = 0; t < n; t++) {
            float s = 0;
            for (int i = 0; i < hd; i++) s += q[i] * (round ? f16(k[t][i]) : k[t][i]);
            att[t] = s / (float) Math.sqrt(hd);
            max = Math.max(max, att[t]);
        }
        for (int t = 0; t < n; t++) { att[t] = (float) Math.exp(att[t] - max); sum += att[t]; }
        for (int t = 0; t < n; t++) for (int i = 0; i < hd; i++) out[i] += att[t] / sum * (round ? f16(v[t][i]) : v[t][i]);
        return out;
    }

    public static void main(String[] args) {
        Random random = new Random(11);
        int hd = 32;
        for (int n : new int[] {4, 64, 1024}) {
            float[] q = new float[hd];
            float[][] k = new float[n][hd], v = new float[n][hd];
            for (int i = 0; i < hd; i++) q[i] = (float) random.nextGaussian();
            for (int t = 0; t < n; t++) for (int i = 0; i < hd; i++) { k[t][i] = (float) random.nextGaussian(); v[t][i] = (float) random.nextGaussian(); }
            float[] exact = attend(q, k, v, false), half = attend(q, k, v, true);
            double maxDiff = 0, maxAbs = 0;
            for (int i = 0; i < hd; i++) { maxDiff = Math.max(maxDiff, Math.abs(exact[i] - half[i])); maxAbs = Math.max(maxAbs, Math.abs(exact[i])); }
            System.out.printf(Locale.ROOT, "%5d positions: largest output %.4f, largest F32-F16 difference %.6f%n", n, maxAbs, maxDiff);
        }
        long perToken = 2L * 4 * 2 * 32;                            // K and V, 4 layers, 2 KV heads, head dim 32
        System.out.printf("tiny model KV cache per token: %d values, %d bytes in F32, %d bytes in F16%n", perToken, perToken * 4, perToken * 2);
    }
}

    4 positions: largest output 1.6992, largest F32-F16 difference 0.000661
   64 positions: largest output 0.3232, largest F32-F16 difference 0.000111
 1024 positions: largest output 0.0922, largest F32-F16 difference 0.000033
tiny model KV cache per token: 512 values, 2048 bytes in F32, 1024 bytes in F16

Each line compares the attention output with exact F32 keys and values against the same computation with both rounded to F16. The largest difference is 0.000661 with 4 positions, 0.000111 with 64, and 0.000033 with 1,024. FP16 keeps about three decimal digits, so each stored value moves by up to about 0.05% of its size, and the output, a weighted average of many rounded values, moves by less as the number of positions grows, because independent rounding errors partly cancel in the average.

The last line is the benefit: the tiny model stores 512 numbers per token (K and V, 4 layers, 2 KV heads, 32 dimensions), 2,048 bytes in F32 and 1,024 in F16. At the full 4,096-token context, that is 8 MiB against 4 MiB. For an 8B model, the same factor of two decides how long a context fits in memory (Memory Math).

The difference matters for Example 11. When llama.cpp and the Java engine use different cache types, their tensors differ in the fourth digit after the first attention layer, which is larger than any difference between two F32 implementations. A tensor-by-tensor comparison must therefore run both sides with the same cache type.

Example 11: layer-by-layer checks against llama-eval-callback

To confirm that the Java forward pass is correct, checking only the final text is weak: a bug in one operation still produces text, only wrong text, and gives no hint where the bug is. A better method is to compare the intermediate results of each operation against a trusted implementation. Here the trusted implementation is llama.cpp.

llama-eval-callback (source in examples/eval-callback/ of the pinned llama.cpp; built with cmake --build build --target llama-eval-callback) runs a prompt through a model and prints every tensor the graph computes: its name, operation, shape, a few values, and the sum of all its values. Comparing the sum of a tensor is a cheap fingerprint. If two tensors have the same sum to six digits, they are almost certainly the same.

The Java program below runs the same prompt through the Llama class of Project 18 (Example 13), with a trace hook (a callback the engine calls with each intermediate vector, named with llama.cpp's tensor names), and compares the sums. It needs GgufModel.java, GgufTokenizer.java, and Llama.java in the same directory and the model file, so it is java,ignore (the build checker skips it). The program has three steps: parse the sums from the llama.cpp log, run the Java engine and sum each traced tensor, then print both sums and their relative difference (|java - llama.cpp| / |llama.cpp|).

import java.nio.file.Files;
import java.nio.file.Path;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Locale;
import java.util.Map;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class CompareLayers {
    record Ref(double sum, int rows) {}

    public static void main(String[] args) throws Exception {
        // 1. The sums llama-eval-callback printed; for a repeated name the last entry wins (ROPE for Qcur and Kcur).
        Pattern header = Pattern.compile("common_debug_cb_eval:\\s+(\\S+) = \\(\\w+\\)\\s+\\w+\\(.*\\) = \\{(\\d+), (\\d+),");
        Map<String, Ref> ref = new HashMap<>();
        String name = null;
        int rows = 0;
        for (String line : Files.readAllLines(Path.of(args[1]))) {
            Matcher m = header.matcher(line);
            if (line.startsWith("common_debug_cb_eval:")) {
                name = m.find() ? m.group(1) : null;                         // views and permutes have longer names
                rows = name != null ? Integer.parseInt(m.group(3)) : 0;
            } else if (name != null && line.trim().startsWith("sum = ")) ref.put(name, new Ref(Double.parseDouble(line.trim().substring(6)), rows));
        }
        // 2. The Java engine on the same prompt, with each traced tensor summed per position.
        try (GgufModel gguf = GgufModel.open(Path.of(args[0]))) {
            GgufTokenizer tok = new GgufTokenizer(gguf.metadata());
            Llama model = new Llama(gguf.metadata(), n -> gguf.tensors().containsKey(n) ? gguf.floats(n) : null, 4096, args[3].equals("f16"));
            List<Integer> ids = tok.encode(args[2], true, true);
            Map<String, double[]> sums = new LinkedHashMap<>();
            int[] pos = {0};
            model.trace = (n, x) -> { double s = 0; for (float v : x) s += v; sums.computeIfAbsent(n, k -> new double[ids.size()])[pos[0]] = s; };
            for (; pos[0] < ids.size(); pos[0]++) model.forward(ids.get(pos[0]), pos[0]);
            // 3. A reference tensor with one row holds only the last position (llama.cpp skips unused rows in the last layer).
            double worst = 0;
            for (var e : sums.entrySet()) {
                Ref r = ref.get(e.getKey());
                if (r == null) continue;
                double[] perPos = e.getValue();
                double java = r.rows() == 1 ? perPos[perPos.length - 1] : Arrays.stream(perPos).sum();
                double rel = Math.abs(java - r.sum()) / Math.max(1e-3, Math.abs(r.sum()));
                worst = Math.max(worst, rel);
                System.out.printf(Locale.ROOT, "%-14s llama.cpp %12.6f  java %12.6f  rel.diff %.1e%n", e.getKey(), r.sum(), java, rel);
            }
            System.out.printf(Locale.ROOT, "largest relative difference %.1e%n", worst);
        }
    }
}

A few lines of the raw llama-eval-callback output show what the parser reads (pinned build, prompt Hello world; the value blocks of the three Qcur-0 entries are left out):

common_debug_cb_eval:                     embd = (f32)   GET_ROWS(token_embd.weight{256, 151936, 1, 1}, inp_tokens{2, 1, 1, 1}}) = {256, 2, 1, 1}
    [
        [
            [      0.0296,      -0.0015,      -0.0031,    ...,      -0.0259,       0.0172,      -0.0251  ],
            [      0.0059,      -0.0195,      -0.0109,    ...,       0.0022,      -0.0078,       0.0169  ],
        ],
    ]
    sum = -0.252000
common_debug_cb_eval:                   Qcur-0 = (f32)    MUL_MAT(blk.0.attn_q.weight{256, 256, 1, 1}, attn_norm-0{256, 2, 1, 1}}) = {256, 2, 1, 1}
common_debug_cb_eval:                   Qcur-0 = (f32)    RESHAPE(Qcur-0{256, 2, 1, 1}, }) = {32, 8, 2, 1}
common_debug_cb_eval:                   Qcur-0 = (f32)       ROPE(Qcur-0{32, 8, 2, 1}, inp_pos{2, 1, 1, 1}}) = {32, 8, 2, 1}

Each header line has the tensor's name, its data type (f32), the operation that produced it, the operation's inputs with their shapes, and the shape of the result. The shape lists up to four dimensions, innermost first, as in GGUF. The tensor embd has shape {256, 2, 1, 1}: two positions (Hello and world) of 256 values. Its first row is the embedding row of token 9707 that lesson 49 read from the file. A name can appear several times: Qcur-0 is first the matrix product, then a reshape into 8 heads of 32, then the RoPE result. The parser keeps the last entry for each name and ignores entries whose name continues with (view) or (permuted).

Two comparisons follow. For a fair test, both sides must use the same number formats. Flash attention is llama.cpp's fused attention kernel, which computes the scores, softmax, and weighted values in one pass and uses lower precision internally. The first comparison turns it off (-fa off) and uses an F32 cache (-ctk f32 -ctv f32), so both sides compute in F32:

llama-eval-callback -m tiny-random-qwen2vocab-f32.gguf -p "Hello world" -n 1 --temp 0 -fa off -ctk f32 -ctv f32 > evalcb-f32.log 2>&1
java CompareLayers.java tiny-random-qwen2vocab-f32.gguf evalcb-f32.log "Hello world" f32

Output (pinned llama.cpp build and the Java engine; both use F32 everywhere):

embd           llama.cpp    -0.252000  java    -0.252000  rel.diff 1.3e-06
attn_norm-0    llama.cpp   -12.788591  java   -12.788590  rel.diff 9.1e-08
Qcur-0         llama.cpp    -4.197163  java    -4.197165  rel.diff 3.9e-07
Kcur-0         llama.cpp    -2.860843  java    -2.860843  rel.diff 2.6e-08
Vcur-0         llama.cpp    -4.147503  java    -4.147501  rel.diff 4.1e-07
kqv_out-0      llama.cpp   -21.478447  java   -21.478440  rel.diff 3.0e-07
attn_out-0     llama.cpp    -1.653525  java    -1.653522  rel.diff 1.7e-06
ffn_inp-0      llama.cpp    -1.905522  java    -1.905522  rel.diff 6.4e-08
ffn_norm-0     llama.cpp   -20.006741  java   -20.006731  rel.diff 5.0e-07
ffn_gate-0     llama.cpp     7.037997  java     7.037991  rel.diff 8.2e-07
ffn_up-0       llama.cpp     8.566229  java     8.566226  rel.diff 3.3e-07
ffn_swiglu-0   llama.cpp    -1.051283  java    -1.051282  rel.diff 1.3e-06
ffn_out-0      llama.cpp     0.640622  java     0.640621  rel.diff 2.0e-06
l_out-0        llama.cpp    -1.264902  java    -1.264901  rel.diff 6.5e-07
attn_norm-1    llama.cpp   -12.841502  java   -12.841502  rel.diff 9.1e-09
Qcur-1         llama.cpp   -11.108747  java   -11.108742  rel.diff 4.8e-07
Kcur-1         llama.cpp    -6.278861  java    -6.278859  rel.diff 3.1e-07
Vcur-1         llama.cpp     4.046376  java     4.046374  rel.diff 5.2e-07
kqv_out-1      llama.cpp    20.619619  java    20.619611  rel.diff 3.8e-07
attn_out-1     llama.cpp     0.573859  java     0.573858  rel.diff 2.5e-06
ffn_inp-1      llama.cpp    -0.691043  java    -0.691044  rel.diff 7.5e-07
ffn_norm-1     llama.cpp    -5.516462  java    -5.516457  rel.diff 9.1e-07
ffn_gate-1     llama.cpp   -19.379686  java   -19.379678  rel.diff 4.1e-07
ffn_up-1       llama.cpp     2.393784  java     2.393798  rel.diff 5.7e-06
ffn_swiglu-1   llama.cpp     1.289918  java     1.289921  rel.diff 2.6e-06
ffn_out-1      llama.cpp    -1.731370  java    -1.731370  rel.diff 1.9e-08
l_out-1        llama.cpp    -2.422411  java    -2.422414  rel.diff 1.1e-06
attn_norm-2    llama.cpp   -17.675365  java   -17.675366  rel.diff 5.6e-08
Qcur-2         llama.cpp     6.191720  java     6.191720  rel.diff 4.3e-08
Kcur-2         llama.cpp     6.572375  java     6.572371  rel.diff 6.3e-07
Vcur-2         llama.cpp    -0.041756  java    -0.041758  rel.diff 5.1e-05
kqv_out-2      llama.cpp     5.088159  java     5.088148  rel.diff 2.2e-06
attn_out-2     llama.cpp     0.331252  java     0.331254  rel.diff 4.7e-06
ffn_inp-2      llama.cpp    -2.091162  java    -2.091160  rel.diff 9.0e-07
ffn_norm-2     llama.cpp   -12.571352  java   -12.571338  rel.diff 1.1e-06
ffn_gate-2     llama.cpp    -1.868727  java    -1.868707  rel.diff 1.1e-05
ffn_up-2       llama.cpp    10.653346  java    10.653338  rel.diff 7.8e-07
ffn_swiglu-2   llama.cpp    -2.989168  java    -2.989170  rel.diff 7.5e-07
ffn_out-2      llama.cpp    -1.071301  java    -1.071299  rel.diff 1.4e-06
l_out-2        llama.cpp    -3.162463  java    -3.162460  rel.diff 1.1e-06
attn_norm-3    llama.cpp   -18.524584  java   -18.524572  rel.diff 6.3e-07
Qcur-3         llama.cpp   -12.913055  java   -12.913060  rel.diff 3.9e-07
Kcur-3         llama.cpp    -2.514845  java    -2.514841  rel.diff 1.5e-06
Vcur-3         llama.cpp    -2.280368  java    -2.280368  rel.diff 2.2e-07
kqv_out-3      llama.cpp    -7.981837  java    -7.981837  rel.diff 5.9e-08
attn_out-3     llama.cpp    -3.408318  java    -3.408320  rel.diff 7.2e-07
ffn_inp-3      llama.cpp    -3.208521  java    -3.208519  rel.diff 7.1e-07
ffn_norm-3     llama.cpp   -17.305937  java   -17.305929  rel.diff 4.8e-07
ffn_gate-3     llama.cpp    -6.535916  java    -6.535914  rel.diff 2.8e-07
ffn_up-3       llama.cpp    -7.370199  java    -7.370196  rel.diff 4.5e-07
ffn_swiglu-3   llama.cpp     1.060880  java     1.060879  rel.diff 6.7e-07
ffn_out-3      llama.cpp    -0.096544  java    -0.096544  rel.diff 4.0e-06
l_out-3        llama.cpp    -3.305064  java    -3.305063  rel.diff 2.8e-07
result_norm    llama.cpp   -17.803524  java   -17.803516  rel.diff 4.4e-07
result_output  llama.cpp   -64.309860  java   -64.310586  rel.diff 1.1e-05
largest relative difference 5.1e-05

The second comparison shows what happens when the settings differ in the way Example 10 described. It uses llama.cpp's defaults (flash attention auto, which is on for this CPU model, and an F16 cache) and the Java engine's F16 cache:

llama-eval-callback -m tiny-random-qwen2vocab-f32.gguf -p "Hello world" -n 1 --temp 0 > evalcb-default.log 2>&1
java CompareLayers.java tiny-random-qwen2vocab-f32.gguf evalcb-default.log "Hello world" f16

Selected lines of the output:

embd           llama.cpp    -0.252000  java    -0.252000  rel.diff 1.3e-06
Vcur-0         llama.cpp    -4.147503  java    -4.147501  rel.diff 4.1e-07
kqv_out-0      llama.cpp   -21.473192  java   -21.475113  rel.diff 8.9e-05
attn_out-0     llama.cpp    -1.652917  java    -1.653175  rel.diff 1.6e-04
Vcur-2         llama.cpp    -0.039596  java    -0.040896  rel.diff 3.3e-02
kqv_out-3      llama.cpp    -7.970373  java    -7.973954  rel.diff 4.5e-04
l_out-3        llama.cpp    -3.303302  java    -3.304381  rel.diff 3.3e-04
result_norm    llama.cpp   -17.794407  java   -17.799951  rel.diff 3.1e-04
result_output  llama.cpp   -64.209961  java   -64.292579  rel.diff 1.3e-03
largest relative difference 3.3e-02

Read the F32 output as a table with one row per tensor: the sums from both sides, then the relative difference. A relative difference of 1e-6 means the sums agree in the first six digits. In the F32 comparison, every tensor agrees to about six significant digits. Over the 55 compared tensors, the median relative difference is 6.3e-07, all but one are below 1.2e-05, and the largest, 5.1e-05, is Vcur-2, whose sum is small (-0.0418) because positive and negative values cancel; its absolute difference is 0.000002. The lines are in the order the engine computes them, so they can be read as a walk through one block:

• embd and attn_norm-0 are the embedding rows and the first RMSNorm. Their sums match to the printed digits.
• Qcur-0, Kcur-0, and Vcur-0 are the projections after RoPE (for Q and K). The match of Qcur and Kcur confirms the normal RoPE layout, the base 500,000, and the position numbering. A run with the layout switched to NEOX gives Qcur-0 a sum of -5.414 instead of -4.197 and Kcur-0 -2.765 instead of -2.861.
• kqv_out-0 is the attention output of all heads, attn_out-0 the output projection, and ffn_inp-0 the first residual sum.
• ffn_norm-0 to ffn_out-0 are the SwiGLU steps, and l_out-0 is the layer output.
• The last layer's lines from ffn_inp-3 onward compare only the second position. llama.cpp computes the rest of the last layer only for the positions whose logits are needed (inp_out_ids and the GET_ROWS calls in src/models/llama.cpp), so its tensors there have one row. The Java engine computes every position, and the parser compares the last one when the reference tensor has a single row.
• result_norm and result_output are the final norm and the logits of the last position, summed over all 151,936 tokens.

The remaining differences come from summation order. ggml splits each dot product across SIMD lanes (a CPU feature that adds several numbers per instruction) and threads, and adds the partial sums in a different order than the Java loop does. Floating-point addition is not associative: (a + b) + c can differ from a + (b + c) in the last bits, because each addition rounds. So the results differ in the last bits of each value, even though both programs are correct. The sums prove more than the final tokens would: had any operation been missing, transposed, or applied with the wrong constant, the first affected tensor would be off in its first or second digit.

With llama.cpp's defaults, the tensors before the first attention (embd to Vcur-0) match as closely as before. From kqv_out-0 on, the relative differences range from 3e-06 to 3.3e-02, with a median of 1.5e-04. The rounding of K and V to F16 is the same on both sides, but llama.cpp's CPU flash attention (ggml_compute_forward_flash_attn_ext_f16 in ggml/src/ggml-cpu/ops.cpp) also converts each query to F16 for the dot products with the F16 keys, and accumulates the weighted values in F16 when the value cache is F16. Those two extra roundings change kqv_out-0 in the fifth digit, and the change propagates through the residual stream into all later layers. The largest relative difference, 3.3e-02 for Vcur-2, is again the small-sum tensor (absolute difference 0.0013). The next example shows that these differences do not change the generated tokens of this run.

Example 12: full greedy generation on the tiny model, compared with llama-completion

The last check runs a full generation, the way a user would use the model. Generation is a loop. The model reads the whole prompt (called prefill, here one token at a time), takes the greedy token from the logits of the last position, appends that token to the text, and feeds it back in to get the next one. It repeats until an end-of-generation token appears or a length limit is reached. Each step adds one entry to the KV cache.

The program below tokenizes a prompt with GgufTokenizer, runs the prompt through the Llama class one token at a time, then repeatedly takes the greedy token, stops at an end-of-generation token or after a maximum count, and feeds the token back. It prints the prompt ids, the generated ids, and the text decoded with the streaming decoder. It uses an F16 cache, like llama.cpp's default. It was run with java Generate.java tiny-random-qwen2vocab-f32.gguf "Hello world" 32.

import java.nio.file.Path;
import java.util.ArrayList;
import java.util.List;

public class Generate {
    public static void main(String[] args) throws Exception {
        int maxTokens = Integer.parseInt(args[2]);
        try (GgufModel gguf = GgufModel.open(Path.of(args[0]))) {
            GgufTokenizer tok = new GgufTokenizer(gguf.metadata());
            Llama model = new Llama(gguf.metadata(), n -> gguf.tensors().containsKey(n) ? gguf.floats(n) : null, 4096, true);
            List<Integer> prompt = tok.encode(args[1], true, true);
            System.out.println("prompt ids " + prompt);
            float[] logits = null;
            for (int p = 0; p < prompt.size(); p++) logits = model.forward(prompt.get(p), p);   // prefill, one token at a time
            List<Integer> out = new ArrayList<>();
            GgufTokenizer.Utf8Stream text = new GgufTokenizer.Utf8Stream();
            StringBuilder sb = new StringBuilder();
            for (int pos = prompt.size(); out.size() < maxTokens; pos++) {
                int next = Llama.argmax(logits);                                               // greedy
                if (tok.isEog(next)) break;
                out.add(next);
                sb.append(text.push(tok.tokenBytes(next, false)));
                logits = model.forward(next, pos);
            }
            System.out.println("generated ids " + out);
            System.out.println("text: " + args[1] + sb);
        }
    }
}

Output of the Java engine, three lines: the prompt's token ids, the generated token ids, and the decoded text.

prompt ids [9707, 1879]
generated ids [146539, 2373, 126339, 56171, 96902, 56171, 8359, 8359, 146539, 2373, 70670, 8359, 118552, 60522, 118552, 60522, 146539, 146539, 146539, 146539, 146539, 96902, 146539, 96902, 146539, 96902, 146539, 96902, 146539, 50227, 146539, 50227]
text: Hello world쾰rame ngo	java collage	java Russia Russia쾰rame Interesting Russia程序员ureka程序员ureka쾰쾰쾰쾰쾰 collage쾰 collage쾰 collage쾰 collage쾰 Yorkshire쾰 Yorkshire

The same prompt through llama.cpp's one-shot completion tool with greedy sampling (--temp 0), default cache and attention settings, and conversation mode off (the model has a chat template, so llama-completion would otherwise start an interactive chat):

llama-completion -m tiny-random-qwen2vocab-f32.gguf -p "Hello world" -n 32 --temp 0 -no-cnv

Standard output of llama-completion (pinned build; the prompt is echoed first):

Hello world쾰rame ngo	java collage	java Russia Russia쾰rame Interesting Russia程序员ureka程序员ureka쾰쾰쾰쾰쾰 collage쾰 collage쾰 collage쾰 collage쾰 Yorkshire쾰 Yorkshire

llama-completion prints text, not ids. To compare ids, the same request went to llama-server started on the same model (llama-server -m tiny-random-qwen2vocab-f32.gguf --host 127.0.0.1 --port 18611), with return_tokens set so the response includes the generated ids:

curl -s 127.0.0.1:18611/completion -d '{"prompt": "Hello world", "n_predict": 32, "temperature": 0, "return_tokens": true}'

Fields of the response (pinned build):

{"tokens": [146539, 2373, 126339, 56171, 96902, 56171, 8359, 8359, 146539, 2373, 70670, 8359, 118552, 60522, 118552, 60522, 146539, 146539, 146539, 146539, 146539, 96902, 146539, 96902, 146539, 96902, 146539, 96902, 146539, 50227, 146539, 50227], "tokens_evaluated": 2, "stop_type": "limit"}

The prompt ids [9707, 1879] are the same two tokens that llama-eval-callback and llama-completion --verbose-prompt list for this prompt, with no BOS, as the Qwen2 vocabulary rules require (lesson 50, Example 7). The server's tokens_evaluated: 2 confirms that it also processed two prompt tokens.

Now compare the three runs. The 32 generated ids are identical to the 32 ids from llama-server, position by position, and the Java text is identical to llama-completion's output: 쾰rame ngo, a tab, java collage, and so on to 쾰 Yorkshire. The text is nonsense because the weights are random; the model repeats tokens such as 146539 (쾰) and 96902 ( collage) because random weights have no notion of language, and greedy decoding follows the same high logits again and again. stop_type: limit means the server stopped at n_predict, not at an end-of-generation token, and the Java loop also ran to its limit of 32.

The match was also checked with two other prompts. For def fib(n): (4 prompt tokens) and for a Qwen-style chat prompt with <|im_start|> and <|im_end|> control tokens (17 prompt tokens), the first 32 greedy ids from the Java engine equal those from llama-server, both with the Java engine's F16 cache and with its F32 cache. The second check is the stronger one: although the default llama.cpp run differs from the Java engine in the fourth or fifth digit of its logits (Example 11), the largest logit is never close enough to the runner-up for those differences to change a greedy choice in these runs.

That is also the limit of this test. Two correct engines can diverge when two logits are within the numerical noise of each other: the first differing token changes the context, and every later token can differ. A comparison of greedy ids is therefore reported for specific prompts and lengths, and the per-tensor sums of Example 11 are the check that does not depend on luck.

The Java engine takes several seconds for this run on the reference machine, most of it in the 151,936-row output projection and in dequantizing the two vocabulary matrices at load time. llama.cpp is much faster on the same model. Quantized Kernels, the Vector API, and Threads speeds up the kernels with the Vector API, quantized dot products, and threads, and Capstone: A Chat Engine in Pure Java measures the remaining gap to llama.cpp.

Example 13: Project 18, the Llama class

Project 18 collects Examples 1 to 10 into one class, Llama, that is a complete inference engine for Llama-style models. Read it from top to bottom as the examples in order. Its parts are:

• Config.from reads the hyperparameters (Example 1).
• The constructor takes the metadata map and a function from tensor name to float[]. The function returns null for a missing tensor, so output.weight can fall back to token_embd.weight (tied embeddings) and the optional Qwen2 biases attn_q.bias, attn_k.bias, attn_v.bias can be absent. A bias is a vector added to the result of a matrix product; Qwen2 has one for Q, K, and V, and Llama has none.
• matvec, rmsnorm, rope, and silu are the operations of Examples 4 to 8.
• forward(token, pos) runs one token through all layers, fills the KV cache at position pos, and returns the logits.
• The trace field receives each intermediate vector under llama.cpp's tensor name, which is how Example 11 compares sums.
• argmax is the greedy choice of Example 9.

The main method is a self-test that needs no model file. It builds a seeded synthetic model of width 16 with 2 layers and generates 10 tokens greedily with the KV cache. At every step it checks the result against a second model instance that recomputes the whole sequence from position 0, without reusing anything. It runs the test once for the llama and once for the qwen2 architecture with the same weights. The output has a configuration line, a result line, and a work-count line for each architecture.

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Random;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;

public final class Llama {
    public record Config(String arch, int dim, int layers, int heads, int kvHeads, int headDim, int ffn, int vocab,
                         int contextLength, float eps, float ropeBase, int ropeDims, boolean neox) {
        static Config from(Map<String, Object> md) {
            String a = (String) md.get("general.architecture");
            Function<String, Number> num = k -> (Number) md.get(a + "." + k);
            int dim = num.apply("embedding_length").intValue(), heads = num.apply("attention.head_count").intValue();
            int kvHeads = md.containsKey(a + ".attention.head_count_kv") ? num.apply("attention.head_count_kv").intValue() : heads;
            int vocab = md.containsKey(a + ".vocab_size") ? num.apply("vocab_size").intValue() : ((List<?>) md.get("tokenizer.ggml.tokens")).size();
            float base = md.containsKey(a + ".rope.freq_base") ? num.apply("rope.freq_base").floatValue() : 10000f;
            int ropeDims = md.containsKey(a + ".rope.dimension_count") ? num.apply("rope.dimension_count").intValue() : dim / heads;
            boolean neox = switch (a) {                    // llama_model_rope_type in src/llama-model.cpp
                case "llama" -> false;                     // LLAMA_ROPE_TYPE_NORM: pairs (2i, 2i+1)
                case "qwen2" -> true;                      // LLAMA_ROPE_TYPE_NEOX: pairs (i, i + n_rot/2)
                default -> throw new IllegalArgumentException("architecture not supported here: " + a);
            };
            return new Config(a, dim, num.apply("block_count").intValue(), heads, kvHeads, dim / heads,
                    num.apply("feed_forward_length").intValue(), vocab, num.apply("context_length").intValue(),
                    num.apply("attention.layer_norm_rms_epsilon").floatValue(), base, ropeDims, neox);
        }
    }

    record Layer(float[] attnNorm, float[] wq, float[] wk, float[] wv, float[] bq, float[] bk, float[] bv, float[] wo,
                 float[] ffnNorm, float[] gate, float[] up, float[] down) {}

    final Config c;
    final float[] tokenEmbd, outputNorm, output;
    final Layer[] layers;
    final float[][] keyCache, valueCache;          // per layer: contextLength * kvHeads * headDim
    final boolean f16Cache;
    BiConsumer<String, float[]> trace = (name, x) -> {};

    public Llama(Map<String, Object> metadata, Function<String, float[]> weights, int contextLength, boolean f16Cache) {
        Config cfg = Config.from(metadata);
        c = new Config(cfg.arch(), cfg.dim(), cfg.layers(), cfg.heads(), cfg.kvHeads(), cfg.headDim(), cfg.ffn(), cfg.vocab(),
                Math.min(contextLength, cfg.contextLength()), cfg.eps(), cfg.ropeBase(), cfg.ropeDims(), cfg.neox());
        this.f16Cache = f16Cache;
        tokenEmbd = weights.apply("token_embd.weight");
        outputNorm = weights.apply("output_norm.weight");
        float[] out = weights.apply("output.weight");
        output = out != null ? out : tokenEmbd;                                  // tied embeddings
        layers = new Layer[c.layers()];
        for (int l = 0; l < c.layers(); l++) {
            String p = "blk." + l + ".";
            layers[l] = new Layer(weights.apply(p + "attn_norm.weight"), weights.apply(p + "attn_q.weight"), weights.apply(p + "attn_k.weight"),
                    weights.apply(p + "attn_v.weight"), weights.apply(p + "attn_q.bias"), weights.apply(p + "attn_k.bias"),
                    weights.apply(p + "attn_v.bias"), weights.apply(p + "attn_output.weight"), weights.apply(p + "ffn_norm.weight"),
                    weights.apply(p + "ffn_gate.weight"), weights.apply(p + "ffn_up.weight"), weights.apply(p + "ffn_down.weight"));
        }
        int kvDim = c.kvHeads() * c.headDim();
        keyCache = new float[c.layers()][c.contextLength() * kvDim];
        valueCache = new float[c.layers()][c.contextLength() * kvDim];
    }

    // out[r] = sum over i of w[r * cols + i] * x[i]; a GGUF tensor with dims [cols, rows] stores row r at r * cols.
    static void matvec(float[] out, float[] w, float[] x, int rows, int cols) {
        for (int r = 0; r < rows; r++) {
            float s = 0;
            int base = r * cols;
            for (int i = 0; i < cols; i++) s += w[base + i] * x[i];
            out[r] = s;
        }
    }

    static void rmsnorm(float[] out, float[] x, float[] weight, float eps) {
        double ss = 0;
        for (float v : x) ss += (double) v * v;
        float scale = (float) (1.0 / Math.sqrt(ss / x.length + eps));
        for (int i = 0; i < x.length; i++) out[i] = x[i] * scale * weight[i];
    }

    // Rotates each head of vec in place. Normal: pairs (2i, 2i+1). NEOX: pairs (i, i + ropeDims/2).
    static void rope(float[] vec, int nHeads, int headDim, int ropeDims, int pos, float base, boolean neox) {
        float thetaScale = (float) Math.pow(base, -2.0 / ropeDims);
        for (int h = 0; h < nHeads; h++) {
            float theta = pos;
            for (int i = 0; i < ropeDims / 2; i++) {
                float cos = (float) Math.cos(theta), sin = (float) Math.sin(theta);
                int a = h * headDim + (neox ? i : 2 * i), b = neox ? a + ropeDims / 2 : a + 1;
                float x0 = vec[a], x1 = vec[b];
                vec[a] = x0 * cos - x1 * sin;
                vec[b] = x0 * sin + x1 * cos;
                theta *= thetaScale;
            }
        }
    }

    static void addBias(float[] x, float[] bias) {
        if (bias != null) for (int i = 0; i < x.length; i++) x[i] += bias[i];
    }

    static float silu(float x) { return x / (1f + (float) Math.exp(-x)); }

    float store(float v) { return f16Cache ? Float.float16ToFloat(Float.floatToFloat16(v)) : v; }

    // One token at position pos: fills the KV cache at pos and returns the logits.
    public float[] forward(int token, int pos) {
        int dim = c.dim(), hd = c.headDim(), kvDim = c.kvHeads() * hd, group = c.heads() / c.kvHeads();
        float[] x = new float[dim], xb = new float[dim], q = new float[dim], k = new float[kvDim], v = new float[kvDim];
        float[] att = new float[pos + 1], hb = new float[c.ffn()], hb2 = new float[c.ffn()];
        System.arraycopy(tokenEmbd, token * dim, x, 0, dim);
        trace.accept("embd", x);
        for (int l = 0; l < c.layers(); l++) {
            Layer w = layers[l];
            rmsnorm(xb, x, w.attnNorm(), c.eps());
            trace.accept("attn_norm-" + l, xb);
            matvec(q, w.wq(), xb, dim, dim);
            matvec(k, w.wk(), xb, kvDim, dim);
            matvec(v, w.wv(), xb, kvDim, dim);
            addBias(q, w.bq());                                                 // Qwen2 has Q, K, V biases; Llama has none
            addBias(k, w.bk());
            addBias(v, w.bv());
            rope(q, c.heads(), hd, c.ropeDims(), pos, c.ropeBase(), c.neox());
            rope(k, c.kvHeads(), hd, c.ropeDims(), pos, c.ropeBase(), c.neox());
            trace.accept("Qcur-" + l, q);
            trace.accept("Kcur-" + l, k);
            trace.accept("Vcur-" + l, v);
            for (int i = 0; i < kvDim; i++) { keyCache[l][pos * kvDim + i] = store(k[i]); valueCache[l][pos * kvDim + i] = store(v[i]); }
            float[] heads = new float[dim];
            for (int h = 0; h < c.heads(); h++) {
                int kvOff = (h / group) * hd;                                    // grouped-query attention
                float max = Float.NEGATIVE_INFINITY;
                for (int t = 0; t <= pos; t++) {
                    float s = 0;
                    for (int i = 0; i < hd; i++) s += q[h * hd + i] * keyCache[l][t * kvDim + kvOff + i];
                    att[t] = s / (float) Math.sqrt(hd);
                    max = Math.max(max, att[t]);
                }
                float sum = 0;
                for (int t = 0; t <= pos; t++) { att[t] = (float) Math.exp(att[t] - max); sum += att[t]; }
                for (int t = 0; t <= pos; t++) {
                    float a = att[t] / sum;
                    for (int i = 0; i < hd; i++) heads[h * hd + i] += a * valueCache[l][t * kvDim + kvOff + i];
                }
            }
            trace.accept("kqv_out-" + l, heads);
            matvec(xb, w.wo(), heads, dim, dim);
            trace.accept("attn_out-" + l, xb);
            for (int i = 0; i < dim; i++) x[i] += xb[i];
            trace.accept("ffn_inp-" + l, x);
            rmsnorm(xb, x, w.ffnNorm(), c.eps());
            trace.accept("ffn_norm-" + l, xb);
            matvec(hb, w.gate(), xb, c.ffn(), dim);
            matvec(hb2, w.up(), xb, c.ffn(), dim);
            trace.accept("ffn_gate-" + l, hb);
            trace.accept("ffn_up-" + l, hb2);
            for (int i = 0; i < c.ffn(); i++) hb[i] = silu(hb[i]) * hb2[i];   // SwiGLU
            trace.accept("ffn_swiglu-" + l, hb);
            matvec(xb, w.down(), hb, dim, c.ffn());
            trace.accept("ffn_out-" + l, xb);
            for (int i = 0; i < dim; i++) x[i] += xb[i];
            trace.accept("l_out-" + l, x);
        }
        rmsnorm(xb, x, outputNorm, c.eps());
        trace.accept("result_norm", xb);
        float[] logits = new float[c.vocab()];
        matvec(logits, output, xb, c.vocab(), dim);
        trace.accept("result_output", logits);
        return logits;
    }

    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;
    }
    // Test: a seeded synthetic model. Greedy generation with the KV cache must equal recomputing every step from position 0.
    public static void main(String[] args) {
        for (String arch : List.of("llama", "qwen2")) {
            int dim = 16, nLayers = 2, ffn = 32, vocab = 24;
            Map<String, Object> md = new HashMap<>(Map.of("general.architecture", arch, arch + ".embedding_length", 16L,
                    arch + ".block_count", 2L, arch + ".feed_forward_length", 32L, arch + ".attention.head_count", 4L,
                    arch + ".attention.head_count_kv", 2L, arch + ".context_length", 64L,
                    arch + ".attention.layer_norm_rms_epsilon", 1e-5f, arch + ".rope.freq_base", 10000f));
            List<String> tokens = new ArrayList<>();
            for (int i = 0; i < vocab; i++) tokens.add("t" + i);
            md.put("tokenizer.ggml.tokens", tokens);
            Random random = new Random(42);                                    // the same weights for both architectures
            Map<String, float[]> w = new HashMap<>();
            java.util.function.BiConsumer<String, int[]> init = (name, rc) -> {
                float[] a = new float[rc[0] * rc[1]];
                for (int i = 0; i < a.length; i++) a[i] = (float) (random.nextGaussian() / Math.sqrt(rc[1]));
                w.put(name, a);
            };
            init.accept("token_embd.weight", new int[] {vocab, dim});
            init.accept("output.weight", new int[] {vocab, dim});
            float[] ones = new float[dim];
            Arrays.fill(ones, 1f);
            w.put("output_norm.weight", ones);
            for (int l = 0; l < nLayers; l++) {
                String b = "blk." + l + ".";
                w.put(b + "attn_norm.weight", ones);
                w.put(b + "ffn_norm.weight", ones);
                init.accept(b + "attn_q.weight", new int[] {dim, dim});
                init.accept(b + "attn_k.weight", new int[] {8, dim});
                init.accept(b + "attn_v.weight", new int[] {8, dim});
                init.accept(b + "attn_output.weight", new int[] {dim, dim});
                init.accept(b + "ffn_gate.weight", new int[] {ffn, dim});
                init.accept(b + "ffn_up.weight", new int[] {ffn, dim});
                init.accept(b + "ffn_down.weight", new int[] {dim, ffn});
            }
            Llama cached = new Llama(md, w::get, 64, false), uncached = new Llama(md, w::get, 64, false);
            Config c = cached.c;
            System.out.printf("%s: dim %d, %d layers, %d heads (%d KV), head %d, ffn %d, vocab %d, rope %s%n", arch, c.dim(), c.layers(),
                    c.heads(), c.kvHeads(), c.headDim(), c.ffn(), c.vocab(), c.neox() ? "NEOX" : "normal");
            List<Integer> seq = new ArrayList<>(List.of(1, 5, 9));
            int cachedCalls = 0, uncachedCalls = 0;
            boolean same = true;
            float[] logits = null;
            for (int p = 0; p < seq.size(); p++, cachedCalls++) logits = cached.forward(seq.get(p), p);   // prefill
            for (int step = 0; step < 10; step++) {
                float[] reference = null;
                for (int p = 0; p < seq.size(); p++, uncachedCalls++) reference = uncached.forward(seq.get(p), p);   // no reuse
                same &= Arrays.equals(logits, reference);
                seq.add(argmax(logits));
                logits = cached.forward(seq.get(seq.size() - 1), seq.size() - 1);
                cachedCalls++;
            }
            System.out.println("  generated " + seq.subList(3, seq.size()) + ", logits bit-identical at every step: " + same);
            System.out.println("  forward calls: " + cachedCalls + " with the cache, " + uncachedCalls + " without");
        }
    }
}

llama: dim 16, 2 layers, 4 heads (2 KV), head 4, ffn 32, vocab 24, rope normal
  generated [16, 18, 1, 3, 3, 21, 1, 1, 1, 1], logits bit-identical at every step: true
  forward calls: 13 with the cache, 75 without
qwen2: dim 16, 2 layers, 4 heads (2 KV), head 4, ffn 32, vocab 24, rope NEOX
  generated [1, 1, 9, 1, 1, 1, 1, 1, 9, 1], logits bit-identical at every step: true
  forward calls: 13 with the cache, 75 without

Each architecture prints its configuration, the generated tokens, the equality check, and the work count.

The configuration lines are identical except for the RoPE layout: normal for llama and NEOX for qwen2. The two runs use the same seeded weights, the same prompt [1, 5, 9], and the same code path, so the different generated sequences ([16, 18, 1, 3, 3, 21, 1, 1, 1, 1] against [1, 1, 9, 1, 1, 1, 1, 1, 9, 1]) come from the RoPE layout alone. A model run with the wrong layout produces different, plausible-looking tokens without any error, which is why the layout is part of Config and not a guess.

logits bit-identical at every step: true compares, at each of the 10 steps, the logits from the cached model with the logits from recomputing the whole sequence. They are equal with Arrays.equals, not only within a tolerance. Both paths run the same operations in the same order: recomputation rewrites cache entries 0..t with the same values the cached run stored earlier, and attention at position t reads only those entries. The cache changes how much work is done, never the result.

The work count shows the difference: 13 forward calls with the cache (3 prompt positions and 10 generated tokens) against 75 without it (3 + 4 + ... + 12, the sequence length at each of the 10 steps). The uncached cost grows with the square of the sequence length, and the cached cost grows linearly.

Llama is used as follows: save the program as Llama.java next to GgufModel.java and GgufTokenizer.java, and create a model with new Llama(gguf.metadata(), n -> gguf.tensors().containsKey(n) ? gguf.floats(n) : null, 4096, true), as Examples 11 and 12 do. The class supports the llama and qwen2 architectures with F32, F16, and Q8_0 weights, and rejects other architectures. The quantized-kernels lesson replaces matvec with Vector API kernels that read Q8_0 blocks from the mapped file and split rows across threads, and Capstone: A Chat Engine in Pure Java builds a chat engine on top: chat template, sampler chain, streaming, and stop handling.

Interview points

• A Llama-style forward pass handles one token in these steps: look up the embedding row; then for each layer apply RMSNorm, the Q/K/V projections, RoPE on Q and K, cached attention, the output projection, a residual add, RMSNorm, SwiGLU, and another residual add; finally apply the last RMSNorm and the output projection to get logits.
• Hyperparameters come from {arch}.* GGUF keys, and four of them have defaults when missing: head_count_kv, rope.dimension_count, rope.freq_base, and the vocabulary size.
• GGUF dimensions [ne0, ne1] mean ne1 rows of ne0 values each, so a linear layer's input size is ne0 and output row r starts at float r * ne0.
• Normal RoPE pairs values (2i, 2i+1) and NEOX pairs (i, i + n_rot/2). llama.cpp chooses by architecture (Llama normal, Qwen2 NEOX), and the converter permutes Llama's Q/K weights once so they fit the normal layout.
• Grouped-query attention maps query head h to KV head h / (heads / kvHeads), which shrinks the KV cache by the group size.
• RMSNorm divides by sqrt(mean(x^2) + eps), and ggml accumulates the sum of squares in double precision.
• Tied embeddings reuse token_embd.weight as the output projection when output.weight is absent.
• Greedy decoding picks the first strict maximum of the logits, which matches llama.cpp's temperature-0 sampler when two logits tie.
• llama.cpp's default KV cache is F16, and its CPU flash attention also rounds the query and the value accumulator to F16, so exact comparisons need -fa off -ctk f32 -ctv f32.
• llama-eval-callback prints every graph tensor with its sum, which lets you localize a porting bug to one operation.
• In the last layer, llama.cpp computes only the rows whose logits are needed (inp_out_ids), so its tensors there have one row.
• A KV cache changes the work from quadratic to linear in sequence length, and it must not change the logits.

Summary

• The tiny random model has head dimension 32, KV dimension 64, 4 query heads per KV head, 705,024 weights per layer, and 80,611,584 weights in total.
• As float[], an 8B model needs 29.92 GiB; as Q8_0, 7.95 GiB.
• The embedding of token t is floats t * dim to t * dim + dim - 1 of token_embd.weight.
• RMSNorm of [1, 2, 3, 4] is [0.365, 0.730, 1.095, 1.461].
• Weights stored as 1 2 3 4 5 6 with dims [3, 2] times [1, 2, 3] give [14, 32]; reading them as columns gives [22, 28].
• The tiny model's output projection costs 38,895,616 multiply-adds per token, more than its four layers together.
• At position 1, normal RoPE maps (1, 2) to (-1.143, 1.922); NEOX maps (1, 5) to (-3.667, 3.543).
• silu(2) = 1.7616, silu(-2) = -0.2384, silu(0) = 0.
• An F16 KV cache halves the tiny model's cache from 2,048 to 1,024 bytes per token and changes an attention output by at most 0.0007 in the test.
• With -fa off -ctk f32 -ctv f32, every traced tensor of the Java engine matches llama-eval-callback within a relative difference of 5.1e-05.
• With llama.cpp's defaults, the differences grow to the order of 1e-4 because of F16 rounding inside flash attention.
• The 32 greedy tokens for Hello world are identical in the Java engine, llama-completion --temp 0, and llama-server.
• With the KV cache, generating 10 tokens after a 3-token prompt takes 13 forward calls; without it, 75, with bit-identical logits.