今天小编分享的科学经验：把注意力计算丢给CPU，大模型解码吞吐量提高1.76~4.99倍，欢迎阅读。
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 部署成本，探索异构计算的可能性。
论文：
https://arxiv.org/abs/2410.16179
项目地址：
www.lsh-ai.com
— 完 —
投稿请发邮件到：
标题注明【投稿】，告诉我们：
你是谁，从哪来，投稿内容
附上论文 / 项目主页链接，以及联系方式哦
我们会（尽量）及时回复你
点这里关注我，记得标星哦～
一键三连「分享」、「点赞」和「在看」
科技前沿进展日日相见 ~
>
