Prefill in LLM Inference

1. What is Prefill?

Prefill is the first phase of every LLM inference call. Before the model generates a single token of output, it must process your entire input prompt in one shot.

Think of it as the model "reading and understanding" everything you said before it starts "writing" a response.

You type: "Explain how a CPU works to a 5-year-old"

  ┌─────────────────────────────────────────────────┐
  │                   PREFILL PHASE                  │
  │                                                  │
  │  Input: all your tokens processed IN PARALLEL    │
  │  Output: KV cache + first generated token        │
  │  Duration: one-time, milliseconds                │
  │  Bottleneck: GPU compute (FLOPS)                 │
  └─────────────────────────────────────────────────┘
                         │
                         ▼
  ┌─────────────────────────────────────────────────┐
  │                   DECODE PHASE                   │
  │                                                  │
  │  Input: one token at a time (autoregressive)     │
  │  Output: next token, then next, then next...     │
  │  Duration: ongoing, seconds                      │
  │  Bottleneck: memory bandwidth (GB/s)             │
  └─────────────────────────────────────────────────┘

2. Real-World Token Simulation

Let's trace a real prompt through prefill. We'll use:

Prompt: "Explain how a CPU works to a 5-year-old"

Step 1: Tokenization

The tokenizer (e.g., Llama's SentencePiece / tiktoken) breaks your text into tokens. Tokens are NOT always full words:

Raw text:  "Explain how a CPU works to a 5-year-old"

Tokenized:
┌────────┬──────────┬─────────────────────────────────┐
│ Index  │ Token ID │ Token Text                      │
├────────┼──────────┼─────────────────────────────────┤
│   0    │  849     │ "Explain"                       │
│   1    │  1268    │ " how"                          │
│   2    │  264     │ " a"                            │
│   3    │  18622   │ " CPU"                          │
│   4    │  4375    │ " works"                        │
│   5    │  311     │ " to"                           │
│   6    │  264     │ " a"                            │
│   7    │  220     │ " 5"                            │
│   8    │  -       │ "-"                             │
│   9    │  3236    │ "year"                          │
│  10    │  -       │ "-old"                          │
└────────┴──────────┴─────────────────────────────────┘

Total: 11 tokens

Step 2: Embedding

Each token ID is looked up in an embedding table to get a high-dimensional vector:

Token "Explain" (ID 849)  →  [0.023, -0.156, 0.891, 0.042, ..., -0.334]
                               ╰──────────── 8192 dimensions ────────────╯
                               (for Llama-3-70B; 4096 for smaller models)

Token " how"    (ID 1268) →  [0.178, 0.045, -0.223, 0.567, ..., 0.112]
Token " a"      (ID 264)  →  [-0.089, 0.334, 0.012, -0.445, ..., 0.267]
  ... and so on for all 11 tokens

This creates an embedding matrix of shape [11, 8192] — 11 tokens, each with 8192 dimensions.

Step 3: Through Every Transformer Layer (The Big Part)

This embedding matrix now passes through every single transformer layer in the model. For Llama-3-70B, that's 80 layers. At each layer:

Layer 1 of 80:
┌──────────────────────────────────────────────────────────────────┐
│                                                                  │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │              SELF-ATTENTION                              │    │
│  │                                                          │    │
│  │  For each of the 11 tokens, compute:                     │    │
│  │                                                          │    │
│  │    Q (Query)  = token_embedding × W_Q                    │    │
│  │    K (Key)    = token_embedding × W_K    ← STORED in KV  │    │
│  │    V (Value)  = token_embedding × W_V    ← STORED in KV  │    │
│  │                                                          │    │
│  │  Then: Attention = softmax(Q × K^T / √d) × V            │    │
│  │                                                          │    │
│  │  Every token attends to every other token (11 × 11)      │    │
│  │  That's 121 attention scores computed simultaneously      │    │
│  └─────────────────────────────────────────────────────────┘    │
│                           │                                      │
│                           ▼                                      │
│  ┌─────────────────────────────────────────────────────────┐    │
│  │              FEED-FORWARD NETWORK (MLP)                  │    │
│  │                                                          │    │
│  │  Each token goes through:                                │    │
│  │    hidden = token × W_gate  (upproject)                  │    │
│  │    hidden = SiLU(hidden) × (token × W_up)                │    │
│  │    output = hidden × W_down (downproject)                │    │
│  │                                                          │    │
│  │  For 70B: 8192 → 28672 → 8192 dimensions                │    │
│  └─────────────────────────────────────────────────────────┘    │
│                                                                  │
│  Output: Updated embeddings [11, 8192] + KV cache for layer 1   │
└──────────────────────────────────────────────────────────────────┘

... repeat for layers 2, 3, 4, ... all the way to layer 80

Step 4: The Attention Matrix (Visualized)

Here's what the 11×11 attention matrix looks like for our prompt at one attention head:

                Explain  how    a    CPU  works   to    a     5     -   year  -old
  Explain      [  1.0   0.0   0.0   0.0   0.0   0.0  0.0   0.0  0.0  0.0   0.0 ]
  how          [ 0.35   0.65  0.0   0.0   0.0   0.0  0.0   0.0  0.0  0.0   0.0 ]
  a            [ 0.15   0.25  0.60  0.0   0.0   0.0  0.0   0.0  0.0  0.0   0.0 ]
  CPU          [ 0.30   0.10  0.05  0.55  0.0   0.0  0.0   0.0  0.0  0.0   0.0 ]
  works        [ 0.10   0.15  0.05  0.40  0.30  0.0  0.0   0.0  0.0  0.0   0.0 ]
  to           [ 0.20   0.05  0.05  0.10  0.15  0.45 0.0   0.0  0.0  0.0   0.0 ]
  a            [ 0.05   0.05  0.10  0.05  0.05  0.30 0.40  0.0  0.0  0.0   0.0 ]
  5            [ 0.10   0.05  0.05  0.05  0.05  0.10 0.10  0.50 0.0  0.0   0.0 ]
  -            [ 0.05   0.03  0.02  0.03  0.02  0.05 0.05  0.35 0.40 0.0   0.0 ]
  year         [ 0.05   0.05  0.05  0.05  0.05  0.05 0.05  0.15 0.10 0.40  0.0 ]
  -old         [ 0.10   0.05  0.05  0.05  0.05  0.10 0.05  0.15 0.05 0.15  0.20]

  ↑ Each row = how much each token "attends to" every previous token.
  ↑ CAUSAL attention — tokens only see tokens before them (upper triangle = 0).
  ↑ "CPU" attends strongly to "Explain" (0.30) — it knows this is about explaining CPUs.
  ↑ "-old" attends to "5" (0.15) and "year" (0.15) — understands "5-year-old" as one concept.

The critical point: During prefill, this ENTIRE 11×11 matrix is computed in ONE shot. All 121 attention scores calculated simultaneously using GPU parallelism. This is why prefill is compute-heavy but fast.

Step 5: Output — First Token + KV Cache

After all 80 layers, the final hidden state of the LAST token goes through a linear layer to produce logits over the entire vocabulary (~128,000 tokens for Llama-3):

Final hidden state of "-old" → Linear projection → Logits [128,256]

Top predictions:
  "A"        → probability 0.12
  "Imagine"  → probability 0.09
  "Think"    → probability 0.08
  "So"       → probability 0.07
  "OK"       → probability 0.06
  ...

Sampled (temperature=0.7): "Imagine"   ← This is the FIRST generated token (TTFT)

Simultaneously, all Key and Value vectors from every layer are stored as the KV cache. This allows decode to skip reprocessing the entire prompt.

3. The Restaurant Analogy (Prefill Edition)

┌─────────────────────────────────────────────────────────────────────┐
│                     🍽️  THE LLM RESTAURANT                         │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  CUSTOMER (User) walks in and places an order:                      │
│  "I'd like a detailed Italian pasta dish, gluten-free,              │
│   with a side salad, no onions, for two people"                     │
│                                                                     │
│  This order = YOUR PROMPT (many specific requirements)              │
│                                                                     │
├─────────────────────────────────────────────────────────────────────┤
│                                                                     │
│  STEP 1: ORDER TICKET (Tokenization)                                │
│  ┌───────────────────────────────────────────────────────────┐      │
│  │ The waiter breaks the order into individual items:        │      │
│  │                                                           │      │
│  │  [Italian] [pasta] [gluten-free] [side salad]            │      │
│  │  [no onions] [for two]                                   │      │
│  │                                                           │      │
│  │ = 6 "tokens" on the ticket                               │      │
│  └───────────────────────────────────────────────────────────┘      │
│                                                                     │
│  STEP 2: KITCHEN READS FULL ORDER AT ONCE (Self-Attention)          │
│  ┌───────────────────────────────────────────────────────────┐      │
│  │ The head chef reads the ENTIRE ticket simultaneously.     │      │
│  │ They understand RELATIONSHIPS between items:              │      │
│  │                                                           │      │
│  │  "Italian" + "pasta" → ah, we're making Italian pasta     │      │
│  │  "gluten-free" modifies "pasta" → use rice noodles        │      │
│  │  "side salad" + "no onions" → salad without onions        │      │
│  │  "for two" → double everything                            │      │
│  │                                                           │      │
│  │ This cross-referencing of every item with every other     │      │
│  │ item = the ATTENTION MATRIX.                              │      │
│  └───────────────────────────────────────────────────────────┘      │
│                                                                     │
│  STEP 3: PREP STATIONS SET UP (KV Cache Generation)                 │
│  ┌───────────────────────────────────────────────────────────┐      │
│  │ Based on understanding the full order, the kitchen sets   │      │
│  │ up PREP STATIONS — pre-measured ingredients, sauces       │      │
│  │ heated, pans selected:                                    │      │
│  │                                                           │      │
│  │  Station 1: Rice noodles (measured for 2)                 │      │
│  │  Station 2: Italian sauce (no onion variant)              │      │
│  │  Station 3: Salad greens (no onions, for 2)              │      │
│  │  Station 4: Garnishes ready                               │      │
│  │                                                           │      │
│  │ These prep stations = KV CACHE                            │      │
│  └───────────────────────────────────────────────────────────┘      │
│                                                                     │
│  STEP 4: FIRST DISH ELEMENT PRODUCED (First Token)                  │
│  ┌───────────────────────────────────────────────────────────┐      │
│  │ With everything prepped, the kitchen produces the first   │      │
│  │ component: the boiling water hits the noodles.            │      │
│  │                                                           │      │
│  │ Time from order to first action = TTFT                    │      │
│  │ (Time To First Token)                                     │      │
│  └───────────────────────────────────────────────────────────┘      │
│                                                                     │
│  NOW: The kitchen moves to DECODE phase — producing each            │
│  dish component one at a time, using the prep stations (KV cache)   │
│  instead of re-reading the order ticket every time.                 │
└─────────────────────────────────────────────────────────────────────┘

4. Step-by-Step Technical Walkthrough

def prefill(prompt: str, model: TransformerModel) -> Tuple[Token, KVCache]:
    """
    Process the entire prompt and return the first generated token + KV cache.
    This runs ONCE per request.
    """

    # ═══════════════════════════════════════════════════════
    # STEP 1: Tokenize
    # ═══════════════════════════════════════════════════════
    token_ids = tokenizer.encode(prompt)
    # "Explain how a CPU works to a 5-year-old"
    # → [849, 1268, 264, 18622, 4375, 311, 264, 220, -, 3236, -]
    seq_len = len(token_ids)  # 11

    # ═══════════════════════════════════════════════════════
    # STEP 2: Embed
    # ═══════════════════════════════════════════════════════
    hidden_states = embedding_table[token_ids]
    # shape: [11, 8192]  (seq_len × hidden_dim)

    # ═══════════════════════════════════════════════════════
    # STEP 3: Process through all transformer layers
    # ═══════════════════════════════════════════════════════
    kv_cache = {}

    for layer_idx in range(80):  # 80 layers for Llama-3-70B

        # --- Self-Attention ---
        Q = hidden_states @ W_Q[layer_idx]  # [11, 8192] × [8192, 8192] → [11, 8192]
        K = hidden_states @ W_K[layer_idx]  # [11, 8192] × [8192, 1024] → [11, 1024]
        V = hidden_states @ W_V[layer_idx]  # [11, 8192] × [8192, 1024] → [11, 1024]

        # ★ STORE K and V — this IS the KV cache
        kv_cache[layer_idx] = (K, V)

        attention_scores = (Q @ K.transpose()) / sqrt(head_dim)  # [11, 11]
        attention_scores = apply_causal_mask(attention_scores)
        attention_weights = softmax(attention_scores)             # [11, 11]
        attention_output = attention_weights @ V                  # [11, 8192]

        # --- Feed-Forward Network ---
        gate = SiLU(attention_output @ W_gate[layer_idx])  # [11, 28672]
        up   = attention_output @ W_up[layer_idx]           # [11, 28672]
        ffn_output = (gate * up) @ W_down[layer_idx]       # [11, 8192]

        hidden_states = layer_norm(hidden_states + ffn_output)

    # ═══════════════════════════════════════════════════════
    # STEP 4: Predict first token
    # ═══════════════════════════════════════════════════════
    last_hidden = hidden_states[-1]                    # [8192]
    logits = last_hidden @ W_vocab                     # [128256]
    first_token = sample(logits, temperature=0.7)      # → "Imagine"

    return first_token, kv_cache

Key observations

1. All 11 tokens processed at once — hidden_states is always [11, 8192]. One big matrix multiplication.
2. KV cache = K and V saved per layer — kv_cache[layer_idx] = (K, V). 80 layers × 2 matrices × 11 tokens.
3. Attention matrix is N×N — 11×11 = 121 scores. For 4096 tokens: 16.7 million scores. This is why prefill is compute-bound.
4. Only the LAST token predicts the next token — the other 10 outputs are discarded (their K,V stay in cache).

5. KV Cache: What Gets Stored

KV Cache for "Explain how a CPU works to a 5-year-old" (11 tokens):

Layer 0:  K: [11, 8, 128]  V: [11, 8, 128]
Layer 1:  K: [11, 8, 128]  V: [11, 8, 128]
... (80 layers total) ...
Layer 79: K: [11, 8, 128]  V: [11, 8, 128]

Total: 80 × 2 × 11 × 8 × 128 = 22,528,000 float16 values = ~43 MB

For a 4096-token prompt:

80 × 2 × 4096 × 8 × 128 × 2 bytes (FP16) = ~1.34 GB per request
50 concurrent users = 67 GB = almost an entire H100's memory

6. Why Prefill is Compute-Bound

Q × K^T  →  shape: [N, N]

For N = 11:     11 × 11   = 121 multiply-adds           → trivial
For N = 512:    512 × 512 = 262,144 multiply-adds       → easy
For N = 4096:   4096 × 4096 = 16,777,216 multiply-adds  → heavy
For N = 32768:  32K × 32K = 1,073,741,824 multiply-adds → MASSIVE

Arithmetic Intensity = FLOPs / Bytes Loaded
Prefill:  ~N² FLOPs / ~N bytes  =  ~N  (HIGH — compute-bound)
Decode:   ~N FLOPs  / ~N bytes  =  ~1  (LOW — memory-bound)

GPU Utilization During Prefill:

CUDA Cores: ████████████████████████████░░░  ~85% utilized (doing matmuls)
Memory Bus:  ████████████░░░░░░░░░░░░░░░░░░  ~35% utilized (not the bottleneck)

This is GOOD — the GPU is doing what it's designed for.

7. Real Numbers: Llama-3-70B

Where the time goes (4096-token prefill, Llama-3-70B):

Self-Attention (QKV + attention + output):   ~55% of time
Feed-Forward Network (gate + up + down):     ~40% of time
Layer norms, residuals, embedding:            ~5% of time

Total: ~180ms on 1× H100, ~45ms with 4-way tensor parallelism

8. Prefill vs Decode Side-by-Side

┌──────────────────────────────────┬──────────────────────────────────┐
│           PREFILL                │           DECODE                 │
├──────────────────────────────────┼──────────────────────────────────┤
│ Processes ALL tokens at once     │ Generates ONE token per step     │
│ Attention: N×N matrix (121 ops)  │ Attention: 1×N vector (11 ops)   │
│ Bottleneck: COMPUTE              │ Bottleneck: MEMORY BANDWIDTH    │
│ GPU cores: ~85% utilized         │ GPU cores: ~12% utilized         │
│ Memory BW: ~35% utilized         │ Memory BW: ~88% utilized         │
│ Runs ONCE per request            │ Runs N times (once per token)    │
│ Duration: milliseconds           │ Duration: seconds                │
│ Output: KV cache + 1st token     │ Output: all remaining tokens     │
│ Metric: TTFT                     │ Metric: ITL, throughput          │
│ Wants: More FLOPS, higher TP     │ Wants: More bandwidth, batching  │
│ Scaling: O(N²) with prompt len   │ Scaling: O(N) per step           │
└──────────────────────────────────┴──────────────────────────────────┘

BEFORE (Aggregated — same GPU does both):

  GPU 0: ██ PREFILL ██░░░ DECODE ░░░░██ PREFILL ██░░░ DECODE ░░░░
         compute-heavy  BW-heavy      blocks decode   wastes compute

AFTER (Disaggregated — dedicated GPUs):

  Prefill GPU:  ██ REQ1 ██ REQ2 ██ REQ3 ██ REQ4 ██  (always computing)
  Decode GPU:   ░░ REQ1 ░░░░░░░░░░░░░░░░░░░░░░░░    (batches users)
  KV Transfer:  Prefill GPU ──NIXL──→ Decode GPU (~1ms NVLink)

Sources

1. Attention Is All You Need — Vaswani et al., 2017 → arxiv.org/abs/1706.03762
2. PagedAttention / vLLM — Kwon et al., 2023 → arxiv.org/abs/2309.06180
3. Splitwise — Patel et al., 2024 → arxiv.org/abs/2311.18677
4. DistServe — Zhong et al., 2024 → arxiv.org/abs/2401.09670
5. NVIDIA Dynamo → github.com/ai-dynamo/dynamo
6. FlashAttention — Dao et al., 2022 → arxiv.org/abs/2205.14135
7. Llama 3 Model Card — Meta, 2024 → github.com/meta-llama/llama3