MagicPIG: LSH Sampling for Efficient LLM Generation

发表时间: 2025-04 · arXiv:2410.16179 (ICLR 2025)

原文: https://arxiv.org/abs/2410.16179

速读

一句话结论 本文提出了基于局部敏感哈希（LSH）采样的异构推理系统 MagicPIG，通过将大模型的 KV Cache 和注意力计算卸载到 CPU，并用无偏采样替代传统的 TopK 稀疏注意力，在极低计算开销下保持了长上下文任务的高精度，并将解码吞吐量提升了最高 5 倍。

要解决什么问题 长上下文大模型在自回归生成时，KV Cache 会随批次大小和序列长度线性增长，导致严重的显存瓶颈和访存受限（Memory-bound），使得 GPU 算力利用率极低。为了缓解这个问题，现有方法通常采用动态稀疏或基于 TopK 的注意力近似机制，只计算得分最高的少数 Token。但这卡在了三个致命机制上：首先，注意力分布并不总是高度稀疏的，尤其是在需要综合全局信息的聚合任务（如词汇提取）中，注意力呈现长尾分布，强制截断的 TopK 会引入严重偏差，导致模型能力断崖式掉点；其次，看似极高的稀疏度往往是由“注意力下沉（Attention Sink）”现象制造的假象，除了初始 Token 外，其余 Token 的注意力分布相对均匀，且 Query 和 Key 在几何空间上分布于方向几乎相反的狭窄锥体中，这导致精确搜索 TopK 的开销极大；最后，现有的稀疏方法大多只能减少计算量，却无法真正缩减 KV Cache 占用的总显存，在显存紧张时依然无法扩大上下文长度或批次大小。

怎么做的 核心思路是放弃有偏的 TopK 截断，将注意力近似转化为一个带有理论保证的无偏采样估计问题。作者发现，按注意力分数的真实分布进行采样，能以极小的计算预算大幅降低估计误差。为了在未知完整注意力分数的前提下实现高效采样，MagicPIG 采用了自归一化重要性采样（Self-normalized Importance Sampling），并利用局部敏感哈希（LSH）来近似采样概率。具体而言，方法包含三个关键设计。第一是基于 LSH 的采样估计：利用 SimHash 算法，通过 $K$ 个随机投影哈希函数和 $L$ 个哈希表，让与 Query 余弦相似度高的 Key 有更大概率发生哈希碰撞并被采样。采样概率 $u_i$ 与碰撞机制定义为：

$$u_i = 1 - (1 - p_i^K)^L - L p_i^K (1 - p_i^K)^{L-1}$$


其中 $p_i = 1 - \frac{1}{\pi} \arccos \frac{q k_i^T}{|q| \cdot |k_i|}$。最终的注意力输出通过采样集合 $S$ 进行重要性加权估计：

$$\bar{o} = \sum_{i \in S} \frac{\exp(\frac{q k_i^T}{\sqrt{d}} - \log u_i)}{\sum_{i \in S} \exp(\frac{q k_i^T}{\sqrt{d}} - \log u_i)} v_i$$
第二是数据预处理（Centering）：由于除 Sink Token 外的 Key 几乎都集中在 Query 的对立面，直接哈希会导致采样概率趋同。MagicPIG 在构建哈希表前对 Key 进行了中心化平移操作，使其在空间中分布更均匀，从而激活了 LSH 的区分能力。第三是 CPU-GPU 异构系统协同设计：为了彻底打破 GPU 显存限制，MagicPIG 将计算密集型的线性投影（如 MLP、QKV 映射）和 LSH 的随机投影计算保留在 GPU 上；而将极度消耗内存的 KV Cache 和庞大的 LSH 哈希表存储在 CPU DRAM 中，并在 CPU 上完成轻量级的哈希查询与注意力加权计算。此外，最关键的 Sink Token 和局部 Token 会被保留在 GPU 显存中直接计算，以进一步降低 CPU 负载。

效果如何 实验覆盖了 Llama-2-7b、Llama-3.1-8B/70B 以及 CodeLlama 等模型，在 A100、L20 和单张 RTX 4090 上进行了测试。评估任务包括 lm-eval-harness 的中等长度任务、LongBench 和 RULER（上下文最高达 256K）的长文本任务。对比的基线方法主要有两个：代表性能上限的 GPU 完整注意力机制（Full Attention），以及代表动态稀疏路线的强基线方法 Quest。量化结果显示，MagicPIG 在仅使用完整注意力 2% 到 5% 计算开销的情况下，在各类任务上的精度损失均不到 2%。在信息聚合类任务上，它的表现甚至超越了精确的 TopK 注意力。在系统效率方面，由于成功将 KV Cache 卸载到 CPU，MagicPIG 能够支持比基线大 12 倍以上的批次大小，从而在不同硬件上将解码吞吐量提升了 1.5 倍到 5 倍。在单张 24GB 显存的 RTX 4090 上，运行 96K 上下文的 Llama-3.1-8B 模型时，单请求解码延迟仅为 54 毫秒。作者也承认了该方法的局限性：它强依赖于 CPU 拥有充足的 DRAM 来存储卸载的 KV Cache 和庞大的哈希表，不适用于内存受限的场景；此外，目前 MagicPIG 仅针对解码阶段进行了优化，尚未在预填充阶段实现该机制。

Zhuoming Chen, Ranajoy Sadhukhan, Zihao Ye, Yang Zhou, Jianyu Zhang, Niklas Nolte, Yuandong Tian, Matthijs Douze, Leon Bottou, Zhihao Jia, Beidi Chen / Carnegie Mellon University, University of Washington, New York University, Meta AI

主要贡献

具有长上下文窗口的大型语言模型（LLMs）受到了广泛关注，但在自回归生成过程中，为了避免重复计算而存储的键值（KV）缓存成为了一个独特的性能瓶颈。KV缓存随着批处理大小和序列长度呈线性增长，占据了大量GPU内存，导致LLM生成极度受限于内存带宽，从而使得GPU计算能力利用率低下。为了利用“注意力具有稀疏性”这一普遍共识，以往的研究提出了各种动态稀疏或基于TopK的注意力近似方法。

本文的核心研究目标是提出一种理想的稀疏注意力近似方法，该方法需在提供理论保证的同时，在多样化的下游任务中保持全精度、具有低成本的KV缓存选择开销，并能有效节省GPU内存。

本文的主要创新点如下：

  1. 揭示TopK注意力的局限性：首次证明了TopK注意力本身在某些需要聚合全上下文的下游任务中会导致严重的质量下降，因为注意力并不总是如预期般稀疏。
  2. 提出基于采样的注意力估计：证明了相较于选择具有最高注意力分数的键和值，具有理论保证的采样方法能够为注意力输出提供更好的无偏估计。
  3. 设计MagicPIG系统：提出了一种基于局部敏感哈希（Locality Sensitive Hashing, LSH）的异构系统——MagicPIG。该系统将LSH哈希表存储在CPU上，并在CPU上运行注意力计算，从而大幅减少了注意力计算的工作量。
  4. 突破硬件限制：通过CPU-GPU协同设计，MagicPIG能够支持更长的上下文和更大的批处理大小。在单张RTX 4090上，针对96k tokens上下文的Llama-3.1-8B-Instruct模型，MagicPIG可实现54ms的解码延迟，并在各种GPU硬件上将解码吞吐量提升高达5倍。
Figure 1 尽管TopK注意力在检索任务中表现良好，但在聚合任务（如单词提取）中严重退化。X轴为TopK注意力使用的注意力键的比例
Figure 1 尽管TopK注意力在检索任务中表现良好，但在聚合任务（如单词提取）中严重退化。X轴为TopK注意力使用的注意力键的比例

背景知识与关键观察

注意力估计问题定义。在LLM解码阶段，自注意力机制通过 $o = \mathrm{Softmax}(\frac{qK^T}{\sqrt{d}})V = wV$ 计算先前值的加权平均。其中 $d$ 为头部维度，$n$ 为上下文大小。目标是找到采样矩阵 $\Pi \in \mathbb{R}^{n \times m}$ 和对角矩阵 $D \in \mathbb{R}^{m \times m}$，以最小化误差 $\delta = ||wV - w\Pi D\Pi^T V||$，且计算预算 $m \ll n$。

局部敏感哈希（LSH）基础。LSH通过将相似输入分配相同哈希码的概率最大化来进行近似最近邻搜索。SimHash【[13], Similarity estimation techniques from rounding algorithms + 2002 + STOC】是基于余弦相似度的LSH变体。对于向量 $x \in \mathbb{R}^d$，SimHash生成随机超平面 $w$ 并返回 $\mathrm{Sign}(w^Tx)$。两个向量 $x, y$ 共享相同符号的概率为 $p = 1 - \frac{\theta}{\pi}$，其中 $\theta = \arccos \frac{x^Ty}{||x||\cdot||y||}$。

TopK注意力的致命弱点。尽管基于搜索的稀疏注意力算法【[63], Quest: Query-aware sparsity for efficient long-context llm inference + 2024 + arXiv】试图逼近TopK，但TopK注意力本身是有偏且不准确的。特别是当注意力分数分布呈现长尾特性且计算预算有限时，TopK在需要利用完整上下文的聚合任务中表现极差。实验表明，前20%的tokens仅能覆盖70%~80%的注意力分数，导致不可忽视的15%~20%的估计误差。

Figure 4 16k tokens的KV缓存的TopK估计误差
Figure 4 16k tokens的KV缓存的TopK估计误差

注意力的几何特征观察。通过研究查询 $q$ 和键 $k$ 的几何结构，得出以下关键观察：

  1. 初始token（即注意力Sink，记为 $k_{sink}$）的键状态对于任意输入几乎保持不变，相互间的余弦相似度大于0.99【[68], Efficient streaming language models with attention sinks + 2023 + arXiv】。
  2. 键状态的中心（即平均键 $k_{avg} = \frac{1}{n}\sum_{i=1}^n k_i$）在不同输入句子中方向保持稳定，相似度超过0.9。
  3. $k_{avg}$ 和 $k_{sink}$ 的方向几乎完全相反，余弦相似度在 $-0.9 \sim -0.8$ 之间。

这些几何特征表明，注意力Sink独立于输入产生高稀疏性，而其他部分分布更均匀。简单的TopK会过度赋予Sink权重，从而丢失上下文信息。此外，$q$ 和 $k$ 的不对齐也导致了搜索困难【[44], Retrievalattention: Accelerating long-context llm inference via vector retrieval + 2024 + arXiv】。

Figure 5 注意力的几何信息：汇聚键、平均键的方向及余弦相似度
Figure 2 左：LLM中长尾分布的例子。中：注意力Sink使分布显得更稀疏。右：注意力的几何特征，Sink键与平均键方向相反

通过采样估计注意力。将注意力输出 $o$ 重写为来自分布 $w$ 的期望值，即 $o = \mathbb{E}_{i \sim w}(v_i)$。定义3.1（Oracle采样估计）：给定采样预算 $\mathcal{B}$ 和归一化注意力分数 $w$，从 $w$ 中独立采样 $\mathcal{B}$ 个元素，输出估计为 $\bar{o} = \frac{1}{\mathcal{B}}\sum_{j=1}^{\mathcal{B}} v_{i_j}$。
定理3.2：Oracle采样估计是无偏的，且协方差的迹随 $\mathcal{B}$ 单调递减。
定理3.3：预期计算预算 $\mathbb{E}(|S|)$ 具有上限 $1 + \mathcal{B}\epsilon$，其中 $\epsilon = 1 - \max_i w_i$。这证明了实际计算成本通常远小于采样预算。

Figure 6 左/中：Oracle采样估计显著降低数值误差。右：下游任务比较，采样在极低计算成本下保持高准确率
Figure 3 TopK与采样的对比，16k总上下文

方法细节

自归一化重要性采样估计。由于获取精确分布 $w$ 需要计算所有的 $qk_i^T$，Oracle采样无法带来实质加速。因此，引入自归一化重要性采样，通过从提议分布 $u$ 中采样索引 $i_1, i_2, ..., i_{\mathcal{B}}$ 来估计未知分布 $w$。由此产生的估算器为 $X^{\mathrm{IS}} = \frac{1}{\widetilde{Z}} \sum_{j=1}^{\mathcal{B}} \frac{\widetilde{w_{i_j}}}{u_{i_j}} v_{i_j}$，其中 $\widetilde{Z} = \sum_{j=1}^{\mathcal{B}} \frac{\widetilde{w_{i_j}}}{u_{i_j}}$，且 $\widetilde{w_i} = e^{\frac{qk_i^T}{\sqrt{d}}}$。该估算器具有 $\mathbb{P}[\lim_{\mathcal{B} \to \infty} X^{\mathrm{IS}} = o] = 1$ 的优良性质。为了最小化方差，提议分布 $u$ 必须满足 $u_i \propto \widetilde{w_i}|v_i - o|$。

利用LSH进行方差缩减。对目标分布进行分解：$\widetilde{w_i}|v_i - o| = \exp(\frac{qk_i^T}{\sqrt{d}} + \log|v_i - o|)$。经验观察表明，$\log|v_i - o|$ 的波动相比 $\frac{qk_i^T}{\sqrt{d}}$ 并不显著，因此将 $u$ 的要求简化为与 $qk_i^T$ 共享相同的峰值。接着，通过向量变换 $r = \max_{1 \le i \le n} |k_i|$，$\bar{q} = [q, 0]$，$\bar{k_i} = [k_i, \sqrt{r^2 - |k_i|^2}]$，将内积转换为余弦相似度。基于这一转换，利用SimHash【[56], Simhash: Hash-based similarity detection + 2007】构建概率分布 $u_i = \mathbb{P}[h(q) = h(k_i)]$，该概率与 $\cos \frac{qk_i^T}{|q| \cdot |k_i|}$ 单调相关。

估算器近似与哈希函数选择。由于哈希提供的概率未归一化，MagicPIG对估算器进行了调整，计算 $X = \frac{\sum_{i \in S} \frac{\widetilde{w_i}}{u_i} v_i}{\sum_{i \in S} \frac{\widetilde{w_i}}{u_i}}$。在哈希函数选择上，利用 $K \times L$ 个随机向量进行SimHash。对于 $L$ 个哈希表中的每一个，保留投影的符号生成 $K$ 位哈希值。仅当键 $k_i$ 在至少两个哈希表中与 $q$ 共享哈希值时，才对其进行采样。对应的采样概率为 $u_i = 1 - (1 - p^K)^L - Lp^K(1 - p^K)^{L-1}$，其中 $p_i = 1 - \frac{1}{\pi} \arccos \frac{qk_i^T}{|q| \cdot |k_i|}$。

数据预处理与居中化。由于键几乎总是集中在查询的一侧（除了初始token），随机投影无法有效区分键，会导致均匀的采样概率。因此，在构建哈希表前，MagicPIG先对 $k_i$ 向量进行中心化处理（$\bar{k_i} = k_i - \frac{1}{n}\sum_{i=1}^n k_i$），使键分布更优且保持计算上的等价性。

算法执行过程。结合上述设计，MagicPIG的执行步骤如下：
首先，计算新查询 $\pmb{q}$ 的哈希码 $\pmb{q}_{\mathrm{code}} = \mathbf{Encode}(\pmb{q}, \pmb{W})$。
再者，查询哈希表以获取采样集合 $S = \mathbf{Query}(HT, \pmb{q}_{\mathrm{code}})$，并提取对应的键值 $\pmb{K}_S = K[S], \pmb{V}_S = V[S]$。
接着，计算 $\pmb{q}$ 与采样 $\pmb{K}_S$ 及静态缓存 $\pmb{K}_T$ 的内积：$\pmb{w}_S = \pmb{q}\pmb{K}_S^T$，$\pmb{w}_T = \pmb{q}\pmb{K}_T^T$。
然后，计算碰撞概率 $\pmb{p} = 1 - \frac{1}{\pi} \arccos(\pmb{w}_S / (||\pmb{q}|| \cdot ||\pmb{K}_S||))$，并由此计算采样概率 $\pmb{u} = 1 - (1 - \pmb{p}^K)^L - L\pmb{p}^K(1 - \pmb{p}^K)^{L-1}$。
最后，结合静态缓存 $T$（如Sink token）计算注意力输出估计：$\bar{o} = \mathbf{Softmax}(\frac{[\pmb{w}_S, \pmb{w}_T]}{\sqrt{d}} - \log([\pmb{u}, \mathbf{1}_t]))[\pmb{V}_S, \pmb{V}_T]$。

Input: K, V, q, random projectors W, hash tables HT, static KV cache K_T, V_T
# Compute hash code for new query
q_code = Encode(q, W)
# Query hash tables to sample S
S = Query(HT, q_code)
K_S = K[S], V_S = V[S]
# Compute inner product for q and sampled K
w_S = q * K_S.T
w_T = q * K_T.T
# Compute collision probability
p = 1 - (1 / pi) * arccos(w_S / (||q|| * ||K_S||))
# Compute sampling probability
u = 1 - (1 - p**K)**L - L * p**K * (1 - p**K)**(L - 1)
# Compute attention output estimation
o_bar = Softmax([w_S, w_T] / sqrt(d) - log([u, 1_t])) * [V_S, V_T]
Return o_bar

系统协同设计。为了突破GPU VRAM的内存瓶颈，MagicPIG利用带宽为 $100-200\mathrm{GB/s}$ 的CPU DRAM作为GPU的聚合内存【[31], Fastdecode: High-throughput gpu-efficient llm serving using heterogeneous pipelines + 2024 + arXiv】。解码工作负载被划分为四个部分：

  1. 参数计算：所有线性投影（MLP, $W_Q, W_K, W_V, W_O$）在GPU上运行。
  2. 随机投影：为每个 $q$ 进行 $K \times L$ 次随机投影以获取哈希码。由于所有头共享相同的随机投影器，内存开销极小（400KB），受计算限制，因此放置在GPU上。
  3. 检索：在CPU上的 $L$ 个预建哈希表中查找 $q$ 的哈希码。哈希表占用大量内存，适合放在CPU。
  4. 注意力计算：注意力核心计算 $o = \mathrm{Softmax}(\dots)V$ 在CPU上运行。

设备端缓存（On-device cache）：由于Sink token（前几个token）和局部token极有可能被采样，为进一步减少CPU工作量，MagicPIG将这些token存储在GPU上，不对其应用LSH采样，并利用递归注意力技术合并CPU和GPU的输出。

Figure 7 左：硬件内存层次结构。右：MagicPIG的工作负载划分。线性投影和哈希函数计算在GPU上，哈希表检索和注意力计算在CPU上
Figure 7 左：硬件内存层次结构。右：MagicPIG的工作负载划分。线性投影和哈希函数计算在GPU上，哈希表检索和注意力计算在CPU上

实验环境

实验结果

1. 准确率保持评估

2. 硬件效率与系统性能评估

Figure 8 MagicPIG在三种服务场景下的评估。左：A100吞吐量提升1.5倍。中：L20吞吐量提升5.0倍。右：RTX 4090实现54ms延迟并提升吞吐量3.3倍
Figure 8 MagicPIG在三种服务场景下的评估。左：A100吞吐量提升1.5倍。中：L20吞吐量提升5.0倍。右：RTX 4090实现54ms延迟并提升吞吐量3.3倍

3. 消融实验

Figure 9 左：有无中心化的准确率对比。中/右：MagicPIG与TopK在聚合任务中的对比，采样方法击败了确切的TopK注意力
Figure 9 左：有无中心化的准确率对比。中/右：MagicPIG与TopK在聚合任务中的对比，采样方法击败了确切的TopK注意力

结论

本文首先揭示了TopK注意力近似在解决长上下文LLM生成的计算和内存挑战时的局限性。随后证明了Oracle采样能够超越TopK，并提出了一种利用LSH采样近似Oracle采样的创新方法——MagicPIG。MagicPIG通过将哈希表和减少后的注意力计算卸载到CPU的系统协同设计，显著降低了注意力计算的工作量，同时在多样化任务中保持了高准确率。实验结果表明，MagicPIG在多种硬件配置下大幅提高了吞吐量并降低了延迟，全面优于传统的TopK注意力机制。其理论的严谨性、鲁棒性和可扩展性为注意力近似方法及算法-硬件协同设计开辟了新的机遇。未来的工作包括将MagicPIG扩展至预填充（Prefilling）阶段、应用更先进的交叉多胞形哈希（Cross-polytope hash）以减小哈希表体积，以及探索适用于高端GPU的纯GPU LSH加速方案。

附录与补充细节

A 定理证明补充。附录提供了定理3.2（Oracle采样无偏性及方差递减）和定理3.3（计算预算上限约束）的完整数学推导，利用Jensen不等式和凸函数性质严格证明了 $\mathbb{E}[|S|] \le 1 + \epsilon \mathcal{B}$。

B Oracle采样策略分析。理论上，保证最低方差的最优采样概率应为 $u'_i \propto w_i ||v_i||$。但MagicPIG未采用此策略，原因有二：首先，Sink token的值范数（Value norm）显著小于其他token（如图11所示），若按此概率采样会降低其被选中的概率，从而影响注意力的功能；其次，该策略会导致概率分布比 $w_i$ 更平缓，从而大幅增加计算成本。

Figure 10 单次解码步骤中 $\log|v_i - o|$ 和 $\frac{qk_i^T}{\sqrt{d}}$ 波动范围的对比，证明 $\log|v_i - o|$ 相对稳定
Figure 11 注意力Sink的值状态范数 $||v_0||$ 显著小于其他Token

D 扩展评估。

E LSH超参数 (K, L) 的选择。

F TopK与采样的直观对比。附录提供了一个动物园食物消耗估计的直观例子：TopK方法（仅选取数量最多的动物种类）会赋予高消耗动物不成比例的权重，导致估计值严重偏高（有偏估计）；而基于分布概率的放回采样机制，不仅能够提供无偏的估计值，还能随着采样预算的增加有效降低方差。这通俗地解释了为什么采样在聚合任务中超越了TopK。