Project scaffold: modules, CLI, and correctness harness#1
Merged
Conversation
Module map, the correctness-oracle workflow (logit parity vs transformers), and a day-by-day roadmap from weight loading to quantization.
safetensors + tokenizers + memmap2 for I/O, rayon for the mat-vec hot loop, half for f16 weights, clap/anyhow/serde for the CLI and config. Release profile tuned for throughput (LTO, single codegen unit).
Reads the fields the decoder needs (dims, head counts, RoPE theta, RMS eps) with GQA-aware helpers for head_dim, kv_dim, and kv_group_size.
The row-major matrix-vector product that dominates a decode step (stubbed for a hand-written, rayon-parallel implementation), plus an in-place residual add.
The element-wise transformer building blocks as small, self-contained kernels.
Rolling per-layer key/value cache (store/advance/clear) so decode is O(n) per token, plus the GQA scaled-dot-product attention entry point.
Per-row symmetric quantization and a fused dequantize+mat-vec that never expands weights back to f32 — the headline memory/bandwidth optimization.
Greedy arg-max implemented; nucleus sampling stubbed behind the same interface.
Loads config, exposes the per-step forward() and a right-sized KV cache; the safetensors loading and full forward pass are the next implementation targets.
clap-based CLI (prompt, model dir, max tokens, sampling) wired to the decode loop, so the control flow is in place ahead of the kernels.
Prints next-token logits from transformers for a fixed prompt so ember's own logits can be diffed against a trusted reference before wiring generation.
There was a problem hiding this comment.
Pull request overview
Establishes the initial Rust project scaffold for ember (CLI, model/config plumbing, module layout) with the core transformer kernels intentionally left as todo!() placeholders, plus a Python “correctness oracle” script for comparing logits against HuggingFace.
Changes:
- Adds the Rust module skeleton for config parsing, KV cache + attention API, sampling, quantization, and tensor/ops kernels.
- Implements a CLI-driven generation loop wiring (
main.rs) and a model load/forward skeleton (model.rs). - Adds docs + quickstart + roadmap in README, and a
scripts/reference_logits.pyparity helper.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 8 comments.
Show a summary per file
| File | Description |
|---|---|
| src/main.rs | CLI + generation loop wiring (currently kernel-stub dependent). |
| src/config.rs | config.json parsing + helper shape/accessor methods. |
| src/attention.rs | KV cache data structure + attention entrypoint stub. |
| src/model.rs | Model load/forward skeleton (weights/forward TODOs). |
| src/ops.rs | Kernel API stubs for RMSNorm/RoPE/SwiGLU/softmax. |
| src/tensor.rs | Dense mat-vec API stub + implemented add_assign. |
| src/quant.rs | Quantization structs + quantize/matvec stubs. |
| src/sample.rs | Sampler enum + greedy sampling (argmax) implementation. |
| scripts/reference_logits.py | HuggingFace-based reference logits printer for parity checks. |
| README.md | Project overview, design map, quickstart, oracle instructions, roadmap. |
| Cargo.toml | Adds core dependencies and release profile tuning. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+48
to
+50
| pub fn keys(&self, layer: usize) -> &[f32] { | ||
| &self.keys[layer] | ||
| } |
Comment on lines
+53
to
+55
| pub fn values(&self, layer: usize) -> &[f32] { | ||
| &self.values[layer] | ||
| } |
Comment on lines
+58
to
+64
| 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); | ||
| } |
Comment on lines
+14
to
+20
| 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") | ||
| } |
Comment on lines
+28
to
+30
| fn argmax(logits: &[f32]) -> u32 { | ||
| let mut best = 0usize; | ||
| let mut best_val = f32::NEG_INFINITY; |
Comment on lines
+57
to
+59
| pub fn head_dim(&self) -> usize { | ||
| self.hidden_size / self.num_attention_heads | ||
| } |
Comment on lines
+62
to
+64
| pub fn kv_group_size(&self) -> usize { | ||
| self.num_attention_heads / self.num_key_value_heads | ||
| } |
Comment on lines
+54
to
+67
| 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, | ||
| } | ||
| }; |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Lays out ember as a compilable Rust skeleton with the transformer math left as clearly-marked
todo!()kernels.What's here
main.rs) — clap args, decode loop wiredconfig.rs) — parses Qwen2.5config.json, GQA-aware helperstensor.rs(mat-vec),ops.rs(RMSNorm/RoPE/SwiGLU/softmax),attention.rs(GQA + KV cache),quant.rs(INT8/INT4)sample.rs) — greedy implemented, top-p stubbedmodel.rs) — config load + forward skeletonscripts/reference_logits.py) — logit parity vs HuggingFaceNo ML framework in the core — the transformer math is the project. See the README roadmap for the day-by-day build order.
Merging this establishes
mainas the working baseline; kernels land on follow-up branches.