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

TL;DR

BitDecoding addresses the memory and bandwidth bottlenecks of long-context LLM decoding by co-designing a system that cooperatively uses Tensor Cores and CUDA Cores to decode low-bit KV caches. It introduces a layout-inducing, warp-aligned architecture with Residual and Packing kernels, plus a query transformation, enabling mixed-precision execution across GQA/MQA variants and multiple GPU generations. The approach delivers up to 8.6x kernel speedups and 3x end-to-end latency reductions on long-context models, while maintaining near-FP16 accuracy with 4-bit KV caches. This work demonstrates strong practical gains for long-context inference and provides open-source tooling to spur further algorithm–system co-design in KV-cache quantization.

Abstract

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.

Figures (16)

• Figure 1: Comparison of mixed-precision matrix multiplication for low-bit weight and low-bit KV cache. (a) Quantized weights can be preprocessed offline. (b) KV cache requires online quantization and packing for each newly generated token.
• Figure 2: Comparison of different low-bit KV cache systems against half-precision FlashAttention. Each system follows the attention formulation $\mathrm{Out} = \mathrm{softmax}(Q\,\mathcal{D}(K'^\top))\,\mathcal{D}(V')$, where $K'$ and $V'$ are low-bit quantized Key and Value tensors, and $\mathcal{D}(\cdot)$ denotes the dequantization function.
• Figure 3: (a) mma.m16n8k16 fragment layout for matrix B. Each thread ($T_i$) is assigned a specific set of values based on the instruction-defined interleaved mapping. (b) For INT4, quantization packs values contiguously per thread. After dequantization, the layout misaligns with the expected interleaved pattern.
• Figure 4: (a) A single warp along $N$ for register-level operations will experience stalls due to dequantization (DQ) (b) Micro-level comparision with and without dequantization.
• Figure 5: Overview of methods for optimizing low-bit layout on Tensor Cores. (1) Fused computation and quantization within Tensor Cores fragments. (2) The low-bit packing data preserves FP16 values. (3) Low-bit Layout matches with the dequantized half-precision layout. (4) Layout remapping for faster dequantization.