veloxquant_mlx.quantizers.zipcache module

ZipCache-adapted quantizer — 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.

What ZipCache adds that the repo did not have: per-token mixed bit-width within the quantized space. Three existing methods use the key-norm proxy for token importance; this is the fourth use and the first with a different decision:

KIVI-Sink — top-k norm tokens → fp16 (binary: quantized vs not) AdaKV-proxy — mean norm per head → head-level budget reallocation Kitty — per-channel variance → channel-level bit allocation ZipCache-adapted → per-token norm → hi_bits vs lo_bits (both quantized)

Every token stays quantized. High-norm tokens (salient, top hi_fraction by key L2-norm) are quantized at hi_bits; the rest at lo_bits. The effective average key rate is hi_frac·hi_bits + (1-hi_frac)·lo_bits bits per element — explicitly between lo_bits and hi_bits.

Honest proxy limitation: the paper’s true saliency signal is the normalized attention score, which is not observable by a cache wrapper. Key L2-norm is a proxy: attention-sink tokens (highest importance) also exhibit large key norms. This proxy is weaker than true attention scores — reported plainly, never hidden. See the KIVI-Sink docs for the proxy derivation.

Adaptation decisions:

1. Per-group min/max quantization (not per-channel as in the paper’s ideal). Reuses the same group-quant logic as KIVI/CacheGen/GEAR for consistency and test-coverage inheritance.

2. Values are quantized uniformly at hi_bits (no saliency routing for values — the paper focuses on keys; quantize_values=True can be set to compress values too at a single bit-width).

3. Saliency is computed once per incoming block. No retroactive reclassification of already-quantized tokens.

This module holds the pure numerics: norm computation, saliency masking, per-group channel quant, full compress/reconstruct, and honest byte accounting. The cache wrapper owns the per-layer prefill/decode state.

class veloxquant_mlx.quantizers.zipcache.ZipCacheState(hi_codes, hi_scales, hi_zeros, lo_codes, lo_scales, lo_zeros, hi_mask, hi_bits, lo_bits, seq_len, head_dim, group_size)

Bases: NamedTuple

A ZipCache-compressed key or value matrix with per-token bit routing.

Variables:

• hi_codes – [n_hi, D] uint8 high-bit key codes for salient tokens.

• hi_scales – [n_hi_groups, D] fp32 per-group scale (hi-bit path).

• hi_zeros – [n_hi_groups, D] fp32 per-group zero-point (hi-bit path).

• lo_codes – [n_lo, D] uint8 low-bit key codes for non-salient tokens.

• lo_scales – [n_lo_groups, D] fp32 per-group scale (lo-bit path).

• lo_zeros – [n_lo_groups, D] fp32 per-group zero-point (lo-bit path).

• hi_mask – [S] bool saliency mask; True = salient (hi-bit) token.

• hi_bits – int bit-width for salient tokens.

• lo_bits – int bit-width for non-salient tokens.

• seq_len – int original token count (= S).

• head_dim – int feature dimension (= D).

• group_size – int token group size used for min/max quantization at compress time; must be reused by zipcache_reconstruct for the dequant grouping to match.

Parameters:

• hi_codes (mlx.core.array)

• hi_scales (mlx.core.array)

• hi_zeros (mlx.core.array)

• lo_codes (mlx.core.array)

• lo_scales (mlx.core.array)

• lo_zeros (mlx.core.array)

• hi_mask (mlx.core.array)

• hi_bits (int)

• lo_bits (int)

• seq_len (int)

• head_dim (int)

• group_size (int)

hi_codes: mlx.core.array

Alias for field number 0

hi_scales: mlx.core.array

Alias for field number 1

hi_zeros: mlx.core.array

Alias for field number 2

lo_codes: mlx.core.array

Alias for field number 3

lo_scales: mlx.core.array

Alias for field number 4

lo_zeros: mlx.core.array

Alias for field number 5

hi_mask: mlx.core.array

Alias for field number 6

hi_bits: int

Alias for field number 7

lo_bits: int

Alias for field number 8

seq_len: int

Alias for field number 9

head_dim: int

Alias for field number 10

group_size: int

Alias for field number 11

veloxquant_mlx.quantizers.zipcache.token_key_norms(keys)

L2 norm of each token’s key vector, [S, D] → [S] fp32.

Return type:

array

Parameters:

keys (mlx.core.array)

veloxquant_mlx.quantizers.zipcache.saliency_mask(norms, hi_fraction)

Top-hi_fraction tokens by key-norm are marked salient (True).

Parameters:

• norms (array) – [S] fp32 key L2-norms.

• hi_fraction (float) – Fraction of tokens to mark as salient (0 → all lo-bit; 1 → all hi-bit).

Return type:

array

Returns:

[S] bool array; True = salient token → hi_bits path.

veloxquant_mlx.quantizers.zipcache.channel_quant(x, bits, group_size=32)

Asymmetric per-group min/max quantization along the token axis.

Groups partition the token (row) axis into blocks of group_size. Within each group, a shared (scale, zero) is fitted per channel.

Parameters:

• x (array) – [N, D] fp32/fp16 input.

• bits (int) – Bit-width (1–8).

• group_size (int) – Tokens per quantization group.

Return type:

tuple[array, array, array]

Returns:

(codes, scales, zeros) where codes is [N, D] uint8 and scales/zeros are [n_groups, D] fp32.

veloxquant_mlx.quantizers.zipcache.channel_dequant(codes, scales, zeros, group_size=32)

Reconstruct [N, D] fp32 from codes + per-group (scales, zeros).

Return type:

array

Parameters:

• codes (mlx.core.array)

• scales (mlx.core.array)

• zeros (mlx.core.array)

• group_size (int)

veloxquant_mlx.quantizers.zipcache.zipcache_compress(x, hi_bits=4, lo_bits=2, hi_fraction=0.2, group_size=32)

Compress [S, D] with per-token saliency-adaptive mixed bit-width.

Parameters:

• x (array) – [S, D] fp16/fp32 key or value matrix for one head.

• hi_bits (int) – Bit-width for salient (high-norm) tokens.

• lo_bits (int) – Bit-width for non-salient tokens.

• hi_fraction (float) – Fraction of tokens routed to hi_bits (by key L2-norm).

• group_size (int) – Token group size for min/max quantization.

Return type:

ZipCacheState

Returns:

ZipCacheState carrying split codes and the saliency mask.

veloxquant_mlx.quantizers.zipcache.zipcache_reconstruct(state)

Reconstruct fp16 [S, D] from a ZipCacheState.

Dequantizes the hi and lo groups and scatters them back into their original token positions using the stored saliency mask.

Return type:

array

Parameters:

state (ZipCacheState)

veloxquant_mlx.quantizers.zipcache.zipcache_bytes(state, group_size=None)

Honest stored size (bytes) of a ZipCacheState.

Counts: packed codes (ceil(n * D * bits / 8)) + fp32 params (scale + zero). The saliency mask is stored as bool (1 byte/token).

group_size defaults to state.group_size (the size the state was compressed with); pass it only to model a different grouping (#641).

Return type:

int

Parameters:

• state (ZipCacheState)

• group_size (int | None)

veloxquant_mlx.quantizers.zipcache.base_only_bytes(S, D, bits, group_size=32)

Uniform baseline size: all S tokens at a single bit-width (fp16 params).

Return type:

int

Parameters:

• S (int)

• D (int)

• bits (int)

• group_size (int)

veloxquant_mlx.quantizers.zipcache.zipcache_quant_dequant(x, hi_bits=4, lo_bits=2, hi_fraction=0.2, group_size=32)

Drop-in quant→dequant: full ZipCache round-trip on [S, D] → fp16.

Return type:

array

Parameters:

• x (mlx.core.array)

• hi_bits (int)

• lo_bits (int)

• hi_fraction (float)

• group_size (int)