Jul 2026 draft
building a tiny inference engine in Java
MiniInfer started in Python where every tensor is inspectable. this is the Java port, and what the type system caught versus what it could not.
why Java
Python was the right first language because iteration is cheap and I could print any intermediate. But Java is my day job language, and I wanted three things Python would not give me: real threads for experimenting with batching later, allocation visibility (every float array is an object you can feel in a profiler), and compile time shape discipline. Not full shape safety, Java generics cannot express "this matrix is 768 wide", but at least named classes like AttentionCache and LogitsProcessor instead of bare nested lists drifting through functions.
scope discipline mattered too. The port covers a toy scale only: a small vocabulary, a couple of transformer layers, greedy sampling. Anything bigger and I would be writing a slow clone of existing runtimes instead of learning. The point was the decode loop, the cache, and sampling, each observable on its own.
the decode loop
the loop itself is five stages and fits in one method: tokenize the prompt, look up embeddings, run each transformer block (attention with the growing KV cache, then the feedforward), project the final hidden state to vocabulary logits, sample one token, append it, repeat. Greedy sampling first (argmax, deterministic, debuggable), temperature and top-k only after greedy output looked sane.
keeping the cache append order explicit was the design decision that paid off. Each block owns its cache lists, and the loop appends the new key and value right after computing attention for the step, never before. Getting that ordering wrong produces fluent nonsense (more below), so making it the one obvious sequence in the code was worth more than any abstraction.
what broke
two bugs, one caught by types and one completely invisible to them. The caught one: I transposed a weight matrix the wrong way and Java refused to compile the matmul helper, since my Matrix record carried its row and column counts and the multiply method checked them. In Python that would have been a runtime shape error three calls deep. Point to the type system.
the uncaught one was softmax overflow. My first softmax exponentiated raw logits, and with logits around 20 to 30, exp() overflowed to infinity and every probability came out NaN. The fix is the standard one line: subtract the max logit before exponentiating. Types cannot save you here because every value involved is a perfectly valid double. The output looked like fluent nonsense for an embarrassing hour before I printed the probabilities and saw the NaNs. Now "subtract the max" is the first thing I check in any sampling code.
what MiniInfer taught me
the Java port confirmed the lesson from the Python experiments: the bottleneck is memory movement, not arithmetic. Generating one token streams all the weights plus the full KV cache for a single vector of useful work, which is why batching changes the economics so much (same weight read, many tokens out). Writing the loops in Java, where each array allocation is explicit, made that streaming cost feel concrete in a way NumPy had hidden.
what surprised me most is how little code a working engine actually is. Tokenizer, embeddings, attention with cache, feedforward, logits, greedy sample: a few hundred lines, most of them plumbing. The model weights do all the knowing. The engine is just a careful loop that refuses to recompute the past, which connects directly back to the KV cache note. Next step is measuring where time actually goes at 512 versus 2048 tokens, then sampling strategies, in that order.
all notes