CPU+GPU，模型 KV 缓存压力被缓解了。来自 CMU、华盛顿大学、Meta AI 的研究人员提出 MagicPIG，通过在 CPU 上使用 LSH（局部敏感哈希）采样技术，有效克服了 GPU 内存容量限制的问题。
MagicPIG: LSH Sampling for Efficient LLM Generation论文链接：
https://arxiv.org/abs/2410.16179代码链接：
https://github.com/Infini-AI-Lab/MagicPIG与仅使用 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（注意力得分需要归一化）, 所以研究实际上试图近似的是自归一化重要性采样。
除了精度下降的问题外，受限的 GPU 显存容量也限制了现有动态 KV 缓存压缩方法（如 Quest 和 Loki）在许多场景中的适用性。与此同时，像 DeepSpeed-Zero-Inference 和 FastDecode 这样的技术展示了将 KV 缓存和注意力计算卸载到 CPU 上的潜力。CPU 的内存带宽大约是 GPU 显存带宽的 10%-20%，这引出了一个自然的问题：能否在不牺牲精度的前提下，将注意力计算中的内存访问量减少 10 倍？通过利用采样算法，例如 MagicPIG 中基于 LSH（局部敏感哈希）的采样技术进行注意力估计，研究大幅降低了内存访问量。这种方法等效地提升了 CPU 的内存带宽，使得在维持精度的情况下实现高效的注意力计算。论文的系统设计扩展了以往的工作，将大语言模型（LLM）的解码分为以下四个部分：注意力计算：涉及公式 ，该部分在 CPU 上运行。随机投影：在生成过程中，对于每个 𝑞 执行 K x L 次随机投影以生成哈希码。由于所有注意力头可以共享相同的随机投影器，内存开销较小（在实际实现中约为 400KB）。实验中 K=9 或 10，而 L 为数百，因此该步骤主要受计算限制，放置在 GPU 上运行。 检索：需要在 L 个哈希表中查找 q 的哈希码。这部分计算开销非常轻量，但预构建的哈希表占用的内存较大，因此更适合放置在 CPU 上运行。通过上述任务分区，可以支持更大规模的 K 和 L 哈希表，而无需担心哈希码计算和哈希表存储的开销。
研究从准确率和推理速度两个方面来评估 MagicPIG 系统的能力。图片中的百分比为实际采样的 KV cache 的数量，对于 MagicPIG 而言，K10L150≈2%, K10L170≈2.5%。以 Llama-3.1-8B-Instruct 为例，MagicPIG 在检索和推理任务中比 Quest（稀疏注意力的 SOTA 基线）实现了更高的下游准确率。
在 L20 + Intel 8563C 上测试吞吐量，MagicPIG 与仅使用 GPU 的注意力机制相比，在各种情况下提高了 1.76~4.99 倍的解码吞吐量。
整体而言，MagicPIG 是将经典的哈希算法和高维向量估计用到 LLM 解码上的尝试。接下来，研究将支持更加高效的局部敏感哈希算法，并希望进一步降低 LLM 部署成本，探索异构计算的可能性。
如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。
总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。
PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学术热点剖析、科研心得或竞赛经验讲解等。我们的目的只有一个，让知识真正流动起来。
📝 稿件基本要求：
• 文章确系个人原创作品，未曾在公开渠道发表，如为其他平台已发表或待发表的文章，请明确标注
• 稿件建议以 markdown 格式撰写，文中配图以附件形式发送，要求图片清晰，无版权问题
• PaperWeekly 尊重原作者署名权，并将为每篇被采纳的原创首发稿件，提供业内具有竞争力稿酬，具体依据文章阅读量和文章质量阶梯制结算
📬 投稿通道：
• 投稿邮箱：[email protected]
• 来稿请备注即时联系方式（微信），以便我们在稿件选用的第一时间联系作者
• 您也可以直接添加小编微信（pwbot02）快速投稿，备注：姓名-投稿
△长按添加PaperWeekly小编
🔍
现在，在「知乎」也能找到我们了
进入知乎首页搜索「PaperWeekly」
