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 feeling of disproportion. 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, or an air-gapped machine, or anywhere "just run Python" isn't free, delivery is the entire problem.

So I wrote the architecture directly in Rust, on top of Candle, which is Hugging Face's minimalist tensor framework. The payoff is something you can hold: cargo build --release gives you one static binary that runs the forward pass itself. No Python, no sidecar, no network calls, and it behaves the same everywhere. The transformer goes back to being a function call.

The part I didn't expect to enjoy was everything after that. When you write every piece yourself, there's no framework magic sitting 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. They aren't marketing numbers. They're a map of 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.

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.

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.

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. True, and also a comfortable place to hide, because a benchmark on random weights can't be wrong in the way that matters. Nothing 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 the honest score for my generic config system was zero out of one checkpoints loadable. 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. Same model, same math, different haiku. Text comparison tells you almost nothing. Logit comparison settles 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.

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 lever was always going to work. It was compiled not to.

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 model of f16 wasn't wrong so much as applied to the wrong regime. 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 doesn't just help prefill less, it makes prefill 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 sentence is the paycheck for the whole benchmark section: the gap between me and llama.cpp isn't architecture, it's kernel quality, and now I know that with numbers instead of folklore.

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, visible with your own eyes, and it's the entire reason sliding-window and flash-style attention exist.

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 right in front of you, which is how the benchmarks turned into a to-do list instead of a mystery.

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.

Why I'd do it again

I went into this for the deployment story, one binary, no Python, and that part delivered: the binary now loads a real instruction-tuned model and chats at 34 tokens a second on a CPU. But it isn’t why I’d do it again.

I’d do it again because of what happened to the word “transformer” in my head. Before this project it was a black box I had opinions about. Now it’s a specific pile of matrix multiplies whose costs I’ve personally paid, one bend in a latency chart at a time. When I read a paper now and it says grouped-query attention, I don’t nod along. I remember the afternoon my attention couldn’t load a checkpoint because of it.

Writing the predictions down before measuring changed what the project was, too. Without that file, the benchmark section would be four numbers and a story that fits them, and neither you nor I could tell my understanding from my hindsight. With it, the record shows what I actually knew: the bandwidth story right, the kernel layer invisible to me until it cost a 4x. I know exactly where my model of this system fails now, and it fails below the dtype, at the level where someone’s hand-written matmul quietly beats my whole precision strategy.

So the binary was the excuse, and the kernels are the next dig site. The real product is that I’m no longer afraid of what’s inside.