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