| # 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 |
|
|
