veloxquant_mlx.cache.zipcache_cache module

ZipCache-adapted KV cache — saliency-adaptive per-token mixed-precision.

Inspired by “ZipCache: Accurate and Efficient KV Cache Quantization with Salient Token Identification” (He et al., NeurIPS 2024, arXiv:2405.14256). Documented as “ZipCache-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

Per-token saliency routing: each incoming K/V block has its key tokens sorted by L2-norm; the top hi_fraction receive hi_bits and the rest receive lo_bits. Both paths remain quantized (not fp16) — this distinguishes ZipCache-adapted from KIVI-Sink (which keeps sinks in fp16).

The wrapper compresses each [B, H, S, D] head matrix with zipcache_compress / zipcache_reconstruct and hands the reconstructed fp16 K/V to the parent mlx_lm cache (no .bits attribute; SDPA stays on the clean fp16 path). Byte accounting tracks the mixed-bit stored size against both fp16 and a uniform-lo-bit baseline.

Honest proxy limitation: the saliency signal is the key L2-norm (a proxy for attention importance). The paper uses normalized attention scores, which are not observable by a cache wrapper. The proxy has been used for KIVI-Sink and AdaKV-proxy in this repo; this is the third use, with a different decision (bit-width routing rather than fp16 protection or head budgeting).

Byte accounting:

compressed_key_bytes / compressed_value_bytes — mixed-bit ZipCache stored size baseline_key_bytes / baseline_value_bytes — uniform lo-bit baseline fp16_key_bytes / fp16_value_bytes — uncompressed cost

class veloxquant_mlx.cache.zipcache_cache.ZipCacheKVCache(config)

Bases: KVCache

KV cache implementing ZipCache-adapted per-token mixed-precision for one layer.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: zipcache_hi_bits (int, default 4) — bit-width for salient tokens, zipcache_lo_bits (int, default 2) — bit-width for non-salient tokens, zipcache_hi_fraction (float, default 0.20) — fraction of tokens at hi_bits, zipcache_group_size (int, default 32) — token group size for quant, zipcache_quantize_values (bool, default True) — apply to values too.

Notes

No .bits attribute — keeps mlx_lm SDPA on the clean fp16 path. Single-layer (no coordinator); for_model propagates the zipcache_* fields automatically via dataclasses.replace. Values are quantized uniformly at hi_bits (saliency routing is key-driven; values follow the hi-bit path as the safer default).

update_and_fetch(keys, values)

Sort key tokens by L2-norm saliency, quantize the top hi_fraction at hi_bits and the rest at lo_bits (values uniformly at hi_bits); return reconstructed fp16 K/V.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (mixed hi/lo-bit saliency-routed codes, all heads/batches).

property compressed_value_bytes: int

Realized stored bytes for the compressed value cache (uniform hi_bits codes, all heads/batches).

property baseline_key_bytes: int

Uniform lo-bit baseline for comparison (no saliency routing).

property baseline_value_bytes: int

Uniform lo-bit baseline for comparison (no saliency routing).

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 compression_ratio: float

fp16 bytes / compressed bytes (key side). >1 means storage savings.

property effective_avg_bits: float

Average key bits/element implied by the stored mixed-bit rate.

property merge