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

> Discover how TurboQuant slashes KV cache memory usage with low-bit integer quantization, packing keys and values into compact codebook indices to save significant memory.

- Repository: [Prince Canuma/mlx-vlm](https://github.com/Blaizzy/mlx-vlm)
- Tags: performance
- Published: 2026-04-05

---

**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`](https://github.com/Blaizzy/mlx-vlm/blob/main/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`](https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/generate.py) (lines 44 3‑44 9). This swaps a regular `KVCache` for `TurboQuantKVCache` when the `kv_bits` parameter is set:

```python
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:

```python

# 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`](https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/turboquant.py) – Core `TurboQuantKVCache` class, integer codecs, and Metal kernels for packing/unpacking
- [`mlx_vlm/generate.py`](https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/generate.py) – Public API via `maybe_quantize_kv_cache`
- [`mlx_vlm/tests/test_turboquant.py`](https://github.com/Blaizzy/mlx-vlm/blob/main/mlx_vlm/tests/test_turboquant.py) – Unit tests validating fractional bit-width support
- [`mlx_vlm/tests/test_kv_cache_quantization.py`](https://github.com/Blaizzy/mlx-vlm/blob/main/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`](https://github.com/Blaizzy/mlx-vlm/blob/main/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`](https://github.com/Blaizzy/mlx-vlm/blob/main/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`](https://github.com/Blaizzy/mlx-vlm/blob/main/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.