arXiv:2503.18773cs.ARcs.AI2025-03被引 7

用低比特缓存加速长文本大模型推理,显著提升显卡性能利用率。

BitDecoding: Unlocking Tensor Cores for Long-Context LLMs with Low-Bit KV Cache

  • 创新布局与并行解量化,让显卡张量核心高效处理低比特键值缓存。
  • 在黑格尔显卡上提速最高达8.6倍,128K上下文推理延迟降低3倍。
  • 适合追求高吞吐、低延迟的长文本生成应用,如文档理解与代码生成。

长上下文大语言模型在自回归解码时因键值(KV)缓存扩大而面临巨大内存与带宽压力。尽管4位或2位精度的量化可减少内存占用,但现有系统仅依赖CUDA核心解码,未能充分利用GPU上的主流计算资源——张量核心。本文提出BitDecoding,首个能高效解码低比特KV缓存的推理系统,通过协同利用CUDA核心与张量核心实现加速。该系统设计了适配张量核心的存储布局,引入波次级解量化并行机制,并通过查询变换、高性能张量与通道级量化及软件流水解量化核,支持混合精度执行。架构感知优化进一步利用黑格尔(Hopper)的波组张量指令和布莱克威尔(Blackwell)的NVFP4(MXFP4)格式。在布莱克威尔、黑格尔和安培系列GPU上评估显示,与FP16 FlashDecoding-v2相比,平均提速7.5倍,黑格尔+NVFP4下最高达8.6倍,相较当前最优方法提速最高4.3倍。在LLaMA-3.1-8B模型上,128K上下文单批推理延迟降低3倍。代码已开源:https://github.com/OpenBitSys/BitDecoding。

原文摘要 · Abstract (English)

The growth of long-context Large Language Models (LLMs) significantly increases memory and bandwidth pressure during autoregressive decoding due to the expanding Key-Value (KV) cache. While accuracy-preserving KV-cache quantization (e.g., 4-bit or 2-bit) reduces memory footprint, existing systems decode inefficiently by relying solely on CUDA cores, underutilizing Tensor Cores-the dominant compute resource on GPUs. We present BitDecoding, the first inference system to efficiently decode low-bit KV caches by cooperatively leveraging CUDA cores and Tensor Cores. BitDecoding smartly induces Tensor-Core-friendly layouts, introduces warp-level dequantization parallelism, and provides unified system support through query transformation, high-performance tensor- and channel-wise quantization, and a software-pipelined dequantization kernel enabling mixed-precision execution. Architecture-aware optimizations further leverage Hopper's warpgroup tensor instructions and Blackwell's NVFP4 (MXFP4) tensor formats. Evaluated on Blackwell, Hopper, and Ampere GPUs, BitDecoding achieves an average 7.5x decoding speedup over FP16 FlashDecoding-v2, up to 8.6x on Blackwell with NVFP4, and up to 4.3x over state-of-the-art approaches. On LLaMA-3.1-8B with a 128K context, BitDecoding reduces single-batch decoding latency by 3x. BitDecoding is open-sourced at https://github.com/OpenBitSys/BitDecoding.

大模型推理张量核心低比特量化长上下文

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。