Phase 1: a working f32 forward pass (generates coherent text) - #2
Merged
Conversation
Row-major matrix-vector product with each output row computed in parallel, plus the in-place residual/bias adds.
RoPE uses the HuggingFace rotate-half convention (dim j pairs with j+head_dim/2) that Llama/Qwen weights expect.
Each query head attends over the causal KV cache using its shared kv head (group = n_heads / n_kv), scaled by 1/sqrt(head_dim).
Greedy arg-max plus top-p with a small dependency-free xorshift RNG.
default-features = false selects the pure-Rust regex path, so no C/C++ toolchain is needed to compile the tokenizer.
Loads every tensor (bf16/f16/f32 -> f32), handles Qwen's tied LM head and QKV biases, and runs embed -> N x (RMSNorm -> GQA -> RMSNorm -> SwiGLU) -> norm -> LM head for one decode step.
Encodes the prompt, prefills the KV cache, then samples/decodes/streams tokens until EOS or the token budget, reporting tokens/sec.
A standalone NumPy mirror of the exact algorithm. Verified against Qwen2.5-0.5B: 'The capital of France is Paris. It is the largest city in Europe and the third'.
There was a problem hiding this comment.
Pull request overview
This PR completes the Phase 1 (correctness-first) pipeline for ember: it loads Qwen2.5 weights from safetensors, runs a full f32 transformer forward pass with a causal GQA KV-cache, and generates text via greedy or top‑p sampling. It also adds a NumPy “oracle” implementation to cross-check the Rust implementation token-for-token.
Changes:
- Implemented core kernels and ops (rayon
matvec, RMSNorm, RoPE rotate-half, SwiGLU, softmax) and wired them into attention + the full multi-layer forward pass. - Added safetensors mmap loading (bf16/f16/f32 → f32), KV-cache sizing cap, and tied-LM-head handling.
- Implemented CLI generation loop with tokenizer prefill + streaming decode, plus a standalone NumPy reference forward pass and updated project docs.
Reviewed changes
Copilot reviewed 10 out of 10 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| src/tensor.rs | Implements rayon-parallel row-major matvec and adds add_bias helper. |
| src/ops.rs | Implements RMSNorm, HF rotate-half RoPE, SwiGLU, and numerically-stable softmax. |
| src/attention.rs | Implements grouped-query attention over a causal KV cache and wires in softmax. |
| src/model.rs | Adds safetensors weight loading + full per-token forward pass across all layers. |
| src/sample.rs | Implements greedy + top‑p sampling and adds a small dependency-free RNG. |
| src/main.rs | Wires CLI + tokenizer to prefill and decode loop with streaming output + timing. |
| scripts/reference_forward.py | Adds a NumPy oracle mirroring Rust forward pass for correctness checks. |
| README.md | Updates status/docs to “Phase 1 complete” and adds an example run + references. |
| PHASES.md | Introduces a phased build plan and documents Phase 1 completion criteria. |
| Cargo.toml | Disables default tokenizers features to avoid requiring a C toolchain. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
26
to
+28
| 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") | ||
| debug_assert_eq!(vec.len(), head_dim); | ||
| let half = head_dim / 2; |
Comment on lines
+96
to
+100
| let hd = config.head_dim(); | ||
| let n_heads = config.num_attention_heads; | ||
| let n_kv = config.num_key_value_heads; | ||
| let group = n_heads / n_kv; | ||
| let kv_dim = config.kv_dim(); |
Comment on lines
19
to
+22
| 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") | ||
| let mut probs: Vec<f32> = logits.iter().map(|&l| l / temperature).collect(); | ||
| softmax(&mut probs); | ||
|
|
Comment on lines
+171
to
+187
| let bytes = t.data(); | ||
| let out = match t.dtype() { | ||
| Dtype::F32 => bytes | ||
| .chunks_exact(4) | ||
| .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]])) | ||
| .collect(), | ||
| Dtype::F16 => bytes | ||
| .chunks_exact(2) | ||
| .map(|c| half::f16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f32()) | ||
| .collect(), | ||
| Dtype::BF16 => bytes | ||
| .chunks_exact(2) | ||
| .map(|c| half::bf16::from_bits(u16::from_le_bytes([c[0], c[1]])).to_f32()) | ||
| .collect(), | ||
| other => bail!("tensor {name}: unsupported dtype {other:?}"), | ||
| }; | ||
| Ok(out) |
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.
Implements the whole Phase 1 pipeline — a hand-written transformer that loads a real Qwen2.5 model and generates text.
What landed
matvec(rayon),rms_norm,rope(HF rotate-half), grouped-queryattention+ causal KV cache,swiglu,softmaxVerified
The NumPy oracle (identical math to the Rust) run against Qwen2.5-0.5B:
Coherent, on-topic output confirms every convention is correct — RoPE pairing, GQA head mapping, tied embeddings, biases.
Next: Phase 2 (quantization + benchmarks).