PaperlayerView exact revision on arXiv

Source-linked AI summary

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

William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, Jonathan Ragan Kelly

arXiv:2405.12981v1cs.LGcs.CL

TL;DR

KV-cache memory can constrain serving at long sequence lengths and large batch sizes. This paper introduces Cross-Layer Attention (CLA), which shares key/value activations across layers, and reports a 2× cache reduction with little perplexity change at 1B and 3B scales.

  • Problem

    KV-cache memory grows with sequence length and batch size, potentially limiting batch sizes for long sequences and increasing the cost of persistent-cache storage.

    Source: §1 · p. 1

  • Method

    Cross-Layer Attention reduces the number of unique layers in the KV cache by sharing key/value activations across layers.

    Source: §1 · p. 2, §2.2 · p. 2

  • Results

    2× KV cache reduction at both 1B and 3B scale, with at worst a less than 1% change in perplexity for typical 64- and 128-dimension MQA heads.

    Source: §4 · p. 10

  • Takeaways & Limitations

    CLA advances the memory-efficient transformer accuracy/memory Pareto frontier at 1B- and 3B-parameter scales.

    Source: §6 · p. 12

    The paper leaves end-to-end inference-efficiency evaluations of large, long-context CLA models for future work.

    Source: §4 · p. 11

Abstract

from arXiv · show

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

1 Introduction

KV-cache storage grows with sequence length and batch size, creating memory pressure for long-context LLMs. The paper introduces cross-layer sharing as a distinct way to reduce cache size and reports improved memory/accuracy tradeoffs.

1 Introduction

  • KV-cache size scales with sequence length and batch size, which can limit batch sizes for long sequences and make persistent-cache storage costly.The paper notes that the memory footprint is increasingly important as applications demand longer sequences.Source: §1 · p. 1
  • CLA targets a distinct cache-reduction dimension: reducing the number of unique layers represented in the KV cache.Earlier approaches include low-precision storage [Hooper et al., 2024, Zhang et al., 2024], eviction [Zhang et al., 2023, Liu et al., 2023], and sharing across query heads [Shazeer, 2019, Ainslie et al., 2023b].Source: §1 · pp. 1–2
  • At 1B and 3B scales, combining CLA with MQA achieves a 2× KV cache reduction versus plain MQA with minimal perplexity degradation.The authors also report that CLA advances accuracy/memory Pareto improvements over MQA and GQA.Source: §1 · p. 2
  • The paper evaluates CLA configurations across architectural hyperparameters, learning rates, and model sizes, and recommends sharing between consecutive layer pairs, especially with MQA.The authors report that CLA should be used between pairs of consecutive layers and delivers the most robust benefits with MQA.Source: §1 · p. 2

2 Cross-Layer Attention

Cross-Layer Attention shares key/value activations across layers, building on within-layer sharing in MQA and GQA. Its cache savings depend on the sharing factor, while its system effects extend beyond inference-cache storage.

2.1 Background: Multi-Query Attention and Grouped-Query Attention

  • MQA and GQA reduce cache storage by sharing key/value heads across query heads; MQA is the single-group case, and neither architecture is best for every use case.In GQA, groups of query heads share a key/value head, so cache size depends on the number of distinct key/value heads.Source: §2.1 · p. 2

2.2 Sharing KV Activations Across Layers

  • CLA computes key/value projections in only a subset of layers, letting other layers reuse earlier activations; only projected layers contribute to the cache.The sharing factor specifies how many layers share each projection’s output; CLA2 shares across adjacent pairs, while CLA3 shares across groups of three.Source: §2.2 · pp. 2–3, Figure 1 · p. 3, Figure 2 · p. 4

2.3 Implications for System Design

  • With other architectural hyperparameters fixed, CLA shrinks KV-cache memory by the sharing factor, or slightly less when that factor does not evenly divide the layer count.It also reduces intermediate KV activation memory during training, though those tensors are typically small relative to hidden states and MLP activations in MQA and GQA models.Source: §2.3 · p. 3
  • CLA is compatible with tensor parallelism, but pipeline-parallel stages must keep layers sharing a cache together or communicate their KV activations.CLA slightly reduces model parameters and FLOPs; it has no direct effect on core attention decoding latency because shared caches must still be reread at each attention layer.Source: §2.3 · p. 4

3 Pretraining Experiments

The pretraining experiments assess CLA’s accuracy–memory tradeoffs against MQA and GQA at 1B and 3B parameters. Across the reported configurations, CLA improves perplexity at matched cache capacity or reduces cache capacity with small perplexity changes, while results vary by scale and configuration.

3 Pretraining Experiments

  • Across the pretraining experiments, CLA offers favorable accuracy/memory tradeoffs versus plain MQA or GQA; sharing factor 2 was more effective than other tested factors, with benefits at both model scales.The authors also report preliminary evidence that CLA models benefit from higher learning rates than comparable non-CLA models.Source: §3 · p. 5

3.2 Experiments at 1B-Parameter Scale

  • All 1B-scale models were trained on 30 billion tokens with consistent data order; exploration included seven non-CLA baselines spanning two orders of magnitude in KV-cache memory requirements.The design-space exploration used LR = 3×10^-4, described as conservative; accuracy was measured by perplexity on a held-out validation set of approximately 4M tokens drawn from the SlimPajama corpus.Source: §3.2 · p. 6, §3.2.1 · p. 7

3.2.1 Design Space Exploration

  • At 1B scale, MQA-CLA2 matches the cache footprint of MQA baselines with dhead = 32, 46, or 64 while improving perplexity by 0.21–0.48 points; larger head sizes improve it by 0.03 points over MQA and GQA2 dhead = 128.The accuracy/memory frontier is shown in Figure 3.Source: Figure 3 · p. 5, §3.2.1 · p. 7
  • In the 1B design-space exploration, CLA3/CLA4 beat matched-cache MQA baselines but trailed CLA2 at equal cache footprints; alternative sharing patterns also worsened perplexity and used slightly more cache.Only GQA2-CLA2 beat its same-footprint baseline among the tested GQA-CLA2 configurations, matching MQA-CLA2 within 0.01 points.Source: §3.2.1 · pp. 7–8

3.2.2 Robustness to Learning Rate Tuning

  • With tuned learning rates, 1B MQA-CLA2 has 0.04 points worse validation perplexity than dhead = 128 MQA, but 0.31 points better than dhead = 64 MQA at equal cache memory.The optimal rates were 1.5 × 10^-3 for H128-MQA and 2.25 × 10^-3 for both H64-MQA and H128-MQA-CLA2.Source: §3.2.2 · p. 8, Table 3 · p. 9
  • At 1B scale, tuned models were within 0.01 points on Wikitext perplexity for CLA2 versus dhead = 128 MQA, while downstream results showed no consistent winner and were within 1–5 percentage points.CLA2 was 0.71 points better in Wikitext perplexity than dhead = 64 MQA.Source: §3.2.2 · pp. 8–9, Table 4 · p. 9

3.3 Experiments at 3B-Parameter Scale

  • In the first 3B experiment set, tuned MQA-CLA2 had substantially better perplexity than both MQA baselines, while dhead = 64 MQA outperformed dhead = 128 MQA despite half the cache capacity.The optimal learning rates and perplexities are reported in Table 5.Source: §3.3 · p. 9, Table 5 · p. 9
  • At 3B scale, MQA-CLA2 with dhead = 64 is 0.05 points worse in perplexity than dhead = 64 MQA, but 0.35 points better than dhead = 32 MQA at equal cache memory.The comparisons are from the second set of 3B-scale experiments.Source: §3.3 · p. 10, Table 7 · p. 10

4 Discussion & Future Work

At 1B and 3B scales, CLA2 substantially reduces KV-cache memory with little perplexity change. Higher sharing factors trade off somewhat worse than CLA2, and end-to-end inference efficiency for large, long-context CLA models remains unevaluated.

4 Discussion & Future Work

  • Across tested regimes, MQA-CLA2 consistently achieves the lowest validation perplexity within 0.01 points for a given KV-cache budget and model size.Source: §4 · p. 10
  • CLA3 and higher sharing factors have slightly worse accuracy/memory tradeoffs than CLA2 with varying head dimension, but still Pareto-dominate plain MQA.Source: §4 · p. 10
  • At 1B and 3B scales, CLA2 cuts KV-cache size by 2× with less than 1% worst-case perplexity change for 64- and 128-dimension heads, sometimes improving perplexity.The authors recommend this as a conservative change to existing MQA architectures.Source: §4 · p. 10
  • End-to-end inference efficiency evaluations of large, long-context models using CLA are left for future work; the authors suspect extremely long-sequence models may gain most.Source: §4 · p. 11

5 Related Work

Transformer memory efficiency can target storage or bandwidth, training or inference, and within- or between-pass memory; this section focuses on persistent inference KV-cache storage. Prior approaches compress or selectively retain cache entries, while architectural approaches alter attention or reduce distinct KV values.

5.1 Decreasing KV cache size Post training

  • Post-training KV-cache compression includes one-to-two-bit quantization, selective token retention during generation, and direct tensor encoding.Examples include KVQuant and Coupled Quantization; H2O, Scissorhands, and FastGen; and Cachegen.Source: §5.1 · p. 11

5.2 Architectural Changes that decrease KV cache size

  • Architectural approaches reduce KV-cache size by attending to fewer tokens, replacing softmax attention, or reducing unique KV values relative to queries.Source: §5.2 · p. 11
  • Local-window methods such as Transformer XL [Dai et al., 2019] and Sparse Attention [Child et al., 2019] reduce effective sequence length; Infini attention compresses prior tokens using linear attention.Source: §5.2 · p. 11
  • Lookup over prior tokens or external datastores can reduce computation, but does not reduce storage unless lookup keys and values are offloaded from GPU memory.Examples include Landmark attention, Memorizing transformers, and retrieval [Guu et al., 2020, Izacard et al., 2024, Borgeaud et al., 2022].Source: §5.2 · p. 11
  • Linear-attention alternatives to softmax can use constant space during token generation and make generation time linear in generated-token count.Examples include GLA, Mamba, and RWKV v6.Source: §5.2 · p. 12
  • MQA and GQA [Ainslie et al., 2023a] share KV pairs across query heads; concurrent methods instead use low-rank KV projections or split local and global attention across model halves.The latter approaches are Multi Latent Attention in Deepseek-V2 and You Only Cache Once.Source: §5.2 · p. 12

6 Conclusion

Across evaluated 1B- and 3B-parameter baselines, Cross-Layer Attention advances the memory-efficiency Pareto frontier, reducing KV-cache storage by 2× with roughly equal perplexity.

6 Conclusion

  • Across evaluated 1B- and 3B-parameter baselines, Cross-Layer Attention advances the memory-efficiency Pareto frontier, reducing KV-cache storage by 2× with roughly equal perplexity.Source: §6 · p. 12

A Learning rate sweeps

The section presents learning-rate sweep results at both the 1B- and 3B-parameter scales.

A Learning rate sweeps

Loading 2405.12981v1…