veloxquant_mlx.cache.kvtc_cache module

KVTC-adapted KV cache — local PCA + DP-optimal per-component bit allocation + entropy coding.

Inspired by “KV Cache Transform Coding for Compact Storage in LLM Inference” (NVIDIA, ICLR 2026, accepted poster, arXiv:2511.01815). Documented as “KVTC-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

Modeled on PALUKVCache (both K and V compressed, true-latent storage — not SVDq’s keys-only, reconstruct-fp16 scope). The mechanism this method adds over Palu/SVDq/SpectralQuant: a DP-optimal, discrete, per-component bit allocation (may drop a component entirely, i.e. assign it 0 bits) instead of a fixed hand-chosen split, followed by a real, measured entropy-coding stage. See quantizers/kvtc.py and allocators/kvtc_dp.py for the mechanism and the honesty crux (local per-sequence PCA vs the paper’s pre-calibrated global basis; analytic distortion proxy vs the paper’s real-activation-fit rate-distortion model; order-0 Huffman coder vs the paper’s possibly more sophisticated scheme; paper’s headline numbers are the paper’s, on trained models, never reproduced here).

NOT PATH-DEPENDENT (contrast with the eviction family)

Unlike H2O/TOVA/MorphKV/KVzip (eviction, path-dependent keep-sets), KVTC is in the Palu/SVDq/SpectralQuant family: the local PCA basis (V, mean) and the DP-derived per-component bit allocation are fit once, at the first update_and_fetch call with S > 1 (prefill), and reused unchanged for every subsequent token (prefill continuation or decode). Feeding the same sequence one token at a time after prefill vs. all at once produces the identical stored basis/allocation/codes/decompressed output — pinned by a determinism test in tests/cache/test_kvtc_cache.py.

Design

Prefill (first call, S > 1):

  1. For keys and values independently: run local PCA (via quantizers/kvtc.py::kvtc_compress) on the prefill batch, at kvtc_bit_budget total bits per token. Store the resulting projection V, mean, and per-component bit_allocation as fixed layer state.

  2. Store the (fixed-basis) latent codes for the prefill batch.

Decode (subsequent calls, any S):

  1. Project the new key/value into the already-fitted basis ((x - mean) @ V), quantize each component at its already-fixed bit allocation, and grow the entropy-coded latent store.

Because true latent (entropy-coded) storage is not a fixed-shape tensor the parent mlx_lm fp16 ring buffer cannot hold it — like PALUKVCache, this class bypasses the parent buffer entirely and manages its own offset, re-running kvtc_compress() over the accumulated raw rows on every call (the DP allocation and PCA basis are frozen after prefill; only the per-token codes grow). Every call returns reconstructed fp16 [B, H, S_total, D] for the downstream SDPA call; storage accounting is based on the realized entropy-coded bytes.

Byte accounting

kvtc_bytes / full_seq_bytes / compression_ratio — key + value combined, mirroring Palu’s combined reporting style. pre_entropy_bytes — fixed-width (pre-entropy-coding) size, for comparison. entropy_coding_gain — pre_entropy_bytes / kvtc_bytes. Since #77, kvtc_compress guarantees the entropy-coded payload+table never exceeds fixed-width packing (falling back to fixed-width when Huffman doesn’t pay for itself — see quantizers/kvtc.py’s module docstring), but this ratio divides by the FULL stored size (kvtc_bytes, which also carries the fixed projection-basis overhead unrelated to entropy coding), so it can still legitimately read < 1 when that basis dominates — it is not itself guaranteed >= 1.

class veloxquant_mlx.cache.kvtc_cache.KVTCKVCache(config)[source]

Bases: KVCache

KV cache implementing KVTC-adapted local-PCA + DP-optimal bit allocation + entropy coding.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: head_dim (D), kvtc_bit_budget (int, default 4 * head_dim — total bits

per token across all principal components, for K and V independently; “4 avg-bits-per-component equivalent scaled by head_dim”),

kvtc_bit_choices (tuple[int, …], default

allocators.kvtc_dp.DEFAULT_BIT_CHOICES),

kvtc_beta (float, default

allocators.kvtc_dp.DEFAULT_BETA).

Notes

Applies to both K and V independently (mirrors Palu’s scope, not SVDq’s keys-only scope) — see module docstring. No .bits attribute (this isn’t a fixed-bit-width quantizer in that sense, and it isn’t an eviction cache either). Single (batch, head) state per layer; B and H handled by looping (mirrors PALUKVCache’s per-head loop). Bypasses the parent fp16 ring buffer (true latent storage), like PALUKVCache — manages its own offset.

property merge
update_and_fetch(keys, values)[source]

Fit the local PCA basis + DP bit allocation on prefill, else append through the frozen basis; return reconstructed fp16 K/V.

Parameters:

  • keys (mlx.core.array)

  • values (mlx.core.array)

property offset: int

Number of tokens accumulated so far (own counter — parent ring buffer is bypassed).

size()[source]

Number of tokens currently stored.

Return type:

int

property state

Last reconstructed fp16 (keys, values), since the true-latent buffers bypass the parent ring buffer mlx_lm reads directly.

is_trimmable()[source]

Whether this cache supports trimming.

Return type:

bool

trim(n)[source]

Drop the n most-recently-added tokens, re-quantizing through the frozen basis; returns the number actually dropped.

Return type:

int

Parameters:

n (int)

property kvtc_bytes: int

Realized total stored bytes (K + V, all heads/batches), including the entropy-coded payload, code table, projection basis, and quant params — never the pre-entropy-coding size.

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (entropy-coded payload + table + basis, all heads/batches).

property compressed_value_bytes: int

Realized stored bytes for the compressed value cache (entropy-coded payload + table + basis, 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 pre_entropy_bytes: int

Fixed-width (pre-entropy-coding) size, K + V, all heads/batches.

property entropy_coding_gain: float

pre_entropy_bytes / kvtc_bytes — realized entropy-coding gain.

A modest, honestly-scoped secondary effect (see quantizers/_entropy_coding.py): NOT the theoretical Shannon-entropy bound. Since #77, the entropy-coding stage itself (payload + table) is guaranteed to never exceed fixed-width packing — but this ratio’s denominator (kvtc_bytes) also includes the fixed projection-basis overhead, which is unrelated to entropy coding and can dominate at small survived-component counts, so this property can still read < 1 for that reason.

property full_seq_bytes: int

Hypothetical fp16 K + V cost if nothing were compressed.

property compression_ratio: float

full_seq_bytes / kvtc_bytes; > 1 means memory savings over fp16.

property bit_budget: int

Configured total bits per token across all principal components (K and V independently).