← Back to Logs
05-08-2026

How KV Cache and FlashAttention Accelerate LLM Inference

llm-inferencekv-cacheflashattentiongpu-programmingvllm
Try the interactive lab for this articleTake the quiz (6 questions)

Serving large language models in enterprise data centres across Frankfurt, Amsterdam, and Zurich presents a fundamental hardware challenge. Modern autoregressive transformer architectures such as Llama-3-70B, Mistral-Large, and Qwen-2.5 exhibit two completely different execution profiles depending on whether a request is in the prompt prefill phase or the token generation decode phase.

During prompt ingestion, execution is bound by GPU FLOPS matrix multiplication capacity. During autoregressive token generation, execution flips to being entirely memory bandwidth bound. At single-batch inference, high-performance tensor cores on an NVIDIA H100 SXM GPU remain idle for over 95 percent of execution cycles, waiting for model weight parameters and attention history to transfer across the High Bandwidth Memory (HBM) bus.

Optimizing production LLM inference serving requires addressing two distinct latency and memory bottlenecks. First, key-value (KV) caching eliminates redundant matrix operations across historical token positions, but introduces enormous VRAM capacity overheads and severe memory fragmentation. Second, standard scaled dot-product attention materializes quadratic intermediate matrices in off-chip DRAM, saturating memory channels.

This post details the hardware arithmetic intensity bounds of transformer inference, the mathematical equations governing KV cache footprint across attention variants, the virtual paging mechanics of PagedAttention, the GPU memory hierarchy tiling algorithms of FlashAttention-1, 2, and 3, and the architecture of continuous batching engines.

The Autoregressive Inference Bottleneck

To understand why LLM generation is memory bandwidth bound, transformer execution must be split into two operational phases: the prefill phase and the decode phase.

+-----------------------------------------------------------------------------------+
|                                 PREFILL PHASE                                     |
| Input Prompt: N tokens  ---> Parallel GEMM computation                            |
| Matrix-Matrix Multiplication (GEMM) -> High Operational Intensity (Compute-Bound) |
+-----------------------------------------------------------------------------------+
                                         |
                                         v
+-----------------------------------------------------------------------------------+
|                                  DECODE PHASE                                     |
| Token t=1 -> Token t=2 -> Token t=3 ... (Serial Step-by-Step Generation)           |
| Vector-Matrix Multiplication (GEMV) -> Low Operational Intensity (Bandwidth-Bound)|
+-----------------------------------------------------------------------------------+

Prefill Phase vs. Decode Phase Mechanics

In the prefill phase, the model server processes the full input prompt of length $N_{\text{prompt}}$ simultaneously. All tokens are packed into a 2D tensor matrix $X \in \mathbb{R}^{N_{\text{prompt}} \times d_{\text{model}}}$. Linear projections for Query, Key, and Value projections are executed as Matrix-Matrix Multiplications (GEMM):

$$Q = X W_Q, \quad K = X W_K, \quad V = X W_V$$

Because all $N_{\text{prompt}}$ rows interact with the weight matrices $W_Q, W_K, W_V \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ concurrently, each loaded weight parameter is re-used across $N_{\text{prompt}}$ separate matrix multiply-accumulate (MAC) operations.

In the decode phase, tokens are generated strictly one at a time. To generate token position $t+1$, the model passes only the single newly sampled token $x_t \in \mathbb{R}^{1 \times d_{\text{model}}}$ into the forward pass. The linear projections reduce to Matrix-Vector Multiplications (GEMV):

$$q_{t+1} = x_t W_Q, \quad k_{t+1} = x_t W_K, \quad v_{t+1} = x_t W_V$$

Each weight parameter loaded from off-chip HBM into on-chip registers is multiplied by a single vector element and immediately discarded.

Roofline Model and Arithmetic Intensity Calculations

The operational behavior of hardware accelerators is governed by the Roofline Model, which maps achievable performance (TFLOPS) against Arithmetic Intensity ($I$), defined as the ratio of floating-point operations performed per byte of memory read from or written to main memory:

$$\text{Arithmetic Intensity } (I) = \frac{\text{Floating Point Operations (FLOPs)}}{\text{Memory Access (Bytes)}}$$

Consider an NVIDIA H100 SXM5 GPU operating at 16-bit floating point precision (FP16 or BF16):

  • Peak FP16/BF16 Tensor Core Performance ($P_{\text{max}}$): $1,979 \times 10^{12} \text{ FLOPs/s} = 1,979 \text{ TFLOPS}$
  • HBM3 Memory Bandwidth ($B_{\text{mem}}$): $3.35 \times 10^{12} \text{ Bytes/s} = 3.35 \text{ TB/s}$

The critical operational intensity ceiling ($I_{\text{roof}}$) where execution transitions from memory bandwidth bound to compute bound is calculated as:

$$I_{\text{roof}} = \frac{P_{\text{max}}}{B_{\text{mem}}} = \frac{1,979 \times 10^{12} \text{ FLOPs/s}}{3.35 \times 10^{12} \text{ Bytes/s}} \approx 590.7 \text{ FLOPs/Byte}$$

If a GPU kernel exhibits an arithmetic intensity $I < 590.7 \text{ FLOPs/Byte}$, the arithmetic logic units (ALUs) will stall waiting for data transfers over the HBM bus. The maximum achievable compute performance in the memory bound regime is constrained by:

$$P_{\text{achievable}} = I \times B_{\text{mem}}$$

Decoding Phase Arithmetic Intensity

Consider generating a single token ($N=1$) for a model with parameter count $P_{\text{params}}$ operating at FP16 (2 bytes per parameter). The forward pass requires evaluating weight projections across all transformer layers.

For a linear weight matrix $W \in \mathbb{R}^{d_{\text{in}} \times d_{\text{out}}}$, the number of parameters is $d_{\text{in}} \times d_{\text{out}}$, corresponding to $2 \cdot d_{\text{in}} \cdot d_{\text{out}}$ bytes transferred from HBM. The number of floating-point operations for a single input vector $x \in \mathbb{R}^{1 \times d_{\text{in}}}$ is $2 \cdot d_{\text{in}} \cdot d_{\text{out}}$ FLOPs (one multiplication and one addition per parameter).

Summing across all model weight matrices, generating one token at batch size $B=1$ yields:

$$\text{FLOPs}{\text{decode}} = 2 \times P{\text{params}}$$

$$\text{Bytes}{\text{decode}} = 2 \text{ bytes/param} \times P{\text{params}} = 2 \times P_{\text{params}}$$

$$\text{Arithmetic Intensity}{\text{decode}} = \frac{2 \times P{\text{params}}}{2 \times P_{\text{params}}} = 1.0 \text{ FLOP/Byte}$$

An arithmetic intensity of $1.0 \text{ FLOP/Byte}$ is far below the H100 threshold of $590.7 \text{ FLOPs/Byte}$. On an H100 SXM GPU, the maximum processing throughput for a single decoding sequence of a 70B parameter model in FP16 ($140 \text{ GB}$ weights, requiring 2x H100 GPUs via Tensor Parallelism) is hard capped by memory bandwidth:

$$\text{Max Throughput}_{\text{single sequence}} = \frac{3.35 \text{ TB/s}}{140 \text{ GB}} \approx 23.9 \text{ tokens/second}$$

Even if the H100 Tensor Cores possessed infinite processing speed, memory bandwidth constraints prevent generating faster than 24 tokens per second for a single client stream. To saturate GPU compute units, servers must batch multiple client requests together, boosting arithmetic intensity proportionally by batch size $B$:

$$\text{Arithmetic Intensity}_{\text{batched decode}} \approx B \times 1.0 \text{ FLOPs/Byte}$$

To reach the compute roofline on an H100 SXM, the server requires a minimum active decode batch size of $B \ge 591$.

KV Cache Architecture

In standard self-attention, the Query vector at position $t$ must compute dot products against Key vectors at all preceding sequence positions $1, \dots, t$, followed by a weighted aggregation over Value vectors:

$$\text{Attention}(q_t, K_{:t}, V_{:t}) = \text{softmax}\left(\frac{q_t K_{:t}^T}{\sqrt{d_k}}\right) V_{:t}$$

Without caching, generating token $t+1$ would require recomputing $K_{:t}$ and $V_{:t}$ for all previous tokens $1 \dots t$ from scratch at every generation step. For a output sequence length of $T$, total computational complexity across generation steps scales quadratically:

$$\sum_{t=1}^{T} O(t^2 d_{\text{model}}) = O(T^3 d_{\text{model}})$$

Storing Key and Value Tensors

The KV cache solves this redundant recomputation by persisting the Key ($k_i$) and Value ($v_i$) projection vectors for all historical tokens in GPU VRAM after they are computed once during prefill or past decode steps.

Step t=1:  [k_1, v_1]  ---> Store in KV Cache
Step t=2:  Compute [k_2, v_2] ---> Concatenate with [k_1, v_1] ---> Store [k_{1..2}, v_{1..2}]
Step t=3:  Compute [k_3, v_3] ---> Concatenate with [k_{1..2}, v_{1..2}] ---> Store [k_{1..3}, v_{1..3}]

During step $t+1$:

  1. Compute $q_{t+1}, k_{t+1}, v_{t+1}$ for the single new incoming token.
  2. Append $k_{t+1}$ to the cached Key matrix: $K_{:t+1} = [K_{:t} ,;, k_{t+1}]$.
  3. Append $v_{t+1}$ to the cached Value matrix: $V_{:t+1} = [V_{:t} ,;, v_{t+1}]$.
  4. Evaluate attention of $q_{t+1}$ against updated cached matrices $K_{:t+1}$ and $V_{:t+1}$.

This drops per-step compute complexity for attention from $O(t^2)$ down to $O(t)$, reducing cumulative generation complexity over sequence length $T$ to $O(T^2)$.

Mathematical Formula for KV Cache VRAM Footprint

While KV caching eliminates redundant GEMM operations, it consumes substantial VRAM.

Let:

  • $L$: Number of transformer layers.
  • $H_k$: Number of Key/Value attention heads per layer.
  • $d_k$: Hidden dimension per head ($d_k = d_{\text{model}} / H_q$, where $H_q$ is Query heads).
  • $S$: Total context sequence length (prompt tokens + generated tokens).
  • $B$: Batch size (number of concurrent client sequences).
  • $P$: Precision byte size (2 bytes for FP16/BF16, 1 byte for INT8/FP8, 0.5 bytes for INT4).

Each token position stores two vectors (Key and Value) per layer. The KV cache VRAM footprint per token across a single sequence is:

$$\text{Bytes}_{\text{token}} = 2 \times L \times H_k \times d_k \times P$$

For a batch of $B$ sequences, each at context length $S$, total memory consumption is:

$$\text{Memory}_{\text{KV total}} = 2 \times B \times L \times S \times H_k \times d_k \times P \text{ bytes}$$

Architectural Variations: MHA vs. MQA vs. GQA Footprint Analysis

The structural configuration of attention heads profoundly affects KV cache growth.

Multi-Head Attention (MHA)     Multi-Query Attention (MQA)    Grouped-Query Attention (GQA)
Key/Value Heads = Query Heads   Key/Value Heads = 1            Key/Value Heads = Groups (e.g. 8)
 
  Q Q Q Q   K K K K   V V V V         Q Q Q Q    K    V             Q Q Q Q   K K   V V
  | | | |   | | | |   | | | |         | | | |    |    |             | | | |   | |   | |
  h1 h2 h3 h4 h1 h2 h3 h4 h1 h2 h3 h4     h1 h2 h3 h4 Shared        h1 h2 h3 h4 g1 g2 g1 g2
  1. Multi-Head Attention (MHA): Every Query head has a dedicated Key and Value head ($H_k = H_q$).
  2. Multi-Query Attention (MQA): All Query heads share a single Key and Value head ($H_k = 1$).
  3. Grouped-Query Attention (GQA): Query heads are split into $G$ groups, with each group sharing one Key and Value head ($H_k = H_q / G$).

To evaluate the operational impact, consider three canonical 70-billion parameter transformer configurations running at FP16 precision ($P = 2$ bytes), with context length $S = 8,192$ tokens, layer count $L = 80$, query heads $H_q = 64$, and head dimension $d_k = 128$:

Case 1: Llama-2-70B using Multi-Head Attention (MHA)

  • $H_k = 64$
  • Footprint per token per sequence: $$\text{Bytes}_{\text{token}} = 2 \times 80 \times 64 \times 128 \times 2 = 2,621,440 \text{ bytes} \approx 2.62 \text{ MB/token}$$
  • Footprint for one sequence ($S = 8,192$): $$\text{Memory}_{\text{seq}} = 2,621,440 \times 8,192 = 21,474,836,480 \text{ bytes} \approx 21.47 \text{ GB}$$
  • At batch size $B = 32$, KV cache alone consumes: $$\text{Memory}_{\text{batch}} = 32 \times 21.47 \text{ GB} = 687.04 \text{ GB}$$ This footprint requires 9x 80GB GPUs dedicated exclusively to storing KV cache memory, excluding model weights.

Case 2: Multi-Query Attention (MQA)

  • $H_k = 1$
  • Footprint per token per sequence: $$\text{Bytes}_{\text{token}} = 2 \times 80 \times 1 \times 128 \times 2 = 40,960 \text{ bytes} \approx 40.96 \text{ KB/token}$$
  • Footprint for one sequence ($S = 8,192$): $$\text{Memory}_{\text{seq}} = 40,960 \times 8,192 = 335,544,320 \text{ bytes} \approx 0.335 \text{ GB}$$
  • At batch size $B = 32$: $$\text{Memory}_{\text{batch}} = 32 \times 0.335 \text{ GB} = 10.74 \text{ GB}$$ MQA yields a 64-fold reduction in KV cache memory relative to MHA, but can degrade model capacity and multi-topic reasoning quality.

Case 3: Llama-3-70B using Grouped-Query Attention (GQA)

  • $H_k = 8$ (Group ratio $G = 8$)
  • Footprint per token per sequence: $$\text{Bytes}_{\text{token}} = 2 \times 80 \times 8 \times 128 \times 2 = 327,680 \text{ bytes} \approx 327.68 \text{ KB/token}$$
  • Footprint for one sequence ($S = 8,192$): $$\text{Memory}_{\text{seq}} = 327,680 \times 8,192 = 2,684,354,560 \text{ bytes} \approx 2.68 \text{ GB}$$
  • At batch size $B = 32$: $$\text{Memory}_{\text{batch}} = 32 \times 2.68 \text{ GB} = 85.76 \text{ GB}$$

GQA achieves an 8-fold reduction over MHA while preserving task accuracy and architectural expressiveness.

+-----------------------------------------------------------------------------------------------+
|                        KV CACHE VRAM FOOTPRINT AT S = 8,192 TOKENS                            |
+----------------------+--------------------+--------------------+------------------------------+
| Architecture         | KV Heads (H_k)     | Per-Token Size     | Single Seq Size (S=8,192)    |
+----------------------+--------------------+--------------------+------------------------------+
| MHA (Llama-2-70B)    | 64                 | 2.62 MB            | 21.47 GB                     |
| GQA (Llama-3-70B)    | 8                  | 327.68 KB          | 2.68 GB                      |
| MQA                  | 1                  | 40.96 KB           | 0.335 GB                     |
+----------------------+--------------------+--------------------+------------------------------+

PagedAttention and Memory Fragmentation

In early serving frameworks (such as HuggingFace Transformers and initial TensorRT-LLM releases), model servers managed KV caches by allocating contiguous memory buffers in GPU VRAM for each active request sequence.

The Contiguous Allocation Problem

When allocating contiguous tensors for a sequence with maximum sequence length $S_{\text{max}} = 8,192$, the framework pre-allocated a contiguous memory block of shape $[L, 2, H_k, S_{\text{max}}, d_k]$ upfront.

This naive allocation pattern suffers from three distinct sources of memory waste:

Naive Static Memory Allocation:
+---------------------------------------+---------------------------------------+
|  Actual Used Context (e.g. 1,024 tok)  |    Reserved Unused Space (7,168 tok)  |
+---------------------------------------+---------------------------------------+
|<---------------------- Pre-allocated S_max = 8,192 -------------------------->|
  1. Reserved Unused Capacity: Most client requests do not generate tokens up to $S_{\text{max}}$. If a client prompt is 512 tokens and output generation stops at 512 tokens, $87.5%$ of the pre-allocated contiguous buffer remains empty, yet cannot be assigned to other client requests.
  2. Internal Fragmentation: Requests allocate memory for their estimated maximum length. Due to variable output lengths, memory slots reserved between current generation step $t$ and $S_{\text{max}}$ sit unutilized throughout execution.
  3. External Fragmentation: Over hours of continuous server operation, requests with varying context sizes arrive, allocate contiguous blocks, complete, and release VRAM. The GPU memory allocator becomes fragmented into non-contiguous gaps. Even if 40 GB of total VRAM is free across thousands of small memory holes, a new request requiring 8 GB of contiguous VRAM will trigger an out-of-memory (OOM) error.

In production deployments, contiguous KV cache allocation wasted between $60%$ and $80%$ of total GPU VRAM, capping serving batch sizes at low values.

PagedAttention Architecture

To resolve VRAM fragmentation, Kwon et al. introduced PagedAttention in vLLM, adapting virtual memory management principles from operating systems to GPU memory architectures.

In PagedAttention, Key and Value tensors are not stored in contiguous VRAM blocks. Instead, VRAM is partitioned into a global pool of fixed-size physical memory blocks (typically sized to hold $B_{\text{block}} = 16$ or $32$ tokens).

LOGICAL MEMORY (Client Request Sequence)
+---------------+---------------+---------------+
| Logical Blk 0 | Logical Blk 1 | Logical Blk 2 |   Tokens 0..47 (Block Size = 16)
+---------------+---------------+---------------+
        |               |               |
        v               v               v
PAGE TABLE MAPPING
  Logical Block 0 ---> Physical Block 7
  Logical Block 1 ---> Physical Block 2
  Logical Block 2 ---> Physical Block 12
 
PHYSICAL GPU VRAM POOL (Non-Contiguous Blocks)
+------------------+------------------+------------------+------------------+
| Physical Block 2 | Physical Block 7 | Physical Block 12| Free Block Pool  |
| (Tokens 16..31)  | (Tokens 0..15)   | (Tokens 32..47)  |                  |
+------------------+------------------+------------------+------------------+

A logical sequence is represented as a series of logical blocks. As generation progresses, new physical blocks are dynamically allocated on-demand from the global pool. When a single token exceeds the boundary of the current block, the vLLM engine fetches one new physical block from the free list and updates the request's Page Table.

Logical-to-Physical Block Address Translation

Let:

  • $i$: Absolute token index within sequence ($0 \le i < S$).
  • $B_{\text{block}}$: Fixed block size (number of tokens per block, e.g., 16).
  • $H_k$: Number of Key/Value heads.
  • $d_k$: Hidden dimension per head.
  • $P$: Precision byte length (2 bytes for FP16).

The logical block index $L_{\text{idx}}$ and token block offset $O_{\text{idx}}$ are derived as:

$$L_{\text{idx}} = \lfloor \frac{i}{B_{\text{block}}} \rfloor$$

$$O_{\text{idx}} = i \bmod B_{\text{block}}$$

The Engine looks up $L_{\text{idx}}$ in the request page table array $\text{PageTable}[L_{\text{idx}}]$ to retrieve the Physical Block ID ($\text{P}_{\text{id}}$).

The physical VRAM memory byte address for the Key tensor element of token $i$ at head $h$ is:

$$\text{Address}{\text{Key}}(i, h) = \text{BaseAddr}(\text{P}{\text{id}}) + \left( O_{\text{idx}} \times H_k \times d_k + h \times d_k \right) \times P$$

Python Simulation of PagedAttention Memory Management

Below is an executable simulation demonstrating logical-to-physical block allocation, address translation, and zero-copy page table branching:

import math
from typing import List, Dict, Optional
 
class PhysicalBlockPool:
    """Manages global non-contiguous GPU VRAM block allocations."""
    def __init__(self, total_blocks: int, block_size: int, num_layers: int, num_heads: int, head_dim: int):
        self.block_size = block_size
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.free_blocks: List[int] = list(range(total_blocks))
        self.ref_counts: Dict[int, int] = {b_id: 0 for b_id in range(total_blocks)}
 
    def allocate(self) -> int:
        if not self.free_blocks:
            raise MemoryError("GPU Out of Memory: No physical KV cache blocks available.")
        block_id = self.free_blocks.pop(0)
        self.ref_counts[block_id] = 1
        return block_id
 
    def free(self, block_id: int):
        self.ref_counts[block_id] -= 1
        if self.ref_counts[block_id] == 0:
            self.free_blocks.append(block_id)
 
    def increment_ref(self, block_id: int):
        self.ref_counts[block_id] += 1
 
 
class PagedKVCacheSequence:
    """Manages virtual block tables for a single client sequence."""
    def __init__(self, seq_id: str, pool: PhysicalBlockPool):
        self.seq_id = seq_id
        self.pool = pool
        self.block_table: List[int] = []
        self.seq_len = 0
 
    def append_token(self):
        """Simulate adding one token to the sequence, allocating blocks as required."""
        if self.seq_len % self.pool.block_size == 0:
            # Current blocks full; allocate new physical block
            new_block_id = self.pool.allocate()
            self.block_table.append(new_block_id)
        
        self.seq_len += 1
 
    def translate_token_address(self, token_index: int, layer_idx: int, head_idx: int) -> Dict[str, int]:
        """Translates logical token position to physical block memory parameters."""
        if token_index >= self.seq_len:
            raise IndexError("Token index out of bounds.")
 
        logical_block_idx = token_index // self.pool.block_size
        block_offset = token_index % self.pool.block_size
        physical_block_id = self.block_table[logical_block_idx]
 
        # Calculate stride within physical block
        bytes_per_element = 2  # FP16
        token_stride = self.pool.num_heads * self.pool.head_dim * bytes_per_element
        head_offset = head_idx * self.pool.head_dim * bytes_per_element
        byte_offset_within_block = (block_offset * token_stride) + head_offset
 
        return {
            "logical_block": logical_block_idx,
            "physical_block_id": physical_block_id,
            "block_offset": block_offset,
            "byte_offset_within_block": byte_offset_within_block
        }
 
    def fork(self, new_seq_id: str) -> 'PagedKVCacheSequence':
        """Zero-copy branch creation for parallel sampling / beam search."""
        child = PagedKVCacheSequence(new_seq_id, self.pool)
        child.block_table = list(self.block_table)
        child.seq_len = self.seq_len
        # Increment reference count for shared blocks
        for block_id in child.block_table:
            self.pool.increment_ref(block_id)
        return child
 
    def release(self):
        """Releases sequence memory allocations."""
        for block_id in self.block_table:
            self.pool.free(block_id)
        self.block_table.clear()
        self.seq_len = 0
 
 
# Execution Proof
if __name__ == "__main__":
    # Initialize GPU pool with 1,000 physical blocks (Block size = 16 tokens)
    pool = PhysicalBlockPool(total_blocks=1000, block_size=16, num_layers=80, num_heads=8, head_dim=128)
    
    # Client sequence starts prompt prefill (40 tokens)
    req1 = PagedKVCacheSequence("client_req_001", pool)
    for _ in range(40):
        req1.append_token()
 
    print(f"Req1 Token Count: {req1.seq_len}")
    print(f"Req1 Physical Block Table: {req1.block_table}")  # Uses 3 physical blocks (16 + 16 + 8)
 
    # Address translation for Token 35, Layer 0, Head 3
    addr_info = req1.translate_token_address(token_index=35, layer_idx=0, head_idx=3)
    print(f"Token 35 Physical Translation: {addr_info}")
 
    # Parallel sampling: fork sequence zero-copy
    req2 = req1.fork("client_req_001_branch_2")
    print(f"Req2 Forked Block Table: {req2.block_table}")
    print(f"Shared Block Ref Counts: {[pool.ref_counts[b] for b in req1.block_table]}")
 
    req1.release()
    req2.release()
    print(f"Free Blocks Remaining after release: {len(pool.free_blocks)}")

Advanced Memory Sharing Features

PagedAttention enables complex decoding patterns without copying memory buffers:

  • Zero-Copy System Prompt Sharing: System prompts (e.g. 3,000-token instructions in agent frameworks) are processed once during prefill. Their physical blocks are retained in memory with elevated reference counts. Subsequent client requests append these physical block IDs directly into their own page tables, reducing prompt ingestion compute and VRAM footprint to zero.
  • Copy-on-Write (CoW) Beam Search: During beam search, multiple candidate sequence generations fork from a shared prefix. Parent and child paths share physical blocks. When a child stream appends a new token that modifies a shared block, the engine clones only that specific 16-token physical block, isolating execution across paths.

PagedAttention reduces VRAM waste from $>60%$ to $<4%$, allowing model servers to scale active decoding batch sizes by $2\times$ to $4\times$.

FlashAttention Kernel Mechanics

While PagedAttention solves VRAM capacity fragmentation, self-attention memory bandwidth bottlenecks remain inside the attention computation itself.

GPU Memory Hierarchy and Memory Access Costs

Modern GPU architectures feature a hierarchical memory structure:

+-------------------------------------------------------------------------+
| SRAM / Register File (On-Chip)                                          |
| Size: ~228 KB per SM (~50 MB total across H100 GPU)                     |
| Bandwidth: ~19 - 33 TB/s  | Latency: ~10 - 20 cycles                    |
+-------------------------------------------------------------------------+
                                     ^
                                     | High Speed Transfer
                                     v
+-------------------------------------------------------------------------+
| High Bandwidth Memory - HBM3 / DRAM (Off-Chip)                          |
| Size: 80 GB / 96 GB                                                     |
| Bandwidth: ~2.0 - 3.35 TB/s | Latency: ~200 - 300 cycles                |
+-------------------------------------------------------------------------+

Standard scaled dot-product attention evaluates three matrix projections $Q, K, V \in \mathbb{R}^{N \times d}$ stored in off-chip HBM:

$$S = \frac{Q K^T}{\sqrt{d_k}} \in \mathbb{R}^{N \times N}$$

$$P = \text{softmax}(S) \in \mathbb{R}^{N \times N}$$

$$O = P V \in \mathbb{R}^{N \times d}$$

In a standard PyTorch CUDA implementation, each step materializes full $N \times N$ matrices in off-chip HBM:

  1. Read $Q, K$ from HBM to SRAM $\to$ Compute $S \in \mathbb{R}^{N \times N}$ $\to$ Write $S$ to HBM.
  2. Read $S$ from HBM to SRAM $\to$ Compute $P = \text{softmax}(S) \in \mathbb{R}^{N \times N}$ $\to$ Write $P$ to HBM.
  3. Read $P, V$ from HBM to SRAM $\to$ Compute $O = P V \in \mathbb{R}^{N \times d}$ $\to$ Write $O$ to HBM.

Total memory reads and writes scale quadratically with sequence length $N$, transferring $O(N^2)$ elements over the slow HBM bus. For $N = 16,384$, the intermediate attention matrix $S$ requires storing $16,384 \times 16,384 \times 2 \text{ bytes} = 536.87 \text{ MB}$ per head per layer. Across 80 layers and 64 heads, materializing intermediate attention scores consumes tens of gigabytes of DRAM bandwidth, causing execution to stall.

Tiling Algorithm and Online Softmax

Tri Dao et al. introduced FlashAttention, an exact attention algorithm that computes attention without materializing intermediate $N \times N$ matrices $S$ and $P$ in off-chip HBM.

FlashAttention partitions the Query, Key, and Value matrices into small blocks that fit entirely within the GPU's fast on-chip SRAM (~228 KB per Streaming Multiprocessor), executing all operations in SRAM tiles before writing final outputs back to HBM.

Mathematical Derivation of Online Softmax

Computing Softmax over a full vector $x = [x_1, \dots, x_N]$ requires a global maximum to prevent exponent overflow:

$$m = \max_{1 \le i \le N} x_i, \quad d = \sum_{i=1}^{N} e^{x_i - m}, \quad \text{Softmax}(x)_i = \frac{e^{x_i - m}}{d}$$

Because $m$ depends on every element in the sequence, standard Softmax requires a full pass over the sequence before normalisation.

To compute Softmax block-by-block in SRAM, FlashAttention uses Online Softmax.

Suppose vector $x$ is split into two blocks $x^{(1)}$ and $x^{(2)}$.

  1. For Block 1 ($x^{(1)}$): $$m^{(1)} = \max_{j} x_j^{(1)}, \quad d^{(1)} = \sum_{j} e^{x_j^{(1)} - m^{(1)}}$$

  2. For Block 2 ($x^{(2)}$): $$m^{(2)} = \max_{j} x_j^{(2)}, \quad d^{(2)} = \sum_{j} e^{x_j^{(2)} - m^{(2)}}$$

  3. To combine Block 1 and Block 2 into a global state: $$m^{(\text{new})} = \max\left(m^{(1)}, m^{(2)}\right)$$

    The updated normalization denominator $d^{(\text{new})}$ is derived by rescaling previous sums: $$d^{(\text{new})} = d^{(1)} \cdot e^{m^{(1)} - m^{(\text{new})}} + d^{(2)} \cdot e^{m^{(2)} - m^{(\text{new})}}$$

  4. Rescaling Output Accumulator: Let $O^{(1)} = \frac{e^{x^{(1)} - m^{(1)}}}{d^{(1)}} V^{(1)}$ be the partial output computed from Block 1. When Block 2 arrives, $O^{(1)}$ is updated to $O^{(\text{new})}$ without re-reading Block 1 data from HBM:

    $$O^{(\text{new})} = O^{(1)} \cdot \left( \frac{d^{(1)} e^{m^{(1)} - m^{(\text{new})}}}{d^{(\text{new})}} \right) + \left( \frac{e^{x^{(2)} - m^{(\text{new})}}}{d^{(\text{new})}} \right) V^{(2)}$$

This recurrence formula allows FlashAttention to incrementally stream Key and Value blocks through SRAM, updating running output accumulators while keeping memory footprint bounded by $O(N)$ SRAM storage.

FlashAttention Tiling Loops

Let $B_r$ be the SRAM block row size for Query blocks, and $B_c$ be the SRAM block column size for Key/Value blocks.

       KEY MATRIX (K) -> Split into Column Tiles B_c
       +--------+--------+--------+
       |  K_1   |  K_2   |  K_3   |
       +--------+--------+--------+
Q      +--------+--------+--------+
U  Q_1 |  S_1,1 |  S_1,2 |  S_1,3 |  ---> Computed inside SRAM
E      +--------+--------+--------+       Online Softmax Rescaling
R  Q_2 |  S_2,1 |  S_2,2 |  S_2,3 |       Final Output O written to HBM
Y      +--------+--------+--------+

C++ / CUDA Pseudocode for FlashAttention Forward Pass

// FlashAttention Forward Pass CUDA Kernel Conceptual Structure
#include <cuda_runtime.h>
#include <algorithm>
 
template<int Br, int Bc, int d>
__global__ void flash_attention_forward_kernel(
    const float* __restrict__ Q, // [N, d] stored in HBM
    const float* __restrict__ K, // [N, d] stored in HBM
    const float* __restrict__ V, // [N, d] stored in HBM
    float* __restrict__ O,       // [N, d] output in HBM
    float* __restrict__ L_sum,   // [N] logsumexp buffer in HBM
    const int N,
    const float scale) 
{
    // Shared Memory Allocation in SRAM per SM
    __shared__ float s_Q[Br][d];
    __shared__ float s_K[Bc][d];
    __shared__ float s_V[Bc][d];
    __shared__ float s_S[Br][Bc];
 
    int block_row_idx = blockIdx.x; // Outer loop over Q tiles
    int tx = threadIdx.x;
 
    // Local Registers in SRAM per thread
    float r_o[d] = {0.0f};
    float r_m = -1e30f; // Running max
    float r_d = 0.0f;    // Running denominator sum
 
    // Load Q tile into Shared Memory SRAM
    if (block_row_idx * Br + tx < N) {
        for (int i = 0; i < d; ++i) {
            s_Q[tx][i] = Q[(block_row_idx * Br + tx) * d + i];
        }
    }
    __syncthreads();
 
    // Loop over Key and Value blocks in HBM
    int num_col_blocks = (N + Bc - 1) / Bc;
    for (int j = 0; j < num_col_blocks; ++j) {
        
        // Load K and V tiles from HBM to SRAM
        if (j * Bc + tx < N) {
            for (int i = 0; i < d; ++i) {
                s_K[tx][i] = K[(j * Bc + tx) * d + i];
                s_V[tx][i] = V[(j * Bc + tx) * d + i];
            }
        } else {
            for (int i = 0; i < d; ++i) {
                s_K[tx][i] = 0.0f;
                s_V[tx][i] = 0.0f;
            }
        }
        __syncthreads();
 
        // 1. Compute Tile Logits S_tile = Q_tile * (K_tile)^T
        for (int c = 0; c < Bc; ++c) {
            float score = 0.0f;
            for (int i = 0; i < d; ++i) {
                score += s_Q[tx][i] * s_K[c][i];
            }
            s_S[tx][c] = score * scale;
        }
 
        // 2. Compute local maximum for current SRAM tile
        float tile_max = -1e30f;
        for (int c = 0; c < Bc; ++c) {
            tile_max = fmaxf(tile_max, s_S[tx][c]);
        }
 
        // 3. Compute updated global max and scaling factor
        float new_m = fmaxf(r_m, tile_max);
        float alpha = expf(r_m - new_m);
        float tile_d = 0.0f;
        for (int c = 0; c < Bc; ++c) {
            tile_d += expf(s_S[tx][c] - new_m);
        }
 
        // Rescale running denominator sum
        float new_d = r_d * alpha + tile_d;
 
        // 4. Update and Rescale Output Accumulator Register
        for (int i = 0; i < d; ++i) {
            float pv = 0.0f;
            for (int c = 0; c < Bc; ++c) {
                pv += expf(s_S[tx][c] - new_m) * s_V[c][i];
            }
            r_o[i] = r_o[i] * (alpha * (r_d / new_d)) + (pv / new_d);
        }
 
        r_m = new_m;
        r_d = new_d;
        __syncthreads();
    }
 
    // Write final accumulated result from registers to off-chip HBM
    int global_row = block_row_idx * Br + tx;
    if (global_row < N) {
        for (int i = 0; i < d; ++i) {
            O[global_row * d + i] = r_o[i];
        }
        L_sum[global_row] = r_m + logf(r_d);
    }
}

Architectural Generations: FlashAttention-1 vs. FlashAttention-2 vs. FlashAttention-3

The FlashAttention architecture has evolved across three key iterations:

+---------------------------------------------------------------------------------------------------+
| FlashAttention-1: Outer Loop over K,V tiles | Inner Loop over Q tiles                             |
| FP16/BF16 Support | Theoretical Peak TFLOPS Utilization: ~30-40%                                  |
+---------------------------------------------------------------------------------------------------+
                                                  |
                                                  v
+---------------------------------------------------------------------------------------------------+
| FlashAttention-2: Reversed Loop (Outer Loop over Q tiles) | Parallelized across Sequence Dimension |
| Improved Warp Layouts | Reduced Non-GEMM Operations | Peak TFLOPS Utilization: ~50-70%            |
+---------------------------------------------------------------------------------------------------+
                                                  |
                                                  v
+---------------------------------------------------------------------------------------------------+
| FlashAttention-3: Optimized for NVIDIA Hopper (H100) Architecture                                 |
| Asynchronous Data Pipeline via Tensor Memory Accelerator (TMA) | FP8 Tensor Core Execution       |
| Warp Specialization (Producer-Consumer Warps) | Peak TFLOPS Utilization: ~75-85%                  |
+---------------------------------------------------------------------------------------------------+

FlashAttention-1

  • Outer loop over Key/Value tiles, inner loop over Query tiles.
  • Required writing atomic synchronization statistics to HBM or performing redundant re-computation during the backward pass.
  • Achieved $2\times$ to $4\times$ speedups over PyTorch standard attention, utilizing $30-40%$ of GPU theoretical peak TFLOPS.

FlashAttention-2

  • Reversed loop order: Outer loop over Query tiles, inner loop over Key/Value tiles.
  • Distributed work across Streaming Multiprocessors (SMs) along the sequence length dimension of $Q$, eliminating inter-warp atomic synchronization overhead.
  • Optimized warp layout within SM blocks, aligning Matrix Multiply-Accumulate (MMA) layouts to maximize Tensor Core instruction throughput.
  • Increased GPU Tensor Core utilization from $35%$ up to $55-70%$ of peak TFLOPS.

FlashAttention-3

  • Tailored for NVIDIA Hopper architecture (H100/H200).
  • Leverages the hardware Tensor Memory Accelerator (TMA) unit to transfer data asynchronously between HBM and SRAM in background hardware channels without consuming SM register execution cycles.
  • Implements Warp Specialization: Splits SM warps into dedicated Producer warps (issuing TMA memory transfer instructions) and Consumer warps (executing Tensor Core GEMM operations), overlapping data loading with matrix math.
  • Native FP8 precision support (E4M3 and E5M2 formats), maintaining numerical stability via block-scale quantization while yielding up to $1.2 \text{ PFLOPS}$ execution throughput per GPU.

Production Inference Serving Strategies

Combining PagedAttention and FlashAttention forms the core foundation of modern LLM serving engines like vLLM, TensorRT-LLM, and TGI. Serving deployments in enterprise data centres apply three additional runtime strategies to maximize system throughput.

Continuous Batching (Iteration-Level Scheduling)

In traditional static batching, incoming requests are grouped into a batch and run together through prefill and decode iterations.

STATIC BATCHING (Coarse-Grained Request-Level)
Req 1 [Prefill][Decode 1][Decode 2][EOS]-----------------------> IDLE VRAM & COMPUTE
Req 2 [Prefill][Decode 1][Decode 2][Decode 3][Decode 4][EOS]---> BATCH COMPLETES HERE
      |<---------------- GPU Blocked for Duration of Req 2 ----------------->|
 
CONTINUOUS BATCHING (Fine-Grained Iteration-Level)
Iteration 1: [Req 1 Step 1] [Req 2 Step 1]
Iteration 2: [Req 1 Step 2] [Req 2 Step 2]  (Req 1 hits EOS -> Evicted from pool)
Iteration 3: [Req 3 Prefill] [Req 2 Step 3]  (Req 3 scheduled immediately in empty slot)

Static batching creates two critical failure modes:

  1. Short-Request Starvation: Short generation requests are held hostage in GPU memory until the longest request completes decoding.
  2. GPU Idle Stalls: As requests finish early, their batch slots sit idle, lowering GPU arithmetic intensity.

Continuous Batching operates at token iteration granularity. At the end of every forward pass step:

  • Any sequence emitting the End-of-Sequence (EOS) token is immediately evicted, returning its physical PagedAttention blocks to the global free pool.
  • New incoming prefill requests are scheduled directly into freed physical memory blocks during the very next iteration step.

Continuous batching boosts cluster throughput by $3\times$ to $8\times$ relative to static batching.

Chunked Prefills

When an agentic system submits a long 32,000-token prompt into an active serving node, processing its prefill phase requires a massive compute burst. In naive engines, executing this large prefill blocks decoding for hundreds of active client sequences, causing tail latency (P99) spikes in user-facing applications.

UN-CHUNKED PREFILL (Causes Decode Stalls)
[32k Token Large Prefill GEMM (Takes 250ms)] ---> Active Client Decodes Blocked for 250ms
 
CHUNKED PREFILL (Piggybacked Decode Iterations)
Step 1: [Chunk 1: 512 Prefill Tokens] + [Active Decode Tokens Batch (Size 64)] -> (15ms)
Step 2: [Chunk 2: 512 Prefill Tokens] + [Active Decode Tokens Batch (Size 64)] -> (15ms)

Chunked Prefill divides incoming large prompts into fixed-size chunks (e.g. 512 tokens). In each step, the scheduler piggybacks one 512-token prefill chunk alongside the active decode token batch into a unified forward pass. This saturates GPU compute capacity while keeping decode latency tight.

Speculative Execution Integration

To bypass the memory bandwidth limit of single-token decoding ($1 \text{ token/step}$ per forward pass), production servers deploy Speculative Decoding.

1. Draft Model (Small, Fast e.g. Llama-3-8B)
   Generates K candidate draft tokens rapidly: [y_1, y_2, y_3, y_4]
 
2. Target Model (Large, Slow e.g. Llama-3-70B)
   Executes ONE SINGLE PARALLEL PREFILL PASS over [y_1, y_2, y_3, y_4]
   
3. Verification Sampling:
   Target verifies candidate tokens in parallel. Accepts [y_1, y_2, y_3], Rejects [y_4]
   Emits replacement token y_4'
   RESULT: Generated 4 verified tokens in 1 Target Model Forward Pass duration.

Speculative execution works as follows:

  1. A compact Draft Model (e.g., Llama-3-8B) generates $K$ candidate tokens sequentially (e.g. $K=4$) at high speed using minimal VRAM bandwidth.
  2. The Target Model (e.g., Llama-3-70B) receives the $K$ candidate tokens and evaluates them in a single parallel prefill forward pass step.
  3. A modified rejection sampling algorithm accepts draft tokens that match the target model's output distribution.

Because evaluating $K$ candidate tokens in parallel during prefill has a higher arithmetic intensity than $K$ individual decoding steps, speculative execution delivers a $2.0\times$ to $2.8\times$ latency reduction without altering the model's output probability distribution.

Summary & Operational Guidelines

To synthesize the operational impact of these techniques on enterprise hardware, consider the architectural trade-offs:

+--------------------------------------------------------------------------------------------------+
| TECHNIQUE          | PRIMARY BOTTLENECK SOLVED    | OPERATIONAL MECHANISM                        |
+--------------------+------------------------------+----------------------------------------------+
| KV Cache           | Autoregressive Redundancy    | Stores historical K, V tensors in VRAM;     |
|                    |                              | converts O(T^3) compute to O(T^2).           |
+--------------------+------------------------------+----------------------------------------------+
| Grouped-Query      | KV Cache VRAM Footprint      | Shares KV heads across Query groups;         |
| Attention (GQA)    |                              | reduces VRAM size by 8x relative to MHA.     |
+--------------------+------------------------------+----------------------------------------------+
| PagedAttention     | VRAM Capacity Fragmentation  | Virtual memory paging for KV blocks;         |
|                    |                              | eliminates external memory waste (<4% waste).|
+--------------------+------------------------------+----------------------------------------------+
| FlashAttention-3   | HBM DRAM Bandwidth Overhead  | SRAM tiling & online softmax; eliminates     |
|                    |                              | materializing N x N intermediate matrices.   |
+--------------------+------------------------------+----------------------------------------------+
| Continuous         | Coarse Request-Level Stalls  | Iteration-level scheduling; evicts EOS       |
| Batching           |                              | sequences dynamically per generation step.   |
+--------------------+------------------------------+----------------------------------------------+
| Speculative        | Decode Memory Bandwidth Bound| Draft model generates K tokens; Target model |
| Decoding           |                              | verifies K tokens in 1 parallel prefill pass.|
+--------------------+------------------------------+----------------------------------------------+

Data Centre Deployment Checklist

When deploying LLM inference pipelines on NVIDIA H100 clusters across data centres in Frankfurt or Amsterdam, infrastructure engineers should enforce the following operational parameters:

  1. Precision Configuration: Deploy models in FP8 or BF16 precision. For FP8 execution on H100, ensure KV cache blocks are stored in FP8 to double active batch capacity.
  2. PagedAttention Block Size Tuning: Set PagedAttention block size ($B_{\text{block}}$) to 16 for short sequence workloads or 32 for long context applications ($>16,384$ tokens). Avoid block sizes $>64$ to prevent internal allocation fragmentation.
  3. FlashAttention Optimization: Verify that FlashAttention-3 kernels are active for Hopper architectures. Set chunked prefill max token limits to match the GPU's arithmetic intensity saturation point ($N_{\text{chunk}} = 512$ or $1,024$).
  4. Memory Allocation Ratio: Set the server GPU memory utilization fraction (gpu_memory_utilization) to 0.90. This reserves $90%$ of total VRAM for model weights and the PagedAttention block pool, while retaining $10%$ for CUDA driver overhead and dynamic runtime activation buffers.