LLM Inference: How KV Caching and PagedAttention Save LLM InferenceLLM 推論：KV 快取、FlashAttention 與 PagedAttention

This is the second article in a three-part series. Our first installment, LLM Pretrain: Anatomy of a Decoder-Only Transformer, explored the mathematical mechanics of how a model turns raw text into next-token probabilities.

Pretraining is only the beginning. Once training finishes and the model weights are frozen, the challenge shifts from learning to execution. In production, the model processes dynamic prompts where every newly generated token must be fed back into the model to extend its growing attention history. Each request generates tokens sequentially and retains its own growing state, transforming the core problem from pure training compute into one of memory bandwidth, hardware efficiency, and real-time state manipulation.

Concern	Pretraining	Inference
Input	Large batches of complete sequences	Dynamic prompts arriving over time
Goal	Minimise next-token loss	Generate useful tokens within memory and latency constraints
State	Weights, gradients, optimiser state, activations	Weights and per-request KV cache
Parallelism	Many sequence positions processed together	Prefill is parallel; decode advances autoregressively
Memory pressure	Parameters, gradients, optimiser, activations	Parameters, KV cache, hardware workspaces
Primary metrics	Training tokens/s, convergence, utilisation	Time to First Token (TTFT), inter-token latency, throughput

(Note: Broader production serving architecture, request scheduling, and Service Level Objectives will be explored in detail in the final article of this series).

In this article, we look at why that shift makes LLM inference fundamentally harder than training, and how specific systems-level breakthroughs—like KV caching, FlashAttention, PagedAttention, and model compression—tame the growing attention state and bandwidth bottlenecks to keep execution fast and efficient.

The content is structured as:

1. The two phases of inference: prefill and decode
2. KV cache and attention efficiency
3. FlashAttention
4. PagedAttention and vLLM
5. Breaking the bandwidth bottleneck: model compression
6. From inference mechanics to a service

1. The two phases of inference: prefill and decode

When a large language model generates a response, the forward pass is split into two distinct phases with entirely different hardware bottlenecks and performance characteristics.

Prefill (prompt processing) is the model's first action after receiving an input. Since the entire prompt is known upfront, the model computes the hidden states for all tokens in parallel (subject to the causal mask). During this massive matrix multiplication, the model extracts and stores the key and value vectors for every prompt token, populating the initial KV cache. Because it maximizes parallel computation, prefill is typically compute-bound and dictates the Time to First Token (TTFT).

Decode (token generation) is the autoregressive loop that produces one new token at a time. At each step, the model must read the saved keys and values for all previous positions and append the newest key and value to the KV cache to generate a single token representation. Because the math is minimal but the data movement is massive, decode is highly sequential and typically memory-bandwidth-bound.

While these are the general rules, hardware specifications, batch sizes, and context lengths can shift either phase into different bottlenecks.

Phase	Input per request	Output	Persistent effect	Bottleneck	Primary metric
Prefill	Full prompt (or chunk)	First-token logits	Populates KV cache	Compute-bound	Time to first token (TTFT)
Decode	One newly selected token	Next-token logits	Extends KV cache	Memory-bandwidth-bound	Inter-token latency

When prefill and decode share the same accelerator, long prompt computations can block the rapid iteration of decode cycles. Modern serving engines mitigate this resource contention using techniques like chunked prefill (limiting how much prompt computation enters a single scheduler iteration) or disaggregated serving (assigning prefill and decode to entirely different hardware clusters and passing the KV cache state between them).

Decoding strategies

At the end of each generation step, the model produces a raw probability distribution over its vocabulary (softmax). However, the model does not actually pick the word—the decoding strategy does. A decoding strategy may always select the most likely token, sample from several likely tokens, or maintain multiple candidate sequences.

Because the chosen token is appended back to the context and permanently alters all future predictions, the decoding algorithm directly dictates the generation's diversity, quality, and computational cost.

Strategy	Selection rule	Typical trade-off
Greedy	Choose the highest-probability token	Highly deterministic; optimal for extraction tasks but can become repetitive in creative text.
Temperature	Divide logits by a scalar before sampling	Low values sharpen focus; high values increase diversity.
Top-k	Sample only from the $k$ highest-probability tokens	Fixed candidate count
Top-p	Sample from the smallest set whose cumulative probability reaches $p$	Candidate count adapts to uncertainty
Beam search	Maintain several high-scoring partial sequences	Yields high quality for constrained tasks, but heavily degrades throughput by multiplying the KV cache memory requirement by the beam width.

Generation parameters are as critical as the model checkpoint itself. For production implementations, altering top-p or temperature can fundamentally shift both output behavior and memory utilization. See the Transformers generation guide for production implementations.

2. KV cache and attention efficiency

During causal autoregressive generation, past token representations cannot attend to future positions. Consequently, once a token's Key ($K$) and Value ($V$) vectors are computed at a given layer, their values remain completely invariant for the remainder of the sequence. Recomputing all previous keys and values at every subsequent decode step would introduce $O(S^2)$ redundant compute overhead per layer.

To eliminate this redundant compute, the inference engine implements a Key-Value (KV) cache. The model computes and stores the $K$ and $V$ vectors for past positions. When generating a new token $t$, the model only computes the Query ($Q$), Key ($K$), and Value ($V$) projections for the newly arrived position $t$. It appends $K_t$ and $V_t$ to the per-layer cache and performs attention by querying $Q_t$ against the aggregated history $K_{\le t}$ and retrieving from $V_{\le t}$.

A simplified per-layer cache tensor has the shape:

$$[B,H_{KV},S,d_{head}]$$

• $B$ (Batch size): The number of concurrent sequences the GPU is processing at the same time.
• $H_{KV}$ (KV heads): The number of attention heads dedicated to Keys and Values. In newer models (using MQA or GQA), this is often smaller than the number of Query heads to save memory.
• $S$ (Sequence length): The cumulative context length (prompt tokens plus generated tokens).
• $d_{head}$ (Head dimension): The hidden size per individual attention head ($d_{model} / H_Q$).

While the KV cache resolves the compute bottleneck, it shifts the problem entirely to device memory (VRAM). Across all $L$ layers of a model, the total memory footprint is:

$$\text{KV Memory (Bytes)} = 2 \cdot L \cdot B \cdot S \cdot H_{KV} \cdot d_{head} \cdot p$$

• The factor of $2$ accounts for storing two separate tensors: Keys and Values.
• $L$ is the number of transformer blocks..
• $p$ is the precision in bytes per value ($p=2$ for 16-bit precision such as BF16 or FP16; $p=1$ for FP8).

Because KV memory scales linearly with both context length ($S$) and concurrent batch size ($B$), it quickly dominates GPU memory capacity.

Concrete Example: Consider a 32-layer model with $H_{KV} = 8$, $d_{head} = 128$, operating at 16-bit precision ($p = 2$). For an 8,192-token context ($S = 8192$), a single sequence consumes:

$2 \times 32 \times 1 \times 8192 \times 8 \times 128 \times 2 \approx 1.07\text{ GB of VRAM}$

If the serving engine processes 100 concurrent requests ($B = 100$), the cache alone requires over 107 GB of VRAM—surpassing the total 80 GB capacity of an NVIDIA A100 GPU before even allocating memory for the static model weights.

Where the cache physically lives?

In high-performance serving environments, model weights and the active KV cache reside directly in high-bandwidth device memory (HBM). Upon initialization, the serving runtime loads the static weights, pre-allocates execution workspaces, and assigns the remaining VRAM to a dynamic KV cache pool.

Accelerator memory (GPU HBM)
├── Model weights (static parameters)
├── Dynamic KV-cache pool
├── Short-lived activations and output logits
└── Kernel execution and graph workspaces

Host memory (CPU RAM)
├── Request queue and scheduler metadata
├── Tokenized input buffers
└── Swapped/offloaded cache blocks (evicted state)

Storage (SSD / disk)
└── Model checkpoints, tokenizer configs, quantized weight artifacts

When models exceed single-device memory limits, distributed serving partitions both the model and the KV cache:

• Tensor Parallelism (vertical split): Shards the attention heads vertically across GPUs; each worker retains the KV cache corresponding only to its assigned subset of heads ($H_{KV} / \text{TP}$).
• Pipeline Parallelism (horizontal split): Shards the model horizontally across layers; each stage stores the KV cache only for its subset of layers ($L / \text{PP}$).

Managing Memory Pressure: Offloading vs. Recomputation

What happens when the VRAM pool gets full? A usual approach is offloading (swapping)-moves cache blocks to host CPU RAM. While this expands effective capacity, host memory bandwidth is substantially slower, and bus transfers introduce high latency tail spikes.

Instead of saving the cache to CPU RAM, prefix recomputation evicts cache blocks and recomputes the prefix representations on demand during future steps. This eliminates transfer overhead and frees VRAM, but trades off computational throughput.

Mitigating Cache Footprints: Grouped-Query Attention (GQA)

When an LLM generates text, it uses "Attention Heads" to look back at previous words. As we learned with the KV cache, the engine must store the Keys and Values for every single token in memory. How a model connects Queries to Keys/Values dictates how much memory it will consume.

The architectural ratio between Query heads ($H_Q$) and Key-Value heads ($H_{KV}$) directly dictates the baseline size of the KV cache.

• Multi-Head Attention (MHA): Maximum Quality, Maximum Memory Every Query head has a dedicated Key and Value head ($H_Q=H_{KV}$).The model is very smart (high "representational capacity"), but the KV cache is enormous.

• Multi-Query Attention (MQA): Maximum Speed, Lower Quality All Query heads share a single Key head and a single Value head ($H_{KV} = 1$). This reduces cache size by a factor of $H_Q$, drastically lowering memory bandwidth demand, but can degrade model expressive capacity on complex tasks.

• Grouped-Query Attention (GQA) Serves as an optimal middle ground ($1 < H_{KV} < H_Q$). Query heads are partitioned into groups, with each group sharing one Key and Value head.

MHA (32 Q, 32 KV)       GQA (32 Q, 8 KV)            MQA (32 Q, 1 KV)
 Q Q Q Q ... Q Q        [Q Q Q Q] ... [Q Q Q Q]     [Q Q Q Q ... Q Q]
 │ │ │ │   │ │           \ /           \ /               \  /
 K K K K ... K K             K             K                 K
 V V V V ... V V             V             V                 V
 (Cache: 1.00×)             (Cache: 0.25×)           (Cache: 0.031×)

If a model has 32 query heads and 8 key-value heads, each KV pair serves 4 query heads, shrinking the cache footprint by 75% relative to standard MHA:

Architecture	Query heads (HQ​)	KV heads (HKV​)	Relative KV-Cache Footprint
Multi-Head Attention (MHA)	32	32	1.00×
Grouped-Query Attention (GQA)	32	8	0.25×
Multi-Query Attention (MQA)	32	1	0.03125×

Because head counts and projection matrices ($W_K, W_V$) are fixed architectural parameters defined during pretraining, GQA cannot be applied purely as a post-training runtime toggle. Modern open-weight foundation models (such as LLaMA 3, Mistral, and Qwen) standardize on GQA out of the box to maximize serving efficiency.

The original GQA work showed a useful quality-speed compromise for its evaluated models: GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints.

3. FlashAttention

In standard attention, every prompt token compares itself with every token it is allowed to attend to. Mathematically, this creates an $S \times S$ attention score matrix, where $S$ is the sequence length.

The primary bottleneck during the prefill phase is not the math itself, but the physical location where the math occurs:

• HBM (High-Bandwidth Memory): The main, large memory pool on the accelerator (e.g., an NVIDIA A100's 80 GB). It holds massive amounts of data but is relatively slow to access.
• SRAM (Static RAM): The tiny memory cache located directly on the computing chip itself. It is incredibly fast but strictly limited in size (often just 20 to 40 MB per GPU).

The Memory Wall

To execute standard attention, the GPU must materialize the entire $S \times S$ matrix. It computes the raw scores ($QK^\top$) in SRAM, writes them out to the slow HBM, and then reads them back into SRAM to calculate the softmax probabilities (because standard softmax requires a global sum across all keys). Finally, it writes the result back to HBM. For long sequences, moving these massive intermediate matrices back and forth across the memory bus completely dominates execution time. This data traffic jam is known as the "Memory Wall."

The Solution: Tiling and Kernel Fusion

FlashAttention fundamentally changes the execution order without changing the final mathematical result. It operates via "tiling" — breaking the queries, keys, and values into small blocks that fit comfortably inside the fast SRAM.

For each query tile, FlashAttention loads the required blocks into SRAM and computes the attention scores, softmax, and value multiplication all at once. To solve the softmax problem without reading the whole row, it uses a mathematical trick called online softmax update, maintaining a running statistical tally directly in SRAM.

Because all operations are fused into a single kernel, the finished output is written back to the slow HBM only once.

Production Impact

FlashAttention is an exact, IO-aware attention algorithm, meaning it produces the exact same result as standard attention (subject only to standard floating-point variance). It is important to note what it does and does not do:

• It does not remove the $O(S^2)$ arithmetic computations (FLOPs), shrink the model weights, or reduce the persistent memory footprint of the decode KV cache.
• It does drastically reduce memory traffic and eliminate the need to allocate HBM for the temporary $S \times S$ attention matrix.

Production serving engines invoke highly optimized, hardware-specific implementations of this algorithm rather than recreating it in application code. By completely eliminating the intermediate memory overhead, FlashAttention is the primary engineering breakthrough that allows modern LLMs to process massive context windows (like entire books or codebases) without suffering catastrophic out-of-memory (OOM) errors.

The method is described in FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness.

4. PagedAttention and vLLM

The Fragmentation Problem

Naive inference engines allocate a contiguous chunk of GPU memory based on a request's maximum possible sequence length. Because most requests finish long before reaching this maximum limit, huge portions of the reserved memory remain completely empty. This wasted space is called "fragmentation." Furthermore, as these large contiguous blocks are allocated and freed, the remaining free memory becomes Swiss cheese (external fragmentation)—preventing the system from accepting new requests even when total available VRAM is theoretically sufficient.

PagedAttention solves this bottleneck by borrowing a foundational concept from computer operating systems: virtual memory paging.

Instead of reserving one massive, contiguous chunk of memory per request, PagedAttention divides the KV cache into small, fixed-size physical blocks (e.g., storing the vectors for 16 or 32 tokens). Each sequence is assigned a logical block table whose entries point to these physical blocks scattered across a shared memory pool.

Request A: logical [0] [1] [2] ──► physical [7] [2] [11]
Request B: logical [0] [1]     ──► physical [4] [9]
Free pool:                           [1] [3] [5] [6] [8] [10] ...

When a request starts the prefill phase, it is only allocated the exact number of blocks it needs. As the model generates text during the decode phase, the system dynamically provisions one new physical block at a time. Because the attention kernel knows how to read the block table map, the physical memory can be scattered randomly across the GPU without issue.

Blocks are allocated strictly as sequences grow and returned immediately when requests finish, meaning waste is limited strictly to the final, partially filled block.

Prefix Sharing and vLLM

Because memory is mapped logically, PagedAttention enables Copy-on-Write (CoW) memory sharing via reference counting. If multiple requests share the exact same system prompt, or if a generation strategy like beam search branches out from a single prefix, the sequences can safely point to the exact same physical blocks for the shared context, saving massive amounts of memory. (Note: These “pages” are a software-level cache abstraction managed by the inference engine, not literal hardware OS pages).

The open-source serving engine vLLM pioneered this technique, pairing paged cache management with a continuous batching scheduler. While PagedAttention maximizes usable cache capacity, the scheduler weaponizes that freed capacity to massively increase concurrent throughput.

For implementation specifics, see Efficient Memory Management for Large Language Model Serving with PagedAttention.

Operational Debugging: Choosing the Right Optimization

Attention and cache optimizations should not be applied blindly; they must be chosen based on measured constraints in the production environment.

Observed symptom	Likely pressure	Candidate interventions
High TTFT for long prompts	Prefill compute, attention IO, or scheduling queue	FlashAttention, chunked prefill, prefix caching, prompt truncation, or disaggregated prefill serving.
KV-cache OOM or low concurrency	Cache capacity or fragmentation limits	PagedAttention, migrating to a model with GQA/MQA, KV-cache quantization (e.g., FP8), or adding device memory.

To summarize the stack: FlashAttention accelerates the physical execution of attention math, GQA permanently shrinks the baseline architectural shape of the cache, and PagedAttention dynamically manages how that cache is allocated in real-time. They are entirely complementary, but none substitutes for strictly monitoring the actual serving workload.

5. Breaking the bandwidth bottleneck: model compression

In addition to managing VRAM capacity and attention IO, engineers can optimize the payload itself to circumvent the strict autoregressive bottleneck. There are several architectural and mathematical techniques to make massive LLMs smaller, faster, and more efficient to serve.

Number formats

Number formats dictate a model's memory footprint, dynamic range, and which hardware kernels are invoked during execution.

Format	Bits	Bytes/Value	Characteristics
FP32	32	4.0	Maximum range and precision; strictly computationally prohibitive for LLM inference.
FP16	16	2.0	Standard inference format; high fractional precision but smaller dynamic range than BF16.
BF16	16	2.0	FP32-like exponent range with fewer fractional bits; highly stable.
INT8	8	1.0	Halves memory footprint; requires scaling metadata to reconstruct values.
INT4	4	0.5	Highly compressed; packs two values per byte before scales and metadata.

Storage precision and computation precision are not always identical. In Weight-Only Quantization (WOQ), a GPU stores the model weights in INT4 to drastically reduce VRAM footprint and memory bandwidth demands. The compression is mostly to solve the data-movement bottleneck (the Memory Wall). However, during execution, the kernel unpacks those INT4 values back into FP16 to perform the actual matrix multiplication.

Quantisation

Quantization is the mathematical process of mapping high-resolution floating-point numbers (FP16) into low-resolution integer spaces (INT8/INT4).

A standard symmetric quantization scheme operates by calculating a scale factor ($s$) and clamping the rounded values:

$$s=\frac{\max|W|}{q_{max}}, \qquad W_q=\operatorname{clamp}(\operatorname{round}(W/s),q_{min},q_{max})$$

To utilize the tensor during inference, the kernel multiplies it by the scale to reconstruct an approximation ($\hat{W} = sW_q$). Conceptually, it divides the tensor's dynamic range into discrete buckets, forces every number into its nearest bucket, and stores only the bucket ID.

The scale $s$ reconstructs an approximation $\hat W=sW_q$. Practical LLM quantisers typically use per-channel or grouped scales, clipping, outlier handling, and packed kernels.

When implementing quantization, several critical distinctions arise:

• Weight-only vs. Weight-and-Activation: Compressing only weights targets the memory bandwidth wall. Compressing activations alongside weights enables specialized, high-throughput integer compute kernels (like INT8 Tensor Cores) but is highly sensitive to outlier values.
• Post-training Quantisation (PTQ) vs. Quantisation-aware Training (QAT): PTQ mathematically compresses a frozen checkpoint; QAT simulates these precision drops during the pretraining phase, allowing the network to adapt to the information loss.
• Component Isolation: The static model weights and the dynamic KV cache exist as separate memory pools and can safely utilize entirely different quantization formats (e.g., FP4 weights with an FP8 KV cache).

Because compression alters the fundamental precision of the model, fewer bits do not automatically guarantee lower latency unless the target hardware possesses an optimized execution kernel for that specific format. Furthermore, aggressive quantization can introduce silent quality regressions. Evaluation must go beyond basic memory profiling to include task accuracy, perplexity, structured-output validity, TTFT, and inter-token latency (ITL) under representative loads. The Hugging Face quantisation documentation summarizes supported production integrations such as AWQ, GPTQ, and bitsandbytes.

If compression produces a quality regression, the correct operational response is rarely more serving optimization. Instead, engineers must revisit the calibration data, adjust the quantization group size, retain highly sensitive layers at a higher precision, or simply restore the previous format. Performance changes must pass the exact same task-level and safety evaluations as an entirely new model release.

Knowledge distillation

While quantization mathematically compresses an existing model, knowledge distillation is a pre-deployment training technique that forces a smaller "student" model to imitate a massive "teacher" model.

Standard pretraining relies on hard targets (e.g., next-token prediction). If the prompt is "The dog chased the...", the dataset simply dictates the answer is "cat". In distillation, the teacher model exposes its entire predicted probability distribution: "I am 80% sure it is 'cat', but 15% sure it is 'squirrel', and 5% sure it is 'ball'."

By calculating the cross-entropy alongside the KL divergence between these softened distributions, the student learns nuanced relationships and generalized logic that it could never infer from hard labels alone:

$$\mathcal{L} = \alpha\mathcal{L}_{hard} + (1-\alpha)T^2 \operatorname{KL} \left(p_T^{teacher}\,\|\,p_T^{student} \right)$$

The resulting student model requires less VRAM, features a smaller hidden width, and generates tokens faster—acting as a permanent architectural optimization before inference even begins.

The foundational method is described in Distilling the Knowledge in a Neural Network.

Speculative decoding

Distillation replaces the large model with a student. Speculative decoding explicitly targets the memory bandwidth wall by trading spare computational power for reduced VRAM reads.

Speculative decoding keeps the target model and uses a cheaper draft process (a tiny, cheap model) to propose multiple tokens. The target scores those proposed positions in one parallel pass and accepts a prefix of them.

Draft (Small Model):   Rapidly proposes candidate sequence [t1, t2, t3, t4]
Target (Large Model):  Verifies all proposed positions simultaneously
Result:                Accepts the valid prefix; corrects upon rejection

Because the target model evaluates the entire drafted sequence in a single forward pass, it only has to read its massive weight matrices from HBM once. For exact sampling, a token $x$ proposed from draft distribution $q$ is accepted with a probability of $\min(1,p(x)/q(x))$.

If the lightweight draft model predicts accurately, the system successfully generates multiple tokens for the IO cost of a single memory read. If it fails, the target model rejects the sequence, corrects the prediction, and seamlessly resumes the autoregressive loop.

See Fast Inference from Transformers via Speculative Decoding.

6. From inference mechanics to a service

The mathematical realities of the KV cache represent the fundamental bridge between model execution and serving capacity. Its VRAM footprint dictates exactly how many concurrent sequences can physically fit onto a GPU; its autoregressive growth forces request scheduling to be dynamic; and the extreme hardware differences between the compute-heavy prefill and the memory-bound decode phases dictate our core performance metrics: Time to First Token (TTFT) and Inter-Token Latency (ITL).

By combining memory management frameworks like PagedAttention with model compression strategies like quantization and speculative decoding, we can fundamentally alter the execution math to survive these hardware bottlenecks.

The final article in this series, Running LLMs in Production: Batching, Quantization, and Service-Level Objectives, takes these optimized inference mechanics and turns them into a production reality. We will explore how serving engines handle continuous request queues, distribute workloads across different GPU clusters, enforce Service-Level Objectives (SLOs), and monitor system performance.

加速器記憶體（GPU HBM）
├── 模型權重（靜態參數）
├── 動態 KV 快取池
├── 短暫存在的激活值與輸出 logits
└── Kernel 執行與計算圖工作空間

主機記憶體（CPU RAM）
├── 請求佇列與排程器中繼資料
├── 已 tokenize 的輸入緩衝區
└── 已置換／卸載的快取區塊（已驅逐的狀態）

儲存空間（SSD／磁碟）
└── 模型檢查點、tokenizer 設定、量化後的權重檔案

MHA（32 個 Q，32 個 KV）  GQA（32 個 Q，8 個 KV）      MQA（32 個 Q，1 個 KV）
 Q Q Q Q ... Q Q        [Q Q Q Q] ... [Q Q Q Q]     [Q Q Q Q ... Q Q]
 │ │ │ │   │ │           \ /           \ /               \  /
 K K K K ... K K             K             K                 K
 V V V V ... V V             V             V                 V
 （快取：1.00×）             （快取：0.25×）           （快取：0.031×）

請求 A：邏輯 [0] [1] [2] ──► 實體 [7] [2] [11]
請求 B：邏輯 [0] [1]     ──► 實體 [4] [9]
空閒池：                     [1] [3] [5] [6] [8] [10] ...

草稿模型（小模型）：     快速提議候選序列 [t1, t2, t3, t4]
目標模型（大模型）：     同時驗證所有被提議的位置
結果：                有效的前綴會被接受；遭拒絕的部分則會被修正