veloxquant_mlx.cache.kvquant_cache module

KVQuant-NUQ KV cache wrapper — non-uniform quantization + outlier isolation.

Inspired by “KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization” (arXiv:2401.18079, NeurIPS 2024). Documented as “KVQuant-adapted (VeloxQuant-MLX implementation)” — implements the four cache-observable pillars (per-channel keys / per-token values, NUQ datatype, per-vector dense-and-sparse outlier isolation, and Attention Sink-Aware quantization) and documents pre-RoPE key quantization as out of scope. See veloxquant_mlx.quantizers.kvquant for the numerics.

Quantization axes (matching KVQuant’s asymmetry, the same axes KIVI uses):

Keys — per-channel: each head-dim channel gets its own non-uniform levels

(sample axis = tokens). Channels have stable, distinct distributions.

Values — per-token: each token gets its own levels (sample axis = channels).

Level lifecycle:

Prefill (S > 1) fits the NUQ levels from the incoming batch and freezes them (kvquant_refit_interval == 0, the default — mirrors SVDq’s frozen-V). Decode tokens quantize against the frozen levels. With a positive refit interval, levels are re-fit every N decode steps from the most recent token. Value levels are per-token (one signpost table per token, sample axis = channels) and are, by design, always fit fresh every call — never frozen, regardless of kvquant_refit_interval (that field only gates the key path). This is inherent to the per-token scheme, not a bug.

Performance (VeloxQuant-MLX#504): the per-head application was originally a Python for b in range(B): for h in range(H): loop, and the Lloyd-Max fit inside it forced a host sync (mx.eval) every iteration — together measured a real 59x mlx_lm.generate() decode slowdown (94.6 -> 2.8 tok/s, Llama-3.2-1B, Apple M4). Both are fixed: application is now batched over the flattened B*H axis (split_dense_sparse_batched(), fit_nuq_levels_batched(), etc.), and the Lloyd-Max loop no longer evaluates per iteration — measured 9.6x real end-to-end recovery (2.8 -> 27 tok/s on the same benchmark). This does not close the full gap to fp16 speed. Unlike AdaKV’s equivalent fix (VeloxQuant-MLX#512), the remaining cost here is not dispatch overhead — it’s the Lloyd-Max fit’s own O(n_iters x N x L x D) arithmetic, paid fresh for values on every decode step by the algorithm’s own per-token- levels design. kvquant_lloyd_iters (default 8) is the existing, already-correctly-wired lever for trading fit quality for speed — lower values measured a further 27 -> 41 tok/s at lloyd_iters=1 on the same benchmark, still well short of fp16 parity. There is no further “fix” to apply here without changing the algorithm’s per-token-value-levels design.

Byte accounting:

compressed_* — NUQ codes + per-(channel/token) level table + fp16 outlier

side-channel (value + position index).

fp16_* — uncompressed cost for the ratio.

Attention Sink-Aware quantization (paper §3.5):

The first kvquant_n_sink tokens of the sequence are restored to exact fp16 and excluded from both the level fit and the outlier thresholds — the model is disproportionately sensitive to quantization error at the sink positions it dumps excess attention onto.

What is NOT implemented (documented):

• Pre-RoPE key quantization (needs a model-forward hook — outside the cache contract; we see post-RoPE keys only).

• Offline calibration-set level fitting (we fit online; zero calibration).

• Fisher-information sensitivity weighting of the Lloyd-Max objective (Eq. 1 weights each squared error by F_ii; needs gradients from a calibration pass, so our fit is the unweighted k-means special case).

class veloxquant_mlx.cache.kvquant_cache.KVQuantKVCache(config)

Bases: KVCache

KV cache implementing KVQuant-NUQ non-uniform quantization.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: kvquant_bits (int, default 3), kvquant_outlier_fraction (float, default 0.01), kvquant_group_size (int, default 32; reserved for grouped fits), kvquant_lloyd_iters (int, default 8), kvquant_refit_interval (int, default 0 = freeze prefill levels).

No .bits attribute — mlx_lm’s SDPA checks hasattr(cache, "bits") to route to its quantized-matmul kernel path, which expects a different cache layout (mx.quantize’s native tuple format) and a .group_size attribute this cache doesn’t have. We expose .nuq_bits instead.

update_and_fetch(keys, values)

Fit (or reuse frozen) per-channel key / per-token value NUQ levels, isolate outliers, and preserve sink tokens in fp16; return dequantized K/V.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

trim(n)

Drop the last n tokens and keep this cache’s own state in step.

The inherited trim only moves offset. Left alone, _n_tokens stays at its old value, so after trimming to empty the next prefill would not get its sink tokens (n_sink applies only at _n_tokens == 0) and would reuse the previous sequence’s frozen key levels and outlier thresholds.

Trimming to empty resets everything (a new sequence starts from scratch). A partial trim keeps the levels fit on the surviving prefix and scales the byte accounting down with the token count (an estimate: per-call table and outlier bytes are not tracked per token).

Return type:

int

Parameters:

n (int)

property merge

property nuq_bits: int

Configured non-uniform quantization bit-width.

property outlier_fraction: float

Configured nominal fraction of elements treated as dense-and-sparse outliers.

property outlier_count: int

Realized count of elements isolated as outliers (K + V), across all calls so far.

property n_sink: int

Configured number of leading (attention-sink) tokens kept in fp16.

property sink_kept: int

Sink tokens actually retained in fp16 (bounded by the prefill length).

property key_outlier_thresh

Frozen per-channel |key| outlier threshold reused during decode.

property key_levels

Frozen per-channel non-uniform quantization levels fit at prefill.

property value_levels

Frozen per-token non-uniform quantization levels fit at prefill.

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (NUQ codes + per-channel level table + fp16 outlier/sink side-channel, all heads/batches).

property compressed_value_bytes: int

Realized stored bytes for the compressed value cache (NUQ codes + per-token level table + fp16 outlier/sink side-channel, all heads/batches).

property fp16_key_bytes: int

Hypothetical fp16 key cost if nothing were compressed.

property fp16_value_bytes: int

Hypothetical fp16 value cost if nothing were compressed.

property effective_bits: float

Effective per-element key bits (codes + table + outliers vs fp16).