Token 稀疏化 + Channel 稀疏化：16 倍 LLM 推理加速

最新推荐文章于 2025-04-26 16:43:32 发布

小马不会过河

最新推荐文章于 2025-04-26 16:43:32 发布

阅读量1.6k

点赞数 13

文章标签：网络 gpt 大数据人工智能 microsoft

本文链接：https://blog.csdn.net/m0_59163425/article/details/142467834

版权

一、背景

我们之前介绍过一些针对长序列场景使用 Token 稀疏化来实现 LLM 推理加速的文章，本文中我们介绍一种新的方案，其额外引入了 Channel 稀疏化，也获得不错的加速。当我们在看这一系列稀疏化或者量化方案的时候可以从如下几个角度来分析相应方案的优劣，本文中我们同样会提及到：

显存占用：是否减少 KV Cache 大小，或者碎片化。
IO 大小：是否减少 KV Cache 的访存量；是否额外引入 CPU 与 GPU（offload）或者 GPU 之间的通信量（分布式 KV）。
访存效率：是否是连续存储，方便高效访问；除此之外，非连续存储也可能导致无法使用 Tensor Core。
计算效率：是否引入额外的计算量（比如，反量化）；是否能充分利用 GPU 的计算单元（比如，Tensor Core）；是否存在无效计算（比如，稀疏矩阵乘或者 padding）。

对应的论文为：[2408.07092] Post-Training Sparse Attention with Double Sparsity

对应的代码库为：https://github.com/andy-yang-1/DoubleSparse/tree/main

二、摘要

LLM 推理过程缓慢，并且会占用大量内存，其最主要的瓶颈是过多的 KV Cache 存储和访问。本文中，作者提出 Double Sparsity，是一种新的训练后稀疏注意力机制，旨在通过减少 KV Cache 访问来缓解上述瓶颈。Double Sparsity 将 Token 稀疏性和 Channel 稀疏性相结合：

Token 稀疏性侧重于利用重要的 Token 来计算 Attention，也就是 Sequence 维度的稀疏。
Channel 稀疏性是使用重要的 Channel 来识别重要的 Token，也就是 Hidden 维度的稀疏。

作者的关键见解是：Channel 的稀疏性通常是相对静态的，可以通过离线校准的方式来执行，以便推理更高效，从而能够准确和高效的识别重要的 Token。此外，这种方法可以与 Offload 结合，以显著减少内存占用。

作者对 LLaMA2-7B、LLaMA2-70B 以及 MIstral-8x7B 模型进行评估实验，其在 wiki-2 上的困惑度、key-value 检索，以及长上下文基准上的影响都很小。在注意力计算上可以获得 14.1x 加速，在 GPU 的端到端推理方面也可以获得 1.9x 加速。通过 Offload，与 SOTA 的方案相比，在 256K 的序列长度下可以实现 16.3x 的 Decoding 加速（PS：这里我们标题党了😂）。

PS：Paper 中很多部分说的不够清晰明了，需要结合整篇文章的内容或者图表去推测，比如其中的 top-k 和 β。

三、方案

3.1 概览

整体的方案比较简单，如下图所示，首先通过离线的方式（Offline Calibration）识别出重要的 Channel（比如，识别出 Channel 0 和 Channel 2 是重要的 Channel）。为了避免不连续内存访问导致的性能下降问题，使用 Label Cache 存储筛选后的重要 Channel。（PS：需要说明的是，这里其实相当于还要存储全量的 Key Cache，Label Cache 是新增的部分，所以反而增加了显存占用）

其代码如下图所示（DoubleSparse/blob/main/models/model.py#L71-L96），额外存储了 Label Cache：

之后，如下图所示，实际计算时分为 3 步：

第一步：根据重要 Channel 的 Index 提取 Query，同样要连续存储。
第二步：使用稀疏但连续存储的 Query 与 Key Cache 计算 Attention Score（PS：近似值）。
第三步：根据 Attention Score 选出 Top-k 的重要 Token（Key/Value Cache），然后使用这些重要的 Token 进行 Attention 计算。（PS：这里使用的 Query，Key，Value 都不是 Channel 稀疏化的，这也就是为什么依旧要存储完整的 Key/Value Cahe）

具体的算法如下图所示，其中的 1，2，3 步可以理解为动态的稀疏性计算，以获得重要的 Token：

3.2 Offline Calibration

如下图所示，借鉴 AWQ 中识别重要 Channel 的方式，使用少量样本数据，通过 Offline 的方式识别哪些 Channel 更加重要（PS：这里不同 Layer 的 Channel 可能是不一样的）。此外，这里有几种不同的策略：

Activation Outlier：根据 Activation 中的 Outlier 来识别重要的 Channel，红色部分。
Q-Outlier：根据 Query 中的 Outlier 来识别重要的 Channel，绿色部分。
K-Outlier：根据 Key 中的 Outlier 来识别重要 Channel，蓝色部分。
QK-Outlier：根据 Query * Key 来识别重要 Channel，紫色部分。（PS：和使用 Attention Score 是等价的，不过没有必要）

Top Channel 的数量与质量密切相关，作者进行消融实验验证，当 Top Channel 的占比达到 0.25 时，通过 Offline 方式确定的 Outlier Channel 与通过 Online 方式获得的 Outlier Channel 的重合度达到了 0.95：

3.3 使用 Label Cache 前向推理

上述实验验证 Top Ratio 为 0.25 即可，此时相当于 Key Cache 只用保留 1/4 的 Channel，此外，还可以进一步使用 INT4 对 Key Cache 进行量化，可以进一步将 Key Cache 降低到 1/16。为了保证连续访问，并提高 L1/L2 Cache 命中率，作者额外使用 Lable Cache 存储稀疏、量化后的 Key Cache。此外，作者也对是否使用 Label Cache 进行了消融实验，如下图 Table 4 所示，使用 Label Cache 相比不使用 Label Cache 可以带来 2-4x 的 Decoding 加速（PS：结合摘要中端到端 1.9x 加速，也就意味着不使用 Label Cache 反而会导致降速）：

3.4 Double Sparsity-Offload

从上面也可以看出，Double Sparsity 方案虽然可以在一定程度上减少显存 IO，但需要额外存储 Label Cache，实际上是增加了显存占用。为了解决这个问题，作者提出了将全量的 Key/Value Cache Offload 到 CPU 内存中，在需要的时候动态的加载到显存中，也就是 Double Sparsity-Offload。

然而，如果每次使用 Label Cache 获得 Top-k 的 Token 后再去加载对应的 Key/Value Cache 到显存中，其延时可能比较大，并且无法与计算 overlap。为了解决这个问题，作者引入了 Double Buffer，其核心思想是：在计算当前 Transformer Layer 的时候可以异步的预取下一个 Layer 的相关 Token 的 Key/Value Cache。为什么可以提前加载下一层的 Key/Value Cache 呢？如下图 Figure 8 所示，作者对比相邻层的 Attention Head 发现，除了前两层以及最后两层外，后续相邻层的相似性都很高，甚至达到 95% 以上（PS：这个与之前的很多工作结论类似），因此也就可以使用当前 Layer 来近似下一个 Layer。（PS：实际实现是从 Layer 2 开始，在计算 Layer 2 的 Attention 时，开始使用对应的 Query 来异步估计并加载下一 Layer 的稀疏 Token Key/Value Cache）

如下图所示，代码实现中的预取也跳过了相关层：

3.5 复杂度分析

作者进一步对比了不同方案中 GPU 上 Cache 大小，以及 IO 大小，其中的 β 就是 Token 的稀疏度 Top-k，α 为 Channel 的稀疏度，因为其他部分考虑了 Key、Value，而 α 这里主要是 Label Cache，只有 Key，所以是 α/2。可以看出，Double Sparsity-Offload 的主要优势就是干掉了 GPU 中的全量 Key/Value Cache；此外，可以将 Token 稀疏度 β 进一步降低到 1/16：

四、实验

4.1 精度评估

4.1.1 Wiki-2 Perplexity

如下图 Table 2 所示，可以看出，Wiki-2 上的困惑度（越低越好）在 1/8 稀疏度下基本无损，在 1/16 的稀疏度下损失还可以接收，但是在 1/32 稀疏度下损失就比较大了：

此外，作者进一步对比了不同 Channel 稀疏度和 Token 稀疏度情况下的困惑度（PS：需要说明的是，图中是 10 - 困惑度，因此越高越好）。可以看出，当 Channel 稀疏度和 Token 稀疏度为 1/16 时获得最好的权衡：

除此之外，作者也进一步验证了不同模型结构（GQA、MoE）及不同的 Channel 稀疏选择下的精度。如下图 Table 3 所示，可以看出，Q-Outlier 在大部分情况下都还不错，QK-Outlier 在 Mixtral-8x7B（MoE）情况下会差点：

4.1.2 长上下文基准

如下图 Figure 3 所示，作者与 H2O，StreamingLLM 对比了长上下文场景的精度，可以看出，在 1/4，1/16 稀疏度的情况下本文的 Double Sparsity 都获得更高精度：

4.1.3 Key、Value 检索

如下图 Figure 4 所示，在长上下文 Key、Value 检索基准上依然获得不错的效果：

4.2 速度评估

4.2.1 实验配置

测试硬件主要是 A10 和 A100 GPU，Batch Size 为 4-32，序列长度为 1K-16K。使用 Pytorch 实现 Attention Score 计算，稀疏 Attention 计算使用 Triton 实现。对于 Double Sparsity-Offload，也实现了异步的 CPU 到 GPU 拷贝。

4.2.2 Attention 算子加速

如下图 Figure 5 所示，在比较大的 Batch Size 下，2K、4K 序列长度时可以获得近似线性加速；即使 Batch Size 或者 Seq 比较小也可以获得几倍加速：

4.2.3 端到端加速

如下图 Figure 16 所示，非 Offload 方式，相比 GPT-Fast 最多可以加速 1.9x；而 Offload 方式，相比 FlexGen 最多可以加速 16.3x 倍：

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述