Large Language Models (LLMs) have revolutionized the AI landscape, but serving them efficiently at scale remains a massive computational challenge. The autoregressive nature of text generation—where tokens are generated one by one—places an immense burden on both computation (FLOPs) and memory bandwidth. To make real-world deployment practical, researchers have developed an array of brilliant optimization techniques.
In this blog post, we will break down the core optimizations that power modern LLM inference, exploring them in this specific order: KV Cache, PagedAttention, FlashAttention, Continuous Batching, Quantization, and Speculative Decoding.
When autoregressive models generate text, each new token requires reading all previously generated tokens. A naive implementation recalculates the hidden representations of these past tokens at every step, creating immense computational waste.
The KV (Key-Value) Cache solves this by storing the previously computed key and value vectors of the past tokens. During generation, the model only passes the newest token into the network and fetches the necessary historical representations directly from the cache.
-
The Impact: This simple caching trick reduces the computational complexity of generating the next token from
$O(N^2)$ to$O(N)$ , saving an enormous amount of floating-point operations (FLOPs) as the sequence length grows. -
The Trade-off: While it saves compute, the KV Cache requires a large, dedicated memory footprint of
$O(N)$ , introducing a new bottleneck in memory capacity.
Because the KV Cache grows dynamically as new tokens are generated, managing its memory is incredibly difficult. Traditional memory management allocates contiguous blocks for each sequence upfront. This causes severe fragmentation and inefficiency, as memory is over-allocated for short sequences or requires costly re-allocation for long ones.
PagedAttention solves this by borrowing a concept from operating system virtual memory.
- How it works: Instead of storing the KV cache in one large, contiguous chunk, PagedAttention splits it into smaller, fixed-size "blocks" or "pages" (e.g., storing the keys and values for 128 or 256 tokens per block). These blocks are allocated dynamically on demand and do not need to be physically adjacent in memory.
- The Impact: This block-based approach nearly eliminates wasted memory due to fragmentation and enables much larger batch sizes, reportedly increasing serving throughput by 2-4x in engines like vLLM.
Standard exact attention is an
FlashAttention is an IO-aware algorithm designed to drastically reduce these HBM memory accesses.
- Tiling: It splits the queries, keys, and values into blocks, loading them from the slow HBM into the ultra-fast on-chip SRAM, and incrementally computing the softmax function.
- Recomputation: Instead of saving the giant intermediate attention matrix for the backward pass during training, it only saves scaling factors and recomputes the attention on-the-fly.
- The Impact: FlashAttention provides a massive speedup (e.g., up to 3x on GPT-2) and reduces the memory footprint to scale linearly rather than quadratically with sequence length.
To maximize throughput, inference engines try to generate tokens for multiple user prompts simultaneously. Traditional batching forces all prompts to have the same length, leading systems to inject "padding" tokens. When one prompt finishes early, the GPU wastes compute generating padding tokens while waiting for the longest prompt to finish, severely harming efficiency.
Continuous Batching reimagines this process to eliminate padding waste completely.
- Ragged Batching: Instead of a padded rectangular tensor, prompts of varying lengths are concatenated into a single flat sequence. Token interaction is strictly controlled using the attention mask so prompts do not interfere with one another.
- Dynamic Scheduling: As soon as one prompt finishes generating its sequence, it is instantly removed from the batch and dynamically replaced with a new prompt waiting in the queue.
- The Impact: Combining prefill and decoding phases without padding keeps the GPU fully utilized, allowing engines to serve thousands of concurrent users efficiently.
Modern LLMs consist of billions of parameters, conventionally stored as 32-bit or 16-bit floating-point numbers. An 8-billion parameter model in 16-bit precision requires roughly 16 GB of memory just to load the weights.
Quantization involves mapping these parameters to lower-precision data types, such as 8-bit or 4-bit integers.
- How it works: By using fewer bits per weight, the overall memory required to run the model is vastly reduced. For example, moving an 8B model from 16-bit to 4-bit shrinks its memory requirement from 16 GB down to just 4 GB.
- The Trade-off: While quantization drastically lowers VRAM consumption, going beyond 8-bits can result in precision loss and degraded model quality. Additionally, smaller bit-widths do not automatically guarantee faster computation, as some quantization methods require overhead to convert numbers back to half-precision during execution.
Autoregressive decoding is painfully sequential: each step requires a full forward pass through the massive model just to generate a single token, leaving the GPU's massive parallel compute power underutilized.
Speculative Decoding accelerates this by shifting the paradigm from "predict one token" to "guess multiple tokens, then verify".
- Draft and Verify: A lightweight, highly efficient "draft" mechanism guesses a sequence of upcoming tokens (e.g., 3 to 12 tokens). The main, massive "target" model then takes these guessed tokens and processes them all simultaneously in a single parallel forward pass.
- Rejection Sampling: The target model checks the draft's probabilities against its own. It accepts the tokens that match its own predictions, and discards the rest (falling back to standard generation at the point of rejection).
- Advanced Methods (EAGLE-3 / Multi-Token Prediction): Instead of running a completely separate draft model, advanced methods like EAGLE-3 use lightweight prediction heads attached to the main model's internal layers to propose draft tokens, further reducing overhead.
- The Impact: Because verification happens in parallel, the user sees text materialize in fast, multi-token chunks, significantly lowering latency without sacrificing a single drop of output quality.
By combining these six techniques—KV Cache, PagedAttention, FlashAttention, Continuous Batching, Quantization, and Speculative Decoding—developers can squeeze every ounce of performance out of their hardware, turning massive, sluggish language models into blazing-fast inference engines capable of serving the world.