INT8 KV-Cache Quantization

The problem, from first principles

At long context the KV cache, not the weights, dominates memory, and every decoding step reads the whole cache from memory. Storing keys and values as 8-bit integers instead of 16-bit floats halves both the memory and the bandwidth. The cache is quantized on write and dequantized on read. A single scale for the whole tensor would be ruined by one outlier, so the standard choice is a scale per token and per head: each (token, head) vector of length d is divided by its own absmax / 127, rounded, and stored as int8 together with that one float scale.

From theory to code

Implement quantize_kv and dequantize_kv.

Constraints

• kv has shape (T, H, d). For every (t, h) the scale is max(|kv[t, h, :]|) / 127, using scale 1.0 where that maximum is zero.
• quantize_kv(kv) returns (q, scale) where q = clip(round(kv / scale), -127, 127) as np.int8 (use np.round, which rounds halves to even) and scale has shape (T, H).
• dequantize_kv(q, scale) returns q * scale[..., None] as float64.

Hints