veloxquant_mlx.cache.turboquant_rvq_cache module

TurboQuantRVQ KV cache wrapper for production mlx_lm integration.

This wraps the existing TurboQuantRVQ quantizer in the same update_and_fetch cache protocol that mlx_lm expects.

Storage: keys are kept packed as two streams of b-bit RVQ codebook indices (stage-1 Gaussian codebook + stage-2 Laplacian residual codebook), bit-packed into uint32 words using the same LSB-first, per-row layout as mlx.core.quantize — see _pack_indices()/_unpack_indices(). Dequantization (codebook gather + inverse rotation) happens on every fetch, matching the accepted cost of mlx_lm’s own native QuantizedKVCache (dequantize-on-fetch, see mlx_lm/models/cache.py). Values pass through as plain fp16, unchanged from earlier versions of this wrapper.

Per-vector key storage at bit-width b and head-dim d:

2 * ceil(d * b / 32) * 4 bytes (two packed uint32 index streams) plus a bf16 per-vector norm (2 bytes; unlike fp16 it cannot overflow past ||k|| > 65504, at the cost of ~0.2% relative scale error).

For d=128, b=1 this is 32 bytes/vector vs 256 bytes fp16 → 8x compression. For d=128, b=2 this is 64 bytes/vector vs 256 bytes fp16 → 4x compression.

Note: this differs slightly from the pre-packed-storage accounting formula (ceil(d * 2 * b / 8) + 2), which measured a hypothetical tight packing. The real packed layout pads each row to a whole number of uint32 words, so actual bytes are equal to or slightly higher than that formula predicts — :pyattr:`compressed_key_bytes` now reports the true packed size, not the formula’s estimate.

Measured (not accounting) impact, b=1, mlx-community/Llama-3.2-1B-Instruct-4bit, 4002-token prompt + 100 decode tokens, mx.get_peak_memory(), gc-controlled across runs: peak memory 1402.3 MB, vs. 1608.4 MB for the stock fp16 mlx_lm.models.cache.KVCache (-12.8%), vs. 2537.2 MB for mlx_lm’s own native QuantizedKVCache(bits=4) at the same prompt (-44.7%) – mlx-lm’s native quantized cache is worse than fp16 at this prompt length because its mx.quantize() call materializes full intermediate buffers per update_and_fetch call; this cache’s storage is written incrementally in packed form so it does not pay that cost. This gap should be expected to shrink or reverse at other context lengths / chunking strategies – treat these as one measured data point with its methodology stated, not a universal claim.

class veloxquant_mlx.cache.turboquant_rvq_cache.TurboQuantRVQKVCache(config)

Bases: KVCache

KV cache backed by TurboQuantRVQ two-pass residual quantization.

Compresses keys via TurboQuantRVQ, storing the two RVQ index streams bit-packed (see module docstring); values pass through unchanged. Does not subclass update_and_fetch’s fp16-tensor storage from mlx_lm.models.cache.KVCache for keys — only step-based growth semantics are inherited/mirrored.

Parameters:

config (Any) – KVCacheConfig (consumes head_dim, bit_width_inlier, seed).

Notes

Reports byte accounting via :pyattr:`compressed_key_bytes` and :pyattr:`fp16_key_bytes` so users can compute the realized compression ratio after a benchmark run. Unlike earlier versions of this class, :pyattr:`compressed_key_bytes` now reflects genuinely packed storage, not a bit-width accounting estimate — see issue #27.

step = 256

update_and_fetch(keys, values)

Normalize keys to unit vectors + norm, residual-VQ-encode and pack the direction, store values in fp16; return dequantized K and raw V.

Return type:

Any

Parameters:

• keys (Any)

• values (Any)

property state: Any

Packed key streams (packed1, packed2, norms, values) — the raw stored (not reconstructed) state, live portion only.

property meta_state: Any

(head_dim, bits, seed, offset) as strings.

Type:

Config fields needed to rebuild derived state after a restore

classmethod from_state(state, meta_state=None)

Reconstruct a cache from .state and .meta_state.

meta_state defaults to None only to keep this an LSP-valid override of _BaseCache.from_state(cls, state) in mlx_lm >=0.32 (which dropped the separate meta_state arg); every actual caller in this codebase (fetch_nearest_cache, this module’s own tests) still passes it explicitly, and omitting it is not a supported call.

Return type:

TurboQuantRVQKVCache

Parameters:

• state (Any)

• meta_state (Any)

is_trimmable()

Whether this cache supports trimming.

Return type:

bool

trim(n)

Drop the n most-recently-added tokens; returns the number actually dropped.

Return type:

int

Parameters:

n (int)

make_mask(*args, **kwargs)

Build the attention mask for the current offset, delegating to create_attention_mask.

Return type:

Any

Parameters:

• args (Any)

• kwargs (Any)

empty()

Whether the cache holds no tokens yet.

Return type:

bool

property nbytes: int

True packed byte count for cached keys and values (not an estimate).

property compressed_key_bytes: int

Cumulative compressed key bytes since construction (true packed size).

property fp16_key_bytes: int

Cumulative fp16-equivalent key bytes (for ratio computation).

property assigned_bits: int

Bit-width used by this layer’s quantizer.

Named assigned_bits (not bits) to avoid colliding with mlx_lm’s quantized-SDPA check, which uses hasattr(cache, "bits") to switch attention kernel paths.

property merge