KV Cache Quantization and Per-Token Scale Placement

The KV cache is the part of a transformer’s memory that grows with the conversation. Weights are a fixed cost per model; the cache is a cost per sequence per token per layer, and it is multiplied by the number of concurrent requests. For long contexts and large batches the cache can exceed the weights in size, which makes its precision a first-order decision rather than an afterthought.

Quantizing it is not the same problem as quantizing weights. Weights are static, calibrated offline, and read once per forward pass. Cache entries are written incrementally, read repeatedly, and their distribution depends on what has been said so far.

What the cache holds and how it is read

Each attention layer stores, for every processed token, a key vector and a value vector of dimension $d_h$ per head. During generation, a new token’s query is compared against every stored key, the scores are normalized by softmax, and the result weights a sum over the stored values. The keys are read once per new token, across the whole history; the values likewise.

The read pattern determines the arithmetic intensity. For a new token at position $t$, the attention operation reads $2 t d_h$ elements per head and performs about $2 t d_h$ multiply-adds, so the intensity is roughly one operation per byte at 8 bits, and two operations per byte at 4 bits. That is far below any processor’s balance, which makes attention during decode a memory-bound operation whose cost is dominated by the cache bytes. Halving the cache precision halves the bytes and, if the kernel is efficient, roughly halves the attention time.

Unlike the weight matrices, the cache is read by a kernel that has no reuse across a large tile: each key is used with exactly one query at a time during decode. There is no equivalent of the batch dimension that rescues weight-only quantization at large batch sizes, because batching adds sequences, not queries against the same key. This is why cache quantization produces a speedup that does not evaporate as concurrency rises, up to the point where the attention arithmetic becomes the limit.

Keys and values are not symmetric

The two tensors enter the computation differently, and the difference dictates how each should be quantized.

Keys are multiplied by the query and the result goes through softmax. The softmax is exponentially sensitive to its input: a change of one unit in a score changes the corresponding weight by a factor of $e$. The scores themselves are the dot product of a query with a key, so an error in the key becomes an error in the score scaled by the query’s norm. A key quantization error that is small relative to the key’s own magnitude can still be significant relative to the spread of scores, because the spread of scores is what softmax amplifies.

Values are summed with weights that sum to one. An error in a value contributes to the output in proportion to its attention weight, and the output is a convex combination. There is no exponential amplification. Value errors average out across the weighted sum in the way that ordinary quantization errors do.

The empirical consequence, reported consistently in cache-quantization work, is that keys tolerate less error than values at the same bit width, and that the sensitivity is concentrated in a small number of channels. A scheme that quantizes both at the same width and the same axis leaves accuracy on the table in one and spends too much on the other.

Per-token scales and where they can live

A cache entry for one token is a vector of $d_h$ values per head. The natural granularity is per token per head: one scale for the $d_h$ elements written at that position. This is a per-token scale in the sense used for activations, and it has the same motivation, which is that the magnitude of a newly written entry varies with the token.

The alternative is per-channel, where the scale is shared across all tokens for a given feature index. Per-channel scales can be computed from a running estimate that updates as the cache grows, but the scale for a channel must then be applied to entries written before the estimate was final, which means either re-quantizing the cache (expensive and disruptive) or storing the entries in floating point until the estimate stabilizes.

Per-token scales avoid that problem entirely: the scale is computed from the vector being written, at write time, and stored alongside it. The cache becomes a sequence of independently quantized vectors, each with its own pair of parameters. Reading is a per-vector dequantize, which the kernel does as it loads.

The cost is the metadata: for a head dimension of 128 and 8-bit entries, a per-token scale of 4 bytes adds 3 percent; a scale and zero point add 6 percent. At 4-bit entries the same 4-byte scale is 6 percent, and at 2 bits it is 12 percent. The metadata fraction is the reason cache-quantization schemes at 2 bits usually store scales at 8 bits relative to a per-channel or per-layer super-scale, or group several tokens under one scale.

There is a second cost that does not appear in the byte count. A per-token scale must be fetched with its vector, so the read is two streams rather than one, and the kernel must apply a scale per $d_h$ elements rather than per tile. On hardware where the attention kernel is written to stream keys into shared memory and process a tile of them at once, per-token scales mean the dequantization cannot be hoisted out of the tile loop. Per-channel scales, applied once per feature across the tile, can be hoisted.

The axis choice is a kernel decision

The two axes produce different kernels, and the difference is not cosmetic.

Per-token quantization means each stored vector is scaled independently. A kernel that processes a tile of $T$ keys loads $T$ vectors and $T$ scales, dequantizes each vector with its own scale, and computes $T$ dot products against the query. The dot product is over the head dimension, so the scale can be applied after the dot product only if the scale is uniform across the head dimension, which it is: a per-token scale multiplies the whole score by a constant. That is convenient, because the score is a scalar and the scale can be folded into it after the reduction. The dot product itself runs on the integer values.

Per-channel quantization means the scale varies along the head dimension, which is the contraction axis of the score dot product. A per-channel scale cannot be pulled out of the sum, so the products must be scaled individually, which forces either floating-point arithmetic in the inner loop or a fixed-point approximation. This is the same structural argument that makes per-output-channel scales the right choice for weights and per-channel scales the wrong choice for activations, appearing here along a different axis.

The consequence is that per-token is the axis the kernel prefers, and per-channel is the axis the statistics prefer for stability. Schemes that use per-channel scales for keys usually apply them in a preprocessing step that rescales the query instead, moving the per-channel factor onto the query where it is applied once per new token rather than once per cached entry. The rescaling is exact if the query is not itself quantized, and it costs a vector multiply per head per token.

def quantize_kv_vector(vec, qmax=127):
    """Symmetric per-token quantization of one head's key or value vector."""
    amax = max(abs(float(v)) for v in vec)
    if amax == 0.0:
        return [0] * len(vec), 1.0
    scale = amax / qmax
    q = [max(-qmax, min(qmax, round(float(v) / scale))) for v in vec]
    return q, scale

def attention_scores(q_vec, keys, scales):
    """Scores with per-token scales folded in after the integer dot product."""
    out = []
    for k, s in zip(keys, scales):
        acc = sum(qi * ki for qi, ki in zip(q_vec, k))
        out.append(acc * s)          # one multiply per score, not per element
    return out

What degrades first

The first casualty is the softmax’s discrimination between similar scores. When two keys are nearly equally relevant, the difference in their scores decides the attention weight, and that difference is a small number computed as the difference of two large dot products. Quantization error in the keys enters both scores and partially cancels, but only partially: the errors are independent across the two keys, so the difference carries the sum of two error terms while the true difference may be much smaller than either score. At 8 bits the effect is usually negligible; at 4 bits it shows up as attention distributions that are flatter than the reference, and at 2 bits as attention that misses the correct position.

The second casualty is the value sum’s low-magnitude contributions. A value vector whose entries are mostly small relative to its per-token maximum loses resolution in exactly the entries that carry the fine detail of the output. Because the output is a weighted sum, the error is spread across the output dimension and does not concentrate, which makes it less damaging per unit of error than a key error, but it is present in every output token.

The third is positional. Cache entries written early in a sequence were quantized with scales computed from the token at that position. If the distribution of key magnitudes drifts as the sequence grows, early entries and late entries are quantized on different effective grids, and the attention scores comparing them are not on a common scale. Per-token scaling handles this correctly by construction, since each entry carries its own scale; a scheme that assumed a stationary distribution across positions does not.

A fourth, more subtle effect concerns the residual stream. The cache is written from the key and value projections of the normalized hidden state, and the outliers discussed in the activation-outlier treatment appear in those projections as well. A per-token scale computed from the maximum of the vector is set by the outlier channel, and the remaining channels lose resolution exactly as they do in the activation case. The remedy is the same family of techniques, applied at the projection that produces the cache rather than at the layer input.

Precision floors and mixed cache formats

Because keys and values have different sensitivity, and because the sensitivity varies with the head and the layer, a uniform cache format is a compromise. Practical schemes assign different widths to keys and values, and sometimes different widths to different layers, keeping the first and last layers at higher precision because their contributions to the residual stream are less redundant.

There is also a floor imposed by the accumulation. The attention scores are accumulated in whatever precision the kernel uses, and a 4-bit key multiplied by an 8-bit query produces a product that fits easily in 16 bits; the sum over a head dimension of 128 needs a few more. The accumulator width is not the constraint. The constraint is that the softmax needs the scores in floating point, so there is a conversion from the integer accumulator to float before the exponential, and that conversion fixes the effective resolution of the score regardless of the key bit width. Quantizing keys below the resolution that survives that conversion buys nothing.

For very long contexts, the cache also becomes a target for eviction and summarization, which interact with precision: an entry that will be evicted soon does not need a precise representation, and an entry that will be retained for thousands of tokens does. A scheme that varies precision by recency is trading bytes against error in a way that a uniform format cannot express, and it is the same trade that layer-level bit allocation makes across a model’s depth, applied along the sequence axis instead.

The cache is also where quantization meets the bandwidth argument most directly. During decode the weights are read once per token and the cache is read once per token, and the ratio between the two determines which dominates. For a short context the weights dominate; for a long context the cache does. A deployment that quantizes weights to 4 bits and leaves the cache at 16 has shifted the bottleneck rather than removed it, and the point at which the shift happens is a function of context length, batch size, and the ratio of cache size to parameter count.