KV cache compression¶

The KV cache grows with every token of context. TurboQuant compresses it with the same recipe it uses for weights (Hadamard rotation plus a Lloyd-Max codebook), at run time. The compressed cache is expanded to FP16 only when attention reads it, so it works with attention sinks, sliding windows and linear-attention layers.

turboquant-generate --model ./model-tq3 --prompt "..." \
    --kv-k-bits 8 --kv-v-bits 3 --kv-min-tokens 128

The same options work on turboquant-serve.

Choosing the bits¶

Keys need more precision than values: softmax magnifies any error in K, while V tolerates aggressive quantization. Mixed K8/V3 is the recommended setting, about 4× smaller than FP16.

--kv-min-tokens 128 keeps the first 128 tokens in FP16. Many models attend heavily to those "sink" tokens.

It saves memory; it rarely saves time¶

Whether compression speeds decoding up or slows it down depends on how big the KV cache is per token. With a large per-token cache, the 4× smaller cache saves more memory bandwidth than decompression costs. With a small one, decompression dominates.

On models with a small cache, the penalty grows with context. Qwen3.6-35B-A3B (measured by a community member on an M5 Pro) decodes 1.14× faster with FP16 at 65 tokens and 6.6× faster at 63K, while compression saves 2.08 GB at 63K. So on small-active MoEs, use KV compression to fit a longer context, not to go faster. Similar-sized models don't all behave alike: Qwen3.5-122B doesn't flip the way GPT-OSS-120B does.

The headline case is GPT-OSS-120B: with weights and KV both at 3-bit, its full 131K-token context fits in 50 GB on a 64 GB Mac, and decoding is faster than with an FP16 cache.

Fused decode+attend (experimental)¶

Most of the decoding cost above is expanding the cache back to FP16 every token. An experimental Metal kernel reads the packed cache directly and runs attention on the fly:

turboquant-serve --model ./model-tq3 --kv-k-bits 8 --kv-v-bits 3 \
    --kv-min-tokens 0 --kv-fused

Output is identical. It covers decode only, batch 1, with no FP16 sink (--kv-min-tokens 0), bit-packed keys and values, and equal K/V head sizes that are a multiple of 32. Any step it can't handle (prefill, sink or sliding-window layers) falls back to the standard path, so it is safe to turn on.

The gain grows with context: a single attention layer is about 10× faster than the standard path at 32K tokens. End to end, it scales with the share of full-attention layers, which is why the hybrid Qwen3.6 gains less than the dense Llama.

From Python¶

from mlx_lm.models.cache import make_prompt_cache
from turboquant_mlx.layers import convert_cache_to_turboquant, enable_fused_attend

cache = make_prompt_cache(model)                 # the right cache type per layer
cache = convert_cache_to_turboquant(
    cache,
    k_bits=8, v_bits=3,                          # keys need the precision
    min_tokens_before_quant=128,                 # first 128 tokens stay FP16
    group_size=64,
)

Layers that don't hold a growing KV cache (RotatingKVCache, ArraysCache) are left as they are. For the fused kernel, call enable_fused_attend(cache) after converting (it needs min_tokens_before_quant=0).