AAAI 2026 AMD论文Spark方法揭秘：查询感知的 KV 缓存通道剪枝

AAAI 2026 AMD论文Spark方法揭秘：查询感知的 KV 缓存通道剪枝

原文作者：Huanxuan Liao, Yixing Xu, Shizhu He, Guanchen Li, Xuanwu Yin, Dong Li, Emad Barsoum, Jun Zhao, Kang Liu

在这篇博客中，我们将讨论SparK，一种无需训练、即插即用的大语言模型（LLMs）KV缓存压缩方法。通过解决特征通道中被忽视的冗余，并采用“剪枝与恢复”的策略，SparK 在保持模型精度的同时，相比传统方法减少了超过 30% 的 KV 缓存存储。它为长上下文推理提供了一个稳健的解决方案，并在非结构化稀疏性方面开辟了新的视角。

SparK 与 AMD ROCm™ 软件栈协同设计，以充分发挥 AMD GPU 的并行计算能力。我们的 KV 缓存剪枝方法能够帮助提升 LLM 在 AMD GPU 上的性能。

欢迎阅读完整论文 [1] 并尝试实现 [2] 。本研究已被AAAI 2026 [3] 接收。

为什么 KV 缓存压缩很重要

在大语言模型（LLMs）的长上下文推理中，越来越受到 KV 缓存瓶颈 的限制：内存使用随着序列长度 线性增长，而注意力计算则随着序列长度 二次方增长。这限制了单个 GPU 上可处理的最大批量大小和序列长度。

图 1：示意性比较 (a) 完整 KV 缓存，(b) 基于淘汰的 KV 压缩，(c) 基于结构化通道剪枝的 KV 缩减，以及 (d) 我们提出的SparK，其采用非结构化通道剪枝，并在注意力分数计算过程中进行后续恢复。

如图 1 所示，现有方法通常通过在时间轴（即 token 维度）上压缩 KV 缓存来解决这一问题。诸如 token 淘汰（移除不重要的 token）或 token 合并 等策略，一直是减少内存开销的标准做法。 然而，这些方法往往忽视了特征维度（通道）中存在的冗余。它们将所有通道视为同等重要，可能会保留一些“无效”或不相关的特征信息，从而占用宝贵的内存资源。

SparK 内部机制：它是如何工作的

图 2：SparK 示意图。SparK 在 prefill 阶段 计算通道级显著性分数，并进行 非结构化剪枝。在 解码阶段，SparK 利用 F 和从缓存分布中采样来重建被剪枝的通道，然后执行标准的完整注意力计算。

SparK（Query-Aware Unstructured Sparsity with Recoverable KV Cache Channel Pruning） 采用了一种不同的方式。它并不仅仅是淘汰 tokens，而是如图 2 所示，直接针对通道级冗余进行处理。

SparK 的核心洞察在于：通道显著性在不同的查询和位置之间存在显著差异。对于某个特定查询，某些特征通道几乎不携带任何信息，而另一些通道则会在相关性上急剧提升。

SparK 基于一个简单但有效的原则运行：

• 查询感知剪枝（Query-Aware Pruning）：在通道层面识别并剪枝那些对当前查询无关的 KV 条目。
• 动态恢复（Dynamic Recovery）：关键在于，它会在注意力分数计算过程中动态恢复被剪枝的条目。

这种“剪枝—恢复”机制使 SparK 能够在应用非结构化稀疏性的同时，不会永久丢失高精度注意力所需的关键信息。值得注意的是，SparK 与现有的压缩技术是正交的。这意味着它可以与量化或token淘汰方法结合使用，在AMD GPU上实现更进一步的内存节省。

在 AMD GPU 上的结果：稳健性与高效性

表 1：在 LongBench 上对 LLaMA-3-8B-Instruct 的性能比较。 SparK (λ) 表示通道级 Key 缓存剪枝比例 λ。 基准测试在 AMDGPU加速器上进行。

SparK 相较于基线的淘汰式方法展现出令人印象深刻的稳健性：

• 存储减少（Storage Reduction）：在相同序列长度下，SparK 相比标准淘汰方法可减少超过30% 的 KV 缓存存储，如图 3 (b) 所示。
• 精度保持（Accuracy Preservation）：通过减少通道级冗余，SparK 能够在相同的内存预算下处理更长的序列。在测试中，相较于基线方法它要么保持，要么提升了模型精度，如表 1 所示。
• 高稀疏容忍度（High Sparsity Tolerance）：即使在激进的80% 剪枝比例下，SparK 的性能仍能保持在基线淘汰方法的水平之上，性能下降不到5%，如图 3 (a) 所示。

图 3：在 LLaMA3-8B-Instruct 上对 SparK 的性能分析。 (a) LongBench 在不同剪枝比例 (λ) 下的平均性能。SparK 在所有压缩水平上均显著优于 ThinK。 (b) 缓存大小与性能的权衡。SparK 相比 ThinK 和 SnapKV 实现了更优的存储—性能平衡。 实验在 AMDGPU加速器上进行。

这些结果突显了 SparK 在处理长上下文场景方面的强大能力，使其成为内存受限环境中的稳健选择。

总结

在这篇博客中，我们探讨了SparK——一种缓解大语言模型（LLMs）中 KV 缓存瓶颈的全新方法。不同于传统的时间维度压缩，SparK 利用通道维度中的非结构化稀疏性。通过剪枝无关通道并在计算过程中动态恢复，它在无需重新训练模型的情况下实现了显著的内存节省。

SparK的突出之处在于，它是一种即插即用的解决方案，能够与现有的 KV 压缩和量化技术兼容，为优化长上下文 LLM 推理提供了灵活的工具。

您可以在我们的论文中深入了解方法论和详细的基准测试，并在 GitHub [2] 上获取实现。我们欢迎研究人员在支持 AMD ROCm 的 GPU 上探索 SparK，并与社区分享反馈。

我们也邀请您体验 AMD Developer Cloud[4]，其中配备了专为 AI 工作流打造的 AMD GPU加速器。如有问题或合作机会，请联系 AMD 团队：amd_ai_mkt@amd.com。请持续关注后续的博客文章、扩展工具以及实践教程，我们将不断推进 KV 缓存剪枝的研究与应用。

参考链接

[1] 原文：https://arxiv.org/abs/2508.15212

[2] AMD-Spark实现：https://github.com/AMD-AGI/AMD-Spark

[3] AAAI 2026：https://aaai.org/conference/aaai/aaai-26/

[4] AMD Developer Cloud：https://amd.digitalocean.com/login

获取更多AMD开发者资源及技术支持，欢迎添加小助手微信:

AMD_Developer（备注：开发者）