把注意力计算丢给CPU，大模型解码吞吐量提高1.76~4.99倍

基于局部敏感哈希采样技术

原标题：把注意力计算丢给CPU，大模型解码吞吐量提高1.76~4.99倍
文章来源：量子位
内容字数：5268字

来自卡内基梅隆大学、华盛顿大学和Meta AI的研究人员提出了一种名为MagicPIG的新方法，有效解决了大语言模型(LLM)推理过程中GPU内存容量限制的问题，该问题主要由KV缓存的巨大内存占用引起。

1. 在长上下文LLM推理中，KV缓存用于存储中间注意力键和值，避免重复计算，但其内存占用随批量大小和序列长度线性增长，严重限制了GPU的批量处理能力。

2. 现有技术如TopK Attention、Quest、H2O和Loki等，试图通过选择注意力得分最高的子集来压缩KV缓存，但这些方法存在精度损失的问题，尤其在需要高上下文利用率的复杂任务中，性能下降严重。

3. 研究人员观察到注意力机制的三个关键特性，解释了TopK方法的局限性：首个输入token的隐藏状态几乎不随输入变化；键状态的中心方向在不同输入句子中保持稳定；键状态的中心与汇聚点token的键状态几乎相反。

1. **基于采样而非搜索:** MagicPIG采用基于采样的注意力估计方法，而非像TopK那样搜索最高的键值对。这通过引入基础分布信息，显著提高了估计精度，并降低了估计误差。

2. **局部敏感哈希(LSH)技术:** MagicPIG利用LSH生成采样概率，并结合重要性采样技术，有效地从已知分布中抽取样本，来估计未知分布的期望，从而降低估计方差。

3. **异构计算:** MagicPIG将注意力计算和哈希表卸载到CPU上，充分利用CPU的内存资源，降低了对GPU显存的依赖。实验表明，这种方法等效地提升了CPU的内存带宽，在维持精度的情况下实现了高效的注意力计算。

1. MagicPIG将LLM解码分为四个部分：参数计算(GPU)、注意力计算(CPU)、随机投影(GPU)和检索(CPU)。这种任务分区有效地利用了CPU和GPU的优势。

2. 实验结果表明，MagicPIG在Llama-3.1-8B-Instruct模型上，相比于Quest，在检索和推理任务中实现了更高的下游准确率。同时，MagicPIG的解码吞吐量提高了1.76~4.99倍。

MagicPIG通过巧妙地结合LSH采样技术和异构计算，有效缓解了LLM推理过程中的KV缓存压力，在保证精度的同时大幅提升了推理速度和吞吐量。这项研究为降低LLM部署成本，探索异构计算的可能性提供了新的思路。

文章来源：量子位
作者微信：
作者简介：追踪人工智能新趋势，关注科技行业新突破

文章版权归作者所有，未经允许请勿转载。

暂无评论...