The Roofline Model & KV Cache Mechanics

Serving large language models presents an architectural paradox: the exact same model running on the exact same GPU behaves like two completely different computational workloads depending on whether it is processing input tokens or generating output tokens.

This chapter details the physics of GPU execution: the Roofline Model, the mathematical separation between the Prefill and Decode phases, the exact memory scaling of the KV Cache, and how PagedAttention eliminates memory fragmentation.

1. The Roofline Model & Operational Regimes

The Roofline Model defines the attainable performance of any kernel or workload on hardware:

Deriving Operational Intensity

For any neural network operation, the Operational Intensity ($I$) is:

$I = \frac{\text{Floating Point Operations (FLOPs)}}{\text{Bytes Read from Memory (Bytes)}}$

Consider multiplying an activation matrix $X \in \mathbb{R}^{B \times d}$ by a weight matrix $W \in \mathbb{R}^{d \times d_{\text{out}}}$:

• Total FLOPs: $2 \times B \times d \times d_{\text{out}}$ (1 multiply and 1 add per element).
• Bytes read (assuming weights dominate): $d \times d_{\text{out}} \times \text{bytes\_per\_elem}$.

Thus, the arithmetic intensity is:

$I = \frac{2 \times B \times d \times d_{\text{out}}}{d \times d_{\text{out}} \times \text{bytes\_per\_elem}} = \frac{2 \times B}{\text{bytes\_per\_elem}}$

The Batch Size Relationship: Notice that $d$ and $d_{\text{out}}$ cancel out! Arithmetic intensity is directly proportional to Batch Size ($B$). At batch size $B=1$ (single user decode), arithmetic intensity is roughly $2 / 2 = 1.0$ FLOP/byte—deep inside the memory-bandwidth-bound region. At batch size $B=64$, arithmetic intensity jumps to $64$ FLOPs/byte, moving close to the compute roof!

2. Prefill vs Decode Phases in LLMs

An LLM request is serviced in two fundamentally distinct phases:

3. The KV Cache Memory Formula

During autoregressive decoding at step $t$, the attention calculation needs to attend to all past tokens $1 \dots t-1$. Rather than recomputing earlier key and value vectors from raw embeddings, they are stored in the KV Cache.

The exact memory footprint for a single request of length $L$ tokens is:

$\text{Memory}_{\text{KV}} = 2 \times n_{\text{layers}} \times n_{\text{kv\_heads}} \times d_{\text{head}} \times L \times \text{bytes\_per\_elem}$

Where:

• The factor of $2$ accounts for both Key and Value tensors.
• $n_{\text{layers}}$ is the number of transformer blocks.
• $n_{\text{kv\_heads}}$ is the number of key/value heads (in Multi-Head Attention, $n_{\text{kv\_heads}} = n_{\text{heads}}$; in Grouped-Query Attention, $n_{\text{kv\_heads}} \ll n_{\text{heads}}$).
• $d_{\text{head}}$ is the dimension of each attention head.
• $L$ is sequence length.

4. PagedAttention: Virtual Memory for LLM Serving

In traditional serving frameworks, memory for a request was allocated contiguously for the user's max_tokens setting.

Explore how PagedAttention partitions KV blocks into non-contiguous physical pages:

The PagedAttention Mechanism

1. Logical to Physical Mapping: The KV Cache of a sequence is split into fixed-size logical blocks (e.g. 16 tokens per block).
2. Block Table: The serving engine maintains a routing table that maps logical block numbers to arbitrary physical memory frames allocated from a global GPU memory pool.
3. Zero External Fragmentation: Because every physical block has the exact same size, the memory allocator never needs to find contiguous blocks.
4. Copy-on-Write for Parallel Sampling: When an agent or user branches into multiple outputs (e.g. Speculative Decoding, beam search, or tree-of-thought), child sequences share the parent prompt's physical blocks until one of them writes a new token.

5. Python Simulation: KV Cache Profiler

Run this live in-browser Python implementation to compute the exact KV Cache memory scaling across sequence lengths and batch sizes: