kv-cachequantizationkivillm-inferencevllm

KIVI: Asymmetric 2-Bit KV Cache Quantization Guide

September 20, 2026

KIVI Architecture: Asymmetric Per-Channel Key and Per-Token Value Cache Quantization with Residual FP16 Window

The Problem: The KV Cache Memory Explosion in Long-Context LLMs

While post-training weight quantization (such as AWQ, GPTQ, or FP4) successfully compresses static model parameters, enterprise AI systems face a far more aggressive memory bottleneck during inference: the dynamic Key-Value (KV) cache.

During autoregressive generation, the model caches the Key and Value representations of all previous tokens to avoid recomputing past context. As context lengths expand to 32,000, 64,000, and 128,000 tokens in retrieval-augmented generation (RAG) and document analysis, the KV cache size grows linearly:

$$\text{KV Cache Size (Bytes)} = 2 \times 2 \times L \times H \times D \times T$$

Where $L$ is the number of transformer layers, $H$ is the number of attention heads, $D$ is the head dimension, and $T$ is the sequence length.

For a 70B parameter model serving a 128k-token sequence in 16-bit precision, the KV cache consumes over 80 gigabytes of GPU VRAM—exceeding the size of the 4-bit quantized model weights (~35GB) for a single concurrent user.

Why Standard Uniform Quantization Fails on KV Caches

Historically, attempts to quantize the KV cache to 4 bits or 2 bits caused catastrophic perplexity degradation. Detailed empirical analysis revealed that Keys and Values exhibit fundamentally conflicting outlier distributions:

Keys Exhibit Channel-Wise Outliers: Specific hidden feature channels (often just 1 or 2 dimensions across the head) persistently maintain extreme numerical magnitudes across the entire sequence. If Keys are quantized per-token, the massive scale of these outlier channels collapses the numerical resolution of all remaining channels for that token, destroying attention selectivity.
Values Exhibit Token-Wise Outliers: In contrast to Keys, Values do not display persistent outlier channels. Instead, specific tokens (such as sentence boundaries, punctuation, or pivotal entity nouns) trigger large activation bursts across all channels. If Values are quantized per-channel, the scale factor is distorted along the entire sequence length, corrupting normal tokens.

Applying a uniform quantization scheme across both tensors forces coarse scaling factors that destroy model reasoning at low bit-widths.

The Architecture: KIVI’s Asymmetric Quantization Strategy

To compress the KV cache down to 2 bits without fine-tuning, the KIVI algorithm introduces an asymmetric quantization framework that aligns quantization axes directly with the physical outlier geometry of the attention tensors.

Rather than treating Keys and Values symmetrically, KIVI decouples their quantization directions:

KIVI Memory Topology: Streaming Sequence Ingestion, Residual FP16 Window, Per-Channel Key Quantization, and Per-Token Value Packing

The Three Structural Innovations of KIVI

Per-Channel Key Quantization: Because Key outliers are confined to specific channel dimensions across tokens, KIVI groups Key elements along the sequence dimension ($T$) for each individual channel ($d$). By computing scale and zero-point parameters per channel, outlier channels retain dedicated dynamic ranges without diminishing the precision of neighboring channels.
Per-Token Value Quantization: Because Value outliers occur at specific token timestamps across hidden dimensions, KIVI groups Value elements along the hidden channel dimension ($d$) for each individual token ($t$). The scale and zero-point parameters encapsulate the token’s unique burst magnitude without corrupting historical sequence states.
The Streaming Residual FP16 Window: Local attention context (the most recent $W$ tokens, typically $W = 32$ to $128$) dominates attention weight distribution during immediate next-token prediction. KIVI maintains a small rolling FIFO buffer in native FP16/BF16. Once newly generated tokens slide past the window, they are grouped and packed into low-bit integer storage, preserving full precision where it matters most.

Comparing KV Cache Management Paradigms

Strategy	Precision	128k Context Memory (70B)	PPL Degradation	Retention Capability
Standard FP16	16-bit float	~80 GB VRAM	0.0% (Baseline)	Full context retention
Uniform INT8	8-bit int	~40 GB VRAM	< 0.1%	Full context retention
Uniform INT4	4-bit int	~20 GB VRAM	Moderate (> 1.5%)	Severe needle-retrieval errors
Sliding Window / Sinks	16-bit float	~4 GB to 8 GB VRAM	High on long reasoning	Evicts middle tokens permanently
KIVI INT2 / INT4	2-bit / 4-bit asym	~10 GB to 20 GB VRAM	< 0.1% (Lossless)	Full context retention (100% NIAH)

Implementation: Simulating KIVI Asymmetric Quantization in PyTorch

Let us implement a complete, standalone Python implementation demonstrating KIVI’s per-channel Key quantization, per-token Value quantization, and residual window management.

Step 1: Implement the Quantization and Packing Primitives

import torch
from typing import Tuple

def quantize_per_channel_keys(
    keys: torch.Tensor,
    num_bits: int = 2
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Quantizes Key tensors along the sequence dimension (dim=2) for each channel (dim=3).
    keys shape: [batch, heads, seq_len, head_dim]
    """
    q_max = (1 << num_bits) - 1 # 3 for 2-bit (0, 1, 2, 3)

    # Compute min and max along the sequence dimension
    min_val = torch.min(keys, dim=2, keepdim=True)[0]
    max_val = torch.max(keys, dim=2, keepdim=True)[0]

    scale = (max_val - min_val) / float(q_max)
    scale = torch.clamp(scale, min=1e-7)
    zero_point = torch.round(-min_val / scale)

    # Quantize: q = round(x / scale) + zero_point
    q_keys = torch.clamp(torch.round(keys / scale) + zero_point, 0, q_max).to(torch.uint8)

    return q_keys, scale, zero_point

def quantize_per_token_values(
    values: torch.Tensor,
    num_bits: int = 2
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """
    Quantizes Value tensors along the head dimension (dim=3) for each token (dim=2).
    values shape: [batch, heads, seq_len, head_dim]
    """
    q_max = (1 << num_bits) - 1

    # Compute min and max along the channel/head_dim dimension
    min_val = torch.min(values, dim=3, keepdim=True)[0]
    max_val = torch.max(values, dim=3, keepdim=True)[0]

    scale = (max_val - min_val) / float(q_max)
    scale = torch.clamp(scale, min=1e-7)
    zero_point = torch.round(-min_val / scale)

    q_values = torch.clamp(torch.round(values / scale) + zero_point, 0, q_max).to(torch.uint8)

    return q_values, scale, zero_point

Step 2: Implement Dequantization for Attention Computation

def dequantize_keys(
    q_keys: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor
) -> torch.Tensor:
    """Dequantizes 2-bit/4-bit Keys back to floating-point during attention scoring."""
    return (q_keys.to(torch.float32) - zero_point) * scale

def dequantize_values(
    q_values: torch.Tensor,
    scale: torch.Tensor,
    zero_point: torch.Tensor
) -> torch.Tensor:
    """Dequantizes 2-bit/4-bit Values back to floating-point during attention scoring."""
    return (q_values.to(torch.float32) - zero_point) * scale

Step 3: Implement the KIVI KV Cache Manager with Rolling Window

class KIVIKVCacheManager:
    def __init__(
        self,
        num_layers: int = 32,
        num_heads: int = 32,
        head_dim: int = 128,
        residual_window: int = 64,
        num_bits: int = 2
    ):
        self.num_layers = num_layers
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.residual_window = residual_window
        self.num_bits = num_bits

        # In-memory storage structures
        self.fp16_key_window: List[torch.Tensor] = []
        self.fp16_val_window: List[torch.Tensor] = []

        # Quantized storage for older tokens
        self.quantized_keys: Optional[torch.Tensor] = None
        self.key_scales: Optional[torch.Tensor] = None
        self.key_zeros: Optional[torch.Tensor] = None

        self.quantized_values: Optional[torch.Tensor] = None
        self.val_scales: Optional[torch.Tensor] = None
        self.val_zeros: Optional[torch.Tensor] = None

    def append_token_step(self, new_key: torch.Tensor, new_val: torch.Tensor):
        """
        Appends a newly generated token representation [batch, heads, 1, head_dim].
        Once window exceeds threshold, packs older tokens into 2-bit quantized store.
        """
        self.fp16_key_window.append(new_key)
        self.fp16_val_window.append(new_val)

        if len(self.fp16_key_window) > self.residual_window:
            # Pop oldest token to quantize
            old_k = self.fp16_key_window.pop(0)
            old_v = self.fp16_val_window.pop(0)

            # In production kernels, tokens are batched into groups of 32 for efficient bit-packing
            q_k, k_scale, k_zp = quantize_per_channel_keys(old_k, self.num_bits)
            q_v, v_scale, v_zp = quantize_per_token_values(old_v, self.num_bits)

            # Concatenate into quantized historical cache
            if self.quantized_keys is None:
                self.quantized_keys, self.key_scales, self.key_zeros = q_k, k_scale, k_zp
                self.quantized_values, self.val_scales, self.val_zeros = q_v, v_scale, v_zp
            else:
                self.quantized_keys = torch.cat([self.quantized_keys, q_k], dim=2)
                self.quantized_values = torch.cat([self.quantized_values, q_v], dim=2)

    def calculate_memory_savings(self, total_seq_len: int) -> Tuple[float, float, float]:
        """Calculates memory usage in megabytes comparing FP16 against KIVI."""
        bytes_per_token_fp16 = 2 * 2 * self.num_layers * self.num_heads * self.head_dim
        fp16_mb = (bytes_per_token_fp16 * total_seq_len) / (1024 * 1024)

        # KIVI: residual window in FP16, rest in 2-bit (0.25 bytes) + scale metadata
        quantized_tokens = max(0, total_seq_len - self.residual_window)
        bytes_per_token_kivi = 2 * 0.25 * self.num_layers * self.num_heads * self.head_dim
        kivi_mb = (
            (bytes_per_token_fp16 * min(total_seq_len, self.residual_window)) +
            (bytes_per_token_kivi * quantized_tokens)
        ) / (1024 * 1024)

        savings_pct = (1.0 - (kivi_mb / fp16_mb)) * 100
        return fp16_mb, kivi_mb, savings_pct

Step 4: Run the Memory Benchmark

cache_manager = KIVIKVCacheManager(
    num_layers=32,
    num_heads=32,
    head_dim=128,
    residual_window=64,
    num_bits=2
)

seq_len = 131072 # 128k context
fp16_mb, kivi_mb, savings = cache_manager.calculate_memory_savings(seq_len)

print(f"[KIVI KV Cache Benchmark: 128k Sequence Length]")
print(f"  -> Baseline FP16 KV Cache: {fp16_mb / 1024:.2f} GB VRAM")
print(f"  -> KIVI 2-Bit Quantized Cache: {kivi_mb / 1024:.2f} GB VRAM")
print(f"  -> Net VRAM Reduction: {savings:.1f}%")

Production Serving and vLLM Deployment Trade-offs in 2026

When deploying KIVI in enterprise inference engines (such as vLLM, SGLang, and TensorRT-LLM):

Zero Fine-Tuning Overhead: Unlike Quantization-Aware Training (QAT), KIVI requires zero training or calibration datasets. It functions as a drop-in architectural modification during the PagedAttention memory allocation phase.
Decoding Latency and Memory Bandwidth Arbitrage: During autoregressive decoding, generation speed is strictly memory-bandwidth bound. Transporting 2-bit integer weights across the GPU bus requires 75% less bandwidth than FP16. Even with on-the-fly dequantization in SRAM, KIVI delivers up to 2.5x higher token generation speeds under high concurrent batch loads.
Hardware Synergy with 4-Bit Weights: Combining 4-bit weight quantization (e.g., AWQ/FP4) with KIVI 2-bit/4-bit KV caching allows engineering teams to deploy a complete 70B parameter model with 128k active context onto a single 48GB GPU workstation (such as an NVIDIA RTX 6000 Ada or dual RTX 4090s), eliminating the need for expensive multi-node GPU clusters.

Conclusion

KIVI breaks through the long-context memory wall by aligning quantization mathematics with the physical geometry of attention tensors. By quantizing Keys per-channel and Values per-token while preserving a rolling FP16 window, KIVI compresses the dynamic KV cache by up to 75%—delivering higher inference throughput and expanding context boundaries with zero degradation in reasoning capability.