tiny-vllm

You're going to build a high performance LLM inference engine with C++ and CUDA - tiny-vllm, a younger and smaller sibling of vLLM

We will learn a lot along the way, make mistakes and derive the ideas and maths from scratch

This repository consists of two things: 1. a full source code of the inference server and 2. a course where I lead you through the process of implementing the engine. Feel invited to use it as a learning tool on your learning path or if you are a lecturer, feel welcome to use it as a teaching resource at your university

The inference engine consists of:

• load a real LLM model from Safetensors (Llama 3.2 1B Instruct)
• full LLM forward pass (prefill + decode)
• all computation with CUDA kernels
• KV cache
• static batching
• continuous batching
• online softmax, FlashAttention-like
• PagedAttention

Make yourself a hot beverage and let's begin

• tiny-vllm
  • Intro: LLM, vLLM, models, inference servers
  • Technical prerequisities
  • Safetensors and your model
  • How floating-point numbers work and why we use bfloat16
  • GPU and CPU memory
  • Single token inference
  • Tokenization
  • Embeddings
  • CUDA kernel engineering - embeddings
  • RMSNorm and parallel reduction in CUDA
  • RoPE
  • Residual connections
  • cublasGemmEx
  • The column-major to row-major transposition trick
  • Prefill vs decode
  • Why KV cache exists
  • Attention
  • GQA
  • SiLU
  • Softmax
  • Causal mask
  • Argmax
  • Feed forward network
  • Buffer reuse
  • Static batching
  • Continuous batching
  • Online softmax
  • Paged Attention
  • Paged KV cache
  • Paged Attention CUDA kernel

Intro: LLM, vLLM, models, inference servers

It's easy to get lost with so much going on recent years. Let's unpack it

LLM is a model. Physically, LLM is a file which contains a lot of float numbers. Conceptually, these numbers represent weights of operations. Weights are learned/discovered/found during training phase. Some of the operations use these weights. Every operation is a function, which takes some data as input, do something with it and produces data as output. Operations and their order are defined by LLM's architecture. Every model has its own architecture, which is designed by engineers and researchers.

The process of going from 0 to LLM writing a text is like this:

0. Design the model - engineers and researchers use high level language like Python with tensor library like PyTorch or tinygrad to design model's architecture. They train small versions of the model, make experiments with different operations, data and hyperparameters (parameters for operations). It's the phase of figuring out the specification
1. Implement the model - Once they decide on final model architecture and prepare the data for training, they write the code that defines the final model. It can be also in PyTorch or similar
2. Train the model - The chosen model archit