veloxquant_mlx.cache.palu_cache module

PALU KV cache — true low-rank latent storage for keys and values.

Inspired by “PALU: Compressing KV-Cache with Low-Rank Projection” (arXiv:2407.21118, ICLR 2025). Documented as “PALU-adapted (VeloxQuant-MLX implementation)” — not a faithful port.

The defining property (and what separates this from the repo’s SVDq cache): the cache stores the latent codes ``[B, H, S, r]`` directly and never materialises full fp16 ``[B, H, S, D]`` keys/values for storage. SVDq reconstructs full fp16 keys, hands them to the parent mlx_lm KVCache, and so wins only on byte-accounting/bandwidth. PALU keeps the cache itself low-rank and reconstructs to fp16 only at attend time, so its peak storage is genuinely smaller.

Design:

Prefill (first call, S > 1):

  1. Partition the H heads into palu_n_head_groups contiguous groups.

  2. For keys (and values, unless palu_quantize_values is False but still low-rank): fit one shared projection per group via group-head SVD.

  3. Project each head into its group’s latent space, mixed-bit quantize the [S, r] latents, and seed the growing latent buffers.

Decode (S == 1):

  1. Project the new key/value into the stored group projections.

  2. Mixed-bit quantize and append to the latent buffers.

Every call returns reconstructed fp16 [B, H, S, D] (latent → dequant → reconstruct) for the downstream SDPA, but storage stays latent.

Because the parent’s fp16 self.keys / self.values ring buffer is bypassed, this class manages offset itself and stores latents in per-group lists, then reconstructs the full sequence on each fetch.

Byte accounting:
compressed_key_bytes / compressed_value_bytes — latent storage at the mixed-

bit effective rate (this is the real stored size, not a reconstruction)

fp16_key_bytes / fp16_value_bytes — what full fp16 would cost projection_bytes — V_g + mu_g (fp32), amortised

Effective bit-width ≈ (r / D) * weighted_avg(hi_bit, lo_bit).

class veloxquant_mlx.cache.palu_cache.PALUKVCache(config)[source]

Bases: KVCache

KV cache implementing PALU true-latent low-rank K/V compression.

Parameters:

config (Any) – KVCacheConfig. Fields consumed: head_dim (D), palu_rank (int | None — explicit rank; None → energy threshold), palu_energy_threshold (float, default 0.90), palu_n_head_groups (int, default 4), palu_hi_bit / palu_lo_bit / palu_hi_fraction (mixed-bit), palu_group_size (int, default 32), palu_quantize_values (bool, default True).

Notes

Does not expose .bits — mlx_lm’s SDPA checks hasattr(cache, "bits") to route to a quantized kernel; we keep that path clean and hand it reconstructed fp16. Exposes .assigned_avg_bits (max of key/value effective bit-widths).

update_and_fetch(keys, values)[source]

Fit group-head SVD projections on prefill, else project + mixed-bit quantize into the latent buffers; 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 from the latent buffers; returns the number actually dropped.

Return type:

int

Parameters:

n (int)

property nbytes: int

Total realized stored bytes (compressed key + value latents).

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (mixed-bit latents + projection basis, all heads/batches).

property compressed_value_bytes: int

Realized stored bytes for the compressed value cache (mixed-bit latents + projection 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 projection_bytes: int

Per-group SVD projection basis (V + mu, fp32) overhead, shared across heads in a group.

property rank: int

Latent rank actually stored (uniform across head groups).

property assigned_avg_bits: float

Effective bit-width — max of the key and value latent rates.

property merge