CHAI: Clustered Head Attention for Efficient LLM Inference

TL;DR

CHAI tackles inference bottlenecks in large decoder-only transformers by exploiting cross-head redundancy in multi-head attention. It performs offline per-model clustering to determine the number of head clusters per layer and uses online, context-aware membership after the first five tokens to assign heads to clusters, computing attention only for representative heads within each cluster. This approach reduces both compute and KV-cache size while preserving accuracy within a few percent across multiple models and tasks, yielding up to $1.73\times$ latency reduction and $21.4\%$ KV-cache savings without fine-tuning. The method is broadly applicable, does not require retraining, and leverages existing hardware kernels, making it a practical improvement for scalable LLM inference.

Abstract

Large Language Models (LLMs) with hundreds of billions of parameters have transformed the field of machine learning. However, serving these models at inference time is both compute and memory intensive, where a single request can require multiple GPUs and tens of Gigabytes of memory. Multi-Head Attention is one of the key components of LLMs, which can account for over 50% of LLMs memory and compute requirement. We observe that there is a high amount of redundancy across heads on which tokens they pay attention to. Based on this insight, we propose Clustered Head Attention (CHAI). CHAI combines heads with a high amount of correlation for self-attention at runtime, thus reducing both memory and compute. In our experiments, we show that CHAI is able to reduce the memory requirements for storing K,V cache by up to 21.4% and inference time latency by up to 1.73x without any fine-tuning required. CHAI achieves this with a maximum 3.2% deviation in accuracy across 3 different models (i.e. OPT-66B, LLAMA-7B, LLAMA-33B) and 5 different evaluation datasets.

Figures (23)

• Figure 1: Accuracy vs Flops: We study various methods of clustering attention heads, and plot the runtime for sequence length of 2048. For random head selection we randomly choose heads to combine in increasing number of 4, 8, 16 and 24. For Static Head Selection, we choose the heads to combine based on activations. CHAI is our proposed method.
• Figure 2: Redundancy across heads for LLaMa--7B
• Figure 3: Clustered Head Attention: Schematic of clustered head attention, comparing it with popular Multi-Head Attention. In clustered head attention, we remove the query and key vectors which produce similar attention scores.
• Figure 4: Activations for OPT-66B and LLaMa-7B for an exemplary sentence: We observe that OPT-66B has several heads which give uniform attention scores to tokens whereas LLaMa-7B does not. However, both models have redundancies across heads, i.e., groups of heads are give similar attention to each token.
• Figure 5: CHAI Flow: In the offline phase, we run clustering and perform elbow plot analysis for each new model. Then, for each new inference request we only perform cluster membership identification based on online performance.