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

一水 2024-12-27 13:59:43 来源：量子位

CPU+GPU，模型KV缓存压力被缓解了。

来自CMU、华盛顿大学、Meta AI的研究人员提出MagicPIG，通过在CPU上使用LSH（局部敏感哈希）采样技术，有效克服了GPU内存容量限制的问题。

与仅使用GPU的注意力机制相比，MagicPIG在各种情况下提高了1.76~4.99倍的解码吞吐量，并在检索和推理任务中实现了更高的下游准确率，优于Quest等现有技术。

概括而言，这项研究主要贡献有两点：

1、相比于其他的稀疏注意力（Sparse Attention），MagicPIG基于采样/估计而非搜索，提升了推理质量。

2、研究把解码阶段注意力模块的计算和哈希表卸载到CPU上，探索了异构计算的可能性，并且提升了吞吐量，有望降低实际模型部署成本。

下面具体来看。

KV缓存限制了GPU高效利用

在长上下文大模型（LLM）的推理过程中，KV缓存（Key-Value Cache）成为关键瓶颈。KV缓存主要用于存储中间的注意力键和值，从而避免重复计算。

然而，其显存占用随着批量大小和序列长度的线性增长而迅速增加，这严重限制了GPU的批量处理能力，导致计算资源无法被充分利用。

以NVIDIA A100-40GB GPU为例，在处理Llama-3.1-8B模型且上下文长度为128k时，仅支持单个请求，且近一半的解码时间都消耗在访问KV缓存上，GPU利用率明显不足。

此外，推理过程中采用的一些策略，如多样性生成（Best-of-N）和长链式推理（Long Chain-of-Thoughts），会进一步增加生成的Token数量，加剧显存压力，导致推理效率进一步下降。

TopK Attention的问题

众所周知，注意力机制本质上具有稀疏性，因此动态稀疏注意力和基于TopK的近似方法得到了广泛研究。

然而，这些方法往往伴随着显著的质量下降问题。

目前已有的KV缓存压缩技术，如Quest、H2O和Loki，主要通过筛选出KV缓存中注意力得分最高的子集来提高效率。然而，尽管这些方法在实践中表现出一定的效果，基于TopK的注意力依然是一种存在偏差的近似方法，且缺乏理论上的严格保障。

这种不足限制了其在高精度场景中的广泛应用。

下图显示，即使是精确的TopK注意力机制也会导致显著的估计误差和下游任务性能下降。

这一问题在需要高上下文利用率的复杂任务中尤为突出，例如聚合任务、常用词提取（CWE）、高频词提取（FWE）以及逻辑推理任务。在这些场景中，基于TopK近似方法的性能下降尤其严重。

以下几点观察揭示了为何TopK注意力机制无法始终有效工作。

这些观察不仅解释了注意力机制的行为，还可能对模型训练具有重要意义：

1、首个输入token（注意力汇聚点，sink）的隐藏状态（包括但不限于键和值状态）几乎不随输入变化而改变。（见左图， 在采样的输入中，其最小相似度均高于0.99）

2、键状态的中心方向在不同输入句子中保持稳定。（见中图， 相似度均高于0.9）

3、键状态的中心与汇聚点token的键状态几乎相反。（见右图， -0.9至-0.8之间）

这些现象为理解注意力机制提供了新的视角，同时也表明传统的TopK近似方法在某些场景下可能存在局限性。

为了解决这一问题，研究提出了一种基于采样而非搜索TopK键值缓存的新方法。

算法：基于采样的注意力估计

与仅依赖注意力分数最高的键值对相比，融入基础分布信息可以显著提高估计的准确性。

研究将这一问题视为采样中的偏差校正问题。在生物学、社会学和机器学习等领域，无偏且高效的采样技术已被广泛研究，并具有坚实的理论保障。

如图所示，基于注意力分数按比例进行采样（即所谓的Oracle Sampling，研究把注意力模块的输出看成value向量的期望值，对应的分布是注意力得分）相比于传统的TopK选择方法，其估计误差要小得多，最多可降低4倍。

这表明采样技术在注意力近似中的潜力。

从注意力得分中采样，在实际中不可行。重要性采样（Importance Sampling）允许从一个已知分布中抽取样本1，2，…，B，来估计未知分布的期望。

最终的输出由下式给出：

重要性采样要求和的峰值对应以降低估计方差，为此，研究使用局部敏感哈希（LSH） 来生成采样概率。

需要指出的是，因为存在Softmax（注意力得分需要归一化）, 所以研究实际上试图近似的是自归一化重要性采样。

系统：将注意力计算和哈希表放在CPU上

除了精度下降的问题外，受限的GPU显存容量也限制了现有动态KV缓存压缩方法（如Quest和Loki）在许多场景中的适用性。

与此同时，像DeepSpeed-Zero-Inference和FastDecode这样的技术展示了将KV缓存和注意力计算卸载到CPU上的潜力。

CPU的内存带宽大约是GPU显存带宽的10%-20%，这引出了一个自然的问题：

能否在不牺牲精度的前提下，将注意力计算中的内存访问量减少10倍？

通过利用采样算法，例如MagicPIG中基于LSH（局部敏感哈希）的采样技术进行注意力估计，研究大幅降低了内存访问量。这种方法等效地提升了CPU的内存带宽，使得在维持精度的情况下实现高效的注意力计算。

论文的系统设计扩展了以往的工作，将大语言模型（LLM）的解码分为以下四个部分：

• 参数计算：包括所有线性投均在GPU上运行。
• 注意力计算：涉及公式，该部分在CPU上运行。
• 随机投影：在生成过程中，对于每个执行K x L次随机投影以生成哈希码。由于所有注意力头可以共享相同的随机投影器，内存开销较小（在实际实现中约为400KB）。实验中K=9或10，而L为数百，因此该步骤主要受计算限制，放置在GPU上运行。
• 检索：需要在L个哈希表中查找q的哈希码。这部分计算开销非常轻量，但预构建的哈希表占用的内存较大，因此更适合放置在CPU上运行。通过上述任务分区，可以支持更大规模的K和L哈希表，而无需担心哈希码计算和哈希表存储的开销。

实验

研究从准确率和推理速度两个方面来评估MagicPIG系统的能力。

图片中的百分比为实际采样的KV cache的数量，对于MagicPIG而言，K10L150≈2%, K10L170≈2.5%。

长文本RULER

以Llama-3.1-8B-Instruct为例，MagicPIG在检索和推理任务中比Quest（稀疏注意力的SOTA基线）实现了更高的下游准确率。

推理速度和吞吐量

在L20 + Intel 8563C上测试吞吐量，MagicPIG与仅使用GPU的注意力机制相比，在各种情况下提高了1.76~4.99倍的解码吞吐量。

整体而言，MagicPIG是将经典的哈希算法和高维向量估计用到LLM解码上的尝试。

接下来，研究将支持更加高效的局部敏感哈希算法，并希望进一步降低LLM部署成本，探索异构计算的可能性。