CHAI: Clustered Head Attention for Efficient LLM Inference

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Agarwal, Saurabh, Acun, Bilge, Hosmer, Basil, Elhoushi, Mostafa, Lee, Yejin, Venkataraman, Shivaram, Papailiopoulos, Dimitris, Wu, Carole-Jean
Format: Preprint
Publié: 2024
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866911856621256704
author Agarwal, Saurabh
Acun, Bilge
Hosmer, Basil
Elhoushi, Mostafa
Lee, Yejin
Venkataraman, Shivaram
Papailiopoulos, Dimitris
Wu, Carole-Jean
author_facet Agarwal, Saurabh
Acun, Bilge
Hosmer, Basil
Elhoushi, Mostafa
Lee, Yejin
Venkataraman, Shivaram
Papailiopoulos, Dimitris
Wu, Carole-Jean
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2403_08058
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle CHAI: Clustered Head Attention for Efficient LLM Inference
Agarwal, Saurabh
Acun, Bilge
Hosmer, Basil
Elhoushi, Mostafa
Lee, Yejin
Venkataraman, Shivaram
Papailiopoulos, Dimitris
Wu, Carole-Jean
Machine Learning
Computation and Language
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.
title CHAI: Clustered Head Attention for Efficient LLM Inference
topic Machine Learning
Computation and Language
url https://arxiv.org/abs/2403.08058