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

—  完  —

投稿请发邮件到：

ai@qbitai.com

标题注明【投稿】，告诉咱们：

你是谁，从哪来，投稿内容‍

附上论文 / 面目主页不时，以及有关步地哦

咱们会（尽量）实时酬金你

点这里� � 关爱我，紧记标星哦～

一键三连「分享」、「点赞」和「在看」

科技前沿进展日日重逢 ~