发表时间: 2025-03 · arXiv:2503.18773 (Microsoft Research & Univ. of Edinburgh)
原文: https://arxiv.org/abs/2503.18773
文章标题:BitDecoding:利用低比特位KV缓存为长上下文大语言模型解锁Tensor Cores
作者/机构:Dayou Du (University of Edinburgh), Shijie Cao (Microsoft Research), Jianyi Cheng (University of Edinburgh), Luo Mai (University of Edinburgh), Ting Cao (Institute for AI Industry Research (AIR), Tsinghua University), Mao Yang (Microsoft Research)
本文针对长上下文大语言模型(LLMs)在自回归解码过程中因KV缓存不断扩大而导致的显著内存和带宽压力问题,提出了一种名为BitDecoding的高效推理系统。
核心问题与研究目标:
现有的解决方案通过KV缓存量化(如4位或2位)来减少内存占用,但这些系统仅依赖CUDA核进行解码,导致GPU中占主导地位的计算资源——Tensor Cores——被严重低估和利用不足。此外,低比特KV缓存的量化和反量化会引入显著开销,尤其是在自回归解码中,每一新生成的令牌都需要在线进行这些操作(如图1所示),这与可以离线处理的静态权重不同。现有系统要么采用分离的内核导致高昂的内存开销,要么采用仅限CUDA核的融合内核而无法利用Tensor Cores的强大算力(如图2所示)。
研究目标是设计一个能够协同利用CUDA核与Tensor Cores的推理系统,以高效解码低比特KV缓存,从而在保持高精度的同时,大幅提升长上下文LLM的解码速度和吞吐量。
主要创新点与贡献:
1. 提出BitDecoding系统:这是首个通过协同利用CUDA核(用于反量化)和Tensor Cores(用于矩阵乘法)来高效解码低比特KV缓存的推理系统。
2. 创新的布局与并行策略:
* Tensor-Core友好的布局诱导:通过硬件指令(如ldmatrix)自动生成与Tensor Cores计算布局兼容的低比特数据布局,避免了昂贵的全局重塑操作。
* Warp级反量化并行:设计了新的Warp布局策略,通过增加并行处理的Warp数量,利用SM调度器隐藏反量化带来的延迟,从而避免Tensor Cores停顿。
统一且高效的系统支持:
Residual Kernel,以及一个采用软件流水线、能够重叠计算与数据移动的Packing Kernel。面向新架构的优化:
warpgroup张量指令和异步执行能力。显著的性能提升:在Blackwell、Hopper和Ampere等GPU上进行了评估,与FP16的FlashDecoding-v2相比,BitDecoding平均解码速度提升7.5倍,在Blackwell上使用NVFP4最高可达8.6倍。与最先进的低比特方法相比,速度提升高达4.3倍。在LLaMA-3.1-8B模型和128K上下文长度下,单批次解码延迟降低了3倍。
图 1: 低比特权重与低比特KV缓存的混合精度矩阵乘法比较。(a) 量化权重可以离线预处理。(b) KV缓存需要为每个新生成的令牌进行在线量化和打包。
图 2: 不同低比特KV缓存系统与半精度FlashAttention的比较。每个系统遵循注意力公式 Out = softmax(Q D(K'⊤)) D(V'),其中 K' 和 V' 是低比特量化的键和值张量,D(·) 表示反量化函数。
LLM推理与低比特KV缓存。LLM推理包含预填充(Prefill)和解码(Decode)两个阶段。KV缓存的大小与批次大小b和序列长度l成线性关系,成为长上下文和大批量工作负载下的内存瓶颈。由于批处理推理中每个序列的上下文独立,KV缓存的访问通常受限于内存带宽。因此,使用更低比特位的KV缓存【【12,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】,【18,Kivi: A tuning-free asymmetric 2bit quantization for kv cache,2024】,【36,Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization,2024】】成为一种有前景的解决方案,它能在保持接近非量化基线精度的同时,减少内存占用并提高吞吐量。
现代GPU上的Tensor Cores与CUDA Cores。在优化LLM推理时,必须同时利用Tensor Cores和CUDA Cores。Tensor Cores专为矩阵运算(如GEMM)设计,提供绝大部分的计算FLOPS,而CUDA Cores则提供更灵活的向量、标量和控制流能力,但峰值FLOPS远低于Tensor Cores。例如,A100上的Tensor Cores在FP16/BF16下可达312 TFLOPS,远超CUDA Cores的19.5 TFLOPS FP32。这一性能差距在Hopper和Blackwell架构中进一步扩大,后者通过支持原生微缩放格式(如MXFP4)可提供高达20 PFLOPS的算力。
利用Tensor Cores进行高效推理。为了充分利用Tensor Cores,最新的LLM【【10,The llama 3 herd of models,2024】,【17,Deepseek-v3 technical report,2024】,【34,Qwen3 technical report,2025】】越来越多地采用多查询注意力(MQA)【【26,Fast transformer decoding: One write-head is all you need,2019】】和分组查询注意力(GQA)【【1,Gqa: Training generalized multi-query transformer models from multi-head checkpoints,2023】】。这些变体通过在多个查询间复用KV头来减少内存带宽,从而提高算术强度和计算效率,这与Tensor Cores以高吞吐量矩阵为中心的设计非常契合。因此,利用Tensor Cores对于长上下文和分组注意力LLM的高效推理至关重要。
现有低比特KV缓存系统的局限性。尽管已有一些支持低比特KV缓存的系统【【16,Qserve: W4a8kv4 quantization and system co-design for efficient llm serving,2024】,【18,Kivi: A tuning-free asymmetric 2bit quantization for kv cache,2024】,【37,Atom: Low-bit quantization for efficient and accurate llm serving,2024】】,但它们往往因GPU利用率不足而性能不佳,主要原因如下:
* 采用分离的低比特KV缓存内核:以Kivi【【18,Kivi: A tuning-free asymmetric 2bit quantization for kv cache,2024】】为代表的方法将混合精度注意力分解为多个独立内核。这种设计灵活,但频繁的内核启动、中间数据存取导致了高昂的启动开销、增加了全局内存流量并破坏了片上数据复用,从而降低了有效吞-吐量。
* 仅在CUDA Cores上融合低比特KV缓存内核:这类方法将反量化和矩阵运算都放在CUDA Cores上执行。虽然优于非融合设计,但它未利用Tensor Cores。CUDA Cores必须处理昂贵的反量化、缩放和逐元素操作,这些任务是内存密集型的,消耗了指令槽、寄存器带宽和缓存容量,降低了占用率,并为计算密集的矩阵乘法留下了更少的资源。因此,在CUDA Cores上同时运行反量化和矩阵乘法会引入显著开销。
本文的核心方案是协同使用Tensor Cores和CUDA Cores来支持低比特KV缓存。设计思路是将密集的矩阵乘法调度到Tensor Cores上执行,而将量化、打包和反量化等非矩阵运算高效地在CUDA Cores上执行。为了实现高效协同,系统需要平衡两者之间的工作负载,并精心编排数据移动,确保反量化过程能无缝地为Tensor Core的GEMM提供数据,从而最小化内存流量并最大化端到端解码吞吐量。
系统设计目标是实现广泛适用性。该系统旨在支持包括MHA、MQA和GQA在内的多种注意力变体,并能跨越多代GPU架构。这要求系统提供一个简洁的接口,能与现有注意力实现集成,并且设计易于适配,能够快速针对不同的GPU后端进行优化,同时保持高解码吞-吐量。
预期收益巨大。例如,通过在FlashAttention-3 (FA-3)【【25,Flashattention-3: Fast and accurate attention with asynchrony and lowprecision,2024】】的基础上构建低比特解码能力,可以利用SM90架构的特定功能(如Warp专业化流水线),相比传统SM80指令实现高达6倍的加速。此外,该设计也为Blackwell架构的原生低精度格式支持做好了准备,有望带来更显著的吞吐量提升。
挑战1:Tensor Cores常因低比特布局不匹配而效率低下。在自回归生成过程中,动态增长的KV缓存数据布局很难与Tensor Cores的要求对齐。低比特KV缓存在反量化后,必须匹配Tensor Cores期望的半精度布局,这面临三个难题:
1. 指令和GPU代际间的片段布局差异:ldmatrix等指令加载到寄存器的数据片段强制了特定的值到线程映射。如图3a所示,mma.m16n8k16指令的布局与其他指令(如mma.m16n8k8)或Hopper架构的wgmma系列指令不同。
2. 低精度比特宽度加剧对齐问题:Tensor Cores的寄存器布局是交错的,低比特数据在量化打包后是连续存储的,解包反量化后无法直接与交错的寄存器布局对齐,导致MMA执行错误(如图3b所示)。即使是Blackwell的原生低精度格式,硬件支持仍然有限,仍需软件仔细处理。
3. 反量化成为执行瓶颈:简单的低比特到FP16类型转换速度很慢【【14,Who says elephants can’t run: Bringing large scale moe models into cloud scale production,2022】】。先前工作如Ladder【【33,Ladder: Enabling efficient {Low-Precision} deep learning computing through hardware-aware tensor transformation,2024】】和Marlin【【9,Marlin: Mixed-precision auto-regressive parallel inference on large language models,2024】】通过为静态权重插入单独的布局转换核来解决此问题,但这会增加显著开销,不适用于动态解码。
图 3: (a) mma.m16n8k16指令的矩阵B片段布局。每个线程(Ti)根据指令定义的交错映射分配一组特定的值。(b) 对于INT4,量化会按线程连续打包值。反量化后,布局与预期的交错模式不匹配。
挑战2:频繁的停顿限制了Tensor Cores的利用率。在高性能注意力内核中,经验性调整的Warp布局和分区策略,在引入低比特KV缓存时往往会无意中降低性能。在FlashAttention原有的Warp分区策略下,额外的反量化(DQ)操作会显著降低吞吐量和Tensor Core利用率。如图4a所示,当一个Warp负责沿N维度的计算时,反量化操作会频繁地使Warp停顿。Nsight Compute性能分析(图4b)证实,增加的DQ开销会增加内存访问停顿,并抑制计算吞吐量和Tensor Cores利用率【【8,Warpdrive: Gpu-based fully homomorphic encryption acceleration leveraging tensor and cuda cores,2025】】。此外,即使使用原生低精度格式,为了执行第二次矩阵乘法PV,也需要在softmax后动态地对概率矩阵P进行重新量化,这同样会引入新的计算瓶颈。
图 4: (a) 沿N维度进行寄存器级操作的单个warp会因反量化(DQ)而经历停顿。(b) 有无反量化的微观层面比较。
挑战3:缺乏针对不同低比特KV缓存方法的可泛化系统优化。流行的KV缓存量化方法采用不同的缩放粒度(如张量级【【12,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】,【37,Atom: Low-bit quantization for efficient and accurate llm serving,2024】】和通道级【【13,Gear: An efficient kv cache compression recipefor nearlossless generative inference of llm,2024】,【18,Kivi: A tuning-free asymmetric 2bit quantization for kv cache,2024】】),这使得构建一个统一的支持系统变得复杂。在线量化和打包需要归约和逐元素转换,增加了运行时开销。辅助元数据(缩放因子和零点)也增加了内存流量。现有的混合精度内核优化【【9,Marlin: Mixed-precision auto-regressive parallel inference on large language models,2024】,【33,Ladder: Enabling efficient {Low-Precision} deep learning computing through hardware-aware tensor transformation,2024】】主要针对静态权重,不适用于KV缓存的动态特性。
本节介绍BitDecoding系统的设计,它实现了Tensor Cores和CUDA Cores在支持低比特KV缓存方面的协同使用。设计主要包含:(i) 优化低比特布局以利用Tensor Cores的新方法和原则,以及(ii) 并行化和协调GPU Warp以最小化反量化停顿的新策略。
图 5: 在Tensor Cores上优化低比特布局的方法概述。(1) 在Tensor Cores片段内融合计算和量化。(2) 低比特打包数据保留FP16值。(3) 低比特布局与反量化后的半精度布局匹配。(4) 为更快的反量化进行布局重映射。
(1)使用硬件指令诱导低比特优化布局。本设计的核心洞见是利用ldmatrix指令加载数据到寄存器时,数据本身就采用了Tensor Core的交错式片段布局。如图5-(2)所示,如果每个线程随后在本地进行量化和打包,那么生成的低比特打包数据会隐式地保留半精度(FP16)的交错布局。在解包和反量化时,这些值已经与Tensor Core寄存器匹配,无需进行全局重塑。这种方法利用硬件指令在计算过程中自动诱导出一个有效的低比特打包布局,实现了零开销的重映射,既高效又与Tensor Cores执行兼容,并避免了额外的数据移动。基于此,我们设计了一个专用的GPU Residual Kernel,它融合了新生成的FP16 KV张量的计算、量化和打包。同时,我们引入了一个Packing Kernel来消费这些缓存,它融合了反量化与计算,并保证使用与Residual Kernel相同的指令配置(ldmatrix和mma变体),从而确保解包后的数据能立即用于矩阵乘法,无需显式布局校正。
(2)对齐Warp与残差KV缓存以饱和Tensor Cores。Tensor Cores执行Warp分块的矩阵运算,要求输入块被完全填充以达到最佳吞吐量。我们的洞见是,通过分配一个大小与Tensor Cores分块能力相匹配的残差缓冲区,可以确保低比特数据与硬件的计算粒度对齐,从而充分利用计算单元的能力。为此,我们引入了一个大小为Nr的半精度残差KV缓存。我们将整个KV缓存X ∈ R^(L×d)划分为:
其中,X_packed是低比特部分,X_res是半精度残差部分。残差块的大小Nr根据低比特量化的比特宽度β、打包存储的字大小ω、沿N维的Warp数Wn以及每个Warp块处理的元素数Pn来计算:
这保证了低比特KV缓存片段与Tensor Core操作的Warp级分块精确对齐,实现了密集的、布局兼容的打包,并最大化了计算单元的占用率。
(3)为更快的反量化重映射布局。尽管布局与Tensor Cores兼容,但直接使用static_cast将低比特值转换为FP16的效率低下。为了缓解这种低效,我们设计了一种基于低级位操作和指令的更快反量化映射方法【【14,Who says elephants can’t run: Bringing large scale moe models into cloud scale production,2022】】。在使用ldmatrix将打包数据加载到寄存器后,我们先将其转换为INT32,然后按照75316420的模式映射到交错的Tensor Core布局。这种布局能够使用lop3指令进行位操作,高效地将INT4/INT2数据转换为FP16,同时与Tensor Core的计算模式对齐。
(4)通过配置设置协调残差和打包核。该设计通过在统一的指令配置下协调Residual和Packing核来执行。首先,根据GPU架构确定硬件指令配置(包括ldmatrix和mma变体)。然后,根据低比特KV缓存的位宽计算残差块大小Nr。如图5所示,Residual核加载高精度KV条目,使用Tensor Cores进行计算,然后融合量化和打包操作存入低比特KV缓存。Packing核使用相同的指令配置加载打包数据,执行高效的反量化,并继续进行Tensor Core计算。
(1)为低精度操作增强Warp并行性。我们引入了一种新颖的Warp布局来并行化处理多个打包数据块。以反量化为例,我们修改了Warp分区策略:将沿M维度的Warp分配限制为Wm = 1(因为解码时查询长度通常很小),并将资源重新分配以增加沿N维度的Warp数量Wn。如图6所示,通过增加Wn,流式多处理器(SM)的Warp调度器【【24,A survey of techniques for warp scheduling in gpus,2015】】可以有效地隐藏反量化停顿,因为多个Warp可以并发执行反量化操作,然后再进行基于Tensor Cores的矩阵乘法。同样,这种并行策略也缓解了原生低精度注意力中动态量化引入的停顿。
图 6: 增强并行性以实现高效Tensor Cores利用。(1) 新的warp布局设计减少了反量化停顿;(2) 协作式softmax利用GPU寄存器和共享内存之间的数据移动进行跨warp归约,开销极小。
(2)利用内存层次结构进行Warp同步。然而,由于计算结果分布在不同的寄存器和Warp中,原有的寄存器级softmax变得不可行。为解决此问题,我们利用寄存器和共享内存的多级内存层次结构,实现softmax计算的跨Warp归约和同步。如算法1所示,我们扩展了现有的高性能注意力算法(如FlashAttention),引入了两个额外的共享内存缓冲区:sTMP ∈ R^(Wn)和sAcc ∈ R^(Tm×Tn)。sTMP用于在计算softmax的行最大值时进行跨Warp归约。sAcc则临时存储在Tensor Core寄存器中计算出的注意力分数P,之后通过ldmatrix重新加载,以确保后续MMA操作的正确对齐。由于Wn通常很小,我们复用sTMP的共享内存指针给sAcc,以最小化内存开销。在Hopper架构上,WGMMA支持直接访问共享内存,无需显式地将数据从共享内存移动到寄存器。
本节描述BitDecoding的实现,如图7所示。它包含三个主要组件:(i) 支持多种注意力变体的查询转换组件;(ii) 执行低成本量化和打包的Residual Kernel;(iii) 具有细粒度流水线以充分利用Tensor Cores和CUDA Cores的Packing Kernel。最后,我们讨论了针对最新GPU架构的特定优化。
图 7: BitDecoding系统概览。(1) 查询转换重构查询张量布局,以在Tensor Cores上高效执行注意力变体。(2) 残差核以最小开销执行量化和打包,支持张量级和通道级缩放。(3) 打包核使用细粒度的异步流水线执行反量化和矩阵乘法,最大限度地利用Tensor Cores和CUDA Cores以及低比特参数。
为了支持GQA和MQA等多种注意力变体,我们执行查询转换。在解码阶段,查询长度Q_len = 1,导致查询张量批次维度很小,直接计算Q·K^T会使Tensor Cores利用率低下。为了解决这个问题,我们将查询张量的形状从 [1, (gq, hkv)] 重塑为 [gq, hkv],其中gq是查询头与KV头的共享比例。如图7(左)所示,这有效地形成了一个更大的Q分块,使得分组的查询头可以作为一个更大的GEMM块并行处理,从而充分填充Tensor Core片段,提高Warp占用率和吞吐量。
为了在不牺牲性能的情况下支持不同的量化算法(特别是张量级和通道级缩放),我们设计了Residual Kernel。该内核包含两个关键优化:
1. 基于残差块大小划分KV缓存:在预填充阶段,我们根据与Tensor Cores对齐的残差块大小Nr(见公式1)来划分KV缓存。大部分条目被量化并打包,剩余部分则存储在半精度的残差KV缓存中。在解码的每一步,新生成的K、V张量被追加到残差缓存中。一旦残差缓存填满达到Nr大小,Residual Kernel就会使用该半精度缓存计算注意力,并选择性地将其量化打包。这种划分方式自然地支持了在残差块内沿序列长度维度进行通道级量化,或沿隐藏维度进行张量级量化。
2. 使用Warp级指令优化归约:如图7(中)所示,计算完半精度KV数据后,它们以Tensor Cores的交错布局保留在寄存器中。为了高效计算量化参数(缩放因子和零点),我们首先执行线程级归约,然后在Warp内使用__shfl_xor_sync PTX指令进行高效的Warp级归约。如果Warp重复因子Wn > 1,则使用一个小的共享内存缓冲区来协调跨Warp的最终归约。计算完参数后,每个线程在寄存器内执行量化和打包,并将缩放因子和零点存储在紧凑的half2格式中,以实现高效的内存访问。
为了处理辅助元数据(缩放因子和零点)带来的内存流量增加,并协调CUDA Cores上的反量化和Tensor Cores上的计算,我们设计了一个细粒度的异步流水线。
1. 优化异步数据移动:
* 从全局内存到共享内存:我们遵循FlashAttention【【6,FlashAttention-2: Faster attention with better parallelism and work partitioning,2024】】的方法,采用分块平铺【【32,Tilelang: A composable tiled programming model for ai systems,2025】】和策略性重计算。我们为量化参数Kp和Vp引入了专用的共享内存缓冲区。所有全局到共享内存的传输都使用cp.async内在函数异步执行。我们根据数据重用性选择不同的缓存策略:cp.async.cg用于Q、Kpack和Vpack(不重用),cp.async.ca用于Kp和Vp(需要更细粒度的内存访问)。在Hopper架构上,我们使用tma.copy指令加载数据。
* 从共享内存到寄存器:我们使用ldmatrix PTX指令高效地将数据从共享内存加载到寄存器。为了消除bank冲突,我们使用了一种交错(sizzling)方案【【5,Cutlass: Cuda templates for linear algebra subroutines and solvers,2024】】,其定义为:
同时,我们重构了Kp和Vp的共享内存布局以进一步减少bank冲突。
i个切片在Tensor Cores上由mma处理时,第(i + 1)个切片同时从共享内存加载(ldmatrix)并进行反量化(Dequant)。这种生产者-消费者流持续进行,提高了指令吞吐量,并最大化了CUDA Cores和Tensor Cores的利用率。(1)解锁Hopper的Warpgroup加速能力。Hopper架构引入的wgmma指令要求矩阵乘法C = AB中的操作数B必须位于共享内存中。为了解决这个问题,我们利用Hopper的STSM PTX指令,将寄存器中反量化后的FP16值高效地存储到共享内存中,供wgmma_SS操作使用。wgmma的异步特性使得存储操作可以与计算重叠,从而优化性能。
(2)使用原生低精度格式加速Blackwell。Blackwell架构原生支持低精度张量操作,从而无需显式反量化。我们直接使用Blackwell的低精度mma指令(如支持mxfp4/nvfp4格式的指令)在打包的4位数据上执行GEMM。尽管这些指令对打包值及其块缩放因子有严格的布局约束,但我们在第四节-A中提出的布局转换策略是布局无关的,能够自动将打包的KV数据与硬件要求的格式对齐,确保与Blackwell原生张量流水线的无缝集成。
硬件配置:
软件配置:
模型与数据集:
实验设置:
Single(单批次长上下文)、Batches(大批次)、Page(使用paged attention的高吞吐量场景)。wgmma和异步内存指令,BitDecoding-v3相比FP16的FlashDecoding-v3实现了高达8.0倍的加速,远超v2版本的4.1倍。Single和Batches设置下,4位和2位版本分别取得了约4倍和超过7倍的加速。在GQA场景下,BitDecoding保持了3倍的加速,而仅依赖CUDA核的QServe性能下降至1.4倍,凸显了利用Tensor Cores的优势。
图 8: 在Blackwell架构上使用mxfp4的内核性能。
图 9: 在Hopper (H100)上的内核性能。
图 10: 在RTX4090上的内核性能。
图 11: 在A100上的内核性能。
Single设置下,当上下文长度达到128K时,BitDecoding实现了高达3.3倍的加速,而Kivi因缺乏分块内核支持而OOM。在Batches设置下,BitDecoding的吞吐量(900-1200 tokens/s)远高于Kivi(低于700 tokens/s)。Page设置下,QServe仅在MHA模型(LLaMA-2-7B)上表现尚可,但在所有GQA模型上性能均出现下降。相比之下,BitDecoding在所有模型和GPU配置下均稳定超越QServe,最大吞吐量是其2倍以上。
与非融合注意力(LLaMA-3.1-8B)的比较
图 12: 与Kivi比较(a)端到端生成时间和(b)解码吞吐量。
图 13: 与Qserve比较解码吞吐量。
开销分析:
性能分解分析 (图16):分解实验表明,BitDecoding的性能增益主要来源于三个方面:首先是自动诱导Tensor Core兼容布局的设计,其次是Warp并行策略带来的显著加速,最后是流水线优化进一步提升了端到端性能。
TABLE I: 低比特KV缓存的效率与准确性权衡。我们使用Llama-3.1-8B-Instruct,序列长度=32K,并在longbench上评估平均准确率。
TABLE II: 推理过程中量化和打包的延迟(ms)比较。
TABLE III: 协作式softmax和warps对性能和有效性的影响。
图 14: 残差KV缓存的运行时开销。
图 15: 反量化开销分析。
图 16: BitDecoding在各代架构上的优化分解。
a) KV缓存量化算法:近期的工作探索了4位、2位甚至1位的KV缓存量化。KIVI【【18,Kivi: A tuning-free asymmetric 2bit quantization for kv cache,2024】】、Gear【【13,Gear: An efficient kv cache compression recipefor nearlossless generative inference of llm,2024】】和KVQuant【【12,Kvquant: Towards 10 million context length llm inference with kv cache quantization,2024】】等方法使用逐通道量化来处理异常值,而RotateKV【【27,Rotatekv: Accurate and robust 2-bit kv cache quantization for llms via outlieraware adaptive rotations,2025】】则应用旋转来平滑通道分布。尽管这些算法在压缩率上有效,但它们缺乏高效的系统实现,导致性能不佳。
b) 混合精度矩阵乘法:低比特权重和低比特KV缓存对混合精度矩阵乘法(mpGEMM)提出了特殊要求。像Ladder【【33,Ladder: Enabling efficient {Low-Precision} deep learning computing through hardware-aware tensor transformation,2024】】和Marlin【【9,Marlin: Mixed-precision auto-regressive parallel inference on large language models,2024】】这样的优化内核通过布局转换和高效反量化提高了性能。然而,这些方法需要预打包和预转换权重,限制了它们在自回归解码中对动态低比特KV缓存的适用性。
c) 低比特KV缓存的系统实现:KIVI【【31,Triton: an intermediate language and compiler for tiled neural network computations,2019】】使用Triton和分离的内核来实现。Atom【【37,Atom: Low-bit quantization for efficient and accurate llm serving,2024】】将量化集成到前置的线性层中,而QServe【【16,Qserve: W4a8kv4 quantization and system co-design for efficient llm serving,2024】】则将量化直接融合到FlashAttention内核中。然而,后两者都依赖于使用FMA指令的GEMV操作,未能利用Tensor Core的加速能力。
BitDecoding通过展示如何利用有原则的系统设计来协同编排CUDA核与Tensor Cores,为高效的低比特KV缓存解码奠定了新的系统基础。其布局诱导和Warp级协调技术可推广到不同的注意力变体、量化方案和GPU代际,并能自然地扩展到如Blackwell等新兴架构乃至未来架构。我们期望BitDecoding能够推动未来在算法-系统协同设计、近无损测试时扩展以及更强大的GPU执行模型等方面的研究,以支持长上下文LLM的推理。
【1】Ainslie, J. et al. “Gqa: Training generalized multi-query transformer models from multi-head checkpoints,” 2023.
【5】NVIDIA Corporation, “Cutlass: Cuda templates for linear algebra subroutines and solvers,” 2024.
【6】Dao, T. “FlashAttention-2: Faster attention with better parallelism and work partitioning,” 2024.
Packing Kernel遵循其通过分块平铺和策略性重计算来管理数据移动的方法。【8】Fan, G. et al. “Warpdrive: Gpu-based fully homomorphic encryption acceleration leveraging tensor and cuda cores,” 2025.
【9】Frantar, E. et al. “Marlin: Mixed-precision auto-regressive parallel inference on large language models,” 2024.
【10】Grattafiori, A. et al. “The llama 3 herd of models,” 2024.
【12】Hooper, C. et al. “Kvquant: Towards 10 million context length llm inference with kv cache quantization,” 2024.
【13】Kang, H. et al. “Gear: An efficient kv cache compression recipefor nearlossless generative inference of llm,” 2024.
【14】Kim, Y. J. et al. “Who says elephants can’t run: Bringing large scale moe models into cloud scale production,” 2022.
【16】Lin, Y. et al. “Qserve: W4a8kv4 quantization and system co-design for efficient llm serving,” 2024.
【17】Liu, A. et al. “Deepseek-v3 technical report,” 2024.
【18】Liu, Z. et al. “Kivi: A tuning-free asymmetric 2bit quantization for kv cache,” 2024.
【24】Sandokji, S. et al. “A survey of techniques for warp scheduling in gpus,” 2015.
【25】Shah, J. et al. “Flashattention-3: Fast and accurate attention with asynchrony and lowprecision,” 2024.
【26】Shazeer, N. “Fast transformer decoding: One write-head is all you need,” 2019.
【27】Su, Z. et al. “Rotatekv: Accurate and robust 2-bit kv cache quantization for llms via outlieraware adaptive rotations,” 2025.
【31】Tillet, P. et al. “Triton: an intermediate language and compiler for tiled neural network computations,” 2019.
【32】Wang, L. et al. “Tilelang: A composable tiled programming model for ai systems,” 2025.
【33】Wang, L. et al. “Ladder: Enabling efficient {Low-Precision} deep learning computing through hardware-aware tensor transformation,” 2024.
【34】Yang, A. et al. “Qwen3 technical report,” 2025.
【36】Zhang, T. et al. “Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization,” 2024.
【37】Zhao, Y. et al. “Atom: Low-bit quantization for efficient and accurate llm serving,” 2024.