Prefill Phase

Unlike the sequential decode phase, the prefill phase processes the entire input prompt in a single, parallel forward pass. This allows the model to leverage the full parallelism of modern accelerators (GPUs/TPUs) but results in high FLOP (Floating Point Operation) consumption proportional to prompt length. The computational complexity is typically O(n²d) for attention, where n is the prompt length and d is the model dimension, making long prompts expensive.

The primary output of the prefill phase is the initial state of the KV cache. For each transformer layer and attention head, the model computes and stores:

• Key (K) tensors: Representations used for matching in attention.
• Value (V) tensors: Representations used to produce the output. This cached state is then read-only during the subsequent decode phase, enabling efficient autoregressive generation without recomputation.

The prefill phase determines the fixed memory allocation for the KV cache for the entire request lifetime. The cache size scales linearly with:

• Batch Size: Number of concurrent requests.
• Context Length: Total prompt tokens.
• Model Configuration: Layers, heads, and hidden dimensions. For a 70B parameter model, the KV cache for a single 4K-token context can require ~10-15 GB of GPU memory. This makes prefill a critical moment for memory planning and potential out-of-memory (OOM) errors.

The time to generate the first output token (Time to First Token - TTFT) is almost entirely dictated by the prefill phase. This latency is sensitive to:

• Prompt Length: Longer prompts increase prefill time linearly or quadratically.
• Hardware Compute: FLOP/s of the accelerator.
• Memory Bandwidth: Speed of loading weights and writing the initial cache. Optimizations like FlashAttention directly target prefill latency by reducing memory I/O during this phase.

Because the prompt length is known upfront, the computation graph for the prefill phase is static and predictable. This allows for aggressive compiler-level optimizations that are not possible during dynamic decoding, such as:

• Optimal Kernel Selection: Using specialized fused kernels for known sequence lengths.
• Operator Fusion: Combining linear, activation, and layer norm operations.
• Memory Planning: Pre-allocating buffers for intermediate activations.

The prefill and decode phases represent two distinct computational regimes:

Understanding this dichotomy is essential for holistic inference optimization.

FlashAttention is an IO-aware exact attention algorithm that dramatically speeds up the prefill phase. It avoids materializing the large, intermediate attention matrix to high-bandwidth memory (HBM) by recomputing parts of the attention operation on-chip. This reduces HBM reads/writes, which are the primary bottleneck. Related kernel fusion techniques combine multiple operations (e.g., softmax, masking, scaling) into a single, optimized GPU kernel to minimize launch overhead and memory traffic.

• Primary Benefit: Reduces memory bandwidth pressure, the key limiter in prefill.
• Impact: Can provide 2-4x speedup for long-context prefill compared to standard attention implementations.

While continuous batching is often associated with the decode phase, it is equally critical for prefill efficiency. An optimized scheduler dynamically groups incoming requests with similar context lengths into a single batch for prefill execution.

• Key Strategy: Grouping sequences of similar length minimizes the padding overhead, which wastes FLOPs and memory bandwidth.
• Dynamic Scheduling: Advanced systems (e.g., vLLM, TGI) use cache-aware scheduling to prefill requests in an order that optimizes subsequent decode-phase memory locality and KV cache utilization.

Multi-Query Attention (MQA) and Grouped-Query Attention (GQA) are architectural modifications that directly reduce the memory and compute cost of the KV cache generated during prefill.

• MQA: All query heads share a single key head and a single value head. This drastically shrinks the size of the KV tensors that must be computed and stored.
• GQA: A tunable middle ground where groups of query heads share a single key/value head.
• Prefill Impact: These variants reduce the total FLOPs and memory volume during the parallel attention computation of the prompt, leading to faster prefill times and a smaller cache to manage.

For very large models or extremely long contexts, the prefill computation must be distributed across multiple devices. Two primary parallelism strategies are used:

• Tensor Parallelism: Splits individual model layers (including attention heads and feed-forward weights) across GPUs. During prefill, the massive matrix multiplications are distributed, reducing compute time per device. This requires careful synchronization and high inter-GPU bandwidth.
• Sequence Parallelism: Splits the input sequence dimension across devices. Each GPU processes a chunk of the prompt tokens, and results are combined via communication (e.g., all-reduce) for the attention operation. This is effective for contexts longer than the model's per-GPU memory limit.

Using lower numerical precision for compute and storage during prefill provides significant speed and memory benefits.

• Compute in BF16/FP16: Performing the heavy matrix multiplications of the prefill phase in 16-bit floating-point (BF16 or FP16) doubles the computational throughput on modern GPUs (e.g., NVIDIA Tensor Cores) compared to FP32.
• KV Cache in FP8/INT8: The newly computed KV cache can be immediately quantized and stored in 8-bit format (e.g., FP8). This reduces the memory footprint written to during prefill and the bandwidth required for subsequent decode-phase reads.
• Requirement: Models must be stable at lower precision, often achieved through quantization-aware training or sophisticated post-training quantization methods.

In a production serving system, the goal is to hide prefill latency and maximize overall throughput. This involves overlapping prefill with other operations.

• Prefill/Decode Overlap: While one request is in its decode phase (reading from its KV cache), the system can schedule the prefill computation for a new request on the same GPU, utilizing idle compute cycles.
• Asynchronous CUDA Graphs: Capturing the entire prefill computational graph as a single, replayable CUDA Graph eliminates Python and kernel launch overhead, leading to more predictable and lower latency.
• Non-Blocking HBM Transfers: Overlapping the transfer of the computed KV cache to its final memory location with the beginning of decode computations for other requests.

The autoregressive token generation stage that follows the prefill phase. During decoding, the model uses the pre-computed KV cache to efficiently attend to the entire context, predicting one token at a time with minimal new computation. This phase is typically memory-bandwidth bound, as performance is limited by reading the cached keys and values.

The memory buffer that stores computed key and value tensors for all previous tokens. The prefill phase's primary job is to populate this cache for the initial prompt. During the decode phase, the model reads from this cache instead of recomputing attention from scratch, which is the key to efficient autoregressive generation.

• Purpose: Eliminates redundant computation.
• Memory Cost: Scales linearly with batch size and sequence length.

A dynamic scheduling technique that maximizes GPU utilization across both prefill and decode phases. Unlike static batching, it allows:

• New requests to join an ongoing batch (prefill).
• Completed sequences to exit, freeing resources.
• It treats the KV cache as a shared resource, managed by systems like PagedAttention, to handle variable sequence lengths efficiently within a single batch.

The maximum sequence length a model can process, a limit directly tied to prefill. The prefill phase must compute the KV cache for all tokens in the prompt, making its cost quadratic with respect to context length for standard attention. Techniques like Sliding Window Attention or StreamingLLM modify this relationship to enable longer contexts by bounding or managing the cache size.

The two primary performance regimes in inference:

• Prefill Phase: Typically compute-bound. Performance is limited by the GPU's FLOPs as it performs dense matrix multiplications for attention across the full prompt.
• Decode Phase: Typically memory-bound. Performance is limited by the bandwidth to read the KV cache for each new token generation. Optimizations target the bottleneck of each phase separately.