MLKV: Multi-Layer Key-Value Heads for Memory Efficient Transformer
  Decoding
Authors: name: Zayd Muhammad Kawakibi Zuhri; name: Muhammad Farid Adilazuarda; name: Ayu Purwarianti; name: Alham Fikri Aji
Published: 2024-06-13 16:33:44

Abstract: Auto-regressive inference of transformers benefit greatly from Key-Value (KV)
caching, but can lead to major memory bottlenecks as model size, batch size,
and sequence length grow at scale. We introduce Multi-Layer Key-Value (MLKV)
sharing, a novel approach extending KV sharing across transformer layers to
reduce memory usage beyond what was possible with Multi-Query Attention (MQA)
and Grouped-Query Attention (GQA). Evaluations on various NLP benchmarks and
inference metrics using uptrained Pythia-160M variants demonstrate that MLKV
significantly reduces memory usage with minimal performance loss, reducing KV
cache size down to a factor of 6x compared to MQA. These results highlight
MLKV's potential for efficient deployment of transformer models at scale. We
provide code at https://github.com/zaydzuhri/pythia-mlkv