README.md

# spec100m — Ultra-fast speculative-decoding language model

A ~480M-parameter transformer + Medusa speculative-decoding architecture optimized for **maximum throughput** on a single GPU. Built from scratch to explore how fast speculative decoding can go when you stop rejecting drafts and accept everything.

## Results

### Inference (A6000, bf16, 480M params)

| Configuration | tok/s | step (ms) | Notes |
|---|---:|---:|---|
| Baseline generation (no n-gram) | 138,607 | 28.9 | Model only, K+1=4096 tokens/step |
| N-gram M=4096 | 266,218 | 28.9 | 1.9x speedup |
| N-gram M=16384 | 491,148 | 40.7 | 3.5x speedup |
| **N-gram M=32768** | **581,646** | 57.3 | **4.2x — exceeds 500k target** |
| **Compile + N-gram M=32768** | **599,716** | 55.6 | **Best result** |
| N-gram potential M=65536 | 686,305 | — | Projected from real-text test |

### Training (A6000, bf16, 271M trainable)

| Config | tok/s | Notes |
|---|---:|---|
| Original (CPU data loading) | 14,267 | Per-step CPU→GPU copy |
| GPU-resident data | 18,068 | Data stays on GPU |
| **GPU data + batch=32** | **38,427** | **2.7x speedup** |

### Long-context prefill (A6000, 20M tokens)

| Version | tok/s | Notes |
|---|---:|---|
| v4 (228M, 1 layer) | 971,092 | Weight-bandwidth bound |
| v5 (480M, 8 layers) | 123,873 | Full stack: MoBA + FP8 + compression |

## Architecture

### Base model (270M params)
- **d_model**: 1024
- **Layers**: 8
- **Heads**: 16 query, 1 KV (MQA)
- **FFN**: SwiGLU, mult=8 (8192 hidden)
- **Vocab**: 50257 (GPT-2 BPE, tiktoken)
- **Max context**: 20M tokens (MoBA block-sparse)
- **Position**: RoPE
- **Norm**: RMSNorm, tied embeddings

### Medusa speculative heads (210M params)
- **K**: 4095 heads (4096 tokens per step)
- **Rank**: 1 (precomputed argmax trick)
- **Batched**: [K, d, 1] and [K, 1, V] weight tensors
- **Precomputed argmax**: `argmax(s * w) = argmax(w) if s > 0 else argmin(w)` — skips the 410MB medusa-out einsum + argmax entirely

### KV cache stack (16x reduction)
1. **MQA**: 1 KV head shared across 16 query heads (4x)
2. **Learned compression**: every 4 tokens → 1 KV entry (4x)
3. **FP8 storage**: E4M3 format (2x)
4. **MoBA**: top-3 block selection via Q @ mean(K_block)
5. **Within-block sparse**: every 4th compressed entry

### N-gram draft extension (0 params)
- 5-gram lookup table trained on 2.4M tokens of Wikipedia
- Self-sustaining: predicts from its own previous predictions
- Adds up to 32768 free tokens per step (no forward pass)
- 4.2x throughput improvement

### Training controller
- Cosine LR schedule with warmup
- Plateau detection (reduce LR on stall)
- Grad norm instability detection
- Convergence prediction (linear extrapolation of loss trajectory)

## Training phases

### Phase 1: Base model (current)
- Train 270M base params only (Medusa frozen)
- Next-token cross-entropy loss
- Goal: produce coherent English
- Status: loss 11.0 → 3.83 in 500 steps on WikiText-2

### Phase 2: Medusa + compression (next)
- Freeze base model
- Train Medusa heads (multi-position loss)
- Train compression params (consistency loss)
- Goal: speculative decoding produces useful drafts

## File layout

| File | Purpose |
|---|---|
| `config.py` | Config dataclass, param counter, VRAM estimator |
| `model.py` | SpecModel: transformer + Medusa + MoBA + compression |
| `inference.py` | V5Engine: compressed MoBA + n-gram + compile + prefill |
| `ngram.py` | NGramModel: 5-gram lookup for free draft tokens |
| `train.py` | Phase 1 training (base only, with controller) |
| `data_pipeline.py` | Download + tokenize WikiText-2 |
| `generate.py` | Generate text from checkpoint |
| `bench_v5_1.py` | Inference benchmark (n-gram + compile) |
| `profile_v5.py` | Step profiling (FFN vs SDPA vs QKVO) |
| `profile_train_speed.py` | Training speed profiling |
| `test_v5_train.py` | Training compatibility test |

## Next steps

1. **Complete Phase 1 training** — run 5000+ steps on WikiText-103 (100M tokens) until loss < 3.0 and output is fluent English
2. **Phase 2: train Medusa heads** — freeze base, train 4095 heads with multi-position loss. Sample 16 heads/step to fit memory. Train compression params with consistency loss.
3. **Evaluate n-gram integration** — after Phase 2, model output should look like real text, so n-gram predictions should be much longer and more accurate. Measure real acceptance rate (not accept-all).
4. **Quality investigation** — the rank-1 Medusa trick limits each head to 2 possible tokens (argmax or argmin). Evaluate whether rank-2 or rank-4 heads are needed for quality, and whether the parameter budget allows it.
5. **Weight quantization** — int8 weights would ~2x memory bandwidth, potentially doubling prefill throughput. Untested.
6. **Larger corpus** — WikiText-2 (2.4M tokens) is too small for a 270M model. Move to WikiText-103 (100M) or OpenWebText for real training.
7. **CUDA graph capture for training** — the training loop has Python overhead that CUDA graphs could eliminate. Requires static shapes (already satisfied).
8. **Fused SwiGLU kernel** — FFN is 65% of inference step time. A custom Triton kernel for SwiGLU could reduce this.

## Environment

- **Remote**: NVIDIA RTX A6000 (48 GB VRAM), PyTorch 2.14.0+cu130, Python 3.10
- **Local**: NVIDIA RTX 3060 (12 GB VRAM), PyTorch 2.6.0+cu124, Python 3.13
- **Dtype**: bfloat16
- **Tokenizer**: tiktoken GPT-2 BPE (50257 vocab)

## License

MIT