Demystifying Transformer Inference: Prefill, Decode, and the Mechanics of KV Caching in PyTorch

Executive Overview

In the lifecycle of a deep learning model, training and inference are often treated as two sides of the same coin. However, when deploying Transformer-based Large Language Models (LLMs), the computational, memory, and structural dynamics of inference diverge sharply from training. While training relies on parallelized computation over fixed-length sequences dominated by compute-heavy matrix multiplications during the backward pass, inference is an autoregressive process driven by sequential forward passes.

This autoregressive nature presents a fundamental engineering challenge: generating text one token at a time requires repeated access to key and value representations of all preceding tokens. Without optimization, generating a response of length $N$ scales quadratically ($O(N^2)$) in compute time due to redundant recalculations of historical attention states.

To overcome this bottleneck, modern inference engines rely on the Key-Value (KV) Cache. By storing the intermediate key and value tensor representations of previously processed tokens in memory, the computational cost of generating subsequent tokens drops from quadratic to linear ($O(N)$) per token. However, this compute speedup introduces a steep memory footprint requirement. Efficiently managing KV cache allocation, memory layout, and phase transitions between initial prompt processing (prefill) and token generation (decode) is central to scaling generative AI workloads.

Detailed Technical Deep-Dive

The Autoregressive Generation Paradigm

Decoder-only Transformer models (such as GPT-4, Llama, or Mistral) generate text sequentially. Given a sequence of input tokens $x_1, x_2, dots, xt$, the model processes them through causal self-attention layers to predict a probability distribution over the vocabulary for the next token, $xt+1$.

$$mathcalP(x_t+1 mid x_1, x_2, dots, x_t) = textSoftmax(textLM_Head(h_t))$$

Where $h_t$ is the final hidden state produced by the network for position $t$. Causal masking ensures that the token representation at position $i$ can only attend to tokens at positions $j le i$.

Once the model produces unnormalized output scores (logits), a sampling strategy—such as greedy decoding (selecting the maximum logit), top-$k$, top-$p$ (nucleus) sampling, or temperature scaling—picks the discrete token $x_t+1$. This new token is concatenated to the input sequence, and the process repeats.

Naive Generation Loop (Without Caching)

In a naive PyTorch implementation, the generation loop re-evaluates the entire concatenated sequence at every iteration:

import torch

@torch.no_grad()
def greedy_decode_naive(model, input_ids: torch.Tensor, max_new_tokens: int) -> torch.Tensor:
    output_ids = input_ids.clone()
    for _ in range(max_new_tokens):
        # Forward pass over the entire accumulated sequence
        logits = model(output_ids)
        # Extract logits corresponding to the final token position
        next_token_logits = logits[:, -1, :]
        next_token = next_token_logits.argmax(dim=-1, keepdim=True)
        # Concatenate the new token to the sequence
        output_ids = torch.cat([output_ids, next_token], dim=1)
    return output_ids

While conceptualizing generation this way is straightforward, it is computationally redundant. If a prompt contains 1,000 tokens and the model generates 200 output tokens, step 1 processes 1,000 tokens, step 2 processes 1,001 tokens, and step 200 processes 1,199 tokens. The attention projections and feed-forward hidden states for the original prompt tokens are recomputed 200 times without modification.

The Two-Phase Execution Architecture: Prefill vs. Decode

To eliminate redundant computation, production inference engines divide generation into two distinct operational phases: Prefill and Decode.

[Prompt Tokens: T1, T2, T3, T4] 
            │
            ▼
┌──────────────────────────────────────┐
│            PREFILL PHASE             │
│  - Processes T1..T4 in parallel      │
│  - Computes & stores K1..K4, V1..V4   │
│  - Generates first new token: T5     │
└──────────────────────────────────────┘
            │
            ▼
┌──────────────────────────────────────┐
│             DECODE PHASE             │
│  - Processes ONLY T5                 │
│  - Computes K5, V5 & appends to cache│
│  - Attends Q5 against [K1..K5]       │
│  - Generates next token: T6          │
└──────────────────────────────────────┘

1. The Prefill Phase

• Input: The full initial user prompt tensor of length $P$.
• Characteristics: Compute-bound operation. Because all prompt tokens are known ahead of time, matrix multiplications across time steps are batch-processed simultaneously.
• Output: Logits for the first generated token, along with the initial Key and Value states ($Ktextprompt, Vtextprompt$) computed across all layers.

2. The Decode Phase

• Input: Exactly one token—the single token generated in the immediately preceding iteration.
• Characteristics: Memory-bandwidth-bound operation. The matrix computations are small (sequence length $= 1$), but the GPU must load all model weights and historical KV tensors from High-Bandwidth Memory (HBM) into SRAM for every single generated token.
• Output: The logits for the next single token, alongside updated KV cache structures.

Mathematical Complexity Analysis

Let $P$ be the prompt length and $G$ be the number of generated tokens. The total length of the sequence reaches $N = P + G$.

• Naive Approach (No Cache):
  For each generated token $i in 1, dots, G$, self-attention computes over sequence length $P + i – 1$.
  $$textTotal Compute Cost = mathcalOleft( sum_i=1^G (P + i)^2 right) = mathcalOleft( P^2 G + P G^2 + G^3 right)$$

• Cached Approach (KV Cache Enabled):

  • Prefill Step: Evaluates $P$ tokens in parallel: $mathcalO(P^2)$.
  • Decode Steps: Each step calculates 1 query against $P + i$ cached keys: $mathcalO(P + i)$ for step $i$.
    $$textTotal Compute Cost = mathcalO(P^2) + mathcalOleft( sum_i=1^G (P + i) right) = mathcalOleft( P^2 + PG + G^2 right)$$

For long-context generation where $P$ or $G$ are large (e.g., 32k+ token contexts), enabling the KV cache reduces latency by multiple orders of magnitude.

Implementation: Building a Cached Causal Transformer in PyTorch

Below is an explicit, modular PyTorch implementation demonstrating how KV states are initialized, stored, updated, and masked during both the prefill and decode phases.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Optional, Tuple, List

class CausalSelfAttentionWithCache(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        assert hidden_size % num_heads == 0, "hidden_size must be divisible by num_heads"
        self.num_heads = num_heads
        self.head_dim = hidden_size // num_heads

        self.qkv_proj = nn.Linear(hidden_size, 3 * hidden_size, bias=False)
        self.out_proj = nn.Linear(hidden_size, hidden_size, bias=False)

    def forward(
        self, 
        x: torch.Tensor, 
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        batch_size, seq_len, hidden_size = x.shape

        # Linear projection -> [batch_size, seq_len, 3 * hidden_size]
        qkv = self.qkv_proj(x)

        # Reshape to separate Query, Key, and Value components
        qkv = qkv.view(batch_size, seq_len, 3, self.num_heads, self.head_dim)
        qkv = qkv.permute(2, 0, 3, 1, 4)  # Shape: [3, batch_size, num_heads, seq_len, head_dim]
        q, k, v = qkv[0], qkv[1], qkv[2]

        # Append to historical KV tensors if present (Decode Phase)
        if past_kv is not None:
            past_k, past_v = past_kv
            k = torch.cat([past_k, k], dim=2)
            v = torch.cat([past_v, v], dim=2)

        current_kv_cache = (k, v)
        total_seq_len = k.size(2)

        # Scaled Dot-Product Attention
        # Query shape: [B, H, seq_len, head_dim]
        # Key transpose shape: [B, H, head_dim, total_seq_len]
        scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)

        # Causal Masking Logic
        past_len = total_seq_len - seq_len
        causal_mask = torch.ones(seq_len, total_seq_len, device=x.device, dtype=torch.bool)
        causal_mask = torch.tril(causal_mask, diagonal=past_len)

        scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), float("-inf"))
        attn_weights = F.softmax(scores, dim=-1)

        # Compute Weighted Context Vector
        context = torch.matmul(attn_weights, v)  # [B, H, seq_len, head_dim]
        context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)

        return self.out_proj(context), current_kv_cache

class TransformerBlock(nn.Module):
    def __init__(self, hidden_size: int, num_heads: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(hidden_size)
        self.attn = CausalSelfAttentionWithCache(hidden_size, num_heads)
        self.ln2 = nn.LayerNorm(hidden_size)
        self.mlp = nn.Sequential(
            nn.Linear(hidden_size, 4 * hidden_size),
            nn.GELU(),
            nn.Linear(4 * hidden_size, hidden_size)
        )

    def forward(
        self, 
        x: torch.Tensor, 
        past_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        attn_out, new_kv = self.attn(self.ln1(x), past_kv=past_kv)
        x = x + attn_out
        x = x + self.mlp(self.ln2(x))
        return x, new_kv

class ModularCausalLM(nn.Module):
    def __init__(self, vocab_size: int = 32000, hidden_size: int = 4096, num_heads: int = 32, num_layers: int = 32):
        super().__init__()
        self.embed = nn.Embedding(vocab_size, hidden_size)
        self.layers = nn.ModuleList([
            TransformerBlock(hidden_size, num_heads) for _ in range(num_layers)
        ])
        self.ln_f = nn.LayerNorm(hidden_size)
        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)

    def forward(
        self, 
        input_ids: torch.Tensor, 
        kv_cache: Optional[List[Tuple[torch.Tensor, torch.Tensor]]] = None
    ) -> Tuple[torch.Tensor, List[Tuple[torch.Tensor, torch.Tensor]]]:
        x = self.embed(input_ids)
        new_cache = []

        if kv_cache is None:
            kv_cache = [None] * len(self.layers)

        for layer, layer_past in zip(self.layers, kv_cache):
            x, layer_new_kv = layer(x, past_kv=layer_past)
            new_cache.append(layer_new_kv)

        logits = self.lm_head(self.ln_f(x))
        return logits, new_cache

Optimized Generation Function Using Cache

@torch.no_grad()
def greedy_decode_cached(
    model: ModularCausalLM, 
    prompt_ids: torch.Tensor, 
    max_new_tokens: int
) -> torch.Tensor:
    output_ids = prompt_ids.clone()

    # --- PHASE 1: PREFILL ---
    # Pass the entire prompt at once to generate initial cache
    logits, kv_cache = model(prompt_ids, kv_cache=None)
    next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
    output_ids = torch.cat([output_ids, next_token], dim=1)

    # --- PHASE 2: DECODE ---
    # Pass ONLY the newest single token in subsequent iterations
    for _ in range(max_new_tokens - 1):
        logits, kv_cache = model(next_token, kv_cache=kv_cache)
        next_token = logits[:, -1, :].argmax(dim=-1, keepdim=True)
        output_ids = torch.cat([output_ids, next_token], dim=1)

    return output_ids

Supporting Context & Metrics

While the KV cache successfully optimizes compute, it creates a significant memory bottleneck. To evaluate system requirements for deployment, serving architects must account for the memory footprint consumed by key and value tensors.

KV Cache Memory Formula

The raw memory required to hold the KV cache for a given batch of requests is governed by the structural dimensions of the model and precision formats:

$$textMemorytextKV = 2 times b times L times htextkv times d_texthead times s times p$$

Where:

• $2$ represents the two separate key and value matrices stored per layer.
• $b$: Batch size (number of concurrent sequences).
• $L$: Number of Transformer layers.
• $h_textkv$: Number of key-value attention heads (equal to total heads in Multi-Head Attention).
• $d_texthead$: Dimension per head ($textHidden Size / textTotal Query Heads$).
• $s$: Sequence length (Prompt length + Generated tokens).
• $p$: Precision size in bytes (e.g., 2 bytes for FP16/BF16, 1 byte for FP8, 4 bytes for FP32).

Practical Analytical Calculation

Consider a standard 7-Billion Parameter Model (e.g., Llama-2-7B) running in 16-bit precision (2 bytes per element):

• $L = 32$ layers
• $h_textkv = 32$ heads
• $d_texthead = 128$
• Precision $p = 2$ bytes (BF16)

Scenario A: Single Request ($b = 1$) with Context Length $s = 4,096$

$$textMemorytextKV = 2 times 1 times 32 times 32 times 128 times 4096 times 2 text bytes$$
$$textMemorytextKV = 2,147,483,648 text bytes approx mathbf2.0 text GiB$$

Scenario B: Server-Scale Deployment ($b = 64$) with Context Length $s = 8,192$

$$textMemorytextKV = 2 times 64 times 32 times 32 times 128 times 8192 times 2 text bytes$$
$$textMemorytextKV = 68,719,476,736 text bytes approx mathbf64.0 text GiB$$

Note: Despite Llama-3-70B being ten times larger in total parameter count than Llama-2-7B, its KV cache overhead per request is lower due to Grouped-Query Attention (GQA), which reduces $h_textkv$ from 64 to 8.

Expert Insights & Architectural Trade-offs

Because memory footprint dominates serving constraints, modern deep learning system design has evolved to optimize both attention mechanisms and memory allocation strategies.

1. Architectural Attention Variants

To prevent the KV cache from overwhelming HBM capacity, modern architectures replace traditional Multi-Head Attention (MHA) with alternative attention mechanisms:

Multi-Head Attention (MHA)   Grouped-Query Attention (GQA)   Multi-Query Attention (MQA)
  Q Q Q Q   K K K K   V V V V      Q Q Q Q   K K   V V           Q Q Q Q    K    V
  │ │ │ │   │ │ │ │   │ │ │ │      │ │ │ │   │ │   │ │           │ │ │ │    │    │
  ▼ ▼ ▼ ▼   ▼ ▼ ▼ ▼   ▼ ▼ ▼ ▼      ▼ ▼ ▼ ▼   ▼ ▼   ▼ ▼           ▼ ▼ ▼ ▼    ▼    ▼
 [8 Q Heads / 8 KV Pairs]        [8 Q Heads / 2 KV Pairs]      [8 Q Heads / 1 KV Pair]
  Cache Size: 100%                Cache Size: 25%               Cache Size: 12.5%

• Multi-Head Attention (MHA): Every query head has an independent key and value head. While highly expressive, it scales memory consumption rapidly.
• Multi-Query Attention (MQA): All query heads share a single key and value head pair. This reduces KV cache size by $h_textheadstimes$, though it can impair model quality on complex reasoning tasks.
• Grouped-Query Attention (GQA): A hybrid approach where query heads are partitioned into $G$ groups, with each group sharing one KV head pair. Adopted by systems such as Llama-3, Mistral, and Command-R, GQA offers a high quality-to-memory trade-off.

2. Physical Memory Management: PagedAttention

The simple naive PyTorch KV implementation presented earlier concatenates tensors dynamically via torch.cat(). This pattern creates severe efficiency problems in enterprise production environments:

• Memory Fragmentation: Dynamic memory allocations lead to contiguous physical memory allocation failures.
• Over-allocation Overhead: Without prior knowledge of the target output length, serving engines must pre-allocate contiguous memory buffers matching the maximum sequence length (e.g., 8192 tokens), wasting memory on ungenerated tail positions.

Engineers at vLLM introduced PagedAttention, borrowing virtual memory partitioning concepts from operating systems.

Logical Blocks (Sequence Space)
[ Block 0: T1..T4 ] ──> [ Block 1: T5..T8 ] ──> [ Block 2: T9..T12 ]
                             │
                             ▼ Physical Page Table Map
Physical Memory (Non-Contiguous DRAM Pages)
[ Page 412 ] <── Block 0
[ Page 087 ] <── Block 2
[ Page 109 ] <── Block 1

By dividing the KV cache of each sequence into fixed-size physical blocks (e.g., blocks of 16 tokens), physical pages can be allocated dynamically as generation progresses. This eliminates internal memory waste and allows non-contiguous memory layouts, drastically increasing GPU concurrency limits.

Future Outlook

As context lengths expand beyond hundreds of thousands of tokens and concurrent LLM usage scales, optimizing inference runtime continues to advance across several fronts:

1. Prefill/Decode Disaggregation (PD Disaggregation)

Because prefill phase tasks are compute-bound (benefiting from dense matrix tensor cores) and decode phase tasks are memory-bound (benefiting from memory bandwidth utilization), single-GPU execution creates resource contention. Next-generation inference systems disaggregate serving pipelines: requests undergo prefill execution on compute-optimized hardware nodes (e.g., NVIDIA H100s) before high-speed Inter-GPU networks stream the generated KV caches to memory-optimized decode nodes.

2. KV Cache Compression & Quantization

To fit longer contexts into limited HBM, systems are shifting from 16-bit floating-point (FP16/BF16) KV storage to quantized lower-precision representations. Utilizing 8-bit (FP8/INT8) and 4-bit (INT4) formats halves or quarters the memory required per token. Additionally, selective eviction algorithms drop low-attention-weight KV tokens from historical caches without degrading context comprehension.

3. Speculative Decoding

To bypass the single-token bottleneck of the decode phase, speculative decoding uses a small, lightweight draft model to rapidly propose sequence continuation candidates (e.g., $K=5$ draft tokens). The primary target LLM verifies all draft candidates simultaneously in a single, parallelized forward pass during its prefill block. Tokens passing validation are accepted instantly, reducing total generation latency without compromising output distribution accuracy.

Summary

• Training vs. Inference: Training relies on compute-bound parallel processing over fixed sequences, whereas autoregressive inference requires sequential, memory-bound forward passes.
• Prefill vs. Decode: Processing begins with a parallelized prefill phase over the input prompt, followed by sequential token-by-token decode iterations.
• KV Cache Mechanics: By caching key and value projections from past steps, time complexity drops from $mathcalO(N^2)$ to $mathcalO(N)$ per generated token.
• Memory Constraints: The KV cache trades compute for memory footprint, making serving engines highly memory-bandwidth constrained.
• Modern Optimizations: Structural architectural changes like Grouped-Query Attention (GQA), memory abstractions like PagedAttention, and lower-precision FP8 quantization are vital techniques for managing system resources effectively.