veloxquant_mlx.quantizers.cachegen module

CacheGen quantizer — delta-locality group quant + entropy-coded byte model.

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

What CacheGen adds that the repo did not have: entropy coding of the quantized KV. Every other method in the suite packs codes at a fixed bit-width; CacheGen exploits the KV cache’s distributional structure to encode the codes into a smaller bitstream. Its three observations:

  1. Token-wise locality — adjacent tokens’ KV vectors are similar, so the delta between consecutive tokens’ quantized codes is concentrated near zero and is far more compressible than the raw codes.

  2. Layer-wise sensitivity — deeper layers tolerate coarser quantization; CacheGen spends fewer bits on later layers.

  3. Arithmetic coding — the delta symbol stream, being low-entropy, is compressed with an entropy coder down toward its Shannon entropy.

  4. Channel/layer grouping — fitting a separate symbol distribution per channel (§5.1.3) yields much lower entropy than pooling all channels together; layer separation is already implicit (each cache instance owns one layer).

Adaptation:

  • The reconstruction path is the existing asymmetric min/max group quant — the values the model sees are identical to KIVI-style quant (no extra loss from the entropy layer; entropy coding is lossless over the codes).

  • We do not ship a per-step arithmetic codec (a serial range coder would bottleneck MLX’s parallel decode and add no quality). Instead we model the entropy-coded byte size from the measured Shannon entropy of the delta-coded symbol stream — an honest lower-ish bound on what a real arithmetic coder achieves, reported through compressed_*_bytes.

  • Layer-wise bit selection is exposed via config (cachegen_bits plus an optional per-depth schedule applied by the builder).

This module holds the pure numerics: group quant exposing integer codes, token-delta transform, and the entropy/byte estimator. The cache wrapper owns the per-layer state and accounting.

class veloxquant_mlx.quantizers.cachegen.CodeStream(codes, scale, zero, n_rows, bits)[source]

Bases: NamedTuple

Quantized codes plus the params needed to dequantize them.

Variables:

  • codes – [n_groups, group_size, D] fp32 integer codes in [0, 2^bits-1].

  • scale – [n_groups, 1, D] fp32 per-group step.

  • zero – [n_groups, 1, D] fp32 per-group min.

  • n_rows – int original (pre-pad) token count.

  • bits – int bit-width.

Parameters:

  • codes (mlx.core.array)

  • scale (mlx.core.array)

  • zero (mlx.core.array)

  • n_rows (int)

  • bits (int)

codes: mlx.core.array

Alias for field number 0

scale: mlx.core.array

Alias for field number 1

zero: mlx.core.array

Alias for field number 2

n_rows: int

Alias for field number 3

bits: int

Alias for field number 4

veloxquant_mlx.quantizers.cachegen.quantize_to_codes(x, bits, group_size=32)[source]

Asymmetric min/max group quant exposing the integer codes.

Parameters:

  • x (array) – [N, D] fp16/fp32 (one head’s keys or values).

  • bits (int) – bit-width.

  • group_size (int) – tokens per group along axis 0.

Return type:

CodeStream

Returns:

CodeStream (codes + dequant params).

veloxquant_mlx.quantizers.cachegen.dequant_codes(stream)[source]

Reconstruct fp16 [n_rows, D] from a CodeStream.

Return type:

array

Parameters:

stream (CodeStream)

veloxquant_mlx.quantizers.cachegen.token_delta(codes_flat)[source]

Token-wise delta of a flattened code matrix [N, D].

Row 0 is kept as-is; row i (i>0) becomes codes[i] - codes[i-1]. Because adjacent tokens’ KV are similar, the deltas concentrate near zero — the locality CacheGen exploits. Reversible: a prefix-sum recovers the codes.

Return type:

array

Parameters:

codes_flat (mlx.core.array)

veloxquant_mlx.quantizers.cachegen.symbol_entropy_bits(symbols)[source]

Shannon entropy (bits/symbol) of an integer symbol array.

This is the per-symbol size an ideal arithmetic coder approaches. Computed over the empirical distribution of the (delta) code values.

Return type:

float

Parameters:

symbols (mlx.core.array)

veloxquant_mlx.quantizers.cachegen.entropy_coded_bytes(stream, use_delta=True, per_channel=True)[source]

Estimate the entropy-coded size (bytes) of a CodeStream’s codes.

Models a real arithmetic coder by measuring the Shannon entropy of the (optionally delta-transformed) code symbols and multiplying by the symbol count. Params (scale/zero) are added at fp16 like every other method.

Parameters:

  • stream (CodeStream) – the quantized CodeStream.

  • use_delta (bool) – apply the token-delta transform before measuring entropy.

  • per_channel (bool) – fit a separate probability distribution per channel before measuring entropy (§5.1.3: grouping by channel/layer gives much lower entropy than pooling all channels together — the paper reports up to 53% smaller bitstreams from this). If False, all channels are pooled into one distribution (the coarser estimate).

Return type:

int

Returns:

Estimated compressed size in bytes (codes via entropy + fp16 params). Capped at the fixed-width packed size: a real arithmetic coder never stores more than the raw codes (it falls back to raw packing when the symbol stream is incompressible), so neither does this estimate.

veloxquant_mlx.quantizers.cachegen.fixed_width_bytes(stream)[source]

Naive fixed-bit-width packed size (bytes) for the same codes — baseline.

Return type:

int

Parameters:

stream (CodeStream)

veloxquant_mlx.quantizers.cachegen.layer_group_bits(n_layers, base_bits, n_groups=3)[source]

Per-layer bit-width schedule from the paper’s layer-wise sensitivity insight (§5.1.2/§5.2).

Splits the n_layers transformer layers into n_groups contiguous groups (earliest first) and assigns progressively fewer bits to deeper groups: base_bits to the first third, base_bits - 1 to the middle third, base_bits - 2 to the last third (floored at 2 bits). This mirrors the paper’s finding that output quality is far more sensitive to precision loss in shallow layers than in deep ones (Figure 4), so shallow layers get conservative (higher-bit) quantization and deep layers get coarser quantization.

Parameters:

  • n_layers (int) – number of attention-bearing layers.

  • base_bits (int) – bit-width assigned to the shallowest layer group (also the ceiling — no layer is assigned more than this).

  • n_groups (int) – number of contiguous layer groups (paper uses 3).

Return type:

list[int]

Returns:

Length-n_layers list of per-layer bit-widths, non-increasing with depth.

veloxquant_mlx.quantizers.cachegen.cachegen_quant_dequant(x, bits, group_size=32)[source]

Drop-in quant→dequant (values identical to plain group quant).

The entropy layer is storage-only; the reconstructed tensor is exactly the group-quant reconstruction, so this is a drop-in for _group_quant_dequant.

Return type:

array

Parameters:

  • x (mlx.core.array)

  • bits (int)

  • group_size (int)

veloxquant_mlx.quantizers.cachegen.quantize_to_codes_batched(x, bits, group_size=32)[source]

Batched quantize_to_codes over a leading G = B*H axis.

Parameters:

  • x (array) – [G, N, D] fp16/fp32 — G independent heads’ keys or values.

  • bits (int) – bit-width.

  • group_size (int) – tokens per group along axis 1.

Return type:

CodeStream

Returns:

CodeStream whose fields carry the extra leading G axis – codes/scale/zero are [G, n_groups, group_size (or 1), D].

veloxquant_mlx.quantizers.cachegen.dequant_codes_batched(stream)[source]

Batched dequant_codes: reconstruct fp16 [G, n_rows, D].

Return type:

array

Parameters:

stream (CodeStream)

veloxquant_mlx.quantizers.cachegen.token_delta_batched(codes_flat)[source]

Batched token_delta over [G, N, D] — delta along axis 1.

Return type:

array

Parameters:

codes_flat (mlx.core.array)

veloxquant_mlx.quantizers.cachegen.entropy_coded_bytes_batched(stream, use_delta=True, per_channel=True)[source]

Batched entropy_coded_bytes over the CodeStream’s leading G axis.

Same estimate as entropy_coded_bytes(), computed for all G heads (and, if per_channel, all D channels) in one vectorized pass with a single final host read, instead of entropy_coded_bytes called once per head with its own internal per-channel Python loop and two .item() calls per channel (G * D host syncs total there).

Return type:

list[int]

Returns:

Length-G list of estimated compressed sizes in bytes (codes via entropy + fp16 params), each capped at that head’s fixed-width size.

Parameters:
veloxquant_mlx.quantizers.cachegen.fixed_width_bytes_batched(stream)[source]

Batched fixed_width_bytes over the CodeStream’s leading G axis.

Return type:

list[int]

Parameters:

stream (CodeStream)