8 Attention optimization

This chapter introduces attention optimization as a practical way to reduce the inference costs of large language models, especially when long contexts make the KV cache and attention computation expensive. It contrasts two complementary directions: compressing the KV cache to lower dynamic memory use, and identifying attention layers that contribute little so they can be removed from the architecture. The main theme is that attention is often the bottleneck during generation, so improving it can increase usable context length, batch size, or efficiency without necessarily retraining a model from scratch.

The first half focuses on KV cache quantization in Hugging Face and vLLM. Using HQQ or similar backends, the book shows how cached key and value tensors can be stored in reduced precision and dequantized when needed, saving memory while preserving the model weights. The experiments reveal that the outcome depends strongly on hardware and runtime: on older GPUs, 4-bit KV cache compression can reduce memory substantially but may slow generation, while on newer GPUs with native FP8 support, capacity can increase with little or no speed penalty. In contrast, weight quantization mainly affects throughput and can degrade text quality more visibly, showing that cache quantization and weight quantization solve different problems.

The second half moves from memory optimization to structural pruning. By measuring cosine similarity between the input and output around the attention phase, the chapter scores attention layers for redundancy and removes the least important ones while keeping the MLPs intact. It explains how to physically patch the model, preserve the modified architecture for saving and reloading, and evaluate the trade-offs in memory, speed, and quality. The results show that modest attention removal can save memory with limited benchmark loss, but aggressive removal quickly harms open-ended generation, especially long-form coherence. The chapter concludes that attention pruning is effective when matched to the task, can be combined with quantization, and may benefit from recovery methods such as knowledge distillation when generation quality matters.

KV cache quantization write-read cycle. At each generation step, the K and V tensors produced by the Attention module are compressed from FP16 to 4-bit before being written to GPU memory. When the next token is processed, the cached values are dequantized back to FP16 before the attention computation. The cycle repeats for every generated token.

Three bar charts comparing static VRAM, dynamic VRAM delta, and throughput across three inference configurations of Llama-3.2-3B on a T4 GPU: FP16 baseline, 4-bit weight quantization (BitsAndBytes), and FP16 with 4-bit KV cache quantization (HQQ).

Three bar charts comparing static VRAM, dynamic VRAM delta, and throughput across three inference configurations of Llama-3.2-3B on an L4 GPU: BF16 baseline, 4-bit weight quantization (BitsAndBytes), and BF16 with 4-bit KV cache quantization (HQQ).

VRAM allocation across inference engines and quantization strategies. In Hugging Face, the KV cache grows dynamically into available memory. In vLLM, a fixed memory budget is divided at startup between model weights and a pre-allocated KV block pool. Reducing weight precision to NF4 shrinks the weight footprint, expanding the KV block pool within the same budget. Quantizing the KV blocks themselves reduces each block's size, fitting more tokens into the same pool.

Comparison of the three configurations run on a T4 GPU: KV cache capacity in tokens (left), percentage gain over baseline (center), and generation throughput in tokens per second (right). On this GPU, FP8 delivers the highest capacity gain (+111.1%) but penalizes throughput, while NF4 triples the baseline throughput.

The same three experiments run on an L4 GPU. On this GPU with native FP8 support, KV cache quantization doubles the available capacity (+105.2%) while keeping throughput virtually identical to the baseline (35.8 vs 35.6 tok/s). NF4 weight quantization offers a different profile: lower capacity gain (+31.8%), but throughput more than twice the baseline (82.5 tok/s). Each technique serves a different goal.

Internal structure of a LlamaDecoderLayer showing the two hook anchor points used to measure attention layer importance. Hook 1 intercepts the hidden state at input_layernorm, capturing X_A before the attention computation. Hook 2 intercepts the tensor at post_attention_layernorm, capturing Y_A after the attention output has been added back through the residual connection.

Attention Module importance in Llama-3.2-3B, sorted by layer index. Layers 0 and 1 stand clearly above the rest. Layers 2 through 17 form a moderately active band. The genuinely low-score region concentrates in layers 18 to 25, where the three removal candidates (marked with a different pattern) fall below the dashed threshold line. Layers 26 and 27 show a slight recovery..

Save and reload pipeline for a structurally modified model. During the save phase, the pruned layer indices and the custom class reference are registered in config.json before serializing the weights. During the load phase, PrunedLlamaForCausalLM rebuilds the patched architecture before the weights are loaded, so attention keys with no target in the skeleton are skipped.

Benchmark scores for Llama-3.2-3B across four configurations. Degradation between 3 and 6 layers removed is negligible across all benchmarks, but Lambada shows progressive sensitivity to attention pruning at all ratios, and collapses almost completely at 15 layers removed, along with the rest of the tasks.

The four pruning strategies explored in He et al. (2024). Each column shows which modules are removed from a Transformer block (crossed out) and which are preserved. The bar at the bottom represents the SDR (Speedup Degradation Ratio). A shorter bar indicates a better trade-off between speed gain and quality loss. Attention Drop achieves the lowest SDR by a wide margin, making it the most efficient strategy despite leaving MLP layers intact.

• The dynamic memory consumption of the attention mechanism scales with context length. Quantizing the KV cache reduces this memory footprint without altering the model weights, allowing longer prompts to be processed and batch sizes to be increased in production engines like vLLM.
• Weight quantization and KV cache quantization are independent levers that act on different areas of GPU memory. Applying one does not interfere with the other, and both can be combined in the same inference pipeline.
• The effectiveness of KV cache quantization depends on both the inference engine and the GPU architecture. On GPUs with native FP8 support, vLLM doubles the available KV token capacity with no throughput penalty; on older architectures like the T4, the same operation introduces a software overhead that penalizes generation speed.
• Attention layers are significantly more redundant than MLP layers. Anchoring PyTorch hooks at input_layernorm and post_attention_layernorm gives finer-grained importance scores than block-level measurement.
• You can modify a model's execution flow without retraining it by creating a custom class that inherits from the original architecture and overrides its forward() method. This technique is reusable across any structural modification that requires bypassing or replacing internal modules.
• Saving and reloading a structurally modified model requires registering the pruned layer indices in config.json and shipping a custom class alongside the weights. Without this mechanism, from_pretrained reconstructs the original architecture and fails at the first forward pass.
• Multiple-choice benchmarks do not reveal the full cost of attention pruning. Benchmarks like Lambada and free-text generation expose a degradation that static, closed-set evaluation hides entirely.
• The optimal pruning strategy depends on the production task. The cosine similarity metric that identifies attention layers as redundant in generative models points to MLP layers as the expendable component in dense retrieval architectures, where the conclusions are precisely inverted.

FAQ