# microgpt

A GPT-style language model written from scratch in PyTorch, with the pieces that matter in
production LLM systems: rotary position embeddings, RMSNorm, SwiGLU, an incremental-decoding
**KV cache**, nucleus sampling, and a test suite that proves each piece is correct.

No `nn.Transformer`, no Hugging Face. Attention, masking, caching and sampling are all
written out explicitly (~250 lines of model code) so every line can be reasoned about.

## What's inside

| Piece | Where | Why it's here |
|---|---|---|
| Causal self-attention with explicit mask | `model.py` | Mask is offset by the cache length, which is the classic place incremental decoding goes wrong |
| RoPE | `model.py` | Relative position without a learned table; cached keys keep their absolute rotation |
| RMSNorm + SwiGLU, tied embeddings | `model.py` | The Llama-style recipe at toy scale |
| KV cache + sliding window | `model.py` | `O(T)` work per generated token instead of `O(T^2)` |
| Temperature / top-k / top-p | `sampling.py` | Written from first principles, tested against analytic distributions |
| Trainer | `train.py` | AdamW with decoupled weight decay (no decay on norms), warmup + cosine LR, grad clipping |
| Benchmark | `scripts/benchmark.py` | Measures the cache speedup instead of asserting it |

## Quickstart

```bash
python3 -m venv .venv && source .venv/bin/activate
pip install -e ".[dev]"
pytest                                  # 21 tests, ~3s on CPU

mkdir -p data && curl -sSL -o data/shakespeare.txt \
  https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt
python -m microgpt.train --data data/shakespeare.txt --steps 1200 --batch-size 24
python -m microgpt.generate --prompt "ROMEO:" --tokens 200 --top-k 40
python scripts/benchmark.py
```

## Correctness: what the tests actually prove

* **No future leakage** — perturbing tokens at positions >= 8 leaves logits at positions < 8 unchanged.
* **KV cache == full forward** — prefill then one-token-at-a-time decoding reproduces the full
  forward pass to `1e-5`; so does multi-token continuation (which exercises the offset causal mask).
* **Greedy generation is identical with and without the cache** (while within the context window).
* **RoPE** preserves vector norms, and `<rope(q, m), rope(k, n)>` depends only on `m - n`.
* **Sampling** — top-k and top-p keep exactly the intended tokens; empirical frequencies match
  the softmax distribution; seeded runs are reproducible.
* **Learning** — the model can overfit a tiny batch (loss drops by more than 70%).

## Results

Measured on an Intel iMac (6 CPU threads, no GPU). Numbers are from this repo's own scripts.

**Training** — 0.81M parameters, 4 layers, 128-dim, character-level Tiny Shakespeare,
1200 steps (~5.6 min on CPU):

| step | train loss | val loss |
|---:|---:|---:|
| 0 | 4.242 | 4.235 |
| 400 | 1.657 | 1.830 |
| 800 | 1.444 | 1.655 |
| 1200 | 1.348 | 1.596 |

(Initial loss 4.24 ≈ ln 65, i.e. uniform over the 65-character vocabulary, as expected.)

**KV-cache decoding** — 2.68M-parameter model, greedy, 16-token prompt:

| new tokens | no cache (tok/s) | KV cache (tok/s) | speedup | identical output |
|-----------:|-----------------:|-----------------:|--------:|:-----------------:|
| 32 | 158.3 | 194.2 | 1.2x | yes |
| 64 | 166.2 | 259.3 | 1.6x | yes |
| 128 | 98.2 | 174.4 | 1.8x | yes |
| 256 | 74.5 | 198.8 | 2.7x | yes |
| 480 | 34.7 | 167.4 | 4.8x | yes |

Cached throughput stays roughly flat as the sequence grows while uncached throughput decays,
which is the `O(T)` vs `O(T^2)` difference showing up in practice. Timings on a shared CPU are
noisy, so treat the small-length rows as approximate.

## Design notes and known limitations

* **Sliding window past `block_size`.** The uncached path always sees the last `block_size`
  tokens. The cached path drops its cache when full and rebuilds it from the latest half-window.
  Outputs are therefore identical only while the sequence fits in the window; beyond that they
  differ by design. (RoPE makes true cache eviction possible, but keys would need re-rotation.)
* **Attention is `O(T^2)` memory** per layer: no FlashAttention, no grouped-query attention.
  Both are natural next steps and would also shrink the KV cache.
* **Character-level, small scale.** The point is correctness and clarity, not sample quality.
  Samples are Shakespeare-flavoured but not coherent.

## Roadmap

- [ ] Grouped-query attention (smaller KV cache) with a parity test against full multi-head
- [ ] Batched generation with per-sequence cache lengths
- [ ] Speculative decoding with a draft model, verified to preserve the target distribution
- [ ] BPE tokenizer