AIOS: Build an LLM Inference Engine from Scratch

Have you ever wondered how ChatGPT generates responses? How a model with billions of parameters actually runs on your GPU? Or why inference optimization matters so much?

AIOS is a hands-on learning project where you build an LLM inference engine from scratch. By the end, you won't just use LLMs—you'll understand how inference works under the hood.

Who This Is For

• Software engineers curious about AI/ML systems
• ML practitioners who want to understand inference beyond training
• System engineers interested in GPU programming and optimization
• Students looking for practical, implementable knowledge

What You'll Build

By completing this project, you will:

• Build a production-grade LLM inference engine from a simple HuggingFace model
• Implement every major optimization: KV cache, paged attention, continuous batching, FlashAttention, CUDA graphs, tensor parallelism
• Go from ~5 tok/s to ~1200+ tok/s — a 240x improvement
• Serve an OpenAI-compatible API

The Key Mental Model: LLM Inference Engine as an "Operating System"

Traditional computing has an OS that bridges applications and hardware. AI computing has an inference engine that bridges LLMs and GPUs:

Performance Progression

Each lesson adds one major optimization. Here's the throughput progression:

Course Roadmap

Foundation (Lessons 0–2)

• Lesson 0: Introduction — Traditional software vs LLMs, why inference matters, course objective
• Lesson 1: LLM Basics — Tokenizer, Transformer architecture, attention, positional encoding, parameters
• Lesson 2: Running Qwen3 with PyTorch — Loading weights, tokenization, forward pass, generation with HuggingFace

Building the Engine (Lessons 3–8)

• Lesson 3: Remove HuggingFace Dependencies — Own your model: pure nn.Module, direct safetensors loading, manual generation loop
• Lesson 4: Prefill/Decode Split — Understanding KV cache: O(n²) → O(n) compute, 5x speedup
• Lesson 5: Pre-allocated KV Cache — Stop memory fragmentation: contiguous cache, position-indexed writes, fused RMSNorm
• Lesson 5: Paged KV Cache — Virtual memory for LLM: block allocator, block tables, slot mapping
• Lesson 7: Batching — Variable-length batching with cu_seqlens, the Context pattern, 13x throughput
• Lesson 8: The Scheduler — Continuous batching, prefill-first scheduling, preemption, engine loop

Optimization (Lessons 9–11)

• Lesson 9: Fused Layers — QKV fusion, gate+up fusion, smart weight loading, packed_modules_mapping
• Lesson 10: CUDA Graphs — Capture and replay decode, eliminate CPU launch overhead, 1.2x speedup
• Lesson 11: Sampling — Gumbel-max trick, per-request temperature, top-k/top-p filtering

Scaling and Serving (Lessons 12–14)

• Lesson 12: Tensor Parallelism — Column/row parallel, vocabulary sharding, NCCL collectives, multi-GPU coordination
• Lesson 13: Prefix Caching (planned) — Radix cache and shared-prefix KV reuse
• Lesson 14: API Server & Benchmarking — OpenAI-compatible API, SSE streaming, throughput/latency benchmarks

Engine Architecture (Final State)

aios/
├── config.py                    # Engine configuration
├── sampling_params.py           # Per-request sampling parameters
├── llm.py                       # User-facing API
├── engine/
│   ├── llm_engine.py            # Orchestrator (model + scheduler + tokenizer)
│   ├── scheduler.py             # Prefill-first continuous batching scheduler
│   ├── model_runner.py          # GPU execution, KV cache, CUDA graphs, TP
│   ├── sequence.py              # Per-request state machine
│   └── block_manager.py         # Paged KV cache + prefix caching
├── models/
│   └── qwen3.py                 # Inference-only Qwen3 (no HF model deps)
├── layers/
│   ├── attention.py             # FlashAttention + Triton KV write
│   ├── linear.py                # TP-aware linear layers
│   ├── layernorm.py             # RMSNorm with fused residual add
│   ├── rotary_embedding.py      # Precomputed RoPE
│   ├── activation.py            # SiluAndMul (fused gate*up)
│   ├── embed_head.py            # Vocab-parallel embedding + LM head
│   └── sampler.py               # Gumbel-max sampling
└── utils/
    ├── loader.py                # Weight loader (safetensors → fused modules)
    └── context.py               # ThreadLocal context for attention metadata

Model Support

Supports all Qwen3 sizes. Exercises are model-size agnostic — config is loaded from config.json:

Quick Start

Prerequisites

• Python 3.10+
• PyTorch 2.0+
• CUDA GPU (recommended: A100/H100 for full benchmarks, any GPU for learning)

Install Dependencies

pip install torch transformers safetensors flash-attn triton xxhash numpy tqdm

Run the Engine

# Simple generation
python generate.py --model /path/to/Qwen3-0.6B

# Benchmark throughput
python benchmark.py --model /path/to/Qwen3-0.6B --num-prompts 32

# Using the Python API
python -c "
from aios import LLM, SamplingParams
llm = LLM('/path/to/Qwen3-0.6B')
outputs = llm.generate(['Hello world'], SamplingParams(max_tokens=64))
print(outputs[0]['text'])
"

Follow the Course

Start from the beginning:

1. Lesson 0: Introduction
2. Lesson 1: LLM Basics
3. Continue through Lessons 2–15

Each lesson includes:

• README.md — Concepts, diagrams, step-by-step guide
• run_lessonN.py — Standalone demo script
• requirements.txt — Dependencies for this lesson

References

This project is inspired by:

• nano-vllm (~1,200 lines) — Minimal vLLM clone with clean architecture
• mini-sglang (~6,400 lines) — Feature-complete SGLang reference

License

Educational project. See individual files for details.