How TurboQuant Reduces Memory Usage for KV Cache in mlx‑vlm

TurboQuant reduces KV cache memory usage by replacing dense 16-bit floating-point tensors with low-bit integer quantization, packing keys and values into compact codebook indices that consume roughly one-fourth the memory of standard caches.

In large vision-language models, the key-value (KV) cache dominates memory consumption during long-context inference. The mlx‑vlm repository implements TurboQuant, a quantization strategy that compresses the KV cache using fractional bit-widths and custom Metal kernels to maintain speed while dramatically cutting memory footprint.

Understanding the KV Cache Bottleneck

Standard transformer inference stores keys and values as 16-bit floating-point arrays (float16 or bfloat16), consuming 2 × D bytes per token per layer for each of the key and value tensors. For long sequences or large batch sizes, this grows into gigabytes of GPU memory. TurboQuant attacks this problem by quantizing these tensors to low-bit integers, typically achieving a 4× memory reduction when configured to 3.5 bits per element.

How TurboQuant Quantizes the KV Cache

TurboQuant implements a custom TurboQuantKVCache class in mlx_vlm/turboquant.py that replaces the standard cache implementation. The mechanism operates through several coordinated steps:

Configuring Fractional Bit-Widths

TurboQuant supports arbitrary bit-widths ≥ 1, including fractional values like 3.5. The _validate_bits function validates the configuration, while _ensure_codecs (lines 48 225‑48 233) separates the setting into distinct bit-widths for keys and values. This flexibility allows fine-grained trade-offs between compression ratio and model accuracy.

Building Integer Codecs with Codebooks

The _build_codec method instantiates either an MSE codec (_TurboQuantMSECodec) or a split codec for non-integer bit-widths (line 48 261). Each codec maintains a codebook—a small array of representative float values—and stores packed indices pointing into this codebook rather than the full-precision tensor values.

Packing Quantized Tensors

When quantizing, key_codec.quantize and value_codec.quantize return a tuple of (norms, packed) where:

• norms stores a single float per token for dequantization scaling
• packed is a uint32 array with width (D·bits+31)//32 (see update_and_fetch, lines 49 502‑49 511)

Instead of storing D float16 values per token, TurboQuant stores one norm plus ⌈bits·D/32⌉ 32-bit words. This reduces storage from D × 2 bytes to approximately D × bits/8 bytes plus overhead.

Fused Kernel Optimization

For maximum efficiency, _try_fused_kv_quantize (lines 48 440‑48 466) attempts to quantize both keys and values in a single Metal kernel dispatch when both use MSE codecs. This minimizes launch overhead while keeping the packed representation compact.

Memory Allocation and On-Demand Dequantization

The _allocate_state_like function (lines 49 504‑49 506) creates the KV cache with the packed shape (new_end tokens, packed width), ensuring the allocated buffers are already reduced in size. During inference, the cache remains in its compact form until the dequantize method (lines 49 436‑49 442) reconstructs the original float tensors on-the-fly using the stored codebooks. Full-precision data is materialized only for the final attention score computation, keeping the memory footprint minimal throughout most of the pipeline.

Enabling TurboQuant in Your Code

You activate TurboQuant through the helper function maybe_quantize_kv_cache in mlx_vlm/generate.py (lines 44 3‑44 9). This swaps a regular KVCache for TurboQuantKVCache when the kv_bits parameter is set:

from mlx_vlm.generate import maybe_quantize_kv_cache

# Enable 3.5-bit quantization for the KV cache

maybe_quantize_kv_cache(
    prompt_cache,        # list of KVCache objects from previous generation

    quantized_kv_start=0,
    kv_group_size=1,
    kv_bits=3.5,              # fractional bit-width

    kv_quant_scheme="mse",    # quantization codec

)

After activation, isinstance(prompt_cache[0], TurboQuantKVCache) returns True, confirming that the cache now uses the quantized backend. During generation, dequantization happens automatically when accessing keys and values:

# Inside a generation step - automatic on-demand reconstruction

keys, values = kv_cache.dequantize()

# Attention computation proceeds with standard float tensors

Source Code Reference

The implementation spans several critical files in the mlx‑vlm repository:

• mlx_vlm/turboquant.py – Core TurboQuantKVCache class, integer codecs, and Metal kernels for packing/unpacking
• mlx_vlm/generate.py – Public API via maybe_quantize_kv_cache
• mlx_vlm/tests/test_turboquant.py – Unit tests validating fractional bit-width support
• mlx_vlm/tests/test_kv_cache_quantization.py – Integration tests for the full generation pipeline

Summary

• TurboQuant reduces KV cache memory by replacing float16 tensors with low-bit integer indices and per-token normalization factors.
• The system supports fractional bit-widths (e.g., 3.5 bits) validated through _validate_bits and processed by _ensure_codecs.
• Quantization uses codebooks and packed uint32 arrays to achieve roughly a 4× memory reduction at 3.5 bits per element.
• Fused Metal kernels (_try_fused_kv_quantize) optimize the quantization process when both keys and values use MSE codecs.
• Memory allocation via _allocate_state_like creates compact buffers upfront, while dequantize reconstructs full-precision tensors only when needed for attention scoring.
• Activation requires only setting kv_bits in maybe_quantize_kv_cache from mlx_vlm/generate.py.

Frequently Asked Questions

What is the typical memory savings when using TurboQuant?

At a bit-width of 3.5, TurboQuant reduces KV cache memory usage by approximately four times compared to standard float16 storage. The exact formula calculates storage as ⌈bits·D/32⌉ × 4 bytes + 4 bytes norm per token instead of D × 2 bytes, yielding significant savings for long sequences.

Can I use non-integer bit-widths like 3.5 bits with TurboQuant?

Yes. TurboQuant explicitly supports fractional bit-widths ≥ 1. The _ensure_codecs function (lines 48 225‑48 233) handles these values by constructing appropriate codecs, and the test suite in mlx_vlm/tests/test_turboquant.py validates this functionality.

How does TurboQuant affect inference speed?

TurboQuant maintains inference speed through custom Metal kernels that quantize and dequantize directly on the GPU. The _try_fused_kv_quantize kernel combines key and value processing into a single dispatch when possible, minimizing overhead. Dequantization occurs on-demand only during attention computation, keeping the pipeline efficient.

Where is the integration point to enable TurboQuant in existing code?

The integration occurs in mlx_vlm/generate.py via the maybe_quantize_kv_cache function (lines 44 3‑44 9). Passing a kv_bits parameter to this function automatically replaces standard KVCache instances with TurboQuantKVCache objects, requiring no changes to the underlying model architecture.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}