8 Attention optimization
This chapter covers
- Quantizing the KV cache for longer contexts
- Scoring attention layers with cosine similarity
- Removing attention modules from the architecture
- Modifying model execution flow with custom classes
- Persisting and reloading modified models
Throughout the book we've worked on modifying model weights, either by removing them through different forms of pruning or by modifying them through Knowledge Distillation or fine-tuning.
In this chapter, we're going to focus on the attention mechanism, the main bottleneck in LLMs during inference. Both its computational demand and its dynamic memory consumption grow as context length increases, limiting the amount of data and the maximum batch size the model can process.
We'll explore two ways to optimize the attention mechanism. The first, purely practical, is KV cache quantization. We'll look at how to apply it with Hugging Face and with vLLM, and what the results tell us about when it makes sense to use each library.
The second part has a much more direct connection to recent research. In it, we'll analyze the redundancy of Attention Modules and apply tensor similarity-based techniques, as you already saw in the depth pruning chapter, to remove those attention layers that don't contribute value to the output.