Inference = Prefill + Decode
LLM Inference is split in 2 components:
- Prefill
- Decode
Prefill (Prompt Encoding)
- Input prompt tokens = processed in parallel through entire Transformer stack.
- Characteristics:
- Compute-bound (large MatMuls, parallelizable).
- KV cache is built (for decode).
Decode (Autoregressive Generation)
Model generates one token at a time, conditioned on KV cache + past tokens. Characteristics:
- Memory-bound (repeated KVcache lookups, smaller MatMul).
- Sequential dependency = less parallelism.
- Dominant phase for long outputs.
Workflow Example
- Step 1 (Prefill):
- User sends a 1k-token prompt = GPU processes all 1k tokens in parallel = builds KV cache.
- Input: [Batch size x Prompt length x Hidden size]
- KV: [Batch size x #head x Prompt length x dim_head]
- Note: #head x dim_head = Hidden size
- Step 2 (Decode):
- Model starts generating tokens (e.g, 200 tokens).
- Each step: query attends to full KV cache.
- New row added to KV per step.
- Latency accumulates linearly with output length.
Why Split Matters?
- Optimization targets differ:
- Prefill → FLOP efficiency (maximize GPU compute).
- Decode → memory efficiency (optimize KV cache movement, avoid stalls).
- Scheduling differs:
- Prefill can handle large batch size.
- Decode needs careful batching (continuous batching, paged attention).
Optimizations
- Prefill: Chunked Prefill, Parallel Prefill/Decode Disaggregation.
- Decode: Continuous Batching, Paged Attention, KV Compression, Speculative Decoding.
