$ cat ./blog/gen-ai/llm-fundamentals/kv-cache-prefill-va-decode.md
seriesHow LLMs Workpart 13

Inference at Scale: KV Cache, Prefill, and Decode

Part 13 of the How LLMs Work series: why serving one user and serving a hundred users on the same GPU are so different. Every request goes through two phases: prefill saturates compute, decode is limited by memory bandwidth. Starting from arithmetic intensity, the post walks through the KV cache, MQA/GQA/MLA, PagedAttention, continuous batching, FlashAttention, and speculative decoding, then shows where long context really gets expensive.

~/diagrams/kv-cache-prefill-va-decode/inference-request-pipeline-en.htmlinteractive
open ↗
FigThe second row is the path of one request: the prompt is tokenized, waits in the queue for a batch slot, runs prefill over the whole prompt at once (compute-bound, sets TTFT), then enters the decode loop that produces one token at a time (memory-bound, sets TPOT), streams tokens until EOS, and finally frees its KV blocks or keeps the prefix for later requests. The third row shows server-side mechanisms: continuous batching admits new requests every step; prefill writes K and V into a block-based KV cache; the attention kernel (FlashAttention) reads K and V at every decode step; speculative decoding (optional) yields several tokens per round. The top and bottom rows are nine common misconceptions, each linked to the step it gets wrong. Section 10 explains each one.
compute step on the GPUschedulerred frame: request in and tokens out; purple frame: KV cache; grey frame: misconception

Part twelve ended with a promise: inference splits into two very different phases, prefill limited by compute and decode limited by memory bandwidth. This post explains why, and works from there to every technique a modern inference server uses.

The question of this post: why is serving one person on a GPU so different from serving a hundred? One user chatting with an 8B model on their own GPU sees a few dozen tokens per second, and the GPU sits mostly idle. A hundred users at once get nearly a hundred times the total throughput, until memory runs out. Understand those two sentences and you understand the KV cache, batching, PagedAttention, FlashAttention, and speculative decoding.

The whole post turns on a single question: how much work does each byte read from memory do? Every technique that follows either reads fewer bytes per token or does more work per byte read. The post follows the order of Figure 1 at the top of the page.


01From a model to a service

Part 2, section 8 drew the big picture: latency and throughput often conflict (8.1), generation is autoregressive and cannot be parallelized over time (8.2), the KV cache trades VRAM for compute (8.3), PagedAttention manages the cache like virtual memory (8.4), continuous batching admits new requests as soon as a slot opens (8.5), and FlashAttention optimizes at the kernel level (8.6). This post does not repeat those parts. It takes up the question part 2 left open: why those techniques work, and how far they go.

Three metrics run through the post:

Metric Meaning Phase that sets it
TTFT (time to first token) from sending the request to receiving the first token queueing + prefill
TPOT (time per output token), also called ITL (inter-token latency) the gap between two consecutive streamed tokens decode
Throughput total tokens generated per second across the whole system decode, times the batch

A chat user feels TTFT (how long before text appears) and TPOT (how fast the text flows). Whoever pays for the GPU cares about throughput. The rest of the post explains why these three numbers depend on very different things.


02Prefill and decode: two phases, two bottlenecks

The two phases of a request

When a request arrives, the model runs two phases:

  • Prefill. The whole n-token prompt goes through the model in one forward pass. Every multiplication with a weight matrix is matrix × matrix: n input vectors times the same weight matrix. K and V for every prompt token are written into the KV cache. The output is the logits at the last position, which give the first token. Prefill sets TTFT.
  • Decode. Each step has only one new token (per sequence). Multiplication with the weights is matrix × vector. The new token computes its own Q, K, and V, appends one row to the KV cache, and attention rereads all the old K and V. This repeats until EOS. Decode sets TPOT.

Both phases do the same kind of arithmetic, but the ratio of computing to reading is completely different.

Counting FLOPs and counting bytes

Part 8, section 6 already did the counting: one forward pass costs roughly 2N FLOPs per token, where N is the parameter count, plus an attention term that grows with context length. This post needs one more count: every forward pass has to read all the weights from GPU memory (HBM, or GDDR6 on consumer cards) onto the chip to multiply them. At fp16 that is 2 bytes per parameter, so 2N bytes.

The ratio of those two numbers is the arithmetic intensity: the FLOPs done per byte read.

  • Decode, batch 1: 2N FLOPs over 2N bytes of weights, about 1 FLOP/byte.
  • Decode, batch B: one read of the weights serves B sequences, so 2N·B FLOPs over 2N bytes, about B FLOPs/byte (ignoring the KV cache for now; it comes later).
  • Prefill, n-token prompt: one read of the weights serves n tokens, about n FLOPs/byte.

The roofline and the RTX 2080 Ti’s ridge point

The roofline model (Williams, Waterman & Patterson 2009) says a kernel is limited by the slower of two things: how fast it can compute, or how fast it can read data. The minimum time is:

t ≥ max( FLOPs / peak_FLOPS ,  bytes / bandwidth )

The crossover between the two regimes is the ridge point: peak_FLOPS / bandwidth. A kernel with arithmetic intensity below the ridge point is memory-bound (the GPU waits for data); above it, it is compute-bound.

Take the RTX 2080 Ti, the card used in earlier parts. According to Table 1 of the NVIDIA Turing Architecture Whitepaper (reference card, not Founders Edition):

RTX 2080 Ti (Turing TU102) spec Value
Memory bandwidth (GDDR6, 352-bit, 14 Gbps) 616 GB/s
Memory 11,264 MB
Peak FP16 Tensor TFLOPS, FP16 accumulate 107.6
Peak FP16 Tensor TFLOPS, FP32 accumulate 53.8
Peak FP32 TFLOPS (no Tensor Cores) 13.4

When multiplying fp16 matrices, PyTorch usually accumulates in fp32 for accuracy, so this post uses 53.8 TFLOPS as the reference. The ridge point:

53.8e12 FLOP/s / 616e9 byte/s ≈ 87 FLOPs/byte     (FP32 accumulate)
107.6e12 / 616e9              ≈ 175 FLOPs/byte    (FP16 accumulate)

For comparison, kipply computes it for the A100: 312 TFLOPS bf16 and 1.5 TB/s, a ratio of about 208 (Transformer Inference Arithmetic, “kv cache” section). Newer GPUs grow compute faster than bandwidth, so the ridge point keeps rising and decode wastes ever more compute.

Decode at batch 1 has an intensity of about 1, nearly 90 times below the ridge point. While decoding a single sequence, the 2080 Ti’s Tensor Cores are idle about 99% of the time. That is what part 12 meant by “decode is limited by memory bandwidth”.

A FLOPs-per-byte table for Qwen2.5-1.5B

Apply this to a model that fits on a 2080 Ti: Qwen2.5-1.5B at fp16 (1.54B parameters, 3.09 GB of weights; 28 layers, 12 query heads, 2 KV heads, head_dim 128, i.e. 28 KiB of KV cache per token, computed as in section 3). The table below is a theoretical roofline calculation, not a measurement: FLOPs include 2N plus attention, bytes include the weights plus the KV cache read.

Phase FLOPs/byte Against ridge point ≈ 87 Lower bound on time per pass
Prefill, 128 tokens ≈ 128 compute-bound ≈ 7.4 ms
Prefill, 2,048 tokens ≈ 2,124 compute-bound ≈ 124 ms
Decode batch 1, context 1,024 ≈ 1.0 memory-bound ≈ 5.1 ms
Decode batch 16, context 1,024 ≈ 14.7 memory-bound ≈ 5.8 ms
Decode batch 64, context 1,024 ≈ 42 memory-bound ≈ 8.1 ms
Decode batch 256, context 1,024 ≈ 79 near the ridge point ≈ 17.2 ms

Two things to read from the table:

  1. Decode is nearly “free” as the batch grows, up to a threshold. From batch 1 to batch 16, the minimum time per step only goes from 5.1 to 5.8 ms, while tokens produced per step grow 16 times. Most of the time goes into reading 3.09 GB of weights, and that read is shared by the whole batch.
  2. The threshold is not only compute but also the KV cache. At batch 256 with context 1,024, the KV cache is already 256 × 1,024 × 28 KiB ≈ 7.5 GB; add 3.09 GB of weights and you are close to the card’s 11 GB. Memory runs out before the ridge point is reached.

Why the KV cache drags intensity down

The weights are shared across the batch; the KV cache is not. Each sequence has its own K and V, and every decode step rereads all of it. The attention part of decode has a fixed intensity that does not grow with batch: each K element read (2 bytes) is shared by g query heads under GQA, each doing one multiply-add (2 FLOPs), so attention during decode runs at roughly g FLOPs/byte, where g is the number of query heads per KV head. Qwen2.5-1.5B has g = 6, Llama 3.1 8B has g = 4, an MHA model has g = 1.

So the longer the context, the deeper decode sinks into memory-bound territory, even at large batch. At the same batch 64 on Qwen2.5-1.5B, the theoretical intensity drops from about 42 at context 1,024 to about 23 at context 4,096 and about 9 at context 32,768 (where the KV cache is already 60 GB, which fits on no card in reach). This is the thread that ties every later technique together: memory available for the KV cache sets the maximum batch, and the batch sets the throughput.

Why weight-only quantization helps decode but not prefill

Part 12, sections 5–6 cited the measurements in the llama.cpp README: at Q4_K_M, text generation is about 2.5 times faster than F16 (29 to 72 tokens/s), while prompt processing gets slightly slower (923 to 822 tokens/s). The roofline explains it immediately: decode is memory-bound, so weights four times smaller mean almost four times fewer bytes read; prefill is compute-bound, and the multiplications still run at 16 bits after dequantization, so it is no faster, and the dequantize step even makes it slightly slower.

Code sample 1: measuring TTFT and time per decode step

The code below measures (a) prefill time by prompt length and (b) time per decode step by batch size, then converts the result into achieved bandwidth to compare with 616 GB/s.

# Illustrative code. Python 3.12, torch 2.14.1+cu130, transformers 5.18.0
import time, torch
from transformers import AutoModelForCausalLM

name = "Qwen/Qwen2.5-1.5B"
model = AutoModelForCausalLM.from_pretrained(name, dtype=torch.float16, device_map="cuda").eval()
cfg = model.config
head_dim = cfg.hidden_size // cfg.num_attention_heads
kv_per_token = 2 * cfg.num_hidden_layers * cfg.num_key_value_heads * head_dim * 2  # bytes, fp16
weight_bytes = sum(p.numel() * p.element_size() for p in model.parameters())

def sync_time():
    torch.cuda.synchronize()
    return time.perf_counter()

@torch.no_grad()
def prefill(ids):
    # logits_to_keep=1: only compute logits at the last position, avoiding a batch × n × vocab tensor
    return model(ids, use_cache=True, logits_to_keep=1)

# (a) Prefill ≈ TTFT when there is no queueing
for n in (128, 512, 2048, 8192):
    ids = torch.randint(1000, 30000, (1, n), device="cuda")
    prefill(ids)                                    # warmup
    t0 = sync_time(); prefill(ids); dt = sync_time() - t0
    print(f"prompt {n:5d}: prefill {dt * 1e3:7.1f} ms  ({n / dt:8.0f} token/s)")

# (b) Each decode step by batch size, starting context 1,024 tokens
ctx, steps = 1024, 32
for B in (1, 4, 16, 64):
    ids = torch.randint(1000, 30000, (B, ctx), device="cuda")
    with torch.no_grad():
        out = prefill(ids)
        past, nxt = out.past_key_values, out.logits[:, -1:].argmax(-1)
        t0 = sync_time()
        for _ in range(steps):
            out = model(nxt, past_key_values=past, use_cache=True)
            past, nxt = out.past_key_values, out.logits[:, -1:].argmax(-1)
        dt = (sync_time() - t0) / steps
    kv_read = B * (ctx + steps // 2) * kv_per_token  # average KV read per step
    print(f"batch {B:3d}: {dt * 1e3:6.2f} ms/step  {B / dt:8.0f} token/s  "
          f"≈ {(weight_bytes + kv_read) / dt / 1e9:5.0f} GB/s")

Not run. Needs an NVIDIA GPU with about 8 GB or more of free VRAM (for example an RTX 2080 Ti 11 GB), CUDA 13.0, Python 3.12, torch 2.14.1+cu130, transformers 5.18.0.

What you should see when you run it: in part (b), ms per step grows very slowly from batch 1 to batch 16, so tokens/s grows almost linearly; the derived bandwidth creeps toward 616 GB/s but cannot exceed it. In part (a), prefill tokens/s is hundreds of times higher than decode at batch 1. The theoretical table above tells you which order of magnitude to expect; real numbers are always worse than the lower bound because of kernel and Python overhead.


03KV cache: the memory math

Bytes per token

The KV cache size formula is in part 2, section 8.3. Written as bytes per token, it is easier to compare across models:

KV bytes per token = 2 × num_layers × num_kv_heads × head_dim × bytes_per_value

The 2 is for K and V. Multiply by the total number of tokens in the batch (all sequences combined) to get the total KV cache. The vLLM paper uses OPT-13B (MHA, 40 layers, hidden size 5,120) as an example: 800 KB per token at fp16, and a single 2,048-token request can take up to 1.6 GB (Kwon et al. 2023, section 3, “Large KV cache” paragraph).

The KV cache stores only K and V, not Q and not the attention scores. Q of an old token is never used again: at a decode step, only the new token’s Q is multiplied with all old Ks (part 7, section 7, in Vietnamese, explains why the causal mask means old tokens never have to “look back” at new ones).

Code sample 2: computing the KV cache from config.json

Instead of looking things up in a table, read config.json directly from the Hugging Face Hub. The code also counts parameters from the safetensors file headers without downloading the weights.

# Python 3.12, huggingface_hub 1.33.0 — only downloads config.json and safetensors headers, runs on CPU
import json, os
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
from huggingface_hub import hf_hub_download, get_safetensors_metadata

MODELS = {
    "Llama 3.1 8B":   "unsloth/Meta-Llama-3.1-8B",  # ungated mirror, config identical to meta-llama
    "Llama 2 7B":     "NousResearch/Llama-2-7b-hf", # MHA: as many KV heads as query heads
    "Qwen2.5-7B":     "Qwen/Qwen2.5-7B",
    "Mistral 7B v0.3": "mistralai/Mistral-7B-v0.3",
    "DeepSeek-V3":    "deepseek-ai/DeepSeek-V3",    # MLA
}
KiB, GiB = 1024, 1024**3

def kv_bytes_per_token(cfg, bytes_per_value=2):
    L = cfg["num_hidden_layers"]
    if "kv_lora_rank" in cfg:                      # MLA: cache the latent c_KV plus the RoPE part of the key
        return L * (cfg["kv_lora_rank"] + cfg["qk_rope_head_dim"]) * bytes_per_value
    n_kv = cfg.get("num_key_value_heads", cfg["num_attention_heads"])
    head_dim = cfg.get("head_dim") or cfg["hidden_size"] // cfg["num_attention_heads"]
    return 2 * L * n_kv * head_dim * bytes_per_value  # 2 = K and V

rows = {}
print(f"{'model':16} {'attn':4} {'L':>3} {'q/kv':>7} {'KiB/token':>9} {'GiB @128k':>9} {'weights GiB':>11}")
for name, repo in MODELS.items():
    cfg = json.load(open(hf_hub_download(repo, "config.json")))
    n_params = get_safetensors_metadata(repo).parameter_count  # counted from headers, no weight download
    per_tok = kv_bytes_per_token(cfg)
    n_q = cfg["num_attention_heads"]
    n_kv = cfg.get("num_key_value_heads", n_q)
    kind = "MLA" if "kv_lora_rank" in cfg else ("MHA" if n_kv == n_q else "GQA")
    weights = sum(n_params.values()) * 2             # assume bf16 for every parameter
    rows[name] = (cfg, per_tok, weights)
    print(f"{name:16} {kind:4} {cfg['num_hidden_layers']:>3} {f'{n_q}/{n_kv}':>7} "
          f"{per_tok / KiB:>9.1f} {per_tok * 131072 / GiB:>9.2f} {weights / GiB:>11.1f}")

# Budget of one 80 GB GPU for Llama 3.1 8B: whatever is left after the weights goes to the KV cache
cfg, per_tok, weights = rows["Llama 3.1 8B"]
budget = 80e9 * 0.9 - weights                       # keep 10% for activations, CUDA context...
n_params = weights / 2
non_embed = n_params - 2 * cfg["vocab_size"] * cfg["hidden_size"]  # drop embedding and lm_head
d_attn = cfg["num_attention_heads"] * cfg["head_dim"]
print(f"\nLlama 3.1 8B: weights {weights/1e9:.1f} GB, KV budget {budget/1e9:.1f} GB, non-embedding N {non_embed/1e9:.2f}B")
print(f"{'context':>8} {'KV/seq GB':>9} {'max seqs':>8} {'prefill TFLOPs':>14} {'attention %':>11}")
for n in (4096, 32768, 131072):
    kv_seq = per_tok * n
    linear = 2 * non_embed * n                               # 2N per token
    attn = 2 * cfg["num_hidden_layers"] * n * d_attn * n     # part 8's context term, n_ctx = n
    print(f"{n:>8} {kv_seq/1e9:>9.2f} {int(budget // kv_seq):>8} {(linear+attn)/1e12:>14.1f} {attn/(linear+attn):>11.0%}")
print(f"crossover: attention = 2N at n ≈ {non_embed / (cfg['num_hidden_layers'] * d_attn):,.0f} tokens")

Output (run on 2026-10-08):

model            attn   L    q/kv KiB/token GiB @128k weights GiB
Llama 3.1 8B     GQA   32    32/8     128.0     16.00        15.0
Llama 2 7B       MHA   32   32/32     512.0     64.00        12.6
Qwen2.5-7B       GQA   28    28/4      56.0      7.00        14.2
Mistral 7B v0.3  GQA   32    32/8     128.0     16.00        13.5
DeepSeek-V3      MLA   61 128/128      68.6      8.58      1275.0

Llama 3.1 8B: weights 16.1 GB, KV budget 55.9 GB, non-embedding N 6.98B
 context KV/seq GB max seqs prefill TFLOPs attention %
    4096      0.54      104           61.6          7%
   32768      4.29       13          738.9         38%
  131072     17.18        3         6333.3         71%
crossover: attention = 2N at n ≈ 53,250 tokens

The original Llama 3.1 8B and Llama 2 7B repos are gated on Hugging Face, so the code uses mirrors (unsloth/…, NousResearch/…). The configuration matches Table 3 of the Llama 3 paper: 32 layers, 32 attention heads, 8 key/value heads for the 8B model (Llama Team 2024, section 3.2).

Reading the table

  • Llama 3.1 8B costs 128 KiB per token. A 128k-token sequence takes 16 GiB of KV cache, about as much as all its bf16 weights (15 GiB). At long context, one request “weighs” as much as a copy of the model.
  • MHA costs four times as much. Llama 2 7B, with 32 KV heads, takes 512 KiB per token. On the same memory budget it serves four times fewer sequences than Llama 3.1 8B, even though the model is smaller.
  • Qwen2.5-7B needs only 56 KiB thanks to 4 KV heads for 28 query heads, so each KV head serves 7 query heads.
  • DeepSeek-V3 has 671B parameters but only 68.6 KiB of KV cache per token, less than Llama 3.1 8B. That is MLA, explained in section 4. The 1,275 GiB weights column assumes bf16; the released model uses FP8, so it is about half that.

The lower part of the output is the budget for one 80 GB GPU: after 16.1 GB of weights and a 10% reserve, about 56 GB remain for the KV cache. At 4k context the GPU holds 104 sequences at once; at 32k, 13; at 128k, only 3. Decode throughput is proportional to the batch, so on the same GPU, a context 32 times longer means about 35 times fewer concurrent requests. The last two columns are used in section 9.

The memory budget of a GPU

GPU memory during serving splits into three parts:

VRAM = weights + KV cache + (temporary activations, CUDA context, kernel buffers)

The weights are fixed. Activations are small at inference because nothing is kept for a backward pass, only for the layer currently running. Everything else goes to the KV cache, and that portion sets the maximum number of concurrent requests. Inference servers like vLLM pre-allocate a fixed fraction of VRAM (the gpu_memory_utilization parameter) and carve whatever is left after the weights into KV blocks.

KV cache quantization

If weights can be quantized as in part 12, so can the KV cache. The simplest approach stores K and V in FP8 instead of fp16/bf16: 1 byte per value instead of 2, which by straightforward arithmetic halves the KV cache. The vLLM docs say quantizing the KV cache to FP8 “can significantly reduce its memory footprint”, letting you store more tokens, which improves throughput and supports longer contexts; the two formats are fp8_e4m3 and fp8_e5m2. Going further, KIVI (Liu et al. 2024) quantizes the KV cache to 2 bits, with an observation that echoes the outlier story of part 12: keys should be quantized per channel and values per token (abstract). On Llama-2-7B with one A100 80 GB, KIVI cuts peak memory 2.6 times (weights included), allows up to 4 times larger batches, and brings 2.35–3.47 times higher throughput (abstract; section 4.2.4, Figure 5). The price is the same as in part 12: you have to re-measure quality on your own task. Hugging Face Transformers ships a QuantizedCache (cache_implementation="quantized") to try it.


04MQA, GQA, and MLA

Looking at the bytes-per-token formula, there are only two knobs that do not touch model depth: the number of KV heads and bytes per value. Section 3 just covered the second. This section covers the first. Part 8, section 6 noted that GQA reduces the parameters of W_K and W_V and the KV cache but does not reduce the FLOPs of QKᵀ; part 1, section 5.10, has a picture of MHA, MQA, and GQA. This section only deals with memory and quality.

MQA: one K/V pair for every head

Multi-query attention (Shazeer 2019) keeps many query heads but has all of them share one K, V pair. The abstract states the goal plainly: reduce the size of the K, V tensors “and hence the memory bandwidth requirements of incremental decoding”. On WMT14 English–German, the decoder went from 46 to 3.8 microseconds per token (Table 2, section 4.3; measured on TPUv2, amortized over a batch of 1,024), while BLEU barely moved: 27.7 versus 27.5 with greedy decoding (Table 1, section 4.2). Pope et al. 2022 show that at PaLM 540B scale, multiquery attention allows context lengths up to 32 times longer within the same memory budget (abstract; Table 1, section 4.2).

The downside: quality does drop, and training is harder to keep stable. With the KV cache shrinking by a factor of n_heads, that is a big trade.

GQA: grouping the heads

Grouped-query attention (Ainslie et al. 2023) sits in between: split the n_heads query heads into n_kv groups, each sharing one K, V pair. n_kv = n_heads is MHA, n_kv = 1 is MQA. The main result is Table 1 (section 3.2) of the paper, on T5-XXL:

Model Inference time per sample Average score
MHA-XXL 1.51 s 47.2
MQA-XXL 0.24 s 46.6
GQA-8-XXL 0.28 s 47.1

GQA with 8 groups has quality close to MHA and speed close to MQA. A practical point: no need to train from scratch. The paper uptrains an MHA checkpoint by merging the K and V heads in each group by mean-pooling (sections 2.1–2.2), then training further with about 5% of the original pretraining compute (abstract; section 3.1 gives α = 0.05). Llama 3 uses GQA with 8 KV heads at all three sizes, 8B, 70B, and 405B (Table 3, section 3.2 of the Llama 3 paper).

MLA: compressing K/V into a latent

Multi-head latent attention (DeepSeek-V2, 2024) goes another way. Instead of fewer heads, MLA projects the input down to a small latent vector c_KV and caches only that latent; each head’s K and V are projected back up from it when needed. The RoPE positional encoding cannot pass through that projection, so MLA splits off a small key part that carries RoPE and caches it separately. Each token caches only d_c + d_h^R values per layer, with d_c = 512 and d_h^R = 64 (section 3.1.2). Table 1 (section 2.1.4) of the paper compares KV cache per token:

Mechanism KV cache per token (elements) Capability per the paper
MHA 2 · n_h · d_h · l strong
GQA 2 · n_g · d_h · l moderate
MQA 2 · d_h · l weak
MLA (d_c + d_h^R) · l ≈ 9/2 · d_h · l stronger than MHA

The table caption says MLA’s KV cache equals GQA with only 2.25 groups, yet its performance is stronger than MHA. Compared with DeepSeek 67B, DeepSeek-V2 cuts the KV cache by 93.3% and raises maximum generation throughput to 5.76 times (abstract). Code sample 2 above applies exactly this formula to DeepSeek-V3: 61 × (512 + 64) × 2 bytes ≈ 68.6 KiB per token.

The cost: MLA is more complex, and to run fast the up-projection has to be “absorbed” into W_Q and W_O, which needs dedicated kernels. Converting an already trained model to MLA is also not as simple as GQA’s mean-pooling.

KV/token vs MHA Quality How you get it
MHA 1× baseline train from scratch
GQA (8 groups, 32 heads) 1/4 close to MHA train from scratch or uptrain with ~5% compute
MQA 1/n_heads slightly lower train from scratch or uptrain
MLA ≈ 9/2 · d_h instead of 2 · n_h · d_h DeepSeek reports stronger than MHA train from scratch

05PagedAttention and prefix caching

The problem: fragmentation

Part 2, section 8.4 covered the core idea of PagedAttention. Here is the number behind it: the vLLM paper measured that in earlier systems only 20.4–38.2% of KV cache memory actually held token states (section 1, Figure 2), meaning roughly 62–80% was wasted. Three sources of waste:

  • Reservation: space set aside for tokens that will be generated later, unused for now but unavailable to anyone else.
  • Internal fragmentation: allocation by max_length, so short requests leave the tail empty.
  • External fragmentation: contiguous regions of different sizes leave gaps between them.

With vLLM, 96.3% of the memory held token states in the same experiment (Figure 2), and throughput rose 2–4 times at the same latency compared with FasterTransformer and Orca (abstract).

The block table

PagedAttention borrows the idea of virtual memory from operating systems:

  • The KV cache is split into fixed-size blocks, by default 16 tokens per block in vLLM (section 7.2).
  • Each request sees a contiguous sequence of logical blocks; a block table maps logical blocks to physical blocks located anywhere in GPU memory.
  • A new block is only allocated when the last one is full. The maximum waste is the empty part of each sequence’s last block.

The attention kernel reads K and V through the block table instead of from one contiguous region. That is why it gets its own name: the kernel had to be rewritten to cope with scattered data.

~/diagrams/kv-cache-prefill-va-decode/paged-kv-cache-en.htmlinteractive
open ↗
FigThe top and bottom rows are the logical blocks of request A and of request B (with its second sample branch B′), in token order; the middle row is the physical blocks in GPU memory, 16 tokens each. The first two blocks of A and B hold the same system prompt, so they point to the same physical blocks 7 and 1 (ref = 2) and are not stored twice. A's own question lives in block 3, B's in block 5. B′ initially shares block 5 with B; when B′ needs to write a token different from B's, block 5 is copied to block 8 (copy-on-write) and only B′'s table points to the new block.
physical block private to one requestlogical block in the block tablepurple frame: shared physical block; grey frame: new block from copy-on-write

Block sharing and copy-on-write

Once the block is the unit of management, sharing becomes cheap. Each physical block has a reference count. Two important cases:

  • Parallel sampling: generating several samples for one prompt (the API’s n parameter). The samples share the prompt’s blocks. When a sample needs to write into a shared block, vLLM copies the block first and then writes, which is copy-on-write at block granularity (section 4.4, Figure 8).
  • Beam search: beams share a common prefix, and the shared part changes constantly as beams are pruned. Part 9, section 3 noted that beam search costs B times the KV cache; with PagedAttention, the overlap is stored once.

The paper measured on OPT-13B with the Alpaca trace: block sharing saves 6.1–9.8% of memory for parallel sampling and 37.6–55.2% for beam search (section 6.3, Figure 15).

Prefix caching

The next step is keeping blocks across requests. Two requests with the same system prompt, the same few-shot examples, or the same conversation history up to the last turn have identical K and V for that prefix. No need to prefill it again.

  • vLLM’s automatic prefix caching: each full block is identified by a hash of itself plus the hash of its parent block, i.e. of the entire prefix up to that block. A new request looks up the hash; on a match it bumps the ref count and skips that part of prefill.
  • SGLang’s RadixAttention (Zheng et al. 2023, section 3) keeps a radix tree mapping token sequences to the KV cache, with one-token pages, and evicts by LRU (the least recently used leaf first). Combined with scheduling that favors requests sharing a prefix, SGLang reports up to 6.4 times higher throughput (abstract).

Prefix caching only helps TTFT and prefill compute; it does not make decode faster, and kept blocks still take KV cache space.

Prompt caching on the API side

The prompt caching offered by commercial APIs is the same idea seen from the invoice: a prefix whose K and V are already cached costs less, because the provider does not have to prefill it again. The common rule is that the prefix must match from the very first token: change one character in the system prompt and you lose the cache for everything after it. So put the fixed part (system prompt, documents, examples) first and the changing part (the user’s question) last. Section 9 covers the price ratios.


06Scheduling: batching, chunked prefill, disaggregation

Continuous batching

Part 2, section 8.5 has an example of static versus continuous batching. The origin is Orca (Yu et al., OSDI 2022), with two ideas:

  • Iteration-level scheduling: the scheduler decides at the level of each step (iteration) instead of each request. A finished request leaves the batch immediately, and a new one joins at the very next step (section 3, solution S1).
  • Selective batching: sequences in a batch have different lengths and cannot be stacked into one rectangular tensor. Orca batches only the operations that do not care about length (the multiplications with weights) and runs attention separately for each sequence (section 3, solution S2). PagedAttention later solved the attention part more neatly.

On GPT-3 175B, Orca reaches 36.9 times the throughput of FasterTransformer at the same latency (abstract; section 6.2, Figure 10b). The number is that large because the baseline has to wait for the longest request in the batch to finish.

From section 2 we know why continuous batching works: it keeps the batch full, and a full batch is the only way to spread the cost of reading the weights over many tokens.

The remaining problem: prefill cutting in

Continuous batching creates a new conflict. When a new request arrives, it needs a prefill. An 8,000-token prefill is hundreds of milliseconds of compute; if it runs in the same step as the sequences being decoded, every user currently streaming stalls for that step. Their TPOT spikes. Conversely, if prefill is postponed to keep TPOT steady, the new request’s TTFT grows.

Chunked prefill

Sarathi-Serve (Agrawal et al. 2024) cuts prefill into chunks of equal compute and runs them across several steps (section 4.1). Each step has a token budget: the decoding sequences stay in the batch, and the rest of the budget goes to one prefill chunk, which the paper calls stall-free batching (section 4.2). Decode is memory-bound and prefill is compute-bound, so mixing them in one step even puts to use the compute that decode leaves idle. Against vLLM, Sarathi-Serve reports 2.6 times higher serving capacity for Mistral-7B on one A100 and up to 3.7 times for Yi-34B on two A100s (abstract). In vLLM V1, chunked prefill is enabled by default whenever possible (vLLM docs, “Optimization and Tuning”).

Splitting prefill and decode onto separate GPUs

The more radical route is disaggregation: prefill runs on one group of GPUs, decode on another, and the KV cache is shipped over the network once prefill is done. The two phases no longer compete, and each group can be configured on its own (parallelism, batch size, even GPU type). DistServe (Zhong et al. 2024) reports serving 7.4 times more requests, or meeting a 12.6 times tighter SLO, while keeping over 90% of requests within latency limits (abstract). Splitwise (Patel et al. 2023) reports 1.4 times higher throughput at 20% lower cost, or 2.35 times higher throughput under the same power and cost budget (abstract). The price is moving the KV cache between machines, which needs fast networking and a much more complex system.

Approach TTFT TPOT Operational complexity
Static batching poor (waits for the whole batch) steady low
Continuous batching good jumps when a prefill cuts in medium
+ chunked prefill good, slightly higher for long prompts steadier medium
Prefill/decode disaggregation good steadiest high: KV over the network, two GPU pools

07FlashAttention: fewer bytes, same FLOPs

Part 2, section 8.6 and part 7, section 10 (in Vietnamese) said FlashAttention computes in blocks inside SRAM, never builds the full attention matrix in HBM, and gives exact rather than approximate results. What they did not say is how softmax can be computed block by block when it needs the sum over the whole row.

Online softmax

Numerically stable softmax needs two things from the whole row: the maximum m (subtracted to avoid overflow) and the sum l of exp(x − m). The naive way reads the row three times: once for the max, once for the sum, once to normalize. Online softmax (Milakov & Gimelshein 2018, Algorithm 3 in section 3) merges the first two passes by keeping a running max and sum, and rescaling the sum whenever the max changes:

m_j = max(m_{j-1}, x_j)
l_j = l_{j-1} · exp(m_{j-1} − m_j) + exp(x_j − m_j)

The factor exp(m_{j-1} − m_j) shrinks the old sum onto the scale of the new max. That paper measured softmax running up to 1.3 times faster just by dropping one memory pass (abstract).

FlashAttention applies the same idea to the output. Walk through K and V block by block; for each block, compute its scores, update m and l, then rescale the accumulated output by the same factor before adding the new part:

for each block j of K, V:
  S   = Q · K_jᵀ / √d                      (only n × block, held in SRAM)
  m'  = max(m, rowmax(S))
  P   = exp(S − m')
  l   = l · exp(m − m') + rowsum(P)
  O   = O · exp(m − m') + P · V_j
  m   = m'
finally: O = O / l

Code sample 3a: online softmax versus standard attention

Here is the algorithm above in plain PyTorch, running on CPU, compared with standard attention (which builds the full n × n matrix) and with PyTorch’s scaled_dot_product_attention. An fp64 version serves as the reference.

# Python 3.12, torch 2.14.1 — runs on CPU, fp32
import torch, torch.nn.functional as F
torch.manual_seed(0)

def attention_standard(q, k, v):
    s = q @ k.transpose(-1, -2) / q.shape[-1] ** 0.5       # builds the full n × n matrix
    mask = torch.ones(s.shape[-2:], dtype=torch.bool).triu(1)
    return torch.softmax(s.masked_fill(mask, float("-inf")), dim=-1) @ v

def attention_online(q, k, v, block=128):
    """Walk K, V in blocks; each query row keeps only a running max m, sum l, and output o."""
    n, d = q.shape[-2], q.shape[-1]
    o = torch.zeros_like(q)
    m = torch.full(q.shape[:-1] + (1,), float("-inf"))
    l = torch.zeros(q.shape[:-1] + (1,))
    rows = torch.arange(n).unsqueeze(1)
    for j in range(0, n, block):
        kj, vj = k[..., j:j + block, :], v[..., j:j + block, :]
        s = q @ kj.transpose(-1, -2) / d ** 0.5            # only n × block
        s = s.masked_fill(torch.arange(j, j + kj.shape[-2]) > rows, float("-inf"))
        m_new = torch.maximum(m, s.amax(-1, keepdim=True))
        p = torch.exp(s - m_new)
        scale = torch.exp(m - m_new)                        # rescale what was summed under the old max
        l = l * scale + p.sum(-1, keepdim=True)
        o = o * scale + p @ vj
        m = m_new
    return o / l

q, k, v = (torch.randn(8, 2048, 128) for _ in range(3))   # 8 heads, 2,048 tokens, head_dim 128
ref64 = attention_standard(q.double(), k.double(), v.double())
std = attention_standard(q, k, v)
onl = attention_online(q, k, v)
sdpa = F.scaled_dot_product_attention(q, k, v, is_causal=True)
for name, x in [("standard fp32", std), ("online fp32", onl), ("SDPA fp32", sdpa)]:
    print(f"{name:14} max |x - ref_fp64| = {(x.double() - ref64).abs().max():.2e}")
print(f"online vs standard: max diff {(onl - std).abs().max():.2e}, "
      f"allclose(atol=1e-5) = {torch.allclose(onl, std, atol=1e-5)}")
n, b = 2048, 128
print(f"score elements held at once per head: standard {n*n:,} · online {n*b:,} ({n*n // (n*b)}× fewer)")

Output:

standard fp32  max |x - ref_fp64| = 8.70e-07
online fp32    max |x - ref_fp64| = 9.05e-07
SDPA fp32      max |x - ref_fp64| = 1.00e-06
online vs standard: max diff 5.96e-07, allclose(atol=1e-5) = True
score elements held at once per head: standard 4,194,304 · online 262,144 (16× fewer)

All three deviate from the fp64 version by about 1e-6, exactly the rounding error of fp32. Online softmax is not an approximation: it is the same computation with a different order of additions, and as part 9, section 7 (tier 2) explained, a different order of additions changes the last digits. The last line shows the payoff: at any moment only an n × 128 block is needed instead of the whole n × n matrix. On a GPU that block fits in SRAM, so the attention matrix never has to be written to HBM.

IO complexity

Theorem 2 (section 3.2) of the FlashAttention paper counts HBM accesses, with N the sequence length, d the head dimension, and M the SRAM size:

standard attention:  Θ(N·d + N²)   HBM accesses
FlashAttention:      Θ(N²·d²/M)    HBM accesses

With d around 64–128 and M around 100 KB, d²/M is many times smaller than 1, so bytes read and written drop sharply. FLOPs do not drop: all of QKᵀ and PV are still computed. Extra memory grows only linearly in N (Theorem 1, section 3.1). In the backward pass, FlashAttention even does more FLOPs: it stores only the output and the two statistics m and l, then recomputes the attention matrix from Q, K, and V when needed (section 3.1, “Recomputation” paragraph). Trading extra compute for fewer bytes pays off because attention is memory-bound. The result: attention on GPT-2 is 7.6 times faster than the PyTorch implementation (Figure 1), GPT-2 training at 1K context is 3 times faster than the HuggingFace implementation, and BERT-large is 15% faster than the MLPerf 1.1 record (abstract).

FlashAttention-2, FlashAttention-3, and Flash-Decoding

  • FlashAttention-2 (Dao 2023) partitions work across thread blocks and warps better and cuts non-matmul operations; about 2 times faster than the original, reaching 50–73% of the A100’s theoretical FLOPs (abstract).
  • FlashAttention-3 (Shah et al. 2024) targets the H100, using Hopper’s asynchrony and FP8: 1.5–2.0 times faster than FlashAttention-2 in the forward pass, up to 740 TFLOPs/s in FP16 (75% utilization) and close to 1.2 PFLOPs/s in FP8 (abstract, section 1).
  • Flash-Decoding (Dao, Haziza, Massa, Sizov 2023) targets the decode phase specifically. During decode there is only one query per sequence, and FlashAttention parallelizes over queries and batch, so at batch 1 it uses less than 1% of the GPU. Flash-Decoding splits K and V into chunks, computes the query’s attention against each chunk in parallel, stores the log-sum-exp of each chunk, then merges them with exactly the rescaling of online softmax. The post reports up to 8 times faster generation for very long sequences.

Hardware note: Turing has no FlashAttention-2

The flash-attn README is explicit: FlashAttention-2 needs an Ampere, Ada, or Hopper GPU (A100, RTX 3090, RTX 4090, H100); for Turing (T4, RTX 2080) you must use the separate flash-attention-turing repo, which supports only a subset of features. In PyTorch, the FLASH_ATTENTION backend of scaled_dot_product_attention is only enabled for sm80 and newer, while the memory-efficient backend (based on xFormers) runs from sm50 (per sdp_utils.cpp in the PyTorch source). The RTX 2080 Ti is sm75, so when you call SDPA, PyTorch picks memory-efficient, not flash. The memory-efficient kernel also works in blocks and does not build the n × n matrix, so the memory benefit remains.

Code sample 3b: comparing the math and memory-efficient backends

# Illustrative code. Python 3.12, torch 2.14.1+cu130
import time, torch, torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend

def bench(backend, n, heads=12, d=128, iters=10):
    q, k, v = (torch.randn(1, heads, n, d, device="cuda", dtype=torch.float16) for _ in range(3))
    torch.cuda.synchronize()
    torch.cuda.reset_peak_memory_stats()
    base = torch.cuda.memory_allocated()
    with sdpa_kernel(backend):                      # force PyTorch to use only this backend
        F.scaled_dot_product_attention(q, k, v, is_causal=True)   # warmup
        torch.cuda.synchronize(); t0 = time.perf_counter()
        for _ in range(iters):
            F.scaled_dot_product_attention(q, k, v, is_causal=True)
        torch.cuda.synchronize()
    dt = (time.perf_counter() - t0) / iters
    return dt, (torch.cuda.max_memory_allocated() - base) / 2**20

for n in (1024, 2048, 4096, 8192, 16384):
    for backend in (SDPBackend.MATH, SDPBackend.EFFICIENT_ATTENTION):
        try:
            dt, peak = bench(backend, n)
            print(f"n={n:6d} {backend.name:20s} {dt * 1e3:8.2f} ms   extra peak memory {peak:8.0f} MiB")
        except torch.OutOfMemoryError:
            print(f"n={n:6d} {backend.name:20s} out of memory")
        torch.cuda.empty_cache()

Not run. Needs an NVIDIA GPU with CUDA (for example an RTX 2080 Ti 11 GB, Turing sm_75), CUDA 13.0, Python 3.12, torch 2.14.1+cu130. On Turing, adding SDPBackend.FLASH_ATTENTION to the loop raises a no-available-kernel error, because that backend needs sm80 or newer.

What you should see when you run it: the peak memory of MATH grows with n² because it builds a heads × n × n score matrix (at n = 16,384 with 12 heads, that fp16 matrix alone is about 6 GiB), while that of EFFICIENT_ATTENTION grows only linearly.


08Speculative decoding

The idea

Section 2 showed that decode wastes compute: it reads all the weights just to compute one token. Batching uses the spare compute by serving many users. Speculative decoding uses it for one user: guess a few tokens ahead, then check them all in one forward pass.

A round has three steps:

  1. A small, fast draft model generates k guessed tokens one after another.
  2. The target model runs one forward pass over the prefix plus the k guesses and gets distributions at all k + 1 positions. That pass is like a short prefill: the weights are read once for k + 1 positions, so when decode is memory-bound it costs about as much as a normal decode step.
  3. A verification step decides how many guesses to keep, then always adds at least one token from the target model.
~/diagrams/kv-cache-prefill-va-decode/speculative-decoding-round-en.htmlinteractive
open ↗
FigThe decode loop sends the current prefix to the draft model; the draft model runs k small decode steps and returns k draft tokens. The decode loop sends the prefix plus k tokens to the target model in one forward pass. The verifier receives the draft distribution q at k positions and the target distribution p at k + 1 positions, then accepts each token with probability min(1, p/q); the first rejected token is resampled from the normalized max(0, p − q); if all k are accepted, one bonus token is sampled from p at position k + 1. At the end of the round the KV of rejected tokens is dropped and the next round starts.
target model forwardresult returned to the decode loopdashed purple: side data and KV cleanup

Why the output keeps the exact distribution

Let p(x) be the target model’s distribution and q(x) the draft model’s at some position. For a draft token x ~ q, the rule from Leviathan et al. (section 2.3) is:

  • If q(x) ≤ p(x): always accept.
  • If q(x) > p(x): accept with probability p(x)/q(x), so in every case acceptance happens with probability min(1, p(x)/q(x)).
  • On rejection: resample from p′(x) = norm(max(0, p(x) − q(x))) and end the round.

A short proof (Appendix A.1 of the paper): the probability of ending up with a particular token x′ is the sum of two paths.

P(accepted, x = x′)       = q(x′) · min(1, p(x′)/q(x′)) = min(p(x′), q(x′))
P(rejected, then x′)      = (1 − Σ_x min(p(x), q(x))) · p′(x′) = p(x′) − min(p(x′), q(x′))
Sum:                      P(x = x′) = p(x′)

The second line holds because Σ_x max(0, p − q) = 1 − Σ_x min(p, q), so the normalizing constant of p′ cancels with the rejection probability. Result: the output tokens are distributed exactly as p, not approximately. Chen et al. 2023 (DeepMind) prove the same result independently (Theorem 1 in the appendix), with the opposite notation: in that paper q is the target and p the draft. With greedy decoding, the rule reduces to: accept a draft token if it matches the target model’s argmax.

There is a subtle point that connects to part 9, section 7 (tier 2): “identical in distribution” is not “identical bit for bit”. The target model computes k + 1 positions in one pass instead of one at a time, so the kernels split the work differently and the logits may differ in the last digits. With greedy decoding, when the top two tokens are nearly tied, speculative output can differ from plain output in a few tokens. That is floating-point arithmetic, not a flaw in the algorithm.

How many tokens per round

Leviathan et al. define the acceptance rate β as the probability of accepting a draft token, and α = E(β) (section 3.1). Assuming acceptances are independent, the expected number of tokens per round is (Equation 1, section 3.1; the paper writes γ for k):

E[tokens per round] = (1 − α^(k+1)) / (1 − α)

For example with k = 4: α = 0.6 gives about 2.3 tokens per round; α = 0.8 about 3.4; α = 0.9 about 4.1. Real speed also has to pay for running the draft model, so Theorem 3.8 (section 3.3) adds the cost ratio c between the draft and target models. The cheaper and more target-like the draft model, the bigger the win; a large k only pays off when α is high.

On T5-XXL, the paper reports 2–3 times faster generation with identical outputs (abstract; Table 2 ranges from 1.4 to 3.4 times). Chen et al. report 2–2.5 times faster sampling from Chinchilla 70B (abstract).

Medusa and EAGLE: no separate draft model

Running two models has two drawbacks: you need a small model with the same tokenizer, and you have to manage two KV caches.

  • Medusa (Cai et al. 2024) adds a few decoding heads on top of the target model itself, each guessing a token further ahead, and checks several candidate branches at once with tree attention. Medusa-1 trains the heads on a frozen backbone and is over 2.2 times faster without hurting quality; Medusa-2 also trains the backbone and reaches 2.3–2.8 times (PDF v3, abstract and section 1); the 3.6 times figure is only the best MT-Bench category. Note that Medusa-2 changes the backbone weights, and its “typical acceptance” scheme keeps quality similar rather than preserving the exact distribution like rejection sampling.
  • EAGLE (Li et al. 2024) guesses at the feature level (the output of the second-to-last layer) rather than the token level, with a lightweight head, and still verifies with rejection sampling, so it preserves the distribution. On LLaMA2-Chat 70B, EAGLE reports a 2.7–3.5 times latency speedup and doubled throughput (abstract; section 4.4, Table 7).

Code sample 4: assisted generation in Transformers

Hugging Face Transformers calls speculative decoding assisted generation: pass a small model as the assistant_model argument of generate(). The docs state it supports greedy search and sampling but not batched inputs. The library does not return the acceptance rate, so the code below counts the target model’s forward passes with a forward hook: tokens generated divided by verification rounds gives the average tokens per round, from which α follows via the formula above.

# Illustrative code. Python 3.12, torch 2.14.1+cu130, transformers 5.18.0
import time, torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tok = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct")
load = dict(dtype=torch.float16, device_map="cuda")
target = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-1.5B-Instruct", **load).eval()
draft = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", **load).eval()  # same tokenizer

calls = 0
def count_call(module, args, output):
    global calls
    calls += 1
target.register_forward_hook(count_call)            # each target forward pass = one verification round

prompts = {
    "code": "Write a Python function that checks whether a number is prime, with a docstring.",
    "vi prose": "Kể một câu chuyện ngắn về một người bán phở ở Hà Nội vào buổi sáng mùa đông.",
}
K = 5
for name, text in prompts.items():
    inputs = tok.apply_chat_template([{"role": "user", "content": text}], add_generation_prompt=True,
                                     return_tensors="pt", return_dict=True).to("cuda")
    gen = dict(max_new_tokens=256, do_sample=False)

    torch.cuda.synchronize(); t0 = time.perf_counter()
    plain = target.generate(**inputs, **gen)
    torch.cuda.synchronize(); t_plain = time.perf_counter() - t0

    calls = 0
    torch.cuda.synchronize(); t0 = time.perf_counter()
    spec = target.generate(**inputs, **gen, assistant_model=draft,
                           num_assistant_tokens=K, num_assistant_tokens_schedule="constant")
    torch.cuda.synchronize(); t_spec = time.perf_counter() - t0

    n_new = spec.shape[1] - inputs["input_ids"].shape[1]
    per_round = n_new / calls                       # the first round also processes the prompt; no separate prefill
    print(f"{name:8s}: same as plain output = {torch.equal(plain, spec)}, "
          f"{t_plain / t_spec:.2f}× faster, {per_round:.2f} tokens per verification round")

Not run. Needs an NVIDIA GPU with about 6 GB or more of free VRAM (for example an RTX 2080 Ti 11 GB), CUDA 13.0, Python 3.12, torch 2.14.1+cu130, transformers 5.18.0. The draft and target models must share a tokenizer (Qwen2.5-0.5B and 1.5B do).

What you should see when you run it: the code prompt usually yields more tokens per round than the Vietnamese prose prompt, because code has many easy-to-guess stretches (indentation, keywords, repeated variable names) that a small model gets right. The “same as plain output” column may be False for some prompts for the floating-point reason above, not because the algorithm is wrong. Also, transformers may stop drafting early when the draft model is not confident enough (assistant_confidence_threshold), so the number of draft tokens per round can be fewer than K.


09Where long context gets expensive

The familiar answer is “attention is O(n²)”. That is only partly right. Splitting the cost of long context into three parts makes it clearer.

Part 1: prefill, where O(n²) actually shows up

Use the FLOPs count from part 8, section 6: each token costs 2N for the matrices plus 2 · n_layer · n_ctx · d_attn for the context-dependent attention term. For a causally masked prefill of n tokens, each token attends to n/2 earlier tokens on average, and both QKᵀ and the multiplication by V cost 2 FLOPs per pair per dimension, so the average attention cost per token equals part 8’s term with n_ctx = n. Attention starts to exceed the matrix part when:

2 · n_layer · n · d_attn > 2N   ⇔   n > N / (n_layer · d_attn)

For Llama 3.1 8B (non-embedding N 6.98B, 32 layers, d_attn 4,096), the crossover is about 53,000 tokens, as the last line of code sample 2’s output shows. At 4k context, attention is only 7% of prefill FLOPs; at 32k, 38%; at 128k, 71%. So O(n²) is a significant cost only at very long context, and it lives in prefill, that is, in TTFT.

Part 2: decode, where the KV cache is the culprit

Every decode step reads the sequence’s entire KV cache. For Llama 3.1 8B that is 0.54 GB at 4k context and 17 GB at 128k, against 16 GB of weights. That means:

  • TPOT grows linearly with context: at 128k, each decode step reads twice as many bytes as at short context, for a single sequence alone.
  • The maximum batch shrinks with context: on an 80 GB GPU, from 104 sequences at 4k to 3 at 128k (code sample 2). Throughput falls with it.

This cost is paid on every generated token, and it is linear, not quadratic. For most real workloads (prompts of a few thousand to a few tens of thousands of tokens, long outputs), this part outweighs part 1. That is why GQA, MLA, KV cache quantization, and PagedAttention, the techniques that cut bytes, matter for long context as much as FlashAttention does.

Summary table for Llama 3.1 8B, bf16, on one 80 GB GPU (numbers from code sample 2’s output):

Context Prefill FLOPs Attention share KV cache per sequence Max sequences
4,096 61.6 TFLOPs 7% 0.54 GB 104
32,768 739 TFLOPs 38% 4.29 GB 13
131,072 6,333 TFLOPs 71% 17.18 GB 3

Part 3: quality

The advertised context window is a technical limit, not the length a model uses well. RULER (Hsieh et al. 2024) evaluates 17 long-context models on 13 tasks: all of them claim contexts of 32K tokens or more, but only half maintain satisfactory performance at 32K (abstract; effective lengths in Table 3). Long context costs money and may not deliver information the model can actually use.

One direction that is often misunderstood is attention sinks and sliding windows. StreamingLLM (Xiao et al. 2023) observes that models dump a lot of attention onto the first few tokens even when they carry no meaning, and calls them “attention sinks”. Keeping the K and V of the first 4 tokens plus a sliding window of recent tokens (section 1, section 3.2) lets a model generate stably for up to 4 million tokens and run up to 22.2 times faster than sliding window with recomputation (abstract). But the paper states clearly that this does not extend the context window or improve long-term memory (Appendix A): a token that has left the window is forgotten. It is a way to keep memory bounded, not a way to remember more.

API prices reflect these three parts

The price lists of the major providers (fetched on 2026-10-08 from the pricing pages of OpenAI, Anthropic, and Google Gemini) share the same shape. This post gives only ratios, since absolute prices change quickly:

  • Output costs about 4–8 times as much as input, most commonly 5–6 times on current flagship models. The reason is section 2: an input token goes through compute-bound prefill, sharing one read of the weights with thousands of other tokens; an output token is a memory-bound decode step that occupies KV cache space for the whole generation.
  • Cached input is cheaper than regular input, typically about 10% of it. There are exceptions: Anthropic’s newest models go down to 2.5–5%, some older OpenAI models are at 25–50%. This is the prefix caching of section 5: a prefix that already has K and V does not need to be prefilled again. Anthropic also charges for cache writes (1.25 times the input price for a 5-minute cache, 2 times for a 1-hour cache), and Gemini charges an hourly storage fee for cached content, because kept blocks occupy GPU memory.
  • Very long prompts can have their own price. Some Gemini Pro models charge more for prompts over 200k tokens, and OpenAI has a long-context tier for inputs over 272K tokens. That is parts 1 and 2 together: prefill is more expensive and each request takes up most of a GPU’s memory.

10Common misconceptions

The nine misconceptions below match the top and bottom rows of Figure 1.

Misconception Reality
“The KV cache also stores Q, or the attention scores” It stores only K and V of processed tokens. Q of an old token is never reused; scores are recomputed every step from the new Q and the old Ks (section 3)
“Decode is slow because the GPU lacks compute” Decode at small batch uses only a few percent of compute; the bottleneck is the bandwidth for reading weights and the KV cache. Intensity is about 1 FLOP/byte against a ridge point of about 87 on a 2080 Ti (section 2)
“The bigger the batch, the better” Throughput grows with batch until the KV cache runs out or compute is reached; each user’s TPOT grows too, and long context makes the maximum batch drop fast (sections 2, 3)
“GQA/MQA make attention cost fewer FLOPs” The number of query heads is unchanged, so QKᵀ FLOPs are unchanged; GQA cuts W_K, W_V parameters and the KV cache, i.e. bytes (part 8, section 6; section 4)
“FlashAttention is approximate attention, or cuts FLOPs” It is exact attention with a different order of additions; FLOPs are unchanged in the forward pass and even higher in the backward pass because of recomputation. It cuts the bytes read from and written to HBM (section 7, code sample 3a)
“Speculative decoding trades quality for speed” With rejection sampling, output tokens follow the target model’s distribution exactly. Any difference is floating-point, as with any kernel change. Medusa-2 and typical acceptance are exceptions with a trade-off (section 8)
“Long context is expensive mainly because attention is O(n²)” O(n²) dominates prefill only beyond about 53k tokens (Llama 3.1 8B). The cost on every output token is reading the KV cache, linear in context, and it also cuts the maximum batch (section 9)
“A 128k context window means the model uses all 128k well” Effective context is often much shorter; RULER found only half of the models hold up at 32K. Attention sinks keep memory bounded but remember nothing extra (section 9)
“Prompt caching caches the answer” It caches the K and V of a prefix that matches from the first token; the answer is generated fresh every time. Change one character at the start and you lose the cache for the rest (section 5)

Summary: how much work does each byte do

This whole post comes down to one ratio: the FLOPs done for each byte read from memory.

  • Prefill processes the whole prompt in one pass, with each byte of weights serving thousands of tokens, so it is compute-bound and sets TTFT. Decode produces one token at a time, rereading all the weights and the KV cache at every step to do about 1 FLOP per byte at batch 1, while an RTX 2080 Ti needs about 87 FLOPs per byte to keep its compute busy.
  • Batching spreads the weight read over many sequences, so throughput grows almost linearly with batch size. But the KV cache cannot be shared: each sequence has its own K and V. So memory available for the KV cache sets the maximum batch, and the batch sets throughput.
  • The KV cache can be computed from config.json: Llama 3.1 8B costs 128 KiB per token, and a 128k sequence weighs as much as the whole set of weights. GQA, MQA, and MLA cut those bytes by sharing or compressing K and V; KV cache quantization cuts the bytes per value.
  • PagedAttention raises the share of KV memory actually used from 20–38% to 96%, and turns prefix sharing into a matter of bumping a ref count. Prefix caching and API prompt caching are the same idea.
  • Continuous batching keeps the batch full; chunked prefill and prefill/decode disaggregation resolve the conflict between new requests’ TTFT and running requests’ TPOT.
  • FlashAttention keeps FLOPs the same and cuts bytes through online softmax and tiling in SRAM; the result is exact, differing only in the order of additions. On Turing, PyTorch uses the memory-efficient kernel instead of flash.
  • Speculative decoding uses decode’s spare compute to verify several guessed tokens in one pass; rejection sampling preserves the target model’s distribution exactly.
  • Long context is expensive mostly because of KV cache memory and bandwidth, paid on every output token; O(n²) dominates prefill only at very long context. API prices reflect exactly that: output costs 4–8 times as much as input, and cached input is usually only about 10% of the input price.

Part 14 covers Mixture of Experts: total weights are huge but each token uses only part of them, so the balance between memory and compute changes completely compared with a dense model, and every roofline calculation in this post has to be redone. Reasoning models, meanwhile, generate very long outputs, so decode dominates their cost; part 15 comes back to that.


References

Papers

Blogs, documentation, and pricing pages

Trương Đức Dũng

AI Engineer - Hanoi

Focused on LLMs, RAG, and AI systems that need to run reliably in production.