Research map / Large language models / Reasoning and chain of thought

CHAI: Clustered Head Attention for Efficient LLM Inference

ICML 2024 · Large language models › Reasoning and chain of thought · approach: Efficient LLM serving · 2 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes CHAI, a training-free runtime method that clusters attention heads with correlated attention patterns (clusters decided from the first five tokens) and computes self-attention only for one representative per cluster, reducing compute and K,V cache for LLM inference.

Problem

Multi-head attention in LLMs is compute- and memory-intensive at inference time (K,V cache growth and repeated attention ops), and existing runtime pruning (DejaVu) and architectural co-design (GQA/MQA) either fail on parameter-efficient models like LLaMA or require retraining.

Approach

Observe that many heads in a layer produce highly correlated attention scores; offline pick the number of clusters per layer via elbow analysis on C4, then per request cluster heads by their attention activations from the first five tokens, compute attention only for one head per cluster, and drop the redundant heads' keys (values kept) from the K,V cache.

Key results

• CHAI reduces K,V cache memory by up to 21.4% and time-to-next-token latency by over 5x at sequence length 2048 on LLaMA-7B; time-to-first-token speedup up to 1.72x.
• On LLaMA-7B, CHAI keeps accuracy within a maximum 3.7% drop from MHA across five tasks, while DejaVu at 30% sparsity drops up to 18.6%.
• On LLaMA-33B, CHAI shows a maximum accuracy degradation of 0.14% versus MHA, and matches MHA on OPT-66B within about 0.15 points on PIQA.

Measured effects reported in the paper

Surprising findings

• DejaVu's contextual sparsity premise (uniform-attention heads) does not hold for LLaMA-7B, which has no such heads, unlike OPT-66B.
• Pruning Q, K and V together (reusing the full head) causes large accuracy losses on LLaMA (e.g., PIQA 77.8 to 61.93 for the QKV variant), so keeping values matters.

Limitations

• Evaluation limited to three models and five multiple-choice tasks; no generation-quality benchmarks.
• Compared to DejaVu only on accuracy; runtime comparison not possible since DejaVu kernels are unreleased.
• Speedup numbers reported on V100 GPUs with inference-time clustering overhead; claims of framework independence not fully tested.

Experimental setup

Datasets: C4

Benchmarks: PIQA HellaSwag ARC-Challenge ARC-Easy BoolQ

Metrics: accuracy accuracy deviation from MHA time to first token latency time to next token latency K,V cache memory reduction

Baselines: MHA DejaVu SpAtten CHAI-static

Builds on: DejaVu Grouped-Query Attention SpAtten xFormers

Closest papers by meaning

• DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming Heads (ICLR 2025 · similarity 0.90)
• SqueezeAttention: 2D Management of KV-Cache in LLM Inference via Layer-wise Optimal Budget (ICLR 2025 · similarity 0.88)
• CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration (ICML 2025 · similarity 0.88)
• SparQ Attention: Bandwidth-Efficient LLM Inference (ICML 2024 · similarity 0.88)
• QUEST: Query-Aware Sparsity for Efficient Long-Context LLM Inference (ICML 2024 · similarity 0.88)
• Scaling Laws Meet Model Architecture: Toward Inference-Efficient LLMs (ICLR 2026 · similarity 0.88)