LLM Inference from Scratch

A from-scratch implementation track covering modern LLM inference optimizations, built as a self-directed deep dive into model serving internals.

The repo has two halves:

1. vlm/ — a complete vision-language model built from scratch (vision encoder, projector, autoregressive decoder), trained end-to-end on image-caption pairs.
2. milestones/ — seven inference-optimization milestones built on top of the VLM's decoder, each isolating one technique that real inference systems use.

Highlights

	what	metric
M5	Fused RMSNorm + Linear Triton kernel	2.15× speedup over PyTorch, max abs diff 2.74e-6 vs fp32 baseline
M7	Speculative decoding (greedy + sampling)	bit-equivalent output to baseline argmax, 3.25 / 4 mean acceptance
M4	Tensor parallelism (Megatron-style)	bit-equivalent to single-rank, max diff 8.94e-7 across 2 GPUs
M6	Prefill/decode disaggregation	token-equivalent across separated GPUs via NCCL send/recv
VLM	From-scratch SigLIP + projector + TinyLM	loss converged 8.89 → 4.02 on 5K image-caption pairs

Stack

PyTorch · Triton · NCCL · multi-GPU torchrun · CUDA (Ampere)

Architecture overview

The VLM's decoder (vlm/tiny_lm.py) is a small Llama-style transformer with RMSNorm, RoPE, multi-head causal attention, SwiGLU, and weight tying. The inference milestones each take this same decoder and add one optimization:

vlm/tiny_lm.py  ←  base model, used by every milestone
    │
    ├── M1: KV cache (past_kvs threading, prefill/decode equivalence)
    ├── M2: roofline analysis (FLOPs / bytes per layer)
    ├── M3: INT8 weight-only quantization (per-channel symmetric)
    ├── M4: tensor parallelism (column + row parallel, NCCL all-reduce)
    ├── M5: fused Triton kernel (RMSNorm + Linear, register-resident)
    ├── M6: prefill/decode disaggregation (NCCL send/recv KV cache)
    └── M7: speculative decoding (small draft + big target verify)

How to run

Each milestone is self-contained. From the repo root:

# VLM training (requires MinimindV pretraining parquet)
python vlm/train_my_vlm.py

# Inference milestones
python milestones/m5_fused_triton/triton_rmsnorm_linear.py
python milestones/m7_spec_decoding/spec_decode.py --mode greedy
torchrun --nproc_per_node=2 milestones/m4_tensor_parallel/run_tp.py
torchrun --nproc_per_node=2 milestones/m6_disagg/run_disagg.py

Milestones add vlm/ to sys.path automatically so the imports resolve.

Notes

The notes/ directory has standalone HTML explainers built alongside the code. They cover the conceptual core of each milestone: GPU memory hierarchy, Triton mental model, RoPE rotation visualizer, tensor-parallelism math derivation, NCCL collectives, KV cache, and speculative decoding. Open in a browser.

Acknowledgments

The VLM scaffolding (TinyLM architecture, MinimindV dataset format) is adapted from MinimindV, an MIT-licensed educational from-scratch VLM project. All milestone implementations are my own.

License

MIT — see LICENSE.