Prefill: Process the Prompt and Build the Cache

In A Decoder Block from Input to Logits, you traced the full forward pass of a decoder block: residuals, Pre-RMSNorm, SwiGLU, and the unembedding projection to logits. You now understand how a single transformer layer transforms token representations. This lesson begins Part IV by revealing a critical fact: LLM generation happens in two completely different phases. We start with the first phase — prefill — where your prompt is processed all at once.

---

The Two Phases of LLM Inference

When you send a prompt to an LLM, the model does not simply “read” it token by token like a human reads a sentence. Instead, generation splits into two distinct phases with fundamentally different computational characteristics:

Phase 1: Prefill (Prompt Processing)

Prefill takes the user prompt of length P (where P is the number of tokens in your prompt), processes all P tokens in parallel, computes their representations, builds the initial Key-Value (KV) cache, and generates the very first new token.

Key properties of prefill:

• All P prompt tokens are known upfront
• Massive parallelism: the GPU multiplies large matrices covering all P tokens simultaneously
• Compute-bound: limited by the GPU’s raw arithmetic processing power (measured in FLOPS — floating-point operations per second)
• Happens exactly once per request

Phase 2: Decode (Token Generation)

Decode takes the newly generated token, reads the cached past keys and values, generates token P+1, appends it to the sequence, and repeats sequentially one token at a time until a stopping condition is met (end-of-sequence token, maximum length, or stop string).

Key properties of decode:

• Only one new token position is processed per step
• Sequential: cannot parallelize across output positions
• Memory-bandwidth-bound: limited by how fast the GPU can read the KV cache from High Bandwidth Memory (HBM)
• Happens once per generated token

Why this distinction matters: The performance bottlenecks, optimization strategies, and latency calculations are completely different for each phase. Understanding prefill first gives you the foundation to analyze decode later.

---

Why Prefill Runs in Parallel

When you submit a prompt (for example, 500 words ≈ 700 tokens), all words are already known. The model does not need to “guess” the prompt — it only needs to process it.

Matrix Multiplication Across All Tokens

Recall from the decoder block that each layer computes projections:

Q = x @ W_Q    # Query projection
K = x @ W_K    # Key projection
V = x @ W_V    # Value projection

During prefill, the input x has shape [B, P, d_model] where:

• B = batch size (usually 1 for a single request)
• P = prompt length in tokens
• d_model = model dimension (e.g., 4096 for Llama 7B)

The weight matrix W_Q has shape [d_model, d_k]. The matrix multiplication x @ W_Q processes all P token rows simultaneously. On a GPU with thousands of CUDA cores, this is a single large GEMM (General Matrix Multiply) operation.

Arithmetic Intensity: The Key to Compute-Bound Behavior

Arithmetic intensity = (FLOPs performed) / (Bytes transferred from memory)

During prefill:

• Weights (e.g., W_Q, W_K, W_V, W_O, FFN weights) are read once from VRAM
• These weights are reused across all P tokens
• The total FLOPs scale with P, but weight bytes transferred stay roughly constant

This means as P grows, arithmetic intensity increases. Large matrix operations achieve high utilization of the GPU’s Tensor Cores. The GPU spends most of its time doing math, not waiting for memory. Therefore, prefill is predominantly compute-bound — limited by the GPU’s peak FLOPS.

FLOPs vs FLOPS reminder:

• FLOPs (floating-point operations) = a countable amount of compute work (e.g., “this layer requires 2 trillion FLOPs”)
• FLOPS (floating-point operations per second) = speed/throughput (e.g., “this GPU delivers 300 TFLOPS peak”)

---

Causal Attention Across the Prompt

Even though all P tokens are processed in parallel, the model must respect causality: token i can only attend to tokens j ≤ i (itself and previous tokens). It cannot “see” future tokens in the prompt.

The Attention Score Matrix

For a prompt of length P, at each attention head:

• Q ∈ [P, d_k] — Query for each position
• K ∈ [P, d_k] — Key for each position
• Score matrix S = (Q @ K^T) / √d_k has shape [P, P]

The causal mask forces all entries where column > row to negative infinity before softmax. After masking and softmax, the attention matrix A is lower-triangular:

A = softmax(mask(S))

Position:    0    1    2    3  ...  P-1
           ┌───────────────────────────┐
        0  │ ●    0    0    0  ...  0  │
        1  │ ●    ●    0    0  ...  0  │
        2  │ ●    ●    ●    0  ...  0  │
        3  │ ●    ●    ●    ●  ...  0  │
        ⋮  │ ⋮    ⋮    ⋮    ⋮   ⋱   ⋮  │
      P-1  │ ●    ●    ●    ●  ...  ●  │
           └───────────────────────────┘
           ● = attends to this position
           0 = masked out (cannot attend)

What This Means in Practice

• Token 0 (first prompt token) attends only to itself
• Token 1 attends to tokens 0 and 1
• Token P-1 (last prompt token) attends to all P tokens
• The final representation at layer L for position P-1 is passed to the unembedding head to produce logits for the very first generated token

Important: Causality is strictly enforced even though all tokens are processed in parallel. The mask is applied inside the softmax, so gradient flow and attention weights respect the same rules as sequential processing.

---

Building the Initial KV Cache

This is the crucial bridge between prefill and decode. At every transformer layer l, the Key tensor K_l and Value tensor V_l are computed for all P prompt tokens:

K_l ∈ [B, H_kv, P, d_head]   # Keys at layer l
V_l ∈ [B, H_kv, P, d_head]   # Values at layer l

Where:

• H_kv = number of KV heads (e.g., 8 for GQA with 32 query heads grouped 4:1)
• d_head = head dimension (typically d_model / H_q)

Crucially, these tensors are not discarded! They are saved in GPU High Bandwidth Memory (HBM) as the initial KV cache.

Why Save Them?

When the decode phase begins:

• The model will never need to recalculate the prompt tokens’ keys and values again
• Each decode step only computes Q, K, V for the single new token
• The new token’s attention reads the full cached history: K_cache ∈ [B, H_kv, P + t, d_head], V_cache ∈ [B, H_kv, P + t, d_head] where t is the number of tokens generated so far

This reuse is why decode is fast per token — but it also means the KV cache grows with every generated token, eventually consuming significant GPU memory.

Cache Shape at End of Prefill

After prefill completes (before any decode step):

KV Cache per layer:
  Keys:   [B, H_kv, P, d_head]
  Values: [B, H_kv, P, d_head]

Total cache across L layers:
  2 × L × B × H_kv × P × d_head  elements

We will calculate the exact byte size in Exercise 2.

---

Time to First Token (TTFT)

Definition: Time to First Token (TTFT) is the elapsed time between when the prompt is received by the inference server and when the first generated token is emitted (available to stream back to the user).

TTFT Components

TTFT ≈ Network Latency + Tokenization + Prefill Compute Time

• Network Latency: Time for the request to reach the server (milliseconds, often negligible on same-region connections)
• Tokenization: Converting raw text to token IDs (microseconds, usually negligible)
• Prefill Compute Time: The dominant term — time to run all L layers over P tokens

Estimating Prefill Latency from FLOPs and FLOPS

Rule of thumb for transformer FLOPs: ~2 FLOPs per parameter per token (1 multiply + 1 add in each matrix multiply).

For a model with N parameters and prompt length P:

Prefill FLOPs ≈ 2 × P × N

If the GPU has peak tensor compute Peak_FLOPS and achieves MFU (Model Flops Utilization) fraction of peak:

Achieved_FLOPS = Peak_FLOPS × MFU
Prefill Time   = Prefill FLOPs / Achieved_FLOPS

Example (illustrative numbers):

• Model: 7B parameters (N = 7 × 10⁹)
• Prompt: P = 1000 tokens
• Prefill FLOPs ≈ 2 × 1000 × 7×10⁹ = 1.4×10¹³ = 14 TFLOPs
• GPU: 300 TFLOPS peak (FP16/BF16 tensor)
• MFU: 50% → Achieved = 150 TFLOPS
• Prefill Time = 14 TFLOPs / 150 TFLOPS = 0.093 s ≈ 93 ms

Key insight: Because prefill FLOPs scale linearly with P, longer prompts directly increase TTFT. This is why very long contexts (e.g., 100K tokens) can take seconds before the first token appears.

---

Hands-On Exercise 1: Prefill Compute (FLOPs) and Ideal Latency

Problem: Calculate the prefill FLOPs and estimated latency for a 7B parameter model with a 1000-token prompt on a GPU with 300 TFLOPS peak FP16/BF16 tensor compute operating at 50% MFU.

Given:

• Model parameters N = 7 × 10⁹
• Prompt length P = 1000
• GPU peak tensor compute = 300 TFLOPS
• MFU = 50% (0.50)

Step-by-step Solution:

Step 1: Calculate prefill FLOPs

Prefill FLOPs ≈ 2 × P × N
              = 2 × 1000 × (7 × 10⁹)
              = 14 × 10¹²
              = 14 TFLOPs

Step 2: Calculate achieved GPU throughput

Achieved FLOPS = Peak FLOPS × MFU
               = 300 TFLOPS × 0.50
               = 150 TFLOPS

Step 3: Calculate prefill compute time

Time = FLOPs / FLOPS
     = 14 TFLOPs / 150 TFLOPS
     = 0.09333... seconds
     = 93.3 ms

Final Answer: ~93 ms of pure compute time for prefill.

Verification: You can verify this by checking 14 / 150 = 0.09333... in any calculator. Real-world TTFT will be higher due to kernel launch overhead, memory transfers, tokenization, and network latency.

Note on MFU: 50% MFU is a reasonable achievable target for well-optimized prefill kernels on modern GPUs (H100, A100). Naive implementations may achieve 10-20%. FlashAttention and fused kernels push toward 50-60%.

---

Hands-On Exercise 2: Initial KV Cache Size Calculation

Problem: Calculate the initial KV cache memory (in bytes, MiB, and MB) after prefill for a prompt of P = 2048 tokens, using a model with L = 32 layers, H_kv = 8 KV heads (GQA), d_head = 128, and BF16 precision (2 bytes per element).

Given:

• P = 2048 (prompt length)
• L = 32 (layers)
• H_kv = 8 (KV heads per layer)
• d_head = 128 (head dimension)
• Precision = 2 bytes (BF16)

Formula:

KV Cache Bytes = 2 × L × H_kv × d_head × P × bytes_per_element

(The factor of 2 accounts for both Keys and Values)

Step-by-step Solution:

Step 1: Calculate elements per layer

Elements per layer (K only) = H_kv × P × d_head
                            = 8 × 2048 × 128
                            = 2,097,152 elements

Elements per layer (K + V) = 2 × 2,097,152
                           = 4,194,304 elements

Step 2: Calculate total elements across all layers

Total elements = L × elements_per_layer
               = 32 × 4,194,304
               = 134,217,728 elements

Step 3: Calculate bytes (BF16 = 2 bytes)

Total bytes = 134,217,728 × 2
            = 268,435,456 bytes

Step 4: Convert to MiB (mebibytes, 2²⁰ = 1,048,576)

MiB = 268,435,456 / 1,048,576
    = 256.0 MiB

Step 5: Convert to MB (megabytes, 10⁶ = 1,000,000)

MB = 268,435,456 / 1,000,000
   = 268.4 MB

Final Answer: 268,435,456 bytes = 256.0 MiB = 268.4 MB

Interpretation: Before generating a single output token, the prompt has already claimed 256 MiB of GPU HBM for the KV cache. For a 24 GB GPU, this is ~1% of memory — but it grows linearly with every generated token during decode.

Check: Using the compact formula: 2 × 32 × 8 × 128 × 2048 × 2 = 268,435,456 ✓

---

Hands-On Exercise 3: Prompt Causal Mask Trace

Problem: For the prompt "AI is fast" tokenized as 3 tokens (P = 3), write out the 3×3 attention connectivity matrix. Explain which tokens each position attends to and why token 2 ("fast") is the one that predicts the first new token.

Step-by-step Solution:

Step 1: Tokenize the prompt

Position 0: "AI"
Position 1: "is"
Position 2: "fast"

Step 2: Construct the causal attention connectivity matrix The matrix shows 1 if row position attends to column position, 0 if masked:

           Attends to →
         ┌───┬───┬───┐
         │ 0 │ 1 │ 2 │  ← Column = Key position
     ┌───┼───┼───┼───┤
   0 │ 1 │ 0 │ 0 │  ← Row 0 (token "AI")
     ├───┼───┼───┼───┤
   1 │ 1 │ 1 │ 0 │  ← Row 1 (token "is")
     ├───┼───┼───┼───┤
   2 │ 1 │ 1 │ 1 │  ← Row 2 (token "fast")
     └───┴───┴───┴───┘
     ↑
     Query position

Step 3: Explain each row

• Row 0 (“AI”): Attends only to position 0 (itself). Cannot see “is” or “fast”.
• Row 1 (“is”): Attends to positions 0 and 1 (“AI” and “is”). Cannot see “fast”.
• Row 2 (“fast”): Attends to all three positions (0, 1, 2). Has full context of the prompt.

Step 4: Why token 2 (“fast”) predicts the first new token During prefill, the model computes representations for all positions in parallel. However, the next-token prediction head (unembedding projection) is applied to the final layer’s hidden state at the last prompt position — position P-1 = 2.

The logits for the first generated token come from:

logits_first = hidden_state[layer_L, position_2] @ W_unembed

Position 2 (“fast”) is the only position that has “seen” the entire prompt through causal attention. Its representation encodes the full context "AI is fast". The model uses this representation to predict what comes next (e.g., ".", "!", "today", etc.).

Final Answer: The connectivity matrix is lower-triangular with ones on and below the diagonal. Token 2 (“fast”) attends to all prompt tokens and its final-layer representation produces the first generated token’s logits.

---

Summary Reference Table: Prefill vs Decode

Aspect	Prefill (Phase 1)	Decode (Phase 2)
Input length	P tokens (full prompt)	1 token (newly generated)
Execution mode	Massively parallel (large GEMMs)	Sequential (one position per step)
Primary bottleneck	Compute-bound (GPU FLOPS)	Memory-bandwidth-bound (HBM reads)
Arithmetic intensity	High (weights reused across P tokens)	Low (KV cache read per token)
KV cache	Written (built from scratch)	Read (appended incrementally)
Latency metric	Time to First Token (TTFT)	Inter-Token Latency (ITL) / Throughput
Scaling with length	FLOPs ∝ P → TTFT ∝ P	Cache size ∝ P + t → ITL grows with context
Parallelism	Across P tokens (batch × sequence)	Across batch requests only (continuous batching)

---

Primary Source Citations

• Vaswani et al. (2017), Attention Is All You Need, Section 3.2 (Scaled Dot-Product Attention, Causal Masking) — arXiv:1706.03762
• Pope et al. (2022), Efficiently Scaling Transformer Inference — arXiv:2211.05102 (Prefill/decode distinction, compute vs memory bounds, KV cache mechanics)
• Dao et al. (2022), FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness — arXiv:2205.14135 (IO-awareness, arithmetic intensity, tiling for prefill)

---

What Comes Next

• Next lesson: Decode: Generate One Position at a Time — Follow each new position through cached attention and explain memory reads, output throughput, and inter-token delay.
• Previous lesson: A Decoder Block from Input to Logits
• Track overview: Start Here: LLM Foundations