Rust Candle LLM Inference: 3x Faster Than Python PyTorch

rustc --version
# rustc 1.77.0 (aedd173a2 2024-03-17) or newer

cargo new candle-inference --bin
cd candle-inference

[package]
name = "candle-inference"
version = "0.1.0"
edition = "2021"

[dependencies]
# Core tensor engine — CPU by default
candle-core = { version = "0.8", features = [] }

# Pre-built model architectures (Phi, Llama, Mistral, Gemma, etc.)
candle-transformers = { version = "0.8" }

# Hugging Face Hub downloader + tokenizer bindings
candle-examples = { version = "0.8" }
hf-hub = { version = "0.3", features = ["tokio"] }
tokenizers = { version = "0.19", default-features = false, features = ["onig"] }

# Async runtime and CLI arg parsing
tokio = { version = "1", features = ["full"] }
anyhow = "1"
clap = { version = "4", features = ["derive"] }

[features]
default = []
# Uncomment to enable CUDA — requires CUDA 12.x toolkit installed
# cuda = ["candle-core/cuda", "candle-transformers/cuda"]
# Metal support for Apple Silicon
# metal = ["candle-core/metal", "candle-transformers/metal"]

# Install the HF CLI if you don't have it
pip install huggingface_hub --break-system-packages

# Download Phi-3 Mini GGUF
huggingface-cli download \
  microsoft/Phi-3-mini-4k-instruct-gguf \
  Phi-3-mini-4k-instruct-q4.gguf \
  --local-dir ./models/phi3

ls -lh models/phi3/
# -rw-r--r-- 1 user user 2.2G Phi-3-mini-4k-instruct-q4.gguf

use anyhow::{Error as E, Result};
use candle_core::{Device, Tensor};
use candle_transformers::models::quantized_phi3::ModelWeights;
use candle_transformers::generation::LogitsProcessor;
use clap::Parser;
use hf_hub::{api::sync::Api, Repo, RepoType};
use std::path::PathBuf;
use tokenizers::Tokenizer;

/// CLI args — add --cuda flag to switch device at runtime
#[derive(Parser, Debug)]
struct Args {
    #[arg(long, default_value = "models/phi3/Phi-3-mini-4k-instruct-q4.gguf")]
    model: PathBuf,

    #[arg(long, default_value = "microsoft/Phi-3-mini-4k-instruct")]
    tokenizer_repo: String,

    #[arg(long, default_value = "Explain the borrow checker in one paragraph.")]
    prompt: String,

    /// Max tokens to generate
    #[arg(long, default_value_t = 200)]
    max_tokens: usize,

    /// Use CUDA if available (requires cuda feature flag)
    #[arg(long)]
    cuda: bool,
}

fn main() -> Result<()> {
    let args = Args::parse();

    // Device selection: CPU is default; CUDA requires the feature flag + hardware
    let device = if args.cuda {
        Device::new_cuda(0).map_err(|e| E::msg(format!("CUDA init failed: {e}")))?
    } else {
        Device::Cpu
    };

    println!("Device: {:?}", device);

    // --- Load tokenizer from HF Hub ---
    // Candle doesn't bundle a Python tokenizer — we pull the JSON config directly
    let api = Api::new()?;
    let repo = api.repo(Repo::new(args.tokenizer_repo.clone(), RepoType::Model));
    let tokenizer_path = repo.get("tokenizer.json")?;
    let tokenizer = Tokenizer::from_file(tokenizer_path).map_err(E::msg)?;

    // --- Load GGUF weights ---
    // GGUF is a binary format that packs quantized weights + metadata in one file.
    // candle-core reads it natively — no conversion needed.
    let start = std::time::Instant::now();
    let mut file = std::fs::File::open(&args.model)?;
    let model_content = candle_core::quantized::gguf_file::Content::read(&mut file)?;
    let mut model = ModelWeights::from_gguf(model_content, &mut file, &device)?;
    println!("Model loaded in {:.2}s", start.elapsed().as_secs_f32());

    // --- Tokenize prompt ---
    let tokens = tokenizer
        .encode(args.prompt.as_str(), true)
        .map_err(E::msg)?;
    let token_ids = tokens.get_ids();
    println!("Prompt tokens: {}", token_ids.len());

    // --- Generation loop ---
    // LogitsProcessor handles temperature + top-p sampling.
    // We use greedy (temp=0) for deterministic benchmark results.
    let mut logits_processor = LogitsProcessor::new(
        42,       // seed — fixed for reproducibility
        Some(0.0), // temperature: 0.0 = greedy, 1.0 = full sampling
        None,      // top_p: None = disabled
    );

    let mut all_tokens: Vec<u32> = token_ids.to_vec();
    let mut generated = 0usize;
    let gen_start = std::time::Instant::now();

    for index in 0..args.max_tokens {
        // Feed only the last token after the first forward pass
        // This is the KV-cache pattern: past context is cached inside the model
        let input = if index == 0 {
            Tensor::new(all_tokens.as_slice(), &device)?.unsqueeze(0)?
        } else {
            let last = *all_tokens.last().unwrap();
            Tensor::new(&[last], &device)?.unsqueeze(0)?
        };

        let logits = model.forward(&input, index)?;
        // Squeeze batch dimension, grab last token's logits
        let logits = logits.squeeze(0)?.squeeze(0)?;

        let next_token = logits_processor.sample(&logits)?;
        all_tokens.push(next_token);
        generated += 1;

        // Phi-3's EOS token ID is 32007 — stop when we hit it
        if next_token == 32007 {
            break;
        }
    }

    let elapsed = gen_start.elapsed().as_secs_f32();
    let tps = generated as f32 / elapsed;

    // Decode only the generated portion (skip the prompt tokens)
    let output = tokenizer
        .decode(&all_tokens[token_ids.len()..], true)
        .map_err(E::msg)?;

    println!("\n--- Output ---\n{output}");
    println!("\n--- Stats ---");
    println!("Generated: {generated} tokens in {elapsed:.2}s ({tps:.1} tok/s)");

    Ok(())
}

# CPU build (default)
cargo build --release

# Run inference
./target/release/candle-inference \
  --prompt "Explain the Rust borrow checker in one paragraph."

# In Cargo.toml, change the cuda feature line:
[features]
cuda = ["candle-core/cuda", "candle-transformers/cuda"]

# Build with CUDA enabled
cargo build --release --features cuda

# Run on GPU 0
./target/release/candle-inference \
  --cuda \
  --prompt "Explain the Rust borrow checker in one paragraph."

# Cargo.toml
[features]
metal = ["candle-core/metal", "candle-transformers/metal"]

cargo build --release --features metal

./target/release/candle-inference \
  --metal \
  --prompt "Explain the Rust borrow checker in one paragraph."

# benchmark_pytorch.py
import torch
import time
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "microsoft/Phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Use float16 to match production serving — bfloat16 on Ampere+
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.float16,
    device_map="cpu",  # cpu to match Candle CPU baseline
)

prompt = "Explain the Rust borrow checker in one paragraph."
inputs = tokenizer(prompt, return_tensors="pt")

start = time.perf_counter()
with torch.no_grad():
    output = model.generate(
        **inputs,
        max_new_tokens=200,
        do_sample=False,  # greedy, same as Candle
    )
elapsed = time.perf_counter() - start

generated = output.shape[1] - inputs["input_ids"].shape[1]
print(f"Generated: {generated} tokens in {elapsed:.2f}s ({generated/elapsed:.1f} tok/s)")

pip install transformers torch --break-system-packages
python benchmark_pytorch.py

// Add to Cargo.toml:
// axum = "0.7"
// serde = { version = "1", features = ["derive"] }
// serde_json = "1"

use axum::{extract::State, routing::post, Json, Router};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
use tokio::sync::Mutex;

#[derive(Deserialize)]
struct InferRequest {
    prompt: String,
    max_tokens: Option<usize>,
}

#[derive(Serialize)]
struct InferResponse {
    output: String,
    tokens_per_sec: f32,
}

// Wrap the model in Arc<Mutex<>> for safe concurrent access
// For high throughput, use a worker pool + channel instead
type ModelState = Arc<Mutex<ModelWeights>>;

async fn infer(
    State(model): State<ModelState>,
    Json(req): Json<InferRequest>,
) -> Json<InferResponse> {
    let mut m = model.lock().await;
    // ... call your inference logic here ...
    Json(InferResponse {
        output: "...".to_string(),
        tokens_per_sec: 29.4,
    })
}

#[tokio::main]
async fn main() {
    let model = Arc::new(Mutex::new(load_model()));
    let app = Router::new()
        .route("/infer", post(infer))
        .with_state(model);

    let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();
    axum::serve(listener, app).await.unwrap();
}

# CPU
time ./target/release/candle-inference --max-tokens 200

# CUDA
time ./target/release/candle-inference --cuda --max-tokens 200

# Linux: watch RSS during inference
/usr/bin/time -v ./target/release/candle-inference --max-tokens 200 2>&1 | grep "Maximum resident"

# macOS
/usr/bin/time -l ./target/release/candle-inference --max-tokens 200 2>&1 | grep "maximum resident"