Deep dive on LLM Inference at Scale — Harshul Jain, Audible & Tanmay Sah, Independent AI Researcher
A workshop by Harshul Jain and Tanmay Sah breaks down LLM inference optimization from first principles, addressing memory, TTFT, and throughput issues. It covers foundational concepts like KV cache mathematics and the prefill/decode phases, then dives into model-side optimizations (quantization, attention variants like GQA, FlashAttention) and serving-side techniques (Paged Attention, continuous batching, prefix caching, KV quantization). The session concludes with benchmarks of popular inference engines like vLLM and SGLang, highlighting their strengths for different workloads, and outlines future research directions.
Harshul Jain, a Senior Software Engineer at Audible, and Tanmay Sah, an independent AI researcher, deliver a workshop on optimizing Large Language Model (LLM) inference. The session aims to provide a first-principles understanding of LLM inference costs and various optimization techniques, targeting both beginner and intermediate audiences in Machine Learning and AI. All workshop materials, including slides and Jupyter notebooks runnable on free RTX 6000 GPUs via Moab, are available in a public GitHub repository.
The Problem Statement: Why LLM Inference is Expensive
LLM inference, encompassing tasks from text generation to medical report analysis, is a multi-billion dollar market. However, it’s inherently expensive due to limited hardware, high compute costs, and the fact that inference is a recurring operational cost that scales with every user and token. This contrasts sharply with training costs, which are typically one-time. The core challenge is to optimize inference solutions to reduce token usage and improve service efficiency.
Three Core Pain Points of LLM Inference
The workshop identifies three critical challenges faced in LLM inference:
- Memory Consumption: GPU memory usage increases directly with context length. For example, a single token of KV cache on Mistral 7B costs 131 KB. A 4,000-token context can consume 0.5 GB, escalating to 2.1 GB for 16,000 tokens. With 80 concurrent users, this can reach 42 GB, quickly overwhelming common 24 GB GPU cards. This often leads to Out-of-Memory (OOM) errors.
- Time to First Token (TTFT): The time taken to generate the first token degrades significantly as the input context length grows. This impacts user experience, especially for interactive applications.
- Throughput Collapse: Naive inference servers process requests sequentially, leading to poor throughput and long completion times when handling multiple users concurrently.
Foundations: Understanding the Inference Pipeline
To address these pain points, the workshop first delves into the LLM inference pipeline:
- Pipeline Overview: Input text is converted into tokens, then embeddings, which are fed into transformer layers. These layers generate new tokens sequentially, which are then appended to the input for subsequent generations.
- Transformer Layers: Approximately 95% of inference compute is spent within these layers. Key components include normalization, attention, and feed-forward layers.
- Attention Layer: The most compute-intensive component, responsible for calculating attention scores for every token relative to all previous tokens. This involves projecting tokens into Key (K), Query (Q), and Value (V) spaces. The number of K and V vectors scales with input size.
- KV Cache: To avoid recomputing K and V vectors for previously processed tokens, they are stored in a Key-Value (KV) cache. This cache, however, is a primary driver of memory consumption, as demonstrated by the 131 KB/token for Mistral 7B.
- GPU Memory Structure: GPU memory is typically divided into fixed model weights, fixed overhead, and dynamic