veloxquant_mlx.cache.svdq_cache module

SVDq KV cache wrapper — sub-2-bit key compression via offline SVD.

Inspired by “SVDq: Singular Value Decomposition-based KV Cache Quantization” (arXiv:2502.15304, Feb 2025, unreviewed preprint). Documented as “SVDq-adapted (VeloxQuant-MLX implementation)” — not a faithful port. See quantizers/svdq.py module docstring for the full list of remaining deviations (per-sequence SVD instead of offline calibration, invented rank heuristic, no model-level accuracy evidence).

Design:

Prefill (first call where S > 1):

  1. Compute truncated SVD of the incoming key batch K ∈ R^{S×D}.

  2. Store the right singular vectors V [D, r] and mean key K̄ [D] as layer state. These are O(D²) and negligible for long sequences.

  3. Project keys into latent space L = (K - K̄) @ V → [S, r].

  4. Apply mixed-precision group quantization to L using the paper’s 8-group fixed bit schedule (Eq. 6): the r latent channels are split into len(bit_schedule) equal-size contiguous groups (largest singular values first), each quantized at its own bit width — 0 bits truncates a group to exactly zero.

  5. Reconstruct fp16 keys for the downstream SDPA call.

  6. Accumulate latents in a growing list (quantized-then-dequantized, so the downstream cache sees fp16, consistent with all other wrappers).

Decode (S == 1 or S < prefill threshold):

  1. Project new key into the already-stored V space.

  2. Quantize and reconstruct fp16.

  3. Pass through to the underlying mlx_lm KVCache.

Values are left at fp16 throughout (the paper notes values have weak low-rank structure; compressing values is left to stacked wrappers).

Byte accounting:

compressed_key_bytes — latent storage at the schedule’s per-group bit rate fp16_key_bytes — what full fp16 would cost (for ratio computation) value_fp16_bytes — values are always fp16 (reported separately)

Effective key bit-width ≈ (r/D) * b̄, where b̄ is the schedule’s mean bit-width (paper Eq. 6: b̄ = mean(bit_schedule)). For the default schedule (8,4,2,1,1,0,0,0), b̄ = 2. At r = 0.5D: effective ≈ 0.5 * 2 = 1.0 bit/element.

class veloxquant_mlx.cache.svdq_cache.SVDqKVCache(config)[source]

Bases: KVCache

KV cache implementing SVDq sub-2-bit key compression.

Parameters:

config (Any) –

KVCacheConfig. Fields consumed: head_dim (D), svdq_rank (int | None — explicit rank; None → energy threshold), svdq_energy_threshold (float, default 0.95), svdq_bit_schedule (Sequence[int], default (8,4,2,1,1,0,0,0) —

paper Eq. 6’s 8-group per-group bit widths, most-significant group first; a 0 truncates that group to exactly zero),

svdq_group_size (int, default 32 — quantization group size

along the token axis, independent of the latent-channel grouping above).

Notes

Does not expose .bits — mlx_lm’s SDPA checks hasattr(cache, "bits") to route to a quantized kernel; we keep that path clean. Exposes .assigned_avg_bits with the effective key bit-width.

update_and_fetch(keys, values)[source]

Fit per-head SVD on prefill (or project into the existing basis on decode), mixed-bit quantize the latents; values pass through fp16 unchanged.

Parameters:

  • keys (mlx.core.array)

  • values (mlx.core.array)

property compressed_key_bytes: int

Realized stored bytes for the compressed key cache (per-head mixed-bit latent codes + amortized projection basis, all heads/batches).

property fp16_key_bytes: int

Hypothetical fp16 key cost if nothing were compressed.

property value_fp16_bytes: int

Actual value cost — values are stored fp16 throughout (key-only method).

property assigned_avg_bits: float

schedule’s mean bit-width scaled by r/D, averaged across heads (each head can land at its own rank/schedule under energy-threshold auto-rank — see _run_prefill_svd).

Type:

Effective key bit-width

property rank: int

Actual SVD rank used after energy-threshold selection, averaged across heads (rounded down) — see assigned_avg_bits for why heads can differ. Use the per-head ranks directly (e.g. via a subclass instance’s internal state) if the per-head breakdown matters.

property merge