Writing an LLM Inference Engine

A step-by-step tutorial: load a quantized GGUF model, run it on the GPU, and generate text — in about forty lines of SuperJ. The sj.gpu stack this tutorial drives is the same engine that runs TinyLlama faster than llama.cpp on the same Apple M4 Max (prefill 5,753 vs 5,624 tok/s, decode 1.24× faster, output byte-identical). You get that engine as a library; this page shows you how to use it, and then how to go beneath it.

Contents

1. What you need
2. The forty-line inferencer
3. The five classes, explained
4. Benchmarking it honestly
5. Going lower: the Tensor layer
6. Why it's fast
7. Knobs and kill switches
8. Correctness: the CPU oracle

1. What you need

• A Mac with Apple Silicon for the Metal backend — or any machine at all: sj.gpu ships a CPU reference backend that runs the identical code (slowly, but bit-for-bit usefully — see §8). A CUDA backend exists for Linux/NVIDIA, but it runs llama-family models only — this tutorial's TinyLlama works on it; the Gemma family described two paragraphs below does not. The quantized matmuls those models need (matmul_q4_0/q5_0/q8_0), gemma3's QK-norm (rmsnorm_heads), gemma4's proportional RoPE (rope_dual_ff_f32) and the whole MoE path (moe_topk, the per-expert matmuls, weighted_row_sum, the indexed gather/scatter) exist in the Metal backend and are absent from sdk/gpu-kernels/cuda/kernels.cu, so a Gemma model cannot dispatch a single forward pass there. Porting them is tracked by #4239.
• A GGUF model file. Everything here uses TinyLlama 1.1B Q4_K_M (636 MB):

curl -L -o tinyllama-q4k.gguf \
  https://huggingface.co/TheBloke/TinyLlama-1.1B-Chat-v1.0-GGUF/resolve/main/tinyllama-1.1b-chat-v1.0.Q4_K_M.gguf

Supported today: the llama family (Llama/Mistral/Qwen-shaped GGUFs), the dense Gemma family (gemma/gemma2/gemma3 — GeGLU, embedding scaling, QK-norm, sandwich norms, alternating sliding-window/global attention with per-layer RoPE theta, all detected from general.architecture and each model's own GGUF headers; #4041), and gemma4 (#4043/#4045/#4046/#4047 — everything gemma2/3 have, plus per-layer attention shape and a REDUCED nKvHead on local/SWA layers, proportional RoPE via a shared rope_freqs.weight, weightless V-norm, "alternating-attention" layers that omit attn_v.weight entirely (V reuses K's pre-norm projection), and a 128-expert top-8 MoE FFN (fused per-expert gate_up/down, per-expert output scale) alongside a dense/shared GeGLU expert, combined through a 5-norm sandwich and a per-layer output scale). gemma4's cross-layer KV-cache sharing and per-layer-embeddings (AltUp) are NOT supported — Model.load fails loudly if a checkpoint's header sets either to a nonzero value (both are 0 for the gemma4:26b checkpoint this was verified against). Tensor dtypes: F32/F16/ Q4_K/Q6_K/Q8_0 (Q4_K_M and Q4_K_S files are the llama sweet spot) plus Q4_0/Q5_0, which small and QAT-quantized Gemma releases commonly use for whole weight matrices, not just embeddings.

A model with a very large declared context_length (gemma4:26b: 262144) can exceed available memory if KvCache.create preallocates its full window — see KvCache.create(dev, m, maxContext)'s 3-arg overload to cap the usable context instead (mirrors llama.cpp's own --ctx-size).

2. The forty-line inferencer

import sj.gpu.GpuDevice;
import sj.gpu.Model;
import sj.gpu.KvCache;
import sj.gpu.TextTokenizer;
import sj.gpu.Tokenizers;

public class tinychat {
    public static void main(String[] args) {
        if (args.length < 2) {
            System.out.println("usage: tinychat <model.gguf> <prompt> [maxTokens]");
            System.exit(1);
        }
        int maxTokens = args.length > 2 ? Integer.parseInt(args[2]) : 200;

        // 1. Open a GPU device (Metal on macOS; falls back to the CPU
        //    reference backend anywhere else — same code, same results).
        GpuDevice dev = GpuDevice.open();

        // 2. Load the model. The GGUF file is mmap'd zero-copy; quantized
        //    weights become GPU tensors in a few milliseconds.
        Model m = Model.load(dev, args[0]);
        TextTokenizer tk = Tokenizers.fromGguf(m.gguf());
        KvCache kv = KvCache.create(dev, m);

        // 3. Prefill: run the whole prompt through the model in one batched
        //    pass. Returns the first generated token.
        int[] promptTokens = tk.encode(args[1]);
        int token = m.prefill(promptTokens, kv);

        // 4. Decode: one token at a time. decodeStep processes `token` at
        //    position `pos` and returns the next token (greedy argmax).
        for (int pos = promptTokens.length; pos < maxTokens; pos++) {
            token = m.decodeStep(token, pos, kv);
            if (token == tk.eos()) break;
            System.out.print(tk.decode(token));
        }
        System.out.println();
    }
}

Compile and run:

superj compile tinychat.sj --sdk-path "$SJ_HOME/sdk" --link --output tinychat
./tinychat tinyllama-q4k.gguf "The capital of France is" 100

That's the whole engine. Everything below is understanding and control.

3. The five classes, explained

GpuDevice — the backend handle. GpuDevice.open() picks the best backend for the machine (Metal → CUDA → CPU); GpuDevice.open("cpu") forces one. Also the door to raw compute: arenas, buffers, kernels (§5).

Model — a loaded transformer. Model.load(dev, path) parses the GGUF header, maps the file, and wraps every weight as a zero-copy tensor view — no 600 MB read; loading takes milliseconds. Introspection: nLayers(), dim(), nHead(), nKvHead(), vocabSize(), maxSeq(). The two verbs:

• prefill(int[] tokens, KvCache kv) -> int — the whole prompt in one batched GPU pass (internally chunked only when a prompt exceeds the scratch-memory budget). Fills the KV cache, returns the first generated token.
• decodeStep(int token, int pos, KvCache kv) -> int — one autoregressive step: processes token at position pos, appends to the cache, returns the greedy next token.

Tokenizers.fromGguf — returns the TextTokenizer the file actually needs, then encode(string) -> int[], decode(int) -> string, isEos(int), eos(), bos().

Do not pick the implementation by hand. A GGUF is either SentencePiece (sj.gpu.Tokenizer — Gemma, Llama 1/2, Mistral) or byte-level BPE (sj.gpu.HfTokenizer — Llama 3, Qwen, most post-2024 releases), and using the wrong one does not fail: it splits every word boundary into byte-fallback tokens and the model answers fluent nonsense (#4240). Tokenizers.fromGguf reads tokenizer.ggml.model and chooses.

Three things about encode that matter as soon as you go past a bare completion prompt (#4241):

• It resolves the vocab's control tokens — <start_of_turn>, <|turn>, <|eot_id|> — to their own single ids, which is what makes a chat template expressible. encode(text, false) turns that off; use it for the user's message, so content containing a literal turn marker cannot forge a turn.
• It adds BOS only if the file says to (tokenizer.ggml.add_bos_token). gemma4 says no, because its template writes <bos> itself.
• isEos(token), never token == eos(). End-of-generation is a set: gemma4's is [<eos>, <turn|>, <|tool_response>], and a chat turn ends with <turn|>. Comparing against eos() alone runs straight past the turn boundary.

demo/sj/demo/ChatDemo.sj is a ~100-line worked example of all of this — tokenizer selection, per-family chat framing, and stopping on the turn boundary. Verified against gemma2:2b and gemma3:1b:

$ ./chatdemo gemma3-1b.gguf "What is the capital of France? Answer in one sentence."
backend=metal arch=gemma3 layers=26 vocab=262144
promptTokens=23

> What is the capital of France? Answer in one sentence.
Paris.

replyTokens=2 endedOnTurnBoundary=true

Two models it does not work on today, both for reasons outside the tokenizer: gemma4 (#4248 — the forward pass is numerically wrong) and any Llama-3.x GGUF (#4258 — rope_freqs.weight is ignored outside gemma4, so RoPE runs unscaled).

KvCache — the attention key/value cache, sized from the model's maxSeq() by default. Create one per conversation; create a fresh one to start over. Two caches over one Model = two independent sessions sharing the weights. KvCache.create(dev, m, maxContext) caps the usable context below the model's declared max instead — needed for a model like gemma4:26b (context_length=262144), where preallocating the full window would need far more memory than most machines have (#4047).

Gguf — the raw file reader (m.gguf()), if you need headers or tensors the Model doesn't surface.

Sampling: decodeStep is greedy (argmax) today. For temperature/top-k you would take logits from the Tensor layer yourself (§5) — a good first extension project.

4. Benchmarking it honestly

Two hard-won rules:

1. Warm up before timing. Apple GPUs boot in a low power state and need ~100 ms of sustained work to reach full clock. A single cold prefill measures the clock ramp, not your code — it reads ~40% slow. Run one throwaway prefill into a scratch KvCache first, then time the real one; report the cold number separately as first-call latency if it matters to your application.
2. Compare same-session. Thermals move every number. Benchmark your engine and any baseline back-to-back, on the same machine, in the same run — a baseline measured an hour ago on a cool machine is not a comparison.

For per-kernel GPU time there are two built-in profilers (environment variables, no rebuild):

• SUPERJ_GPU_PROFILE2=1 — streams each dispatch in its own command buffer at full clock and prints a per-kernel table at exit. Trust it for big kernels (matmuls, attention); it overstates swarms of tiny kernels, because per-buffer overhead dominates a 13 µs dispatch.
• SUPERJ_GPU_CBSTATS=1 — dispatch/barrier counts plus true GPU busy time and occupancy from hardware timestamps.

5. Going lower: the Tensor layer

Model is a few hundred lines of plain SuperJ over public primitives — the sj.gpu sources ship with the SDK, and reading Model end-to-end is the fastest way to learn the whole forward pass (browse the API with superj doc --list / superj doc sj.gpu.Tensor). To build a different architecture, or custom sampling, use the same pieces:

GpuArena a = dev.scratch();                       // bump allocator, reset per step
Tensor w  = Tensor.view(dev, buf, DType.Q4_K, rows, cols);  // zero-copy weight
Tensor x  = Tensor.alloc(dev, DType.F32, n, dim);           // activations

Tensor xn  = x.rmsNorm(normWeight);               // fused normalize
Tensor q   = wq.matmul(xn);                       // GEMV (n==1) or tiled GEMM (n>1)
Tensor.ropeDualInPlace(q, k, ropeSpec, pos);      // rotary embedding, q+k in one dispatch
Tensor att = Tensor.attention(q, keys, vals, attnSpec, seqLen);  // flash attention
wo.matmulAccumInto(att, x);                       // x += Wo·att, fused residual
int next   = logits.argmax();                     // GPU argmax, 4-byte readback

Everything dispatches by batch size automatically: one activation row takes the memory-bound GEMV path, many rows take the tiled GEMM path. A gotcha that bites everyone once: arena is a reserved word in SuperJ — don't name a variable that.

6. Why it's fast

Every item below is measured and was verified output-identical when it landed:

• Zero-copy weights — one no-copy GPU buffer over the mapped file; load in ~5 ms.
• Batched command encoding + rolling commits — one command stream per token; the GPU executes early layers while the CPU encodes later ones.
• GEMV/GEMM split — decode streams quantized weights (memory-bound, the 630 MB/token wall); prefill runs a tiled GEMM that reuses each weight tile across all prompt rows (compute-bound).
• f16 weight twins — prefill weights are pre-dequantized to f16 at load (~2 GB resident for TinyLlama), deleting dequantization from the GEMM inner loop: 14.3 of the measured 15.3 TFLOP/s hardware ceiling. Decode keeps the quantized originals — it's bandwidth-bound, the opposite trade.
• Fused epilogues — rotary embedding is applied inside the q/k GEMM's epilogue, and the k/v projections write directly into the KV cache: whole pipeline stages (and their synchronization stalls) disappear.
• Single-pass prefill + causal skip — a pp512-class prompt runs as one batched pass, and the attention kernel skips key/value tiles beyond each query block's causal horizon instead of computing and masking them.
• Token replay — steady-state decode re-encodes each token's entire dispatch stream natively with two patched scalars.

7. Knobs and kill switches

Every optimization has an environment kill switch (no rebuild) — the first tool to reach for when bisecting a suspected miscomputation or regression:

8. Correctness: the CPU oracle

The CPU backend runs the same Model code with scalar kernels and is the ground truth: run any generation with SUPERJ_GPU=cpu and diff the text. The GPU path is held to token-identical output — and several of its optimizations are bit-identical by construction: they replicate the exact floating-point evaluation order of the path they replace. When you extend the engine, inherit the discipline: give every change a kill switch, A/B it against itself, and diff against the oracle. Fast and wrong is just wrong.