解决多卡推理中注意力头内存分配不均导致的性能瓶颈
FairKV: Balancing Per-Head KV Cache for Fast Multi-GPU Inference
- 通过公平复制机制在多卡间均衡内存负载
- 在LLaMA 70b和Mistral 24b上提升1.66倍吞吐量
- 适合大规模模型多卡推理场景使用
Transformer模型中的键值缓存（KV cache）技术通过减少重复计算来加速推理，但显著增加了内存占用，因此KV缓存压缩成为重要研究方向。现有先进压缩方法采用非均衡的逐头分配策略，在单卡场景下表现优异。然而，我们在多卡部署中发现，这种非均衡压缩会导致严重的负载不均：部分GPU过载而其他GPU闲置。本文提出FairKV，旨在多卡系统中实现注意力头间公平的内存使用。其核心技术为公平复制（Fair-Copying），利用数据并行将少数内存密集型注意力头的副本分发至各GPU，缓解负载失衡。在主流模型如LLaMA 70b和Mistral 24b上的实验表明，与标准张量并行推理相比，FairKV将吞吐量提升1.66倍。代码将在论文接受后开源。
原文摘要 · Abstract (English)
KV cache techniques in Transformer models aim to reduce redundant computations at the expense of substantially increased memory usage, making KV cache compression an important and popular research topic. Recently, state-of-the-art KV cache compression methods implement imbalanced, per-head allocation algorithms that dynamically adjust the KV cache budget for each attention head, achieving excellent performance in single-GPU scenarios. However, we observe that such imbalanced compression leads to significant load imbalance when deploying multi-GPU inference, as some GPUs become overburdened while others remain underutilized. In this paper, we propose FairKV, a method designed to ensure fair memory usage among attention heads in systems employing imbalanced KV cache compression. The core technique of FairKV is Fair-Copying, which replicates a small subset of memory-intensive attention heads across GPUs using data parallelism to mitigate load imbalance. Our experiments on popular models, including LLaMA 70b and Mistral 24b model, demonstrate that FairKV increases throughput by 1.66x compared to standard tensor parallelism inference. Our code will be released as open source upon acceptance.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。