veloxquant_mlx.cache.cachegen_cache module

CacheGen KV cache wrapper — entropy-coded byte accounting over group quant.

Inspired by “CacheGen: KV Cache Compression and Streaming for Fast LLM Serving” (Liu et al., SIGCOMM 2024, arXiv:2310.07240). Documented as “CacheGen-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

The reconstructed K/V the model sees are identical to plain asymmetric min/max group quant (KIVI-style). CacheGen’s contribution here is the byte accounting: it measures the Shannon entropy of the token-delta-transformed code stream and reports compressed_*_bytes from that, modelling an ideal arithmetic coder. This captures the token-wise-locality storage win without shipping a serial range codec that would bottleneck MLX decode.

Per-layer bit-width is configurable (cachegen_bits); KVCacheBuilder.for_model applies a real layer-wise schedule (deeper layers get fewer bits — CacheGen’s layer-sensitivity observation, §5.1.2) via layer_group_bits, injected per layer as cachegen_resolved_bits (see KVCacheBuilder._build_cachegen).

Byte accounting:

compressed_key_bytes / compressed_value_bytes — entropy-coded estimate fixed_width_key_bytes / fixed_width_value_bytes — naive packed baseline fp16_key_bytes / fp16_value_bytes — uncompressed cost for the ratio

class veloxquant_mlx.cache.cachegen_cache.CacheGenKVCache(config)

Bases: KVCache

KV cache implementing CacheGen entropy-coded compression for one layer.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: cachegen_bits (int, default 4) — base/uniform bit-width. cachegen_resolved_bits (int or None) — per-layer bit-width injected

by KVCacheBuilder.for_model (§5.1.2 layer-wise schedule); falls back to cachegen_bits when None (uniform behaviour).

cachegen_group_size (int, default 32), cachegen_use_delta (bool, default True — token-delta transform), cachegen_per_channel (bool, default True — group entropy estimate

by channel per §5.1.3 instead of pooling all channels).

Notes

No .bits attribute — keeps mlx_lm SDPA on the clean fp16 path.

update_and_fetch(keys, values)

Group-quantize K/V (KIVI-style), reporting entropy-coded byte estimates in addition to the fixed-width baseline; return dequantized K/V.

Parameters:

• keys (mlx.core.array)

• values (mlx.core.array)

property compressed_key_bytes: int

Entropy-coded (ideal arithmetic-coder) estimate for the key cache, modelling Shannon entropy of the token-delta code stream.

property compressed_value_bytes: int

Entropy-coded (ideal arithmetic-coder) estimate for the value cache, modelling Shannon entropy of the token-delta code stream.

property fixed_width_key_bytes: int

Naive fixed-width packed baseline for the key cache (no entropy coding).

property fixed_width_value_bytes: int

Naive fixed-width packed baseline for the value cache (no entropy coding).

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

Fraction of the fixed-width code size saved by entropy coding (key side).

property assigned_avg_bits: float

Effective key bit-width after entropy coding (vs fp16=16).