Building a High-Performance LLM Inference Engine from Scratch with tiny-vLLM

The complexity of modern Large Language Model (LLM) inference engines often hides the fundamental mathematics and hardware interactions that make them work. While high-level libraries like PyTorch make model design accessible, the actual process of serving these models at scale requires a deep understanding of GPU memory management, CUDA kernel engineering, and linear algebra.

tiny-vLLM is an educational project designed to demystify this process. Rather than providing a black-box implementation, it serves as a course in C++ and CUDA, guiding developers through the implementation of a high-performance inference engine—a "smaller sibling" to the industry-standard vLLM. By building the engine from the ground up, developers can derive the necessary math and logic from scratch, moving from raw weight files to a functioning server.

The Anatomy of an Inference Server

At its core, an LLM is physically just a file containing a vast array of floating-point numbers (weights). These weights are the product of a costly training phase. However, a weight file is not an executable; it is a blueprint. To make it functional, an inference server must:

  1. Load the weights from a format like Safetensors.
  2. Implement the architecture (the sequence of mathematical operations) defined by the model (e.g., Llama 3.2).
  3. Execute these operations efficiently on hardware, typically using GPUs to handle the massive matrix multiplications involved.

Mastering the Hardware: CUDA and Memory

To achieve high performance, tiny-vLLM leverages CUDA, NVIDIA's parallel computing platform. A critical part of this journey is understanding the distinction between Host (CPU) and Device (GPU) memory.

Because the GPU cannot directly access the system's DRAM, data must be explicitly moved. The typical workflow involves allocating memory on the GPU via cudaMalloc and transferring data using cudaMemcpy. The goal for any high-performance engine is to minimize these transfers and reuse allocated buffers wherever possible to avoid the overhead of frequent allocations.

Deep Dive: Implementing the Inference Pipeline

1. Loading Safetensors

Safetensors is a popular format because it is fast and safe. A Safetensors file consists of a header size, a JSON header (containing tensor names, shapes, and offsets), and the raw tensor data. tiny-vLLM demonstrates how to map these offsets directly to GPU memory pointers, allowing the engine to retrieve specific weights (like the K-projection for a specific layer) with minimal overhead.

2. The Precision Trade-off: BF16

Most modern LLMs use bfloat16 (BF16). Unlike standard float16, which balances exponent and fraction bits, BF16 uses an 8-bit exponent (the same as float32) and a smaller 7-bit fraction. This trade-off is crucial for LLMs because it prevents overflow and underflow issues (range) while accepting a slight loss in precision, which empirical evidence shows has a negligible impact on model accuracy.

3. CUDA Kernel Engineering

One of the most valuable aspects of tiny-vLLM is its focus on writing custom CUDA kernels.

  • Embedding Gather: The first step is mapping token IDs to vectors. Because NVIDIA GPUs typically limit blocks to 1024 threads, but embeddings (like Llama's) are 2048 elements, the project teaches a common optimization: having each thread process two elements to stay within hardware limits.
  • RMSNorm and Parallel Reduction: Root Mean Square Layer Normalization requires calculating the sum of squares across a vector. To do this efficiently without race conditions, tiny-vLLM implements parallel reduction (tree reduction) using __shared__ memory and __syncthreads() to synchronize threads as they sum values in a logarithmic pattern.
  • RoPE (Rotary Positional Embeddings): To give the model a sense of token order, RoPE applies a rotation to the hidden states. This involves complex number math implemented via sine and cosine transformations in a CUDA kernel.

4. Matrix Multiplication and the cuBLAS Trick

Matrix multiplication is the engine's heaviest lift. While one could write a custom kernel, the industry standard is cuBLAS. However, cuBLAS expects column-major format, while most LLM weights are distributed in row-major format.

To avoid the costly process of physically transposing the data in memory, tiny-vLLM employs a mathematical trick: by manipulating the transposition flags (CUBLAS_OP_T and CUBLAS_OP_N) and rearranging the multiplication order ($C^T = B \times A^T$), the engine can treat row-major data as column-major without moving a single byte.

Advanced Inference Concepts

As the engine evolves from single-token generation to a production-ready server, several advanced concepts are introduced:

  • Prefill vs. Decode: The first token (prefill) requires processing the entire prompt. Subsequent tokens (decode) only require processing the most recently generated token.
  • KV Cache: To avoid recomputing the Key (K) and Value (V) projections for every previous token during the decode phase, the engine stores them in a cache. This transforms the complexity of generating each new token from $O(n^2)$ to $O(n)$.
  • Continuous Batching: To maximize throughput, the engine doesn't wait for an entire batch of requests to finish. Instead, it uses "slots," filling a new request into a slot as soon as a previous one completes.
  • PagedAttention: Inspired by virtual memory in operating systems, PagedAttention manages the KV cache in non-contiguous blocks (pages), drastically reducing memory fragmentation and allowing for larger batch sizes.

Conclusion

By breaking down the "magic" of LLM inference into discrete, implementable lessons, tiny-vLLM provides a roadmap for anyone wanting to move from using AI to building the systems that power it. As one community member noted, the lesson-style approach makes the codebase approachable even for those who have never touched CUDA before, turning a complex engineering challenge into a structured learning experience.

Sources