mlsanityData to 7 Oct 2026

CHAI: Clustered Head Attention for Efficient LLM Inference

arXiv 2403.08058 ↗First posted 12 Mar 2024cs.LG cs.CLcited by 13

On the map: Efficient LLM inference and serving → KV cache optimization

Abstract (arXiv)

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.

Key facts. Cited by 13 papers in the map (data through 7 Oct 2026). Placed in KV cache optimization, one of 12 directions in Efficient LLM inference and serving. Builds on 51 of its 63 references that are in the map. First posted on arXiv 12 Mar 2024.

Builds on

51 of 63 references.

Closest ancestors · 1

Same direction, newest first

  1. 2023Scissorhands: Exploiting the Persistence of Importance Hypothesis for LLM KV Cache Compression at Test Time26 May

Other references · 50

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2019Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformercited by 14,491
  4. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  5. 2023Llama 2: Open Foundation and Fine-Tuned Chat Modelscited by 12,865
  6. 2020Scaling Laws for Neural Language Modelscited by 6,411
  7. 2018Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challengecited by 4,710
  8. 2019DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lightercited by 4,243
  9. 2019HellaSwag: Can a Machine Really Finish Your Sentence?cited by 4,018
  10. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,611
  11. 2022OPT: Open Pre-trained Transformer Language Modelscited by 3,464
  12. 2022Training Compute-Optimal Large Language Modelscited by 2,927
38 more
  1. 2019PIQA: Reasoning about Physical Commonsense in Natural Languagecited by 2,818
  2. 2023FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioningcited by 2,551
  3. 2018The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networkscited by 2,250
  4. 2019BoolQ: Exploring the Surprising Difficulty of Natural Yes/No Questionscited by 2,140
  5. 2022GPTQ: Accurate Post-Training Quantization for Generative Pre-trained Transformerscited by 1,838
  6. 2020Reformer: The Efficient Transformercited by 1,658
  7. 2022Fast Inference from Transformers via Speculative Decodingcited by 1,615
  8. 2022SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Modelscited by 1,338
  9. 2022LLM.int8(): 8-bit Matrix Multiplication for Transformers at Scalecited by 881
  10. 2019Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Prunedcited by 844
  11. 2019Are Sixteen Heads Really Better than One?cited by 803
  12. 2023Distilling Step-by-Step! Outperforming Larger Language Models with Less Training Data and Smaller Model Sizescited by 710
  13. 2019Fast Transformer Decoding: One Write-Head is All You Needcited by 700
  14. 2019Reducing Transformer Depth on Demand with Structured Dropoutcited by 434
  15. 2019Q-BERT: Hessian Based Ultra Low Precision Quantization of BERTcited by 355
  16. 2020The Lottery Ticket Hypothesis for Pre-trained BERT Networkscited by 275
  17. 2021I-BERT: Integer-only BERT Quantizationcited by 250
  18. 2022The case for 4-bit precision: k-bit Inference Scaling Lawscited by 247
  19. 2023Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Timecited by 247
  20. 2020SpAtten: Efficient Sparse Attention Architecture with Cascade Token and Head Pruningcited by 228
  21. 2019Structured Pruning of Large Language Modelscited by 201
  22. 2021Block Pruning For Faster Transformerscited by 176
  23. 2022Speculative Decoding: Exploiting Speculative Execution for Accelerating Seq2seq Generationcited by 169
  24. 2019Drawing Early-Bird Tickets: Towards More Efficient Training of Deep Networkscited by 160
  25. 2020The Hardware Lotterycited by 160
  26. 2020On the Effect of Dropping Layers of Pre-trained Transformer Modelscited by 118
  27. 2022The Optimal BERT Surgeon: Scalable and Accurate Second-Order Pruning for Large Language Modelscited by 115
  28. 2020When BERT Plays the Lottery, All Tickets Are Winningcited by 114
  29. 2020A3: Accelerating Attention Mechanisms in Neural Networks with Approximationcited by 97
  30. 2020EarlyBERT: Efficient BERT Training via Early-bird Lottery Ticketscited by 79
  31. 2020EdgeBERT: Sentence-Level Energy Optimizations for Latency-Aware Multi-Task NLP Inferencecited by 69
  32. 2021Pufferfish: Communication-efficient Models At No Extra Costcited by 61
  33. 20158-Bit Approximations for Parallelism in Deep Learningcited by 56
  34. 2023Predictive Pipelined Decoding: A Compute-Latency Trade-off for Exact LLM Decodingcited by 40
  35. 2023Brainformers: Trading Simplicity for Efficiencycited by 31
  36. 2022An Algorithm-Hardware Co-Optimized Framework for Accelerating N:M Sparse Transformerscited by 28
  37. 2023Cuttlefish: Low-Rank Model Training without All the Tuningcited by 25
  38. 2020Principal Component Networks: Parameter Reduction Early in Trainingcited by 9

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 6; 2025: 4; 2026: 3.

13 papers.

Sort:

By area:

  1. 2024A Survey on Large Language Model Acceleration based on KV Cache Managementcited by 125
  2. 2024SampleAttention: Near-Lossless Acceleration of Long Context LLM Inference with Adaptive Structured Sparse Attentioncited by 39
  3. 2024UncertaintyRAG: Span-Level Uncertainty Enhanced Long-Context Modeling for Retrieval-Augmented Generationcited by 14
  4. 2024UNComp: Can Matrix Entropy Uncover Sparsity? -- A Compressor Design from an Uncertainty-Aware Perspectivecited by 14
  5. 2025Scaling Inference-Efficient Language Modelscited by 14
  6. 2025FSA: An Alternative Efficient Implementation of Native Sparse Attention Kernelcited by 7
  7. 2025SPD: Sync-Point Drop for Efficient Tensor Parallelism of Large Language Modelscited by 6
  8. 2024Cross-layer Attention Sharing for Pre-trained Large Language Modelscited by 5
  9. 2025MILLION: Mastering Long-Context LLM Inference Via Outlier-Immunized KV Product Quantizationcited by 3
  10. 2026Towards Efficient Large Language Model Serving: A Survey on System-Aware KV Cache Optimizationcited by 2
  11. 2024RWKV-edge: Deeply Compressed RWKV for Resource-Constrained Devicescited by 1
  12. 2026Linear Predictability of Attention Heads in Large Language Modelscited by 1
1 more
  1. 2026Towards Efficient Large Vision-Language Models: A Comprehensive Survey on Inference Strategiescited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2024Beyond KV Caching: Shared Attention for Efficient LLMs
  2. 2025HeadInfer: Memory-Efficient LLM Inference by Head-wise Offloading
  3. 2024DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads
  4. 2025Progressive Sparse Attention: Algorithm and System Co-design for Efficient Attention in LLM Serving
  5. 2025LLMs Know What to Drop: Self-Attention Guided KV Cache Eviction for Efficient Long-Context Inference
  6. 2026LycheeDecode: Accelerating Long-Context LLM Inference via Hybrid-Head Sparse Decoding
  7. 2025HGCA: Hybrid GPU-CPU Attention for Long Context LLM Inference
  8. 2024NoMAD-Attention: Efficient LLM Inference on CPUs Through Multiply-add-free Attention