GitHub - jmaczan/tiny-vllm: Build your own high performance LLM inference engine in C++ and CUDA - a smaller version of 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: 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 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 Train the model - The chosen model architecture is initialized with dummy weights. They write a script which again uses PyTorch or similar to run learning algorithm like backpropagation on a lot of hardware, like GPUs and TPUs. This phase burns a lot of energy, money and computational power. The product of training phase is a file with model weights, in some format, like Safetensors format. So, the training phase is finding such a set of weights which produces good text using the given architecture Serve the model (we are here) - The file with weights can't be ran on a computer. It's not an executable. It's a lot of numbers. The architecture can't be ran either - it's just a plan, a blueprint, a description of computation. To actually run the model, we need a program that turns the architecture and its operations into executable code and uses file with model weights to load the weights into the architecture. Once you write a program that implements the operations and once the program loads the weights (weights are loaded in the runtime of the program, at the startup), you can finally send prompts to the model and get a meaningful response. Generating an output from a model is called inference. That's why what we build here is called an inference server or inference engine Knowing the reason behind a need for an inference server, let's think why we build it in C++ and CUDA. It's because we want to maximize efficient use of the hardware and get high performance. It means that we want to get responses fast and we want to be able to handle multiple prompts at the same time. CUDA is the whole ecosystem, but also a language that you use to write code that runs on GPUs. We need to write code on GPUs, because many operations inside LLM are multiplying and adding multiple numbers.

If you need to do small amount of math, CPU enough. If a lot, GPU better. LLMs are mostly about multiplying the matrices, which boils down to computing dot products of two vectors, for many numbers and for many vectors. The math of LLMs is simple, we will need basics of linear algebra and you can learn while coding and fill the gaps on the go. I find this way of JIT learning the most effective and perhaps you will like it too My take on a relationship between AI and computation which you maybe find useful is that the intelligence comes from a lot of parameters of the model and a lot of computation of input values using these parameters. There is no a single element, that you can point to and say: "this is what makes the model intelligent or useful". Every part of the model you can replace with a different one and get different tradeoffs in return, like trade accuracy for complexity. I hope I won't forget to get back to this topic later, when we touch the math of attention. Because - the default attention mechanism is very computationally complex (O(n^2*d)). And this complexity can be challenged and in fact people do it and figure out alternative attention mechanisms, like linear attention. If more people will find this course useful, I will think about another one, about ML compilers (a practical one in Python or C++ + some SSA theory) or about alternative attention mechanisms (math + CUDA kernels). If you are interested, please let me know! If you will find this course valuable, please let other people know about it Out of scope: The training phase of an LLM is something we don't do in this course. We take a trained LLM and write a program which will run this LLM fast on NVIDIA GPU for multiple requests in parallel. If you want to train your own LLM, I strongly recommend sensei Karpathy repositories like nanoGPT and llm.c and his YouTube channel. Similarly, we don't design the model, but the tensor libraries are also fascinating topic and worth understanding from scratch. George Hotz's tinygrad is a project which implements a tensor library with a very little amount of code, so if you want to get inspired and learn the internals, it's a good place to do it (also their Discord is nice)! There is also a bit older and smaller version by Andrej Karpathy - micrograd.

And since I brought the Discord, I want to recommend you Mark Saroufim's GPU MODE. Many great people hanging out there! And if you feel lost with what is going on here, and you are new on your AI/ML journey, start with Jeremy Howard and Rachel Thomas fastai book. I conveniently omit the data science and engineering part here, because I don't know much about it. Probably Kaggle can be a good place to start with it and learn on-hands. Last but not least, we're going to code in C++ and CUDA and use cuBLAS where applicable. You can learn on the go. NVIDIA official resources are good and helpful Technical prerequisities You can build and run it on any platform, with minor changes, assuming you have a NVIDIA GPU. You might need to adjust some paths, like CUDA or GCC in c_cpp_propertiesjson or NVCC in CMakeLists.txt I suggest you to fork this repo and make the necessary adjustments so it works on your machine and then create a pull request to jmaczan/tiny-vllm and upstream your changes for benefit of another readers The exact setup on which I develop and test it: Linux (6.19.8 x64_64) CUDA Toolkit (13.1) C++ 17 GCC (15.2.1) The only external dependency you will pull in is JSON parser nlohmann/json 3.12.0, which is a single header file include/json.hpp AMD CPU (Ryzen 7 9800X3D) NVIDIA GPU (RTX 5090) I used Llama 3.2 1B Instruct from Hugging Face (commit hash 898999bd25b40516fce5a5b8f0948f4c81c650bc), you need just model.safetensors file from this repository Install the dependencies and run the program with ./test.sh - it will build it and immediately execute it If you fail to build or run it and your AI of choice won't be able to help, please open an Issue on GitHub - I will try to help.

Make sure to provide all useful context Safetensors and your model First thing you need to do is to download a LLM which we will use to run inference on. I choose Llama 3.2 1B Instruct, because it's easy, small, tuned for dialogs and good enough for us. From perspective of us, the engineers who build an inference server, the model is just a single file containing weights. The model is in Safetensors format. There exist other formats, like Pickle and Parquet. Safetensors is just very popular and widely used, and the model we picked is hosted in Safetensors Let's stop for a moment and understand the Safetensors format before we move on. A safetensor file consists of 3 sections, always in this order: header size, header and tensors data. Header size is always 8 bytes. These 8 bytes are an unsigned 64-bit integer, which says how many bytes the actual header takes. std::ifstream safetensors_file("model.safetensors", std::ios_base::binary); uint64_t header_size; safetensors_file.read(reinterpret_cast<char *>(&header_size), 8); The header is a JSON that contains about all the tensors inside the file. JSON is just a group of pairs <key, value>, where key is a unique string with a tensor name and value is another JSON object, containing info about this tensor. Every key in this JSON is a name of the tensor, except a single key which is called __metadata__, probably for some additonal info when necessary (we won't use it, specs say it's a "special key for storing free form text-to-text map). Every value is a JSON containing three keys - dtype, shape and offsets. dtype says what data type the tensor is stored in. shape says the dimensions of a tensor and offsets say where the tensor is stored, within the tensors data section. Every shape is a list of ints of unknown length and every offsets value is a vector of exactly two ints. First element says where the tensor begins and last element says where the tensor ends. You face the first design decision now. Do you want to make your inference server architecture independent, so it can run any arbitrary model, as long as you implement the operations it needs, or do you want to start simple and focus on our model of choice?