← Back to portfolio

Transformers from Scratch in Rust

I wrote every Transformer piece by hand in Rust on Candle: all four attention variants, RoPE, RMSNorm, and a KV cache. Then I loaded a real Qwen2.5 checkpoint into it, verified it against the reference implementation, and benchmarked it on CPU, with my predictions written down before the numbers came in.

Repo: github.com/lblommesteyn/rust-transformers

Why would you write a transformer in Rust?

It started with a mismatch. Say you have a Rust service. An indexer, a log pipeline, some little API that already compiles to one binary and runs anywhere. You want to add one transformer-shaped thing to it: semantic search embeddings, or a small model that drafts text. Underneath, that model is a stack of matrix multiplies. It should just be a function call.

Instead it asks you to adopt a whole second runtime. A Python interpreter, a virtualenv, a few hundred megabytes of framework, CUDA libraries with opinions about your driver version, and a quiet hope that whatever worked on your laptop also works in prod. The transformer is the simple part. Delivery is the hard part, and on an edge device, an air-gapped machine, or anywhere "just run Python" isn't free, it is most of the work.

So I wrote the architecture directly in Rust, on top of Candle, which is Hugging Face's minimalist tensor framework. The result: cargo build --release gives you one static binary that runs the forward pass itself. There is no Python, sidecar or network call, and it behaves the same on every machine.

What I didn't expect was how much I'd enjoy the rest. When you write every piece yourself, nothing sits between the tokens going in and the vector coming out. If something is slow, you can see exactly which term is slow. If the workload is weird, very long log lines, sequences longer than anything you trained on, you pick the attention mechanism that fits instead of taking whatever the library defaults to. That's why the benchmarks at the end exist: to show where the time goes.

What's in here

From-scratch encoder / decoder / encoder-decoder transformers on Candle, plus a causal LM with KV-cached generation.
Four attention paths, multi-head, flash-style, ALiBi, and sliding-window, chosen by config rather than by swapping libraries.
RoPE / sinusoidal / learned position encodings, LayerNorm or RMSNorm, and GELU / SiLU / Swish / Mish feed-forwards.
A Qwen2 module that loads real Hugging Face checkpoints: grouped-query attention, the gated SwiGLU MLP, and weights verified against the reference implementation to 3.4e-5.
Three weight precisions, f32 / f16 / q8_0, selectable at load time, so the same binary can trade accuracy for speed.
A tests/ suite that checks attention against an independent reference, and verifies that KV-cached generation reproduces a full forward pass exactly.
Real CPU benchmarks below on a real model, produced by cargo run --release --example qwen_bench.

What a transformer actually is

Tokens go in. A vector comes out, or a distribution over the vocabulary. In between is the shape every transformer shares: embed the tokens, then run a stack of identical layers, where each layer does attention and then a feed-forward, with a residual add and a normalization wrapped around each. The two output heads are what decide whether you've built an embedding model or a generator.

One shape, two uses: pool the last hidden state into an embedding for search, or project it through the LM head for generation.

Inside one layer

Attention. Multi-head attention is the baseline. From there the config picks a variant: a flash-style kernel that streams the softmax to keep memory off the hot path, ALiBi linear biases for extrapolating past your training length, or a sliding window (with optional global tokens) so long sequences don't pay the full quadratic bill.

Positioning. Rotary embeddings (RoPE) with a configurable theta are the default on the decoder side. Sinusoidal and learned tables are there too. ALiBi is the odd one out: instead of adding a position vector, it biases the attention scores directly with \(b_{ij} = m_h \cdot (j - i)\), using a per-head slope.

Normalization. LayerNorm is available, but RMSNorm is what I reach for on longer sequences. One fewer statistic to compute, and steadier in practice.

Pooling. For embeddings, masked mean pooling is the default. CLS and attention pooling are selectable if the model was trained for them.

What the attention variants change

The cheapest way to see the difference is to look at which query-key pairs each one lets talk. Full attention is dense, so everything sees everything. Sliding-window keeps a band around the diagonal. ALiBi still attends everywhere, but tilts the scores by distance.

Same query/key grid, three policies. Darker = stronger pull. Sliding-window drops the off-band pairs entirely; ALiBi keeps them but fades them with distance.

Building one

Configuration is a fluent builder. Pick a model type and a shape, flip on the features you want, and back it with a Candle variable store. This example uses random init; you'd swap in .safetensors for real weights.

use rust_transformers::models::transformer::Transformer;
use rust_transformers::utils::config::{ModelType, TransformerConfig};
use candle_core::{DType, Device, Tensor};
use candle_nn::{VarBuilder, VarMap};

let config = TransformerConfig::builder()
    .model_type(ModelType::Encoder)
    .hidden_size(768)
    .num_attention_heads(12)
    .num_hidden_layers(12)
    .intermediate_size(3072)
    .max_position_embeddings(512)
    .use_rotary_embeddings(true)
    .build();

let device = Device::Cpu;
let var_map = VarMap::new();
let vb = VarBuilder::from_varmap(&var_map, DType::F32, &device);
let model = Transformer::new(&config, vb)?;

let input_ids = Tensor::new(&[[1u32, 2, 3, 4]], &device)?;
let out = model.forward(&input_ids, None, None, None, None, None, false, None)?;
println!("{:?}", out.last_hidden_state.shape()); // [1, 4, 768]

Attention variants are config flags, not separate code paths you have to wire up:

// ALiBi linear-bias attention.
let config = TransformerConfig::builder().build().with_alibi(true, Some(8.0));
// Flash-attention code path.
let config = TransformerConfig::builder().use_flash_attention(true).build();
// Sliding-window attention.
let config = TransformerConfig::builder().build().with_sliding_window(256);

Turning that into an embedding

For retrieval, masked mean pooling over the last hidden state is the workhorse:

use candle_core::{DType, Tensor};
use anyhow::Result;

fn mean_pool(last_hidden: &Tensor, mask: &Tensor) -> Result<Tensor> {
    // mask: [B, T] with 1 for real tokens, 0 for padding.
    let mask = mask.to_dtype(DType::F32)?;                  // [B, T]
    let sum = (last_hidden * mask.unsqueeze(2)?)?.sum(1)?;  // [B, D]
    let count = mask.sum(1)?;                               // [B]
    Ok((sum / count.unsqueeze(1)?)?)                        // [B, D]
}

The KV cache

The decoder side ships a CausalLM wrapper with a language-model head and autoregressive generation: greedy, temperature, top-k and nucleus (top-p) sampling. But the engineering I cared about is the key/value cache. Do it naively and generating token n means running a forward pass over all n tokens, which makes producing a sequence \(O(T^2)\) work that re-derives the same keys and values over and over. The cache fixes that. Prefill the prompt once, store each layer's K and V, then every decode step feeds in only the single newest token and reads the rest back out of the cache.

The tests verify this isn't just faster but exact: KV-cached decoding reproduces a full forward pass token-for-token.

use rust_transformers::models::causal_lm::{CausalLM, GenerationConfig};

let lm = CausalLM::new(&config, vb)?;

// Greedy, deterministic.
let tokens = lm.generate(&[10, 42, 7], &GenerationConfig::greedy(20))?;

// Or sample with temperature / top-k / top-p.
let cfg = GenerationConfig { max_new_tokens: 20, temperature: 0.8,
    top_k: Some(50), top_p: Some(0.95), eos_token_id: Some(2), seed: 0 };
let tokens = lm.generate(&[10, 42, 7], &cfg)?;

Loading a real model

Everything above ran on random weights for a long time, and I defended that: random weights measure the architecture's raw cost, and the cost is the same whether the numbers mean anything. That was true, and also convenient, because a benchmark on random weights never checks whether the model works. So I loaded Qwen2.5-0.5B-Instruct, small enough to iterate on a CPU, real enough that a bug shows up as garbage text instead of a slightly different number.

My scaffolding could not load it. Three things were missing, and none of them appear in the tutorial picture of a transformer:

Grouped-query attention. Qwen2.5-0.5B has 14 query heads sharing 2 key/value heads. My attention assumed the counts match. GQA is how modern models shrink the KV cache, and my four attention variants didn't include the one every current checkpoint actually uses.
The gated MLP. I had six activation functions. The checkpoint wants none of them alone: it wants down(silu(gate(x)) * up(x)), three projections with a multiplicative gate. Configurable activations were the wrong axis of freedom.
The other RoPE. My rotary embeddings rotated interleaved pairs. Hugging Face checkpoints rotate the first half of each head against the second half. Both are correct implementations of the same paper, but the weights only work with the layout they were trained under, and the paper doesn't tell you which one you have.

So my generic config system could load zero of the one checkpoint I tried. I wrote a dedicated Qwen2 module instead, reusing what survived (RMSNorm, the KV cache pattern, the sampling loop) and building GQA, the gated MLP, and half-split RoPE from the spec.

Verifying it

The first generation came out coherent: asked for the capital of France, it said Paris. Then I ran three prompts against the Hugging Face implementation and two of the three diverged mid-sentence. My haiku had different lines, and for a while I assumed a bug.

The right check was never the text. On identical token ids, my implementation matches Hugging Face to a maximum logit difference of 0.000034 across all 151,936 vocabulary entries, top ten identical. The generations diverge because greedy decoding is chaotic: when two tokens sit within rounding error, whichever implementation's f32 noise lands on top wins, and the sentences never reconverge. Comparing text told me almost nothing. Comparing logits settled it.

How fast is it?

What these numbers are. Qwen2.5-0.5B-Instruct on CPU only, 16 logical cores, built with -C target-cpu=native. Decode is 64 KV-cached greedy tokens after a 32-token prompt, median of 5 with a warmup discarded. Prefill is one full pass over the prompt. The llama.cpp row is build b10520's CPU binary on the official Q8_0 GGUF of the same model, via llama-bench, same machine. Reproduce mine with cargo run --release --example qwen_bench.

Before measuring anything I wrote four predictions into the repo, because any number can be explained after you've seen it. Condensed: int8 buys 2.5x at decode, because decode is bound by streaming weights through memory. f16 loses to f32, because this CPU has no fast f16 arithmetic and the convert eats the bandwidth saving. int8 helps prefill much less than decode, because prefill is compute-bound. And llama.cpp stays 1.5-2.5x ahead of my int8, because its kernels have years on mine.

Weights	Decode tok/s	vs f32	Prefill 128 (ms)	Prefill 512 (ms)
f32 (mine)	14.0	1.00x	955	3449
f16 (mine)	17.6	1.26x	1338	4116
q8_0 (mine)	34.0	2.42x	2129	8846
llama.cpp Q8_0	52.9	3.77x	~308	~1285

Quality held at every precision: q8 produced token-identical greedy output to f32 on my test prompts, and f16 drifted 0.02 in logits with the same top five.

The prediction that hit, after trying to fail

int8 decode came in at 2.42x against a predicted 2.5x. But the first measurement was 0.26x. Four times slower than f32, with correct output. The cause: Candle's quantized kernels choose their SIMD path at compile time, and a default Rust build enables no AVX2, so the kernel that ran was scalar. The f32 path never showed the problem because its gemm library picks instructions at runtime. One build flag turned 3 tok/s into 34, and lifted f32 itself from 11.3 to 14.0 on the way. The int8 path was fine. The build had switched off its fast kernel.

The prediction that missed

I predicted f16 would lose, 0.4-0.8x. It won, 1.26x at decode. At batch size 1 the weights dwarf everything else, so halving their bytes won anyway; the conversion cost I bet on shows up at prefill, where f16 does lose (0.71x). My reasoning about f16 was right for prefill and wrong for decode. I'm glad the prediction is on the record, because after the fact I would absolutely have "known it all along", in whichever direction the number came out.

What llama.cpp is telling me

The other two predictions were right in sign and wrong in size, and both errors have the same shape. int8 helps prefill less than decode, and in fact makes it 2.2x slower, because Candle's quantized matmul loses to tuned f32 gemm the moment the matmul is big and compute-bound. And llama.cpp landed inside my predicted band at decode, 1.56x ahead of my q8, but is about 7x ahead at prefill, where its hand-written kernels do the same quantized math I do at a completely different speed.

My mental model priced the dtype and ignored the kernel. Whose matmul you run matters more than what dtype it runs in. That is what the benchmark section taught me. The gap between my code and llama.cpp comes from kernel quality, and now I have numbers for it.

The architecture cost, for the record

One chart from the random-weights era survives, because what it shows doesn't depend on the weights meaning anything. Mini model, batch 8. Doubling the sequence from 256 to 512 tokens more than doubles the latency. That bend is attention's quadratic term, and it is the reason sliding-window and flash-style attention exist.

2× the tokens, ~2.4× the time. Linear in tokens would predict 1708 ms at seq 512; the measured 2084 ms is the attention term showing up.

With the KV cache doing its job, decode is linear in length and the cost per token is just the model. Without it, each of those 64 tokens would re-run the full prompt forward, and none of the numbers above would be usable at all.

Project layout

src/
  attention/
    multi_head.rs       # Standard MHA (+ RoPE)
    flash.rs            # Memory-efficient, streamed-softmax attention
    alibi.rs            # Linear-bias attention for length extrapolation
    sliding_window.rs   # Local windows (+ optional global tokens)
  embeddings/
    token.rs            # Token (+ token-type) embeddings
    positional.rs       # Sinusoidal / learned
    rotary.rs           # RoPE
  models/
    transformer.rs      # Top-level encoder / decoder / enc-dec
    encoder.rs          # Encoder stack
    decoder.rs          # Decoder stack (self + cross attention)
    causal_lm.rs        # LM head + KV-cached generation / sampling
    qwen2.rs            # Real Qwen2.5 checkpoints: GQA, SwiGLU, f32/f16/q8_0
    layer.rs            # Encoder / decoder layers
  utils/
    activations.rs      # GELU / SiLU / Swish / Mish
    masking.rs          # Causal + padding masks
    tensor_ops.rs       # Hot-path tensor helpers
    tensor_ext.rs       # Tensor extension traits
    config.rs           # Architecture + hyperparameters
  main.rs               # Runnable forward-pass + generation example
examples/
  bench.rs              # Random-weight architecture benchmark
  qwen_generate.rs      # Chat with a real checkpoint at any precision
  qwen_bench.rs         # The decode/prefill benchmark above
  qwen_logits.rs        # Logit parity checks against Hugging Face
tests/
  forward.rs            # Forward-pass shape / smoke tests
  generation.rs         # KV-cache == full forward, sampling behavior
  correctness.rs        # Attention / RoPE / softmax / norm vs reference

Why Rust, in the end

One binary, anywhere. The reason the whole thing works. You can ship inference to an edge device or into existing Rust and Go infrastructure without a Python runtime tagging along.
Predictable performance. No GC pauses, explicit memory, and the same behaviour across environments instead of laptop-versus-prod surprises.
No language boundary. Embedding and ranking often live inside services that are already Rust. Deleting the FFI hop removes a whole class of bugs and a chunk of latency with it.
You see the whole stack. Every kernel is in the repo, so the benchmarks pointed at specific code to fix.

Things that bit me

Tokenizer lockstep. Keep the tokenizer JSON and the weights matched. Drift between them is the quietest way to get garbage embeddings.
Correctness first. The reference tests caught more than one subtle masking and RoPE bug. Numerical checks are cheap insurance when you're hand-writing kernels.
"RoPE" names two incompatible layouts. Interleaved pairs and half-split are both faithful to the paper, and a checkpoint only works with the one it was trained under. Nothing errors when you pick wrong; the model is just subtly broken.
SIMD can be a compile-time decision. Candle's quantized kernels pick their instruction set when you build, not when you run. Without -C target-cpu=native my int8 path ran scalar and lost to f32 by 4x while producing perfectly correct output. Slow and right is the hardest bug to notice.
Compare logits, not text. Two correct implementations of the same model will write different sentences, because greedy decoding amplifies rounding noise at near-ties. Text divergence told me nothing; a 3.4e-5 max logit diff settled it.

The next step is the kernel layer, where a hand-written matmul beats the whole precision strategy.