arXiv:2503.18773v1 类型: cross 摘要: 长上下文大语言模型(Large Language Models, LLMs)的逐步采用给自回归解码带来了显著的内存和计算挑战，因为关键值-值(Key-Value, KV)缓存的不断扩大。KV缓存量化已 emerged 作为一种有前景的解决方案，先前的工作表明，使用4位甚至2位量化可以在减少内存成本的同时保持模型精度。然而，尽管有这些优点，低位KV缓存的初步实现由于量化和反量化开销以及张量核心利用不足，难以实现预期的加速效果。在本文中，我们提出了BitDecoding，这是一种专为GPU优化的框架，能够利用张量核心实现低位KV缓存下的高效解码。由于在每次解码步骤中KV缓存生成的动态性质，高效利用低位KV缓存的张量核心具有挑战性。BitDecoding 通过一种以张量核心为中心的位融合方案解决了这些挑战，确保数据布局兼容性，从而能够充分利用张量核心。此外，BitDecoding 还包含了一个战程高效并行解码内核和一个细粒度的异步流水线，最小化了反量化开销并提高了计算效率。实验结果显示，与FP16的FlashDecoding-v2相比，BitDecoding 在RTX 4090上的加速比达到了7.5倍，在A100上的加速比达到了4.8倍，在H100上的加速比达到了8.9倍。它还比最先进的低位KV缓存实现(QServe)在某些情况下提高了4.3倍的表现。在LLaMA-3.1-8B模型下，序列长度为128K时，BitDecoding 将单批次解码延迟减少了3倍，展示了其在长上下文生成场景中的有效性。源代码可在 https://github.com/DD-DuDa/BitDecoding 获取。