Writing an LLM inference engine from scratch in Rust

I've been building inferno, a CPU inference engine for Llama-style language models, written in Rust. It doesn't use PyTorch, Candle or GGML. The matrix-multiply kernels, attention, KV cache, thread pool, quantization and sampler are all hand-written, about 3,500 lines in total. The only dependencies are crates to parse the weight file (safetensors), convert f16 values (half), tokenize text (Hugging Face's tokenizers) and generate random numbers (rand).

It runs TinyLlama-1.1B-Chat. The first time I benchmarked it against Ollama (which runs llama.cpp underneath), inferno generated text 5 to 7 times slower, and reading a 1,500-token prompt took almost a minute where Ollama needed about 2.5 seconds.

Here is where it ended up. Both engines run the same 4-bit (Q4_0) model on the same machine (Intel Core Ultra 7 265KF, under WSL2), limited to 8 threads, with Ollama in CPU-only mode:

Prompt length	Prefill: inferno	Prefill: Ollama	Decode: inferno	Decode: Ollama
50 tokens	58 ms	94 ms	111 tok/s	93 tok/s
500 tokens	698 ms	902 ms	104 tok/s	85 tok/s
1500 tokens	2.47 s	3.11 s	89 tok/s	66 tok/s

Prefill is the time to process the prompt. Decode is how fast the reply is generated after that.

A caveat before going further: this is one small model on one machine. Ollama supports hundreds of models and runs almost anywhere. inferno supports exactly one model, on x86 CPUs with AVX2. Beating it here says more about what a narrowly specialized engine can get away with than about llama.cpp.

This post covers how the gap closed: what the profiler showed, which optimizations paid off and which didn't. It also covers two features built on top: speculative decoding that needs no second model, and a chat app that starts processing your message while you're still typing it.

How a Llama model generates text

Everything later builds on this, so here's a short recap.

TinyLlama has 22 transformer layers. Each token is represented as a vector of 2,048 numbers, and each layer does the following:

Normalize the vector (RMSNorm) and project it into queries, keys and values. There are 32 query heads but only 4 key/value heads, so groups of 8 query heads share one key/value head. This is called grouped-query attention (GQA).
Rotate the queries and keys by an angle that depends on the token's position (RoPE).
Attention: compare the query against the keys of every earlier token, which are stored in the KV cache, and take a weighted average of their values.
Apply an output projection, normalize again, and run the feed-forward block (SwiGLU, 2,048 → 5,632 → 2,048).

After the last layer, one more projection produces a score for each of the 32,000 tokens in the vocabulary, and the sampler picks one.

There are two phases:

Prefill processes the prompt. All prompt tokens are known up front, so they can go through each layer together as a batch.
Decode generates the reply one token at a time. Every new token needs a complete pass through all 22 layers.

Nearly all the time goes into multiplying vectors by weight matrices, which is where the first lesson came from.

Lesson 1: decode speed is set by memory bandwidth

In bf16, TinyLlama's weights take 2.2 GB. Generating one token reads every one of those weights once. inferno's bf16 path now decodes about 31 tokens per second, which works out to about 68 GB/s of reads. That is roughly all the bandwidth my RAM has. Faster arithmetic won't help. The only way to generate faster is to read fewer bytes per token.

Quantization does exactly that. inferno uses Q4_0, with the same block size and bit width as llama.cpp's Q4_0. Weights are stored in blocks of 32, each block holding 4-bit integers plus one f16 scale. That's 4.5 bits per weight instead of 16, so about 3.5 times fewer bytes to read.

Lesson 2: my first 4-bit version was slower than 16-bit

My first Q4_0 implementation made prefill about 1.8 times slower than bf16 and barely changed decode speed. The kernel unpacked each 4-bit value with plain scalar code, converted it to a float and then multiplied. That was slow enough that the CPU's arithmetic became the bottleneck, so the savings in memory traffic didn't matter.

The fix was a fused SIMD kernel that unpacks and multiplies in registers without ever writing the unpacked weights to memory. It uses the same trick as llama.cpp: a 16-bit shift right by 4 followed by a 0x0F mask pulls the high nibbles out of 16 bytes at once, byte-interleave instructions put the values back in order, and they're widened straight to floats and fed into fused multiply-adds (FMAs). In a single-threaded microbenchmark this was 10 times faster (4.2 → 43 GFLOP/s). End to end, decode roughly doubled at short context and prefill got 1.7 to 3.2 times faster.

Just before that, I'd fixed a smaller problem in the bf16 kernel. It added every product into a single SIMD register, so each FMA had to wait for the previous one to finish:

// Simplified. Before: one accumulator, so every FMA waits on the last one.
let mut acc = zero();
for i in (0..n).step_by(8) {
    acc = fma(w(i), x(i), acc);
}

// After: four independent chains that the CPU can run in parallel.
let (mut s0, mut s1, mut s2, mut s3) = (zero(), zero(), zero(), zero());
for i in (0..n).step_by(32) {
    s0 = fma(w(i),      x(i),      s0);
    s1 = fma(w(i + 8),  x(i + 8),  s1);
    s2 = fma(w(i + 16), x(i + 16), s2);
    s3 = fma(w(i + 24), x(i + 24), s3);
}
let sum = (s0 + s1) + (s2 + s3);

An FMA takes about four cycles to produce its result, but the core can start two every cycle. With a single chain, most of that capacity sits idle. In isolation, four accumulators ran 2.5 times faster.

The two fixes had very different results end to end. The bf16 fix gave only 1.16 to 1.31 times. The Q4_0 fix gave about 2 times. The bf16 path was already limited by memory bandwidth, so a faster kernel mostly just waited on RAM sooner. The Q4_0 path was limited by arithmetic, so a faster kernel helped directly. Before optimizing a kernel, it helps to know which of the two limits it's hitting, because that predicts how much of the improvement survives outside a microbenchmark.

Lesson 3: 155 thread-pool calls per token add up

The matmuls originally ran in parallel with rayon. Every linear layer made its own parallel call, which comes to about 155 calls per generated token (7 projections × 22 layers, plus the final one). Each call does only tens to hundreds of microseconds of real work.

I measured rayon's overhead on its own, with a parallel call that did almost nothing: about 115 µs per call at 20 threads. Across 155 calls, that's roughly 17 ms of overhead per token, at a time when a whole token took 25 to 55 ms.

So I replaced rayon with a small pool designed for this workload:

Workers stay running between calls. They spin on an "epoch" counter, and the next job starts as soon as the counter changes. After 1 ms without work they go to sleep on a condition variable, so an idle server doesn't keep every core busy.
Work is claimed one chunk at a time from an atomic counter, instead of each thread getting a fixed share. My CPU has 8 fast cores and 12 slower ones. With fixed shares, every call waits for the slowest core to finish.
The calling thread also works on chunks. Nested parallel calls run inline, and a panic in a worker is re-raised on the caller.

At equal thread counts, decode was 1.35 to 1.5 times faster than with rayon. Later, I merged the Q, K and V projections into one matrix and the MLP's gate and up projections into another, which cut the number of parallel calls per token from 155 to 89.

One side effect of spinning workers: running them on every logical core is a bad idea. At 18 to 20 threads on this 20-thread CPU, decode speed varied between 23 and 88 tokens per second from one run to the next. If the OS pauses a spinning worker in the middle of its chunk, the whole call waits until that worker runs again. The default now leaves about a sixth of the cores free.

Lesson 4: profile by wall time, not by operation count

After the thread pool, inferno was 2 times slower than Ollama on short prompts and 4 times slower on long ones. The gap grew with context length, so I added timers around the main parts of the forward pass:

Decode, per token	Total	Matmul	Attention
50 tokens of context	12.3 ms	10.3 ms (84%)	1.4 ms (11%)
1500 tokens of context	46.0 ms	12.5 ms (27%)	32.9 ms (72%)

My own code had a comment saying attention was under 1% of the compute. By number of operations, that was true. By time, attention took almost three quarters of every token at long context. It ran single-threaded and without SIMD, and allocated new buffers for every head. The matmuls alone at 1,500 tokens took about as long as Ollama's entire token.

I changed attention over several rounds:

Spread the (query, head) pairs across the thread pool, use AVX2 for the dot products, and reuse buffers. Decode at 1,500 tokens went from 19.5 to 40.4 tokens per second.
Compute the softmax online over 64-position blocks, keeping a running maximum and sum as FlashAttention does. The full score matrix is never built.
Process query heads that share a key/value head together. The cache is then read twice per key/value head instead of eight times.
Store the KV cache in f16 instead of f32, as llama.cpp does by default.
Tile 4 queries × 4 heads per work item. Each 64-position block of the cache is converted from f16 to f32 once and stays in L1 cache while all 16 query-heads use it.

Profiling turned up a few more surprises:

RoPE. It showed up in the "other" bucket. My implementation computed a powf and a sin_cos for every pair of dimensions, in every head, in every layer: about 25,000 trigonometric calls per token. But the angles depend only on the token's position. Now they're computed once per position and reused by all heads.
A hidden transpose. In a 1,500-token prefill, 1.9 of 5.2 seconds went to reshaping batched matmul output from one row per weight to one row per token, on a single thread. The kernels now write directly into each token's output.
Scalar activation quantization. Converting activations to int8 (see Lesson 6) took about 2 ms per token. A SIMD version that produces exactly the same bytes removed most of that, and decode sped up by about 20%.

Lesson 5: batching needs a real matrix multiply

Prefill can process all prompt tokens together, so every weight row read from memory can be reused for many tokens. My first batched kernel didn't actually get that benefit. It unpacked each 4-bit weight row once per token, not once per batch, so a prefill cost 11.7 ms per token against 10.3 ms per token for plain decode.

The rewrite works like a blocked matrix multiply. Each thread unpacks a group of weight rows into a 256 KB scratch buffer, sized to fit in L2 cache. It then loops over tokens on the outside and rows on the inside, handling up to 3 tokens per weight load. Prefill got 2.6 to 3.3 times faster.

Lesson 6: use the same instructions as your competitor

After all that, prefill was still about 3 times slower than Ollama. I checked which CPU build Ollama had loaded: the variant for Alder Lake and newer, with AVX_VNNI = 1 in its log. Disassembling that library found 461 vpdpbusd instructions. My code had none.

vpdpbusd multiplies 32 unsigned bytes by 32 signed bytes, then adds each group of four products into one of eight 32-bit sums, all in one instruction. For that to work, the activations have to be int8 as well, so before each matmul they're quantized into Q8_0 blocks (32 int8 values with one scale), as llama.cpp does.

My first attempt gained nothing. Each 32-bit lane held part of a single row's dot product, so every block still needed a horizontal sum, an int-to-float conversion and a scale for every row. That was about as many instructions as the f32 path.

The version that worked rearranges the weights at load time into tiles of 8 rows, laid out so that lane r of the accumulator always belongs to row r. One block of 32 weights across 8 rows takes 8 vpdpbusd instructions (256 multiply-adds), and the conversion to float is shared by all 8 rows. No horizontal sums are needed:

for m in 0..4 {
    let v = load_256(&blk[32 * m..]);
    let w_lo = and(v, mask_0f);                   // low nibbles of 8 rows
    let w_hi = and(shift_right_16(v, 4), mask_0f); // high nibbles
    lo = dpbusd(lo, w_lo, broadcast_4_bytes(&x.qs[4 * m..]));
    hi = dpbusd(hi, w_hi, broadcast_4_bytes(&x.qs[16 + 4 * m..]));
}
// Q4_0 stores each weight as w + 8 (range 0..15), so remove the bias:
//   sum((w + 8) * x) - 8 * sum(x) = sum(w * x)
let isum = (lo + hi) - 8 * x.sum;                  // x.sum is precomputed per block
acc = fma(to_f32(isum), w_scales * x.scale, acc);  // 8 rows at once

The stored bias turned out to be convenient. vpdpbusd needs one unsigned and one signed operand, and Q4_0's stored nibbles are already unsigned (0 to 15). The correction is one subtraction per block, using the sum of the activation block, which is computed once when the activations are quantized.

This took prefill from 249 ms to 110 ms at 50 tokens, and decode from 54 to 81 tokens per second. The cost is memory: the other code paths still use the original layout, so the 4-bit weights are stored twice, about 0.6 GB extra.

The whole path in one table

Everything at 8 threads, Q4_0. This starts after the thread-pool change; earlier numbers were measured at a different thread count and aren't comparable. Run-to-run noise is a few percent.

Change	Prefill (50 / 500 / 1500 tokens)	Decode, tok/s (50 / 500 / 1500)
Custom thread pool	766 ms / 9.9 s / 51.3 s	50.4 / 33.6 / 19.5
Parallel SIMD attention	653 ms / 6.9 s / 33.5 s	52.3 / 49.1 / 40.4
Blocked batched matmul	249 ms / 2.70 s / 10.1 s	54.5 / 47.7 / 40.4
Int8 VNNI tiles	110 ms / 1.36 s / 5.66 s	80.8 / 70.0 / 52.3
RoPE, vectorized `exp`, grouped attention heads	93 ms / 1.12 s / 5.00 s	84.6 / 77.5 / 67.8
Fused projections, f16 KV cache	99 ms / 1.08 s / 4.99 s	90.5 / 79.9 / 71.7
No transpose, SIMD activation quantization	60 ms / 718 ms / 2.83 s	109.6 / 91.1 / 81.5
Query-tiled attention	58 ms / 698 ms / 2.47 s	111.4 / 103.8 / 89.2
Ollama, CPU, 8 threads	94 ms / 902 ms / 3.11 s	92.9 / 84.8 / 66.1

Speculative decoding without a draft model

Because decode is limited by memory bandwidth, checking several tokens in one batched forward pass costs little more than generating one. The weights are read once either way. Before building anything, I measured this: one batched pass over k tokens was 2 to 3.6 times cheaper than k separate passes, for k = 2, 4 and 8.

Speculative decoding usually gets its guesses from a second, smaller "draft" model. inferno uses prompt lookup instead. It takes the last few tokens, finds where the same sequence appeared earlier in the context, and guesses that whatever followed it last time will follow again:

fn prompt_lookup_candidates(ids: &[i64], max_ngram: usize, min_ngram: usize, k: usize) -> Vec<i64> {
    let n = ids.len();
    for ngram in (min_ngram..=max_ngram.min(n)).rev() {
        let suffix = &ids[n - ngram..];
        if let Some(start) = ids[..n - ngram].windows(ngram).rposition(|w| w == suffix) {
            let match_end = start + ngram;
            let end = (match_end + k).min(n);
            if end > match_end {
                return ids[match_end..end].to_vec();
            }
        }
    }
    Vec::new()
}

Up to 4 guessed tokens are checked in one batched pass. Guesses are kept up to the first one the model disagrees with, and that position is replaced with the model's own token. With greedy sampling, the output is identical to plain decoding, token for token, and there's a test that checks exactly that.

That test only passes because the batched kernels and the single-token kernels give bit-identical results. Floating-point addition isn't associative, so summing in a different order changes the last bits, which can occasionally change which token wins. Every AVX2 kernel in inferno therefore sums in the same fixed order: 8-wide group g goes into accumulator g mod 4, and the accumulators combine as (a0 + a1) + (a2 + a3). Keeping that invariant constrained every later kernel change, but it means correctness can be checked with a plain equality test instead of a tolerance.

The first end-to-end results were disappointing: 1.05 times faster on code and 0.92 times on a repetition prompt, so actually slower. Two changes fixed most of it:

Require at least a 2-token match. With 1-token matches, a single common word like "the" was enough to trigger a guess. On ordinary prose, that caused 6 wasted batched passes per 60 generated tokens, none of them accepted. With a 2-token minimum, there were no wasted attempts on the prose test. Code improved too, from 1.38 to 1.65 times faster, with acceptance going from 56% to 75%, because the same noise had been diluting genuine matches.
An acceptance gate. inferno tracks acceptance over the last 10 attempts. If it drops below 30%, speculation pauses for 5 steps. Each consecutive pause doubles in length, up to 80 steps.

On Q4_0 with the default settings, code-like output is about 1.34 times faster and a repetitive-text prompt about 1.13 times faster. On the benchmark prompts above, speculation reaches 213 tokens per second at 500 tokens and 182 at 1,500, about twice plain decode. Those prompts are built from a repeated paragraph, though, so the model falls into a loop and every guess is accepted: treat that as the best case. At 50 tokens there's almost nothing earlier in the context to copy, and speculation is about 16% slower than plain decoding (93 vs 111 tokens per second).

I decided not to add a real draft model. It would need its own memory and forward passes, and on genuinely new text, which is where prompt lookup does nothing, a draft model would help only a little.

A chat app that reads while you type

On top of the engine there's a small chat app: an axum backend and a React frontend, with replies streamed over server-sent events. Each conversation keeps its own session with its KV cache, so each new turn only processes the new message, never the whole history.

Since the cache can be extended and cut back cheaply, the backend can also start on your message while you're typing it. The frontend sends the draft to a /compose endpoint 200 ms after you stop typing, and the backend prefills it. When you press Enter, only the part that changed since the last update still needs processing.

A few things made this harder than it sounds:

Edits. Backspacing or changing text mid-message invalidates part of the cache. On each update, the backend re-tokenizes the whole draft, finds how many tokens still match the cache, truncates the cache to that point and processes the rest. I'd planned incremental tokenization, but re-tokenizing takes under 2.3 ms even at the full 2,048-token context, so I didn't need it. Truncation is just Vec::truncate on each layer's cache.
Unstable token boundaries. With BPE tokenization, typing more characters can change how the previous characters are split. "under" might be its own token until you type "standing". So the last token is held back until the text ends in whitespace or punctuation:
```
fn safe_commit_len(candidate_tail_len: usize, text_ends_at_safe_boundary: bool) -> usize {
    if candidate_tail_len <= 1 || text_ends_at_safe_boundary {
        candidate_tail_len
    } else {
        candidate_tail_len - 1
    }
}
```
Earlier turns. Syncing takes an anchor position, and nothing before it is ever truncated, so editing a draft can't damage the conversation history.
Lost requests. A compose request might fail or arrive late. The send path always re-syncs against the final text, so compose requests can only save work, never change the result.

The system prompt gets similar treatment. It's prefilled once, and every new conversation starts from a copy of that cache instead of processing it again.

I tested this end to end by typing a typo into the draft character by character, correcting it and sending. The answer came back correct, with no trace of the deleted tokens.

Is the 4-bit model still good?

All the Q4_0 work so far had been about speed, so I added a scoring function and measured perplexity on a 605-token prose passage (lower is better):

Weights	Perplexity	Top-1 agreement with bf16
bf16	8.59	n/a
Q4_0, my first quantizer	9.22 (+7.4%)	87.4%
Q4_0, llama.cpp-style scale	8.96 (+4.4%)	89.3%

My first quantizer set each block's scale to its largest absolute value divided by 8. The 4-bit range is -8 to +7, so when the largest value was positive, it got clipped to +7, a 12.5% error on the block's largest weight in about half of all blocks. llama.cpp takes the value with the largest magnitude, keeping its sign, and divides by -8, so that value maps exactly to -8. Switching cut the quality loss from quantization by about 40%.

While testing, the 4-bit model kept ending its reply immediately on one prose prompt: its first token was end-of-sequence. It looked like a kernel bug, but it wasn't. That prompt didn't use the chat template, and even the bf16 model gave end-of-sequence a 13% probability there. Quantization noise raised that to about 28%, enough to make it the top choice. Ollama's own Q4_0 TinyLlama does the same thing on the same prompt. With the chat template, the 4-bit model writes a normal story.

Keeping the benchmark fair

Benchmarking against another engine turned out to be about as much work as the optimizing:

Match everything except the engine. Same model, same Q4_0 format, greedy decoding, byte-identical prompts and the same thread count. At 8 threads, Ollama was already close to its memory-bandwidth limit: in an earlier run, letting it use its default thread count made it only about 7% faster.
Watch out for caches. Ollama reuses cached work when a new prompt starts with the same text as the previous one. The prompts for each length are built from unrelated paragraphs, so they share no prefix and both engines do a full prefill every time.
Check which Ollama you're talking to. For a while, the Ollama answering on localhost inside WSL was actually the Windows app, reached through WSL's port forwarding. The Linux service was stuck in a restart loop because the port was taken. The numbers in this post are from Ollama 0.34 running inside WSL2, after I uninstalled the Windows app.
Don't trust single runs. One early one-shot benchmark showed 50-token decode as slower than 500-token decode. For thread-pool comparisons I switched to an A/B harness with a warmup, 7 trials of 32 tokens from an identical cache state, and the median.
Remaining differences: Ollama's tokenizer produces about 3% more tokens for the same text, and inferno quantizes the weights itself, so the weights aren't bit-identical to Ollama's GGUF file.
Context: on this machine's RTX 5070 Ti, Ollama generates 500 to 600 tokens per second. If raw speed is what you need, use the GPU.

Two things also crashed my machine along the way. cargo test runs tests in parallel, and every integration test loads the 2.2 GB model. Several at once reached about 14 GB, and the Linux out-of-memory killer took down the whole WSL virtual machine. Later, the PC itself crashed right after a series of long benchmark runs that kept every core busy. Tests now run one at a time, and all benchmarks run at 8 threads via .cargo/config.toml.

What I learned

Work out whether you're limited by bandwidth or by arithmetic before optimizing. It predicts whether a faster kernel will make any real difference.
Measure a change in isolation and end to end. A 2.5× kernel became 1.2× overall. A 10× kernel became 2×.
Profile by time, not by operation count. Attention was under 1% of the operations and 72% of the time. RoPE was hiding under "other". A transpose took over a third of a long prefill.
Look at what the fast implementation actually does. Disassembling Ollama's CPU library showed me which instruction I was missing.
Make batched and single-token results bit-identical. It cost some flexibility in every kernel, and in return speculative decoding can be tested with a plain equality check.

Limitations and next steps

TinyLlama only. Layer count, head count, head size and a few other constants are hard-coded instead of read from config.json. That's the first thing to fix.
The fast attention path only handles a head size of 64. Most 7B models use 128 and would fall back to scalar code.
No RoPE scaling or attention biases, so Llama 3 and Qwen don't work yet.
The whole weight file is read into memory, and 4-bit weights are stored twice (see Lesson 6).
CPU only, tuned for x86-64 with AVX2 (with a scalar fallback for other CPUs).

Try it

The code is on GitHub at github.com/AtharvaPingale/inferno. Download the model and run it:

huggingface-cli download TinyLlama/TinyLlama-1.1B-Chat-v1.0 --local-dir model
cargo run --release

Or use it as a library:

use inferno::{Engine, SamplerConfig};

let engine = Engine::load_quantized("model/model.safetensors", "model/tokenizer.json")?;
let mut session = engine.new_session();
engine.feed(&mut session, "<|user|>\nWhat is Rust?</s>\n<|assistant|>\n", true)?;

let sampler = SamplerConfig { temperature: 0.0, ..Default::default() }; // greedy
loop {
    let (tokens, is_eos) = engine.generate_speculative_step(&mut session, &sampler, 4);
    for t in tokens {
        print!("{}", engine.decode_token(t));
    }
    if is_eos { break; }
}

The repository includes all the benchmarks from this post: bench_vs_ollama for the Ollama comparison, eval_quality for perplexity, and the speculative decoding sweeps.