Autoregressive inference and the KV cache

A decoder-only transformer generates text one token at a time. At step t, the model takes the full sequence [x_0, x_1, ..., x_(t-1)] as input, computes attention across all positions, and produces a probability distribution over the vocabulary for position t. The highest-probability token (or a sample from the distribution) is appended to the sequence, and the process repeats for step t+1.

The computational cost of this loop grows quadratically. At step t, the model computes Q * K^T for a sequence of length t — an O(t * d) operation per layer. Summing across all generation steps: sum from t=1 to N of O(t * d) = O(N^2 * d). For a 4096-token generation, this means processing roughly 4096^2 / 2 ≈ 8.4M token-positions of attention computation. Most of this work is redundant — at step t, the Q, K, and V vectors for positions 0 through t-2 are identical to what was computed at step t-1.

Eliminating redundant computation

The KV cache stores the key and value projections for all previously processed tokens. At each generation step, only the new token's Q, K, and V are computed. The new K and V are appended to the cache. Attention is then computed between the new token's Q vector (a single row) and the full cached K and V matrices.

python

import torch
import torch.nn as nn
import torch.nn.functional as F

class CachedSelfAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.num_heads = num_heads
        self.d_head = d_model // num_heads
        self.W_q = nn.Linear(d_model, d_model, bias=False)
        self.W_k = nn.Linear(d_model, d_model, bias=False)
        self.W_v = nn.Linear(d_model, d_model, bias=False)
        self.W_o = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x, kv_cache=None):
        B, S, D = x.shape
        q = self.W_q(x).view(B, S, self.num_heads, self.d_head).transpose(1, 2)
        k = self.W_k(x).view(B, S, self.num_heads, self.d_head).transpose(1, 2)
        v = self.W_v(x).view(B, S, self.num_heads, self.d_head).transpose(1, 2)

        if kv_cache is not None:
            k = torch.cat([kv_cache[0], k], dim=2)
            v = torch.cat([kv_cache[1], v], dim=2)

        new_cache = (k, v)

        # q: (B, H, S_new, D_head), k: (B, H, S_total, D_head)
        scores = torch.matmul(q, k.transpose(-2, -1)) / (self.d_head ** 0.5)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)
        out = out.transpose(1, 2).contiguous().view(B, S, D)
        return self.W_o(out), new_cache

Without the cache, generating 1000 tokens from a 1000-token prompt requires computing K and V projections for 1000 + 1001 + 1002 + ... + 1999 ≈ 1.5M token-positions across all layers. With the cache, it requires 1000 (the prompt, computed once) plus 1000 (one new token per step) = 2000 token-positions. The saving is approximately 750x fewer projection computations.

The cache trades compute for memory. Instead of recomputing K and V, the system stores them — and for large models at long context, this storage dominates GPU memory consumption. Llama 3 8B's KV cache at 8K context takes ~1 GB; Llama 3 70B at the same length takes ~2.5 GB. At 128K context, the KV cache for Llama 3 70B exceeds the model weights. The compute-memory tradeoff is almost always worth it — recomputation is far more expensive than storage — but cache memory is the binding constraint for serving throughput (covered in detail in Lessons 8 and 9).

Two phases of inference

Autoregressive inference has two distinct phases with fundamentally different performance characteristics.

Prefill

The entire prompt is processed in a single forward pass. All tokens are available simultaneously, so the computation is a dense matrix multiplication — compute-bound on the GPU. For Llama 3 8B processing a 4096-token prompt on an A100 (80 GB), prefill takes approximately 80-120 ms. This phase populates the KV cache: after prefill, every layer's K and V for all prompt tokens are stored in GPU memory.

The time spent in prefill determines time to first token (TTFT) — the latency a user perceives before the model starts responding. TTFT is directly proportional to prompt length: a 32K-token prompt takes roughly 8x longer to prefill than a 4K prompt.

Decode

Tokens are generated one at a time. Each step computes Q, K, V for a single token, appends K and V to the cache, and computes attention between the new Q and the full cached K/V. The computation per step is small — a single vector-matrix product per head — but the model must read the entire KV cache from GPU memory for every step.

This makes decode memory-bandwidth-bound, not compute-bound. The GPU's arithmetic units are mostly idle, waiting for data to be fetched from HBM. An A100's compute throughput is 312 TFLOPS (FP16), but its memory bandwidth is 2 TB/s. At decode time, the ratio of computation to memory access is extremely low — the GPU moves data rather than computing on it.

Decode speed determines tokens per second — the throughput a user experiences during generation. For Llama 3 8B on an A100, single-request decode runs at approximately 40-60 tokens/sec, bottlenecked by memory bandwidth rather than FLOPS.

Why TTFT and tokens/sec measure different things

TTFT depends on prompt length and compute throughput. A faster GPU (more TFLOPS) reduces TTFT proportionally. Tokens/sec depends on KV cache size and memory bandwidth. A GPU with higher memory bandwidth (A100 vs. A10) improves tokens/sec even if raw FLOPS are similar.

This is why the H100 (3.35 TB/s bandwidth, ~990 TFLOPS dense FP16) is dramatically better for LLM serving than the A100 (2 TB/s, 312 TFLOPS) — the bandwidth improvement matters more than the FLOPS improvement for the decode-bound workloads that dominate inference time.

A useful metric for understanding this: the arithmetic intensity (FLOPS per byte of memory accessed) during decode is extremely low. Reading the full KV cache (say 2.5 GB for Llama 3 70B at 8K) to compute one token's attention scores involves ~2.5 GB of memory reads and a few hundred million FLOPs — an arithmetic intensity of roughly 0.1 FLOP/byte. The A100's peak efficiency requires ~156 FLOP/byte (312 TFLOPS / 2 TB/s). The GPU is operating at less than 0.1% of its compute capacity during decode, entirely limited by how fast it can stream data from HBM.

This asymmetry between prefill and decode has driven an entire branch of systems research into disaggregated serving: using compute-optimized hardware for prefill and bandwidth-optimized hardware for decode. Splitwise (Patel et al. 2024) and DistServe (Zhong et al. 2024) both exploit this observation to improve end-to-end serving efficiency.

Implementing autoregressive generation with KV caching

python

import torch
import torch.nn.functional as F
import time

def generate_with_cache(model, prompt_ids, max_new_tokens, temperature=1.0):
    """Generate tokens autoregressively with KV caching."""
    device = next(model.parameters()).device
    input_ids = prompt_ids.to(device).unsqueeze(0)  # (1, prompt_len)

    # Phase 1: Prefill — process entire prompt, populate cache
    t0 = time.perf_counter()
    logits, kv_caches = model(input_ids, kv_caches=None)
    prefill_time = time.perf_counter() - t0

    next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
    generated = [next_token.item()]

    # Phase 2: Decode — generate one token at a time using cache
    t0 = time.perf_counter()
    for step in range(max_new_tokens - 1):
        logits, kv_caches = model(next_token, kv_caches=kv_caches)
        probs = F.softmax(logits[:, -1, :] / temperature, dim=-1)
        next_token = torch.multinomial(probs, num_samples=1)
        generated.append(next_token.item())
    decode_time = time.perf_counter() - t0

    print(f"Prefill: {prefill_time*1000:.1f} ms ({len(prompt_ids)} tokens)")
    print(f"Decode:  {decode_time*1000:.1f} ms ({max_new_tokens-1} tokens)")
    print(f"TTFT:    {prefill_time*1000:.1f} ms")
    print(f"Decode throughput: {(max_new_tokens-1)/decode_time:.1f} tokens/sec")

    return generated

Measuring the speedup

python

def generate_no_cache(model, prompt_ids, max_new_tokens):
    """Naive generation: recompute everything at each step."""
    device = next(model.parameters()).device
    sequence = prompt_ids.to(device).unsqueeze(0)

    for step in range(max_new_tokens):
        logits, _ = model(sequence, kv_caches=None)
        next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
        sequence = torch.cat([sequence, next_token], dim=1)

    return sequence[0, len(prompt_ids):].tolist()

# Comparison on a small model:
# With cache:    generate 256 tokens from 512-token prompt → ~200 ms
# Without cache: generate 256 tokens from 512-token prompt → ~8000 ms
# Speedup: ~40x

The speedup grows with sequence length. At 4K context generating 1K tokens, the naive approach recomputes K/V for ~4.5M cumulative token-positions. The cached approach computes 5K. That is a 900x reduction in projection computation, offset slightly by the cost of reading the cache from memory at each step.

Cache structure across layers

In a full transformer, the KV cache is a list of (K, V) pairs, one per layer. Llama 3 8B has 32 layers, so the cache is a list of 32 (K, V) tuples. Each K and V tensor has shape (batch, num_kv_heads, seq_len, d_head). During decode, the new token's K and V at each layer are concatenated along the seq_len dimension of the respective cached tensors. The cache grows by one position per step per layer.

Because every layer reads and writes its own slice of the cache, the memory cost scales linearly with depth. A 32-layer model's cache is exactly 32x the cost of a single layer's cache. This is why deeper models (Llama 3 70B with 80 layers, Llama 3 405B with 126 layers) have proportionally larger KV caches even when the per-head dimensions are identical.

← Previous