Optimizing LLM Inference: Compression, Attention, Batching

Efficiently serving LLMs requires separating what pertains to computation from what is constrained by memory. Between quantization, attention variants, and batching policies, several techniques provide tangible gains in throughput and cost. Some, like FlashAttention, can be applied without retraining; others, like MQA, require tailored training.

Reducing Model Size with Quantization

Model compression allows for a decrease in size and acceleration of inference, with a trade-off on precision that can remain limited if the method is well chosen. Quantization involves reducing the numerical precision of weights, and sometimes activations. A model trained in 16-bit floating point occupies about 2 bytes per parameter; in 8-bit integers, 1 byte; and in 4 bits, 0.5 bytes. For a model with 7 billion parameters, moving from FP16 to INT4 reduces the required memory from about 14 GB to 3.5 GB. This reduction can enable the use of less expensive or more accessible hardware. Techniques like GPTQ and AWQ have enabled the creation of 4-bit quantized models.

Adjusting Attention to Limit Memory Traffic

Attention represents the largest share of computation during transformer inference. The classic multi-head attention (MHA) mechanism uses distinct key and value heads for each attention head. With MQA, all query heads share the same set of keys and values, which reduces the amount of data transferred during the decode phase, without altering the computation, but at the cost of a slight loss in precision; the model must be trained or adjusted with MQA to benefit during inference. GQA groups key and value heads into clusters, providing most of the memory gains of MQA while maintaining superior modeling capacity compared to MQA alone. FlashAttention, on the other hand, reorganizes the order of computations and uses tiling to keep intermediate values in fast SRAM memory, significantly reducing memory traffic without changing the mathematical outcome. This optimization applies directly, without requiring model retraining.

Orchestrating Requests: From Static to On-the-Fly Batching

When processing occurs request by request, the GPU is not fully utilized, as the model weights must be loaded into memory for each forward pass, regardless of the number of tokens. Increasing the number of requests processed simultaneously allows for distributing this cost and increasing throughput. Static batching, which waits for a fixed batch to be formed, is inefficient when output lengths vary: the batch does not progress until the longest request is completed, slowing everything down. Dynamic batching groups requests arriving within a short window, with a maximum size and a timeout; it starts as soon as a threshold is reached, balancing latency and batch size. However, once started, all requests advance together, which can block shorter ones. Continuous or on-the-fly batching immediately evicts completed sequences and inserts new requests without waiting for others to finish, thus maintaining high GPU utilization even with highly variable outputs. This strategy is now standard in production environments like vLLM and TensorRT-LLM.

Mastering KV Cache Memory with PagedAttention and Shared Prefixes

KV caching involves storing the keys and values of previous tokens in GPU memory to speed up the decode phase, at the cost of increased memory consumption. The memory requirements for the cache grow with batch size and sequence length; for a 7 billion parameter model in 16 bits, it can consume several gigabytes per request at moderate lengths. With long contexts or large batches, memory becomes the primary limit on the number of simultaneous requests. A naive allocation at maximum length leads to significant internal fragmentation. PagedAttention segments the cache into fixed-size blocks, allocated on demand during generation, without pre-reservation, which reduces waste and allows for increased batch sizes and throughput at constant hardware. vLLM activates this approach by default. When multiple requests share the same prefix, prefix caching allows for calculating and reusing this common segment just once; in RAG pipelines or with long system prompts, this can eliminate a significant portion of redundant computations.

Reading Metrics According to Prefill and Decode

Inference of decoder-only LLMs is divided into two distinct phases: prefill and decode, each with its own constraints. The prefill phase, highly parallelized and limited by computation, prepares the keys and values for the first token. Then, the decode phase generates tokens one by one, without intra-sequence parallelism, and primarily faces memory bandwidth limitations, with the GPU spending a lot of time waiting for data transfers. Effective optimizations are those that reduce the volume of data moved or increase useful work per memory access. Time to first token (TTFT) measures prefill performance, while tokens per second (TPS) after the first token reflect decode performance. Improving one of these phases does not necessarily improve the other; the choice of techniques depends on the priority metric for the intended use.

Other Levers: Speculative Decoding and Multi-GPU

Strategies combining memory management, batching, and attention optimizations can directly improve throughput and latency. Speculative decoding is mentioned as a way to reduce inference costs. Multi-GPU parallelism can increase inference capacity for large or demanding workloads.