Reducing Transformer Key-Value Cache Size with Cross-Layer Attention

TL;DR

This paper introduces Cross-Layer Attention (CLA), a novel method that further reduces transformer KV cache size by sharing key-value heads across layers, enabling longer sequences and larger batch sizes with minimal accuracy loss.

Contribution

The paper proposes Cross-Layer Attention (CLA), a new attention mechanism that shares key-value heads between layers, significantly reducing memory usage in transformer models.

Findings

CLA reduces KV cache size by 2x compared to MQA.

CLA maintains nearly the same accuracy as unmodified MQA.

Experiments show improved memory/accuracy tradeoffs for large language models.

Abstract

Key-value (KV) caching plays an essential role in accelerating decoding for transformer-based autoregressive large language models (LLMs). However, the amount of memory required to store the KV cache can become prohibitive at long sequence lengths and large batch sizes. Since the invention of the transformer, two of the most effective interventions discovered for reducing the size of the KV cache have been Multi-Query Attention (MQA) and its generalization, Grouped-Query Attention (GQA). MQA and GQA both modify the design of the attention block so that multiple query heads can share a single key/value head, reducing the number of distinct key/value heads by a large factor while only minimally degrading accuracy. In this paper, we show that it is possible to take Multi-Query Attention a step further by also sharing key and value heads between adjacent layers, yielding a new attention design we call Cross-Layer Attention (CLA). With CLA, we find that it is possible to reduce the size of the KV cache by another 2x while maintaining nearly the same accuracy as unmodified MQA. In experiments training 1B- and 3B-parameter models from scratch, we demonstrate that CLA provides a Pareto improvement over the memory/accuracy tradeoffs which are possible with traditional MQA, enabling inference with longer sequence lengths and larger batch sizes than would otherwise be possible