candle-mi

Mechanistic Interpretability for the Rust of us.

Note: v0.1.1 — the API may change between minor versions. See the CHANGELOG.

Supported model families

Most HuggingFace transformer models work out of the box via auto-config — no code changes needed. See BACKENDS.md for details and how to add new architectures.

Hardware: candle-mi runs on a single consumer GPU (developed on an RTX 5060 Ti, 16 GB VRAM). Models up to ~7B fit in 16 GB at F32 precision — no H100 cluster required. CPU-only works for small models and tokenizer-only workflows.

Table of Contents

• What is this?
• What can you do with it?
• See it in action
• Quick start
• Design philosophy
• Paper replications
• Feature flags
• Documentation
• License
• Development

What is this?

Mechanistic interpretability (MI) is the study of how a language model arrives at its predictions — not just what it outputs, but what happens inside. By inspecting and manipulating the model's internal activations (attention patterns, residual streams, MLP outputs), researchers can understand which components drive specific behaviors.

candle-mi is a Rust library that makes this possible. It re-implements model forward passes with built-in hook points — type-safe, named locations in the computation graph (e.g., HookPoint::AttnPattern(5) for the post-softmax attention at layer 5) where you can:

• Capture activations (e.g., "what does the attention pattern look like at layer 5?")
• Intervene on activations mid-forward-pass (e.g., "what happens if I knock out this attention edge?" or "what if I steer the residual stream toward a concept?")

This is the Rust equivalent of Python's TransformerLens, built on candle for GPU acceleration. The hook system is type-safe (typos are caught at compile time, not silently ignored at runtime) and zero-overhead (an empty hook spec adds no allocations or clones to the forward pass).

Why Rust? Running published MI experiments — such as Anthropic's Scaling Monosemanticity or Planning in poems — quickly hits the limits of CPU-only Python. Cloud GPUs are always an option, but not a frugal one. With a consumer-grade GPU, memory and runtime become the real bottleneck. candle solves both: Rust's zero-cost abstractions minimize memory overhead, compiled code runs faster, and candle provides direct CUDA/Metal access without Python's runtime tax. That's how candle-mi started: let's bring MI to local hardware. (See the Figure 13 replication for a concrete example — Planning in poems reproduced on a consumer GPU.)

What can you do with it?

candle-mi is (to our knowledge) the only MI toolkit with hook points for recurrent architectures — RwkvState, RwkvDecay, and RwkvEffectiveAttn enable mechanistic analysis of RWKV-6/7 models, a frontier that most MI tooling ignores entirely.

See it in action

The logit lens — what does the model "think" at each layer?

cargo run --release --features transformer --example logit_lens -- "meta-llama/Llama-3.2-1B"

This loads in ~2 seconds and runs in ~112ms on an RTX 5060 Ti (16 GB VRAM) or ~3 seconds on CPU, revealing how factual recall emerges across layers. The same prompt ("The capital of France is") tells three different stories:

• Llama 3.2 1B: "Paris" appears at layer 11 (69% depth) — typical early factual resolution.
• Gemma 2 2B: "Paris" appears at layer 25 (the very last layer, rank 8) — the model hedges until the end.
• StarCoder2 3B: "Paris" never appears as a single token — its BPE tokenizer splits it into " Par", which dominates from layer 22 at 33%→74%. The model knows the answer but its code-oriented vocabulary hides it.

See the full comparison with per-layer tables.

The flagship — Anthropic's circuit tracing on consumer hardware

This library was built to replicate Anthropic's circuit-tracing work on consumer hardware. Here is Figure 13 from "On the Biology of a Large Language Model", running on a single GPU:

cargo run --release --features clt,transformer --example figure13_planning_poems

This uses Llama 3.2 1B with a 524K-feature Cross-Layer Transcoder to suppress natural rhyme features and inject an alternative ("that" → P=0.98), sweeping injection position across all prompt tokens. See the full output and the 15 examples covering logit lens, attention knockout, steering, activation patching, CLT circuits, SAE encoding, RWKV inference, and more.

Quick start

use candle_mi::{HookSpec, MIModel};

fn main() -> candle_mi::Result<()> {
    // 1. Load a model (auto-detects architecture from HuggingFace config)
    let model = MIModel::from_pretrained("meta-llama/Llama-3.2-1B")?;
    let tokenizer = model.tokenizer().unwrap();

    // 2. Tokenize a prompt
    let tokens = tokenizer.encode("The capital of France is")?;
    let input = candle_core::Tensor::new(&tokens[..], model.device())?.unsqueeze(0)?;

    // 3. Run a forward pass (HookSpec::new() = no hooks, zero overhead)
    let cache = model.forward(&input, &HookSpec::new())?;
    let logits = cache.output();  // [1, seq_len, vocab_size]

    // 4. Decode the top prediction
    let last_logits = logits.get(0)?.get(tokens.len() - 1)?;
    let token_id = candle_mi::sample_token(&last_logits, 0.0)?;  // greedy
    println!("{}", tokenizer.decode(&[token_id])?);  // " Paris"
    Ok(())
}

With hooks — capture attention patterns

use candle_mi::{HookPoint, HookSpec, MIModel};

fn main() -> candle_mi::Result<()> {
    let model = MIModel::from_pretrained("meta-llama/Llama-3.2-1B")?;
    let tokenizer = model.tokenizer().unwrap();

    let tokens = tokenizer.encode("The capital of France is")?;
    let input = candle_core::Tensor::new(&tokens[..], model.device())?.unsqueeze(0)?;

    // Capture the post-softmax attention pattern at layer 5
    let mut hooks = HookSpec::new();
    hooks.capture(HookPoint::AttnPattern(5));

    let cache = model.forward(&input, &hooks)?;

    // Retrieve the captured tensor — [1, heads, seq, seq]
    let attn = cache.require(&HookPoint::AttnPattern(5))?;
    println!("Layer 5 attention shape: {:?}", attn.shape());
    Ok(())
}

Here is what an end-to-end run looks like (auto-config loading LLaMA 3.2 1B — config detection, forward pass, and top-5 predictions):

Design philosophy

candle-mi makes a deliberate trade-off: full-sequence recompute at every generation step (no KV cache). This is slower than production inference engines, but it means:

• Maximum observability. Hooks can re-observe how earlier positions change under intervention at every step. Interventions "just work" without KV cache invalidation.
• Interventions compound. When steering the residual stream during autoregressive generation, each new token is generated with the intervention re-applied across the full context. This is why candle-mi's recurrent feedback rescues +2 rhyming couplets (out of 15) where a KV-cached approach gets +1 — the intervention is observed at every step, not just once during prefill.

This is a research-first design: MI analyses need to see everything, and the performance cost is acceptable when the alternative is missing causal effects. candle-mi is not an inference engine — for production serving, see candle-vllm, vllm.rs, or vLLM (Python). It is optimized for observability, not throughput.

Paper replications

Feature flags

Documentation

License

MIT OR Apache-2.0

Development

Exclusively developed with Claude Code (dev) and Augment Code (review). Git workflow managed with Fork.