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

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

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

附上论文 / 技俩主页聚拢，以及关连方式哦

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

点这里� � 热心我，紧记标星哦～

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

科技前沿进展日日再会 ~