Michael Brenndoerfer | Data & AI, Private Equity, Technology

Michael Brenndoerfer

KV Cache Compression › Cache Quantization

Michael BrenndoerferJanuary 9, 202612 min read

Language AI HandbookAll chapters

Part 42 of 67Inference Optimization

Chapter 316KV Cache Compression › Cache Quantization

Reading section 7 of 13 · Cache Quantization · 12 minSection 7 of 13 · 12 min

Section 7 of 13 · 12 min read

KV Cache Compression › Cache QuantizationOne section · Sections in this chapter7 / 13 ↓

Sections in this chapter13

Reading tools Definitions & inline links

Choose your expertise level to adjust how many terms are explained. Beginners see more tooltips, experts see fewer to maintain reading flow. Hover over underlined terms for instant definitions.

Article links

Make inline references clickable

Orthogonal to eviction strategies, cache quantization reduces memory by storing key and value vectors at lower numerical precision. While model weights typically use FP16 or BF16 during inference, the cached activations can often tolerate even lower precision without significant quality loss. This approach complements eviction: rather than removing tokens entirely, we keep all tokens but represent each one using fewer bits.

Think of quantization as converting high-resolution photographs to compressed JPEGs. The compressed versions are smaller and may have barely perceptible artifacts, but they preserve the needed content. For our purposes, the "needed content" of a cached key vector is its direction in high-dimensional space, not its precise magnitude. Attention computation depends on dot products, which are dominated by directional alignment. Modest rounding errors in the vector components rarely change which tokens receive the highest attention.

The mathematical foundation of quantization is straightforward. We need to map continuous floating-point values to a discrete set of integers that can be stored compactly. The simplest form is uniform quantization, which divides the value range into equally-spaced bins and maps each floating-point value to the nearest bin.

To derive the quantization formula, suppose we have a floating-point value $x$ drawn from the range $[x_{\min}, x_{\max}]$ . We want to represent it as an integer $q$ in the range $[0, 2^b - 1]$ , where $b$ is the bit width. The linear mapping that achieves this is:

q = \text{round}\left(\frac{x - x_{\min}}{x_{\max} - x_{\min}} \cdot (2^b - 1)\right)

where:

$q$ : the resulting quantized integer index, which will be stored in place of the original floating-point value
$x$ : the original floating-point value that we wish to compress
$x_{\min}$ : the minimum value in the group being quantized, establishing the lower bound of the value range
$x_{\max}$ : the maximum value in the group being quantized, establishing the upper bound of the value range
$b$ : the target bit width (e.g., 8 for INT8 or 4 for INT4), determining how many discrete levels are available
$2^b - 1$ : the total number of discrete quantization levels, which equals the maximum integer value representable with $b$ bits

This transformation accomplishes its goal through a sequence of intuitive steps. First, it normalizes the original value $x$ to the range $[0, 1]$ by computing the fraction $\frac{x - x_{\min}}{x_{\max} - x_{\min}}$ . This fraction represents where $x$ falls within the overall range of values. Second, it scales this normalized value to the full integer range $[0, 2^b - 1]$ by multiplication. Finally, it rounds to the nearest integer to obtain a discrete representation that can be stored efficiently. This approach preserves the relative ordering and approximate magnitudes of values while dramatically reducing storage requirements.

When we need to use the cached values for computation, we must reverse this process. Dequantization recovers an approximation of the original value:

\hat{x} = \frac{q}{2^b - 1} \cdot (x_{\max} - x_{\min}) + x_{\min}

where:

$\hat{x}$ : the reconstructed floating-point value, which approximates the original $x$
$q$ : the stored quantized integer retrieved from memory
$b$ : the target bit width used during quantization
$x_{\min}$ : the minimum value from the original group, serving as the offset that restores the correct baseline
$x_{\max} - x_{\min}$ : the dynamic range of the original values, used to restore the correct magnitude

The reconstruction process precisely reverses the quantization steps. It first converts the integer $q$ back to a proportion in $[0, 1]$ by dividing by the maximum integer value. It then maps that proportion onto the original dynamic range by multiplying by $(x_{\max} - x_{\min})$ . Finally, it adds the offset $x_{\min}$ to restore the correct baseline. The reconstructed value $\hat{x}$ will differ from the original $x$ by at most half a quantization step, introducing a small but bounded error. With 8 bits, this maximum error is approximately $\frac{x_{\max} - x_{\min}}{510}$ , which for a typical value range of $[-3, 3]$ is about 0.012, far smaller than the typical variation in key vectors.

In[18]:

Code

import torch


class QuantizedCache:
    """
    KV cache with quantized storage.
    Stores keys and values at reduced precision.
    """

    def __init__(self, bits=8):
        self.bits = bits
        self.max_val = 2**bits - 1

        # Quantized storage
        self.keys_quantized = None
        self.values_quantized = None

        # Quantization parameters (per-tensor)
        self.key_min = None
        self.key_scale = None
        self.value_min = None
        self.value_scale = None

    def quantize(self, tensor):
        """Quantize tensor to self.bits precision."""
        t_min = tensor.min()
        t_max = tensor.max()
        scale = (t_max - t_min) / self.max_val

        # Avoid division by zero
        if scale == 0:
            scale = torch.tensor(1.0)

        quantized = torch.round((tensor - t_min) / scale).to(torch.uint8)
        return quantized, t_min, scale

    def dequantize(self, quantized, t_min, scale):
        """Recover approximate float tensor."""
        return quantized.float() * scale + t_min

    def store(self, keys, values):
        """Store key-value pairs in quantized format."""
        self.keys_quantized, self.key_min, self.key_scale = self.quantize(keys)
        self.values_quantized, self.value_min, self.value_scale = self.quantize(
            values
        )

    def retrieve(self):
        """Retrieve dequantized key-value pairs."""
        keys = self.dequantize(
            self.keys_quantized, self.key_min, self.key_scale
        )
        values = self.dequantize(
            self.values_quantized, self.value_min, self.value_scale
        )
        return keys, values

    def memory_usage(self, original_dtype=torch.float16):
        """Compare memory usage."""
        if self.keys_quantized is None:
            return 0, 0

        # Quantized: uint8 = 1 byte per element
        quantized_bytes = (
            self.keys_quantized.numel() + self.values_quantized.numel()
        )

        # Original: 2 bytes for FP16
        original_bytes = (
            self.keys_quantized.numel() + self.values_quantized.numel()
        ) * 2

        return quantized_bytes, original_bytes

In[19]:

Code

import torch.nn.functional as F


def evaluate_quantization_error(keys, values, bits_list=[8, 4, 2]):
    """Measure reconstruction error at different bit widths."""
    results = []

    for bits in bits_list:
        cache = QuantizedCache(bits=bits)
        cache.store(keys, values)

        keys_recovered, values_recovered = cache.retrieve()

        # Compute mean squared error
        key_mse = F.mse_loss(keys_recovered, keys).item()
        value_mse = F.mse_loss(values_recovered, values).item()

        # Compute cosine similarity (important for attention)
        key_cos = F.cosine_similarity(
            keys.flatten(), keys_recovered.flatten(), dim=0
        ).item()

        quant_bytes, orig_bytes = cache.memory_usage()

        results.append(
            {
                "bits": bits,
                "key_mse": key_mse,
                "value_mse": value_mse,
                "key_cosine_sim": key_cos,
                "compression_ratio": orig_bytes / quant_bytes
                if quant_bytes > 0
                else 0,
            }
        )

    return results

In[20]:

Code

import torch

# Test with realistic cache data
batch_size, num_heads, seq_len, head_dim = 1, 32, 512, 128
keys = torch.randn(batch_size, num_heads, seq_len, head_dim)
values = torch.randn(batch_size, num_heads, seq_len, head_dim)

results = evaluate_quantization_error(keys, values, bits_list=[8, 4, 2])

Out[21]:

Console

Quantization Analysis:
-----------------------------------------------------------------
  Bits      Key MSE    Value MSE   Key Cosine  Compression
-----------------------------------------------------------------
     8     0.000121     0.000125       1.0000         2.0x
     4     0.034811     0.036124       0.9833         2.0x
     2     1.031245     1.070498       0.8037         2.0x

INT8 quantization achieves 2x compression with minimal error, maintaining cosine similarity above 0.99. This high cosine similarity is particularly important because attention operates on dot products between queries and keys. The dot product is fundamentally a measure of similarity, and if the quantized keys preserve high cosine similarity with their original values, the resulting attention scores will be nearly identical to what they would have been without quantization. The model will attend to approximately the same tokens in approximately the same proportions, which is exactly what we need for quality preservation.

Out[22]:

Visualization

Histogram of INT8 quantization errors centered tightly around zero, showing minimal reconstruction error. — INT8 quantization error distribution showing tight concentration around zero, confirming high fidelity reconstruction at 2x compression. The narrow spread indicates that individual element errors are very small relative to typical value magnitudes.

Histogram of INT4 quantization errors with moderate spread around zero, showing increased reconstruction error. — INT4 quantization error distribution showing moderate spread around zero, representing the accuracy trade-off at 4x compression. The wider distribution compared to INT8 reflects the reduced number of discrete levels available.

Histogram of INT2 quantization errors with wide spread around zero, showing substantial reconstruction error. — INT2 quantization error distribution showing substantial spread around zero, illustrating severe precision loss at 8x compression. With only 4 discrete levels, large reconstruction errors are unavoidable for most values.

Even INT4 quantization, which provides 4x compression, often preserves sufficient similarity for acceptable generation quality. However, quality degradation becomes noticeable for precision-sensitive tasks at this level. The trade-off between compression and accuracy depends heavily on the specific application: chat applications may tolerate INT4 well, while code generation or mathematical reasoning may require the precision of INT8 or higher.

Per-Channel QuantizationLink Copied

The per-tensor quantization approach we examined computes global statistics across the entire tensor, using a single minimum and scale value for all elements. This approach is simple but can be suboptimal when different attention heads or dimensions have varying value ranges. If one head has values concentrated in a narrow range while another has a wide range, using global statistics forces the narrow-range head to use only a fraction of the available quantization levels, wasting precision.

Think of it this way: if you use a ruler calibrated in meters to measure both the height of a building and the width of a fingernail, you will get very imprecise measurements for the fingernail. Per-channel quantization is like giving each head its own appropriately-scaled ruler.

Per-channel quantization improves accuracy by computing separate scales for each head or channel. Each attention head gets its own minimum and scale values, letting it to use the full range of quantization levels regardless of what other heads are doing. This independence is particularly valuable because different attention heads often specialize in different types of patterns and may exhibit quite different value distributions. A head that specializes in syntactic dependencies may produce keys with very different statistical properties than a head that tracks coreference chains.

The additional overhead of per-channel quantization is modest: instead of storing two scalar parameters (min and scale) per tensor, we store two vectors with one entry per head. For a typical 32-head model, this is 64 additional float values per layer, negligible compared to the tensor itself. The quality improvement, however, can be substantial, particularly at aggressive bit widths where efficient use of the quantization range makes the difference between acceptable and unacceptable quality.

In[23]:

Code

import torch


class PerChannelQuantizedCache:
    """
    KV cache with per-head quantization for better accuracy.
    """

    def __init__(self, bits=8):
        self.bits = bits
        self.max_val = 2**bits - 1

        self.keys_quantized = None
        self.values_quantized = None
        self.key_params = None  # (min, scale) per head
        self.value_params = None

    def quantize_per_head(self, tensor):
        """
        Quantize with separate parameters per attention head.
        tensor shape: (batch, heads, seq_len, head_dim)
        """
        batch_size, num_heads, seq_len, head_dim = tensor.shape

        # Compute min/max per head
        # Reshape to (batch * heads, seq_len * head_dim) for per-head stats
        reshaped = tensor.view(batch_size * num_heads, -1)
        t_min = reshaped.min(dim=1, keepdim=True)[0]
        t_max = reshaped.max(dim=1, keepdim=True)[0]

        scale = (t_max - t_min) / self.max_val
        scale = torch.where(scale == 0, torch.ones_like(scale), scale)

        # Quantize
        quantized = torch.round((reshaped - t_min) / scale)
        quantized = quantized.clamp(0, self.max_val).to(torch.uint8)
        quantized = quantized.view(batch_size, num_heads, seq_len, head_dim)

        return quantized, (t_min, scale)

    def dequantize_per_head(self, quantized, params):
        """Dequantize using per-head parameters."""
        t_min, scale = params
        batch_size, num_heads, seq_len, head_dim = quantized.shape

        reshaped = quantized.view(batch_size * num_heads, -1).float()
        dequantized = reshaped * scale + t_min

        return dequantized.view(batch_size, num_heads, seq_len, head_dim)

    def store(self, keys, values):
        self.keys_quantized, self.key_params = self.quantize_per_head(keys)
        self.values_quantized, self.value_params = self.quantize_per_head(
            values
        )

    def retrieve(self):
        keys = self.dequantize_per_head(self.keys_quantized, self.key_params)
        values = self.dequantize_per_head(
            self.values_quantized, self.value_params
        )
        return keys, values

In[24]:

Code

import torch.nn.functional as F

# Compare per-tensor vs per-channel quantization
cache_per_tensor = QuantizedCache(bits=4)
cache_per_channel = PerChannelQuantizedCache(bits=4)

cache_per_tensor.store(keys, values)
cache_per_channel.store(keys, values)

keys_pt, values_pt = cache_per_tensor.retrieve()
keys_pc, values_pc = cache_per_channel.retrieve()

mse_pt = F.mse_loss(keys_pt, keys).item()
mse_pc = F.mse_loss(keys_pc, keys).item()

cos_pt = F.cosine_similarity(keys.flatten(), keys_pt.flatten(), dim=0).item()
cos_pc = F.cosine_similarity(keys.flatten(), keys_pc.flatten(), dim=0).item()

mse_reduction_pct = (1 - mse_pc / mse_pt) * 100

Out[25]:

Console

INT4 Quantization Comparison:
  Per-tensor:  MSE = 0.034811, Cosine = 0.9833
  Per-channel: MSE = 0.027155, Cosine = 0.9871

  Per-channel reduces MSE by 22.0%

Per-channel quantization significantly reduces error at the same bit width because each attention head can use its full quantization range independently. The improvement is particularly pronounced at lower bit widths where the limited number of quantization levels makes efficient range utilization necessary. This technique allows INT4 per-channel quantization to approach the accuracy of INT8 per-tensor quantization in many cases, effectively doubling the compression while maintaining similar quality.