Project scaffold: modules, CLI, and correctness harness

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"

README.md

@@ -1,5 +1,81 @@
-# ember
+<div align="center">
+
+# 🔥 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)
+
+</div>
+
+`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)

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()

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<Vec<f32>>,
+    values: Vec<Vec<f32>>,
+    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")
+}

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<Path>) -> Result<Self> {
+        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()
+    }
+}