veloxquant_mlx.cache.xquant_cache module¶
XQuant KV cache wrapper — cross-layer KV cache reuse.
Inspired by “XQuant: Achieving Ultra-Low Bit KV Cache Quantization with
Cross-Layer Compression” (arXiv:2510.11236, EMNLP 2025). Documented as
“XQuant-adapted (VeloxQuant-MLX implementation)” — faithful to the cross-layer
reuse core, adapted at the integration boundary via a shared
XQuantCoordinator rather than a modified attention forward pass.
- Per-layer roles (assigned at build time by
pair_layers): - Anchor layer — quantizes K/V with asymmetric min/max group quant, publishes
the integer codes to the coordinator, returns the fp16 reconstruction.
- Reuse layer — fetches the paired anchor’s codes for the same token range,
fits its own per-group scale/zero to those codes (correcting the small cross-layer drift), and reconstructs. Stores only its params (+ optional low-bit residual) — never a full code tensor. That is the byte win.
Both keys and values are compressed (XQuant is a both-tensor method; values are
typically quite correlated across layers too). Set xquant_base_bits for the
anchor; reuse layers inherit the same bit-width plus an optional residual.
- Byte accounting:
Anchor: full
compressed_key_bytes/compressed_value_bytes(codes + params). Reuse: onlyreuse_param_bytes(scale+zero per group) + optional residual.fp16_*always reflects the uncompressed cost for the ratio.
Degenerate case: with no coordinator (single isolated layer) the cache behaves as a plain anchor — useful for unit-testing the anchor path in isolation.
Performance (VeloxQuant-MLX#529): _quantize_anchor/_reconstruct_reuse
used to loop for b: for h: over the batch, calling the scalar
quantize_codes/compute_reuse_params/dequant_with_params/
quantize_residual primitives once per head and restacking — the same
unbatched-per-head-loop pattern already fixed in ChunkKV, CurDKV, GEAR and
RocketKV (#525-#528). Those primitives’ group-quant math (reshape into
groups, min/max, round, along the token axis) already vectorizes over any
number of leading axes, so this fix adds *_batched variants in
quantizers/xquant.py and calls each once per update_and_fetch instead
of B*H times. Verified bit-identical to the old loop; isolated
B=1,H=32,S=128,D=128 anchor-path benchmark: ~3.55ms -> ~1.9ms mean per step
(~1.9x faster).
- class veloxquant_mlx.cache.xquant_cache.XQuantKVCache(config, role='anchor', group_id=0, coordinator=None, n_readers=1)[source]¶
Bases:
KVCacheKV cache implementing XQuant cross-layer reuse for one layer.
- Parameters:
config (
Any) –KVCacheConfig. Fields consumed:xquant_base_bits(int, default 2),xquant_residual_bits(int, default 4 – see VeloxQuant-MLX#380;0 is unsafe on real models, adjacent-layer correlation is too weak for pure reuse to reconstruct coherently),
``xquant_group_quant_size``(int, default 32).
role (
str) –"anchor"or"reuse"(default"anchor").group_id (
int) – Cross-layer group this layer belongs to.coordinator (
XQuantCoordinator|None) – SharedXQuantCoordinator(None → degenerate anchor).n_readers (
int) – Number of reuse layers in this anchor’s group (group_size - 1). Only meaningful forrole="anchor"; tells the coordinator how many fetches to expect before a published segment can be reclaimed.
- update_and_fetch(keys, values)[source]¶
Anchor: group-quantize K/V and publish codes to the coordinator. Reuse: fetch the paired anchor’s codes and fit this layer’s own scale/zero (+ optional residual).
- Parameters:
keys (mlx.core.array)
values (mlx.core.array)
- property compressed_key_bytes: int¶
full codes+params for an anchor, params-only (+residual) for a reuse layer.
- Type:
Realized stored bytes for the compressed key cache
- property compressed_value_bytes: int¶
Realized stored bytes for the compressed value cache (anchor codes+params; reuse layers do not separately compress values here).
- property reuse_param_bytes: int¶
Bytes charged to a reuse layer for its own fitted scale/zero params (+ optional residual) — 0 for an anchor.
- property merge¶