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.

What you need
The forty-line inferencer
The five classes, explained
Benchmarking it honestly
Going lower: the Tensor layer
Why it's fast
Knobs and kill switches
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.
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 architecture family, with F32/F16/Q4_K/Q6_K/Q8_0 tensors (Q4_K_M and Q4_K_S files are the sweet spot).

2. The forty-line inferencer

import sj.gpu.GpuDevice;
import sj.gpu.Model;
import sj.gpu.KvCache;
import sj.gpu.Tokenizer;

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]);
        Tokenizer tk = Tokenizer.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.

Tokenizer — built from the model's own vocab: fromGguf(m.gguf()), then encode(string) -> int[], decode(int) -> string, eos(), bos().

KvCache — the attention key/value cache, sized from the model's maxSeq(). Create one per conversation; create a fresh one to start over. Two caches over one Model = two independent sessions sharing the weights.

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:

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.
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:

Variable	Effect
`SUPERJ_GPU=cpu`	force the CPU reference backend
`SUPERJ_GPU_PREFILL_F16=0`	disable the f16 weight twins
`SUPERJ_GPU_ROPE_FUSED=0`	unfuse rotary embedding / KV-write from the GEMMs
`SUPERJ_GPU_EMBED_BATCH=0`	per-token embedding lookups
`SUPERJ_GPU_REPLAY=0`	disable decode token replay
`SUPERJ_GPU_PROFILE2=1`	full-clock per-kernel GPU profile at exit
`SUPERJ_GPU_CBSTATS=1`	dispatch/barrier counts + GPU occupancy at exit

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.