Prepacking: A Simple Method for Fast Prefilling and Increased Throughput in Large Language Models

TL;DR

Prefilling for autoregressive LLM inference wastes compute when prompts within a batch vary in length due to padding. The authors propose Prepacking, which uses a bin-packing strategy to concatenate prompts into fewer, compact sequences and applies independent masking with restart positional encodings to compute multiple per-prompt KV caches in a single forward pass. This yields substantial improvements in prefilling speed and memory efficiency, with up to 6x TTFT speedups and up to 16x larger practical batch sizes in memory-constrained settings, across diverse datasets and model scales. The approach is architecture-agnostic and easy to implement in PyTorch, offering practical impact for real-world LLM serving and suggesting extensions to generation-heavy workflows.

Abstract

During inference for transformer-based large language models (LLM), prefilling is the computation of the key-value (KV) cache for input tokens in the prompt prior to autoregressive generation. For longer input prompt lengths, prefilling will incur a significant overhead on decoding time. In this work, we highlight the following pitfall of prefilling: for batches containing high-varying prompt lengths, significant computation is wasted by the standard practice of padding sequences to the maximum length. As LLMs increasingly support longer context lengths, potentially up to 10 million tokens, variations in prompt lengths within a batch become more pronounced. To address this, we propose Prepacking, a simple yet effective method to optimize prefilling computation. To avoid redundant computation on pad tokens, prepacking combines prompts of varying lengths into a sequence and packs multiple sequences into a compact batch using a bin-packing algorithm. It then modifies the attention mask and positional encoding to compute multiple prefilled KV-caches for multiple prompts within a single sequence. On standard curated dataset containing prompts with varying lengths, we obtain a significant speed and memory efficiency improvements as compared to the default padding-based prefilling computation within Huggingface across a range of base model configurations and inference serving scenarios.

Figures (13)

• Figure 1: Left: The standard full batching approach (e.g., used in HuggingFace) pads shorter prompts to maximum prompt length in the batch. Each prompt has its own causal attention mask. Right: Prepacking combines multiple prompts into a single sequence using a bin-packing algorithm, and applies independent masking and restart positional encodings (numbers inside token boxes) to avoid prompts attending to other prompts. Both strategies are equivalent at decoding time, but prepacking is more compute efficient during prefilling.
• Figure 2: Prefilling latency scaling with batch size $k$ highlights GPU parallelization limits. Results averaged over 100 runs.
• Figure 3: An actual example of one batch sampled from the MMLU dataset shows how a batch of 16 prompts is packed into a compact packed batch. Restart indices denote the point at where independent mask and position encoding are reset to preserve the semantics of each individual prompt.
• Figure 4: Average inference time per batch for various language models using prepacking and Full Batching, with a batch size of 16. The comparison is conducted across multiple datasets with two metrics, Prefilling Time and TTFT (Time-to-First-Token). Error bars represent the standard deviation of inference times across batches and seed runs. The results show that prepacking consistently leads to reduced inference times compared to Full Batching and exhibits reduced variability, as evidenced by smaller standard deviation errors, indicating more reliable and predictable inference times when adopting prepacking.
• Figure 5: Peak GPU memory usage comparison across models and datasets on a single GPU. Absent data points indicate out-of-memory errors. Prepacking achieves lower peak GPU memory usage and allows for up to 16x larger batch sizes in prefilling computations than Full Batching and Length-Ordered Batching.