diff --git a/Cargo.toml b/Cargo.toml index 3347ffb..4a5410f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,5 +7,23 @@ description = "A from-scratch LLM inference engine in Rust β€” run a real LLM on license = "MIT" repository = "https://github.com/codewithfourtix/ember" readme = "README.md" +keywords = ["llm", "inference", "transformer", "quantization", "cpu"] +categories = ["command-line-utilities", "science"] [dependencies] +# Plumbing only β€” the transformer math is hand-written. +anyhow = "1" # ergonomic error handling +clap = { version = "4", features = ["derive"] } # CLI parsing +half = "2" # f16 weight support +memmap2 = "0.9" # mmap the weights file +rayon = "1" # parallelise the mat-vec hot loop +safetensors = "0.4" # load model weights +serde = { version = "1", features = ["derive"] } # config.json deserialization +serde_json = "1" +tokenizers = "0.20" # BPE tokenizer (tokenizer.json) + +[profile.release] +opt-level = 3 +lto = true +codegen-units = 1 +panic = "abort" diff --git a/README.md b/README.md index 7155795..94e98d3 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,81 @@ -# ember +
+ +# πŸ”₯ ember **Run a real LLM on your CPU β€” a from-scratch inference engine in Rust.** -Status: early scaffolding. +[![Rust](https://img.shields.io/badge/Rust-2021-CE422B?style=flat-square&logo=rust&logoColor=white&labelColor=0d0e11)](https://www.rust-lang.org/) +[![License: MIT](https://img.shields.io/badge/license-MIT-blue?style=flat-square&labelColor=0d0e11)](LICENSE) +[![status](https://img.shields.io/badge/status-scaffolding-f0ad4e?style=flat-square&labelColor=0d0e11)](#roadmap) + +
+ +`ember` loads a **Qwen2.5** model and generates text with a hand-written transformer +forward pass β€” RMSNorm, RoPE, grouped-query attention with a KV cache, and a SwiGLU +MLP β€” then runs the heavy matrices through custom **INT8/INT4 quantization** so a +0.5–1.5B model fits in a laptop's memory and decodes fast. **No `candle`, `burn`, +`tch`, or `ndarray` in the core:** the transformer math is the project. + +> **Status β€” scaffolding.** The architecture, module layout, and CLI are in place; +> the numeric kernels are marked `todo!()` and are being implemented (see the [roadmap](#roadmap)). + +## Why build this + +Writing an inference engine is the clearest way to understand β€” and to demonstrate +understanding of β€” how modern LLMs actually run: attention, the KV cache, the +memory-bandwidth wall of CPU decoding, and quantization. The only dependencies here +are for plumbing (weight loading, tokenization, threading); every kernel is hand-written. + +## Design + +| Module | Responsibility | +|---|---| +| [`config.rs`](src/config.rs) | Parse the model's `config.json` (Qwen2.5 / Llama-style). | +| [`tensor.rs`](src/tensor.rs) | The hot loop β€” hand-written, `rayon`-parallel mat-vec. | +| [`ops.rs`](src/ops.rs) | RMSNorm, RoPE, SwiGLU, softmax. | +| [`attention.rs`](src/attention.rs) | Grouped-query attention + the rolling KV cache. | +| [`quant.rs`](src/quant.rs) | Row-wise INT8/INT4 quantization + fused dequant mat-vec. | +| [`sample.rs`](src/sample.rs) | Greedy / temperature / top-p sampling. | +| [`model.rs`](src/model.rs) | safetensors weight loading + the full forward pass. | +| [`main.rs`](src/main.rs) | CLI and the generation loop. | + +## Quickstart + +```bash +# 1. Rust toolchain (https://rustup.rs) +rustup default stable + +# 2. Fetch a small model (needs: pip install huggingface_hub) +huggingface-cli download Qwen/Qwen2.5-0.5B-Instruct \ + model.safetensors config.json tokenizer.json --local-dir ./weights + +# 3. Build & run +cargo run --release -- --prompt "The capital of France is" --model ./weights +``` + +## Correctness oracle + +Before trusting the generation loop, match a **single forward pass** against +HuggingFace `transformers`: + +```bash +pip install torch transformers +python scripts/reference_logits.py --model Qwen/Qwen2.5-0.5B-Instruct +``` + +`ember`'s next-token logits for the same prompt should agree to ~`1e-3`. Reach parity +here first and every remaining bug is in the loop, not the model. + +## Roadmap + +- [ ] **Day 1** β€” weight loading (safetensors) + tokenizer + embedding β†’ LM head +- [ ] **Day 2** β€” RMSNorm, RoPE, attention, SwiGLU; logit parity vs `transformers` +- [ ] **Day 3** β€” generation loop + sampling β†’ coherent text +- [ ] **Day 4** β€” KV cache; benchmark tokens/sec +- [ ] **Day 5** β€” INT8/INT4 quantization; benchmark memory + speed +- [ ] **Day 6** β€” streaming CLI, benchmark table, demo +- [ ] **Day 7** β€” write-up + +## License + +MIT Β© [Ali Zulfiqar](https://github.com/codewithfourtix) diff --git a/scripts/reference_logits.py b/scripts/reference_logits.py new file mode 100644 index 0000000..d63bf4e --- /dev/null +++ b/scripts/reference_logits.py @@ -0,0 +1,44 @@ +#!/usr/bin/env python3 +"""Print the next-token logits for a fixed prompt, using HuggingFace transformers. + +This is ember's correctness oracle. Run it, then compare against ember's own +logits for the same prompt and token ids β€” they should agree to ~1e-3. Reach +parity here before trusting the generation loop: after that, every bug is in the +loop, not the model. + + pip install torch transformers + python scripts/reference_logits.py --model Qwen/Qwen2.5-0.5B-Instruct +""" + +import argparse + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--model", default="Qwen/Qwen2.5-0.5B-Instruct") + ap.add_argument("--prompt", default="The capital of France is") + ap.add_argument("--topk", type=int, default=10) + args = ap.parse_args() + + tok = AutoTokenizer.from_pretrained(args.model) + model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.float32) + model.eval() + + ids = tok(args.prompt, return_tensors="pt").input_ids + with torch.no_grad(): + logits = model(ids).logits[0, -1] # logits for the *next* token + + print(f"prompt : {args.prompt!r}") + print(f"input ids: {ids[0].tolist()}") + print(f"logits[:5]: {[round(v, 4) for v in logits[:5].tolist()]}") + print(f"top-{args.topk} next tokens:") + top = torch.topk(logits, args.topk) + for score, idx in zip(top.values.tolist(), top.indices.tolist()): + print(f" {idx:>6} {score:+8.4f} {tok.decode([idx])!r}") + + +if __name__ == "__main__": + main() diff --git a/src/attention.rs b/src/attention.rs new file mode 100644 index 0000000..d90aecc --- /dev/null +++ b/src/attention.rs @@ -0,0 +1,101 @@ +//! Grouped-query self-attention with a KV cache. +//! +//! At decode time each new token attends over every previous token, so the keys +//! and values for the whole sequence are cached and only the *new* token's +//! Q/K/V are computed each step β€” the difference between O(nΒ²) and O(n) work per +//! token. Qwen2.5 uses grouped-query attention: several query heads share one +//! key/value head, so the cache stores only `num_key_value_heads` streams. + +use crate::config::Config; + +/// Per-layer rolling cache of past keys and values. +/// +/// Each layer's buffer is laid out as `[pos * kv_dim + i]`, where +/// `kv_dim = num_key_value_heads * head_dim`. +pub struct KvCache { + keys: Vec>, + values: Vec>, + kv_dim: usize, + max_seq: usize, + len: usize, +} + +impl KvCache { + /// Allocate a cache sized for `config` and a `max_seq`-token context. + pub fn new(config: &Config, max_seq: usize) -> Self { + let kv_dim = config.kv_dim(); + let layers = config.num_hidden_layers; + Self { + keys: (0..layers).map(|_| vec![0.0; max_seq * kv_dim]).collect(), + values: (0..layers).map(|_| vec![0.0; max_seq * kv_dim]).collect(), + kv_dim, + max_seq, + len: 0, + } + } + + /// Number of tokens currently cached. + pub fn len(&self) -> usize { + self.len + } + + /// Whether the cache holds no tokens yet. + pub fn is_empty(&self) -> bool { + self.len == 0 + } + + /// The key stream for `layer` (`len * kv_dim` valid entries). + pub fn keys(&self, layer: usize) -> &[f32] { + &self.keys[layer] + } + + /// The value stream for `layer`. + pub fn values(&self, layer: usize) -> &[f32] { + &self.values[layer] + } + + /// Write this step's key/value for `layer` at the current position. + pub fn store(&mut self, layer: usize, key: &[f32], value: &[f32]) { + debug_assert_eq!(key.len(), self.kv_dim); + debug_assert_eq!(value.len(), self.kv_dim); + let off = self.len * self.kv_dim; + self.keys[layer][off..off + self.kv_dim].copy_from_slice(key); + self.values[layer][off..off + self.kv_dim].copy_from_slice(value); + } + + /// Advance the write cursor by one token, after every layer has stored its + /// key/value for this step. + pub fn advance(&mut self) { + debug_assert!(self.len < self.max_seq, "KV cache overflow"); + self.len += 1; + } + + /// Reset for a fresh sequence without reallocating. + pub fn clear(&mut self) { + self.len = 0; + } +} + +/// Compute self-attention for one token at `pos` in a single layer. +/// +/// `q`/`k`/`v` are this token's projections; `k`/`v` are stored into `cache`, +/// then the token attends over positions `0..=pos` and the weighted values are +/// written into `out`. +#[allow(clippy::too_many_arguments)] +pub fn attention( + q: &[f32], + k: &[f32], + v: &[f32], + cache: &mut KvCache, + layer: usize, + pos: usize, + config: &Config, + out: &mut [f32], +) { + // 1. cache.store(layer, k, v) + // 2. for each query head h (mapped to kv head h / group_size): + // score[t] = (q_h Β· k_t) / sqrt(head_dim) for t in 0..=pos + // softmax(score); out_h = Ξ£_t score[t] * v_t + let _ = (q, k, v, cache, layer, pos, config, out); + todo!("grouped-query scaled-dot-product attention over the KV cache") +} diff --git a/src/config.rs b/src/config.rs new file mode 100644 index 0000000..7ccc25c --- /dev/null +++ b/src/config.rs @@ -0,0 +1,70 @@ +//! Model hyper-parameters, parsed from a HuggingFace `config.json`. +//! +//! Only the fields the inference path actually needs are read; everything else +//! in the file is ignored. Defaults match the Qwen2.5 family. + +use std::path::Path; + +use anyhow::{Context, Result}; +use serde::Deserialize; + +/// The subset of `config.json` needed to run a Llama-style decoder. +#[derive(Debug, Clone, Deserialize)] +pub struct Config { + /// Model (embedding) dimension. + pub hidden_size: usize, + /// Number of stacked transformer blocks. + pub num_hidden_layers: usize, + /// Number of query heads. + pub num_attention_heads: usize, + /// Number of key/value heads (`< num_attention_heads` β‡’ grouped-query attention). + pub num_key_value_heads: usize, + /// Hidden dimension of the SwiGLU feed-forward network. + pub intermediate_size: usize, + /// Token vocabulary size. + pub vocab_size: usize, + /// Context length the RoPE tables are built for. + pub max_position_embeddings: usize, + /// RoPE base frequency (ΞΈ). + #[serde(default = "default_rope_theta")] + pub rope_theta: f32, + /// RMSNorm epsilon. + #[serde(default = "default_rms_norm_eps")] + pub rms_norm_eps: f32, + /// Whether the LM head shares weights with the input embedding. + #[serde(default)] + pub tie_word_embeddings: bool, +} + +fn default_rope_theta() -> f32 { + 1_000_000.0 +} + +fn default_rms_norm_eps() -> f32 { + 1e-6 +} + +impl Config { + /// Load and parse a `config.json` from disk. + pub fn from_file(path: impl AsRef) -> Result { + let path = path.as_ref(); + let text = std::fs::read_to_string(path) + .with_context(|| format!("reading config from {}", path.display()))?; + serde_json::from_str(&text).context("parsing config.json") + } + + /// Dimension of a single attention head. + pub fn head_dim(&self) -> usize { + self.hidden_size / self.num_attention_heads + } + + /// How many query heads share each key/value head (the GQA group size). + pub fn kv_group_size(&self) -> usize { + self.num_attention_heads / self.num_key_value_heads + } + + /// Combined width of the key (or value) projection: `num_key_value_heads * head_dim`. + pub fn kv_dim(&self) -> usize { + self.num_key_value_heads * self.head_dim() + } +} diff --git a/src/main.rs b/src/main.rs index ac1d085..70c5bf2 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,3 +1,85 @@ -fn main() { - println!("ember β€” a from-scratch LLM inference engine (scaffolding)"); +//! `ember` β€” a from-scratch LLM inference engine. +//! +//! Loads a Qwen2.5 model and generates text on the CPU. The heavy lifting lives +//! in the sibling modules; this file wires the CLI to the generation loop. +//! +//! The numeric kernels are still `todo!()`, so running the binary today parses +//! args and loads the config, then panics at the first kernel β€” that panic is +//! the map of what to implement next (see the roadmap in the README). + +// Kernels are wired but not yet called from every module during scaffolding. +#![allow(dead_code)] + +mod attention; +mod config; +mod model; +mod ops; +mod quant; +mod sample; +mod tensor; + +use std::path::PathBuf; + +use anyhow::{Context, Result}; +use clap::Parser; + +use model::Model; +use sample::Sampler; + +/// Run a local LLM on your CPU. +#[derive(Parser, Debug)] +#[command(name = "ember", version, about)] +struct Args { + /// Prompt to complete. + #[arg(short, long)] + prompt: String, + + /// Directory with `config.json`, `model.safetensors`, and `tokenizer.json`. + #[arg(short, long, default_value = "weights")] + model: PathBuf, + + /// Maximum number of tokens to generate. + #[arg(short = 'n', long, default_value_t = 128)] + max_tokens: usize, + + /// Sampling temperature (0 β‡’ greedy). + #[arg(short, long, default_value_t = 0.0)] + temperature: f32, + + /// Nucleus (top-p) cutoff, used when `temperature > 0`. + #[arg(long, default_value_t = 0.95)] + top_p: f32, +} + +fn main() -> Result<()> { + let args = Args::parse(); + + let model = Model::load(&args.model) + .with_context(|| format!("loading model from {}", args.model.display()))?; + + let sampler = if args.temperature <= 0.0 { + Sampler::Greedy + } else { + Sampler::TopP { + temperature: args.temperature, + top_p: args.top_p, + } + }; + + // Day 1: load `tokenizer.json`, encode `args.prompt`, prefill the cache with + // the prompt tokens, then continue the loop below from the last one. + let mut cache = model.new_cache(); + let mut pos = 0usize; + let mut token: u32 = 0; // placeholder until the tokenizer is wired + + for _ in 0..args.max_tokens { + let logits = model.forward(token, pos, &mut cache); + cache.advance(); + token = sampler.sample(&logits); + pos += 1; + // Day 3: decode `token`, stream it to stdout, and stop on the EOS id. + } + + println!("(generation loop wired β€” implement the kernels to bring it to life)"); + Ok(()) } diff --git a/src/model.rs b/src/model.rs new file mode 100644 index 0000000..fb5cf31 --- /dev/null +++ b/src/model.rs @@ -0,0 +1,48 @@ +//! The model: weights loaded from safetensors, and the forward pass that ties +//! every kernel together into a single decode step. + +use std::path::Path; + +use anyhow::{Context, Result}; + +use crate::attention::KvCache; +use crate::config::Config; + +/// A loaded model: hyper-parameters plus the raw weight tensors. +/// +/// Weights are held as `f32` for the initial (correctness-first) build; the +/// quantized path in [`crate::quant`] later replaces the hot matrices. +pub struct Model { + pub config: Config, + // Once loading is implemented this holds the per-layer and global tensors: + // token embedding, per-block attention (q/k/v/o) and MLP (gate/up/down) + // projections, the two RMSNorm weights per block, the final norm, and the + // LM head (or a reference to the tied embedding). +} + +impl Model { + /// Load a model from a directory containing `config.json`, + /// `model.safetensors`, and `tokenizer.json`. + pub fn load(dir: impl AsRef) -> Result { + let dir = dir.as_ref(); + let config = Config::from_file(dir.join("config.json")) + .with_context(|| format!("loading model from {}", dir.display()))?; + // Day 1: mmap `model.safetensors`, resolve every tensor by name, and + // keep the buffers this struct needs for the forward pass. + Ok(Self { config }) + } + + /// Run one decode step: embed `token` at absolute position `pos`, run every + /// transformer block (updating `cache`), and return logits over the vocab. + pub fn forward(&self, token: u32, pos: usize, cache: &mut KvCache) -> Vec { + // embed β†’ N Γ— (RMSNorm β†’ GQA attention β†’ residual β†’ RMSNorm β†’ SwiGLU β†’ + // residual) β†’ final RMSNorm β†’ LM head. + let _ = (token, pos, cache); + todo!("full transformer forward pass") + } + + /// A KV cache sized for this model's maximum context length. + pub fn new_cache(&self) -> KvCache { + KvCache::new(&self.config, self.config.max_position_embeddings) + } +} diff --git a/src/ops.rs b/src/ops.rs new file mode 100644 index 0000000..e9f8a80 --- /dev/null +++ b/src/ops.rs @@ -0,0 +1,35 @@ +//! Element-wise transformer building blocks: normalisation, positional +//! encoding, activation, and softmax. Each is a small, self-contained kernel. + +/// RMSNorm: `out = x / sqrt(mean(xΒ²) + Ξ΅) * weight`. +pub fn rms_norm(x: &[f32], weight: &[f32], out: &mut [f32], eps: f32) { + debug_assert_eq!(x.len(), weight.len()); + debug_assert_eq!(x.len(), out.len()); + let _ = (x, weight, out, eps); + todo!("RMSNorm: normalise by the RMS of x, then scale by `weight`") +} + +/// Apply rotary position embeddings (RoPE) in place to a single head's query or +/// key vector at absolute position `pos`. +pub fn rope(vec: &mut [f32], pos: usize, head_dim: usize, theta: f32) { + debug_assert_eq!(vec.len() % head_dim, 0); + // For each (even, odd) dimension pair `i`, rotate by angle + // pos / theta^(2i / head_dim). + let _ = (vec, pos, head_dim, theta); + todo!("RoPE rotation of (even, odd) dimension pairs") +} + +/// SwiGLU feed-forward activation: `out = silu(gate) βŠ™ up`, where +/// `silu(z) = z * Οƒ(z)`. +pub fn swiglu(gate: &[f32], up: &[f32], out: &mut [f32]) { + debug_assert_eq!(gate.len(), up.len()); + debug_assert_eq!(out.len(), up.len()); + let _ = (gate, up, out); + todo!("SwiGLU: silu(gate) elementwise-times up") +} + +/// Numerically-stable softmax over `x`, in place (subtract max, exp, normalise). +pub fn softmax(x: &mut [f32]) { + let _ = x; + todo!("stable softmax over the attention scores") +} diff --git a/src/quant.rs b/src/quant.rs new file mode 100644 index 0000000..21bd428 --- /dev/null +++ b/src/quant.rs @@ -0,0 +1,44 @@ +//! Weight quantization β€” the headline optimization. +//! +//! Weights dominate both the memory footprint and the memory *bandwidth* that +//! bottlenecks CPU decoding. Storing them as INT8/INT4 with a per-row scale +//! shrinks the model ~2Γ—/~4Γ— and, because decode is bandwidth-bound, speeds it +//! up too. The forward path dequantizes on the fly inside the mat-vec, so the +//! weights are never materialised back to `f32`. + +/// Quantization width. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum QuantBits { + /// 8-bit signed, one value per byte. + Int8, + /// 4-bit signed, two values packed per byte. + Int4, +} + +/// A weight matrix stored as row-wise-quantized integers plus per-row scales. +pub struct QuantMatrix { + /// Packed quantized weights (`i8`, or two 4-bit values per byte for INT4). + pub packed: Vec, + /// One dequantization scale per output row. + pub scales: Vec, + pub rows: usize, + pub cols: usize, + pub bits: QuantBits, +} + +/// Quantize a row-major `[rows Γ— cols]` `f32` matrix with one symmetric scale +/// per row: `scale = max|w_row| / qmax`, `q = round(w / scale)`. +pub fn quantize_rowwise(w: &[f32], rows: usize, cols: usize, bits: QuantBits) -> QuantMatrix { + debug_assert_eq!(w.len(), rows * cols); + let _ = (w, rows, cols, bits); + todo!("row-wise symmetric quantization") +} + +/// Fused dequantize + mat-vec: `y = dequant(q) Β· x`, without ever expanding the +/// weights to `f32`. +pub fn matvec_quant(q: &QuantMatrix, x: &[f32], y: &mut [f32]) { + debug_assert_eq!(x.len(), q.cols); + debug_assert_eq!(y.len(), q.rows); + let _ = (q, x, y); + todo!("fused dequantize + mat-vec") +} diff --git a/src/sample.rs b/src/sample.rs new file mode 100644 index 0000000..ee36001 --- /dev/null +++ b/src/sample.rs @@ -0,0 +1,38 @@ +//! Turning a logit vector into the next token id. + +/// Sampling strategy selected from the CLI. +#[derive(Debug, Clone, Copy)] +pub enum Sampler { + /// Always take the arg-max (deterministic). + Greedy, + /// Temperature scaling followed by nucleus (top-p) sampling. + TopP { temperature: f32, top_p: f32 }, +} + +impl Sampler { + /// Pick the next token id from `logits`. + pub fn sample(&self, logits: &[f32]) -> u32 { + match *self { + Sampler::Greedy => argmax(logits), + Sampler::TopP { temperature, top_p } => { + // Divide logits by `temperature`, softmax, keep the smallest set + // of tokens whose probability mass β‰₯ `top_p`, renormalise, draw. + let _ = (temperature, top_p); + todo!("temperature + nucleus (top-p) sampling") + } + } + } +} + +/// Index of the maximum logit. +fn argmax(logits: &[f32]) -> u32 { + let mut best = 0usize; + let mut best_val = f32::NEG_INFINITY; + for (i, &v) in logits.iter().enumerate() { + if v > best_val { + best_val = v; + best = i; + } + } + best as u32 +} diff --git a/src/tensor.rs b/src/tensor.rs new file mode 100644 index 0000000..cd5d607 --- /dev/null +++ b/src/tensor.rs @@ -0,0 +1,28 @@ +//! The numeric core: dense `f32` linear algebra. +//! +//! A decode step is dominated by matrix–vector products against the weight +//! matrices, so [`matvec`] is the single hottest routine in the engine. It is +//! written by hand (and parallelised with `rayon`) rather than pulled from a +//! BLAS / `ndarray` crate β€” implementing it is the point. + +/// Row-major matrix–vector product `y = W Β· x`. +/// +/// `w` is `[out_dim Γ— in_dim]` in row-major order, `x` is `[in_dim]`, and the +/// `[out_dim]` result is written into `y`. +pub fn matvec(w: &[f32], x: &[f32], y: &mut [f32], in_dim: usize, out_dim: usize) { + debug_assert_eq!(w.len(), in_dim * out_dim); + debug_assert_eq!(x.len(), in_dim); + debug_assert_eq!(y.len(), out_dim); + // Day 1: split `y`/`w` into rows and dot each row with `x`, in parallel: + // y[o] = Ξ£_i w[o*in_dim + i] * x[i] + let _ = (w, x, y, in_dim, out_dim); + todo!("hand-written, rayon-parallel row-major mat-vec") +} + +/// In-place element-wise add: `a += b` (transformer residual connections). +pub fn add_assign(a: &mut [f32], b: &[f32]) { + debug_assert_eq!(a.len(), b.len()); + for (ai, bi) in a.iter_mut().zip(b) { + *ai += *bi; + } +}