Transformers from Scratch in Rust

A from-scratch implementation of Transformer architectures on Candle: every attention variant, RoPE, RMSNorm, and a KV cache written by hand — plus honest CPU benchmarks that show where the time actually goes.

Repo: github.com/lblommesteyn/rust-transformers

The story: the model that ships as one file

Here is the moment that started this. You have a Rust service — an indexer, a log pipeline, a little API that already compiles to a single binary and runs anywhere. You want to add one transformer-shaped thing: semantic embeddings for search, or a small model that drafts text. The model is, underneath, a stack of matrix multiplies. It should be a function call.

Instead it asks you to adopt a second runtime. A Python interpreter, a virtualenv, a few hundred megabytes of framework, CUDA libraries that have opinions about your driver version, and a prayer that the versions that worked on your laptop also work on the box in production. The transformer is simple; the delivery is the hard part. And on an edge device, an air-gapped machine, or anything where "just run Python" isn't free, that delivery cost is the whole problem.

So I wrote the architecture directly in Rust on top of Candle, Hugging Face's minimalist tensor framework. The payoff is the thing you can hold: cargo build --release gives you one static binary that does the forward pass itself. No Python, no sidecar, no network, deterministic across environments. The transformer goes back to being a function call.

The second half of the story is the part I didn't expect to enjoy. Writing every piece by hand — each attention variant, the rotary embeddings, the norm layers, the KV cache — means there is no framework magic between the tokens going in and the vector coming out. When something is slow, you can see exactly which term is slow. When a workload is weird (very long log lines, sequences longer than anything you trained on), you can pick the attention mechanism that fits instead of taking whatever the library defaults to. The benchmarks near the end of this note exist because of that: they aren't marketing numbers, they're a map of where the time goes.

TL;DR

• 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 — selected by config, not by swapping libraries.
• RoPE / sinusoidal / learned position encodings; LayerNorm or RMSNorm; GELU / SiLU / Swish / Mish feed-forwards.
• A tests/ suite that checks attention against an independent reference and verifies KV-cached generation reproduces a full forward pass exactly.
• Real CPU benchmarks below, produced by cargo run --release --example bench.

The architecture, end to end

Tokens come in, a vector (or a distribution over the vocabulary) comes out. In between is the same shape every transformer shares: embed, then a stack of identical layers that each do attention and a feed-forward, with a residual add and a normalization around each. The two output heads are what make it an embedding model or a generator.

Inside a 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 HBM hot path, ALiBi linear biases for length extrapolation, 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 for the decoder; sinusoidal and learned tables are also there. ALiBi biases attention scores directly with \(b_{ij} = m_h \cdot (j - i)\), a per-head slope rather than an added position vector.

Normalization. LayerNorm is available, but RMSNorm is the one 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 or attention pooling are selectable when the model was trained for them.

The attention variants, as masks

The cheapest way to see what these mechanisms actually change is to look at which query-key pairs each one lets talk to each other. Full attention is dense; sliding-window keeps a band around the diagonal; ALiBi attends everywhere but tilts the scores by distance.

Building a model

Configuration is a fluent builder. Pick a model type and the shape, flip on the features you want, and back it with a Candle variable store (here random init; 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 different code paths 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);
    

Pooling to 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]
}
    

Generation: why the KV cache earns its keep

The decoder side ships a CausalLM wrapper with a language-model head and autoregressive generation — greedy, temperature, top-k, and nucleus (top-p) sampling. The interesting engineering is the key/value cache. Naively, generating token n means running a forward pass over all n tokens, so producing a sequence is \(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 from 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)?;
    

Benchmarks

What these are. CPU-only (Candle Device::Cpu), f32, single process, random weights, 16 logical cores. No GPU, no quantization, no fused kernels. They measure the architecture's raw cost — a baseline to optimize against, not a tuned serving stack. Reproduce with cargo run --release --example bench; warmup ×3, median of 15.

Forward latency vs sequence length

This is the chart I care about most, because it shows the quadratic term in attention with your own eyes. Mini model, batch 8: doubling the sequence from 256 to 512 tokens more than doubles latency. That bend is exactly why sliding-window and flash-style attention exist.

Generation throughput

With the KV cache, even on CPU, the mini model decodes at a usable clip; the 110M base model is roughly 4× slower, tracking its parameter count.

Full grid

What the numbers say. Three things fall out. (1) Attention is super-linear in sequence length — the seq 256→512 jump is the clearest tell. (2) Model size dominates everything else: base is ~4× mini across the board, roughly its FLOP ratio. (3) The KV cache is what makes generation tractable at all; without it each of those 64 tokens would re-run the full prompt forward.

And what they don't. These are deliberately un-optimized. The obvious next moves — an int8 weight path, a GPU Device, length-bucketed batching to cut padding, and leaning on the flash-style kernel for the long-sequence cases — are exactly the levers the chart above points at. That's the whole argument for building it by hand: the benchmark doesn't just grade the code, it tells you where to dig.

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
    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              # The benchmark harness above
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, and why not just Python

• One binary, anywhere. The reason the story works: ship inference to edge or existing Rust/Go infra without a Python runtime tagging along.
• Predictable performance. No GC pauses, explicit memory, and the same behavior across environments instead of laptop-vs-prod surprises.
• No language boundary. Embedding and ranking often live inside services that are already Rust; deleting the FFI hop removes a class of bugs and a chunk of latency.
• You see the whole stack. Every kernel is in front of you, which is how the benchmarks turned into a to-do list instead of a mystery.

Notes & gotchas

• 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 for hand-written kernels.
• CPU f32 is the floor, not the ceiling. Read the benchmarks as "what does the architecture cost before any optimization," then pick the lever that matches your workload.