veloxquant_mlx.quantizers.xquant module¶

XQuant quantizer — cross-layer KV cache reuse primitives.

Inspired by “XQuant: Achieving Ultra-Low Bit KV Cache Quantization with Cross-Layer Compression” (arXiv:2510.11236, EMNLP 2025, Yang et al.). Documented as “XQuant-adapted (VeloxQuant-MLX implementation)” — faithful to the cross-layer-reuse core, adapted at the integration boundary (a shared coordinator object rather than a modified attention forward pass).

Core idea:

Adjacent transformer layers produce highly similar K/V tensors. Instead of every layer storing its own quantized cache, layers are grouped into anchor/reuse groups. The anchor layer of a group quantizes K/V normally and publishes the integer codes. The reuse layers borrow those codes and store only their own per-group dequantization parameters (scale/zero), which correct for the small cross-layer magnitude/offset drift. Across a group this drives the effective per-element bit-width well below the anchor’s bit-width.

This module holds the pure, side-effect-free numerics:

pair_layers — assign anchor/reuse roles by index stride.
quantize_codes — anchor-side: min/max group quant → (codes, params).
compute_reuse_params — reuse-side: fit this layer’s scale/zero to shared codes.
dequant_with_params — reconstruct fp16 from codes + params.
cross_layer_similarity — diagnostic (MSE / cosine) for tests + benchmark.

All quantization uses the same asymmetric min/max group scheme as _quant_utils._group_quant_dequant (groups along the token axis), so XQuant composes with the rest of the suite’s correctness guarantees.

class veloxquant_mlx.quantizers.xquant.GroupParams(scale, zero, n_rows, bits)[source]¶

Bases: NamedTuple

Per-group asymmetric quantization parameters.

Variables:

scale – [n_groups, 1, D] fp32 — per-group step size.
zero – [n_groups, 1, D] fp32 — per-group minimum (zero-point offset).
n_rows – int — original (pre-pad) number of token rows.
bits – int — bit-width these params were computed at.

Parameters:

scale (mlx.core.array)
zero (mlx.core.array)
n_rows (int)
bits (int)

scale: mlx.core.array¶: Alias for field number 0

zero: mlx.core.array¶: Alias for field number 1

n_rows: int¶: Alias for field number 2

bits: int¶: Alias for field number 3

veloxquant_mlx.quantizers.xquant.pair_layers(n_layers, group_size)[source]¶

Assign each layer an anchor/reuse role and a group id.

Layers are chunked into contiguous groups of group_size. The first layer in each group is the anchor; the rest reuse it.

Parameters:

n_layers (int) – Number of attention-bearing layers.
group_size (int) – Layers per group (>=2). 2 → pairs; 3 → anchor + 2 reusers.

Return type:

list[tuple[str, int]]

Returns:

Length-n_layers list of (role, group_id) where role is "anchor" or "reuse". A trailing partial group whose anchor is its only member is still a valid (degenerate) anchor.

veloxquant_mlx.quantizers.xquant.quantize_codes(x, bits, group_size=32)[source]¶

Anchor-side quantization: return integer codes and the group params.

Mirrors _group_quant_dequant but exposes the codes (for cross-layer sharing) instead of immediately dequantizing.

Parameters:

x (array) – [N, D] fp16 or fp32 (one head’s keys or values).
bits (int) – Bit-width.
group_size (int) – Tokens per group along axis 0.

Return type:

tuple[array, GroupParams]

Returns:

codes – [n_groups, group_size, D] fp32 integer codes in [0, 2^bits - 1]. params: GroupParams for dequantization.

veloxquant_mlx.quantizers.xquant.compute_reuse_params(x, codes, bits, group_size=32)[source]¶

Reuse-side: fit this layer’s per-group scale/zero to shared codes.

Given the anchor’s integer codes, find the scale/zero that best reconstructs this layer’s keys/values. Because the codes are fixed, the optimal affine fit is the one mapping the code range [0, levels] onto this layer’s group min/max — i.e. a fresh asymmetric min/max calibration on this layer’s data, reusing the anchor’s code assignment.

Parameters:

x (array) – [N, D] this layer’s fp16/fp32 keys or values.
codes (array) – [n_groups, group_size, D] anchor integer codes (for shape only; the affine fit is data-driven, the codes carry the bin assignment).
bits (int) – Bit-width the codes were produced at.
group_size (int) – Tokens per group.

Return type:

GroupParams

Returns:

GroupParams (scale, zero) calibrated to x.

veloxquant_mlx.quantizers.xquant.dequant_with_params(codes, params)[source]¶

Reconstruct fp16 [N, D] from integer codes + group params.

Parameters:

codes (array) – [n_groups, group_size, D] fp32 integer codes.
params (GroupParams) – GroupParams (scale, zero, n_rows).

Return type:

array

Returns:

Reconstructed [n_rows, D] fp16.

veloxquant_mlx.quantizers.xquant.quantize_residual(x, recon, bits, group_size=32)[source]¶

Quantize the reuse-layer residual (x - recon) and return its contribution.

A low-bit correction applied on top of the shared-code reconstruction. With bits == 0 the caller skips this entirely (pure reuse).

Parameters:

x (array) – [N, D] this layer’s true keys/values.
recon (array) – [N, D] the shared-code reconstruction before correction.
bits (int) – Residual bit-width (>= 1).
group_size (int) – Tokens per group.

Return type:

array

Returns:

[N, D] fp16 quantized residual to add back to recon.

veloxquant_mlx.quantizers.xquant.quantize_codes_batched(x, bits, group_size=32)[source]¶

Batched quantize_codes() over an arbitrary leading batch shape.

Parameters:

x (array) – [..., N, D] fp16 or fp32 (e.g. [B, H, S, D]).
bits (int) – Bit-width.
group_size (int) – Tokens per group along axis -2.

Return type:

tuple[array, GroupParams]

Returns:

codes – [..., n_groups, group_size, D] fp32 integer codes. params: GroupParams (scale/zero shaped [..., n_groups, 1, D]).

Bit-identical to calling quantize_codes() once per leading index and stacking the results — the reshape/min/max/round path already vectorizes over any number of leading axes.

veloxquant_mlx.quantizers.xquant.compute_reuse_params_batched(x, codes, bits, group_size=32)[source]¶

Batched compute_reuse_params() over an arbitrary leading batch shape.

Parameters:

x (array) – [..., N, D] this layer’s fp16/fp32 keys or values.
codes (array) – [..., n_groups, group_size, D] anchor codes (shape only).
bits (int) – Bit-width the codes were produced at.
group_size (int) – Tokens per group.

Return type:

GroupParams

Returns:

GroupParams calibrated to x, shaped [..., n_groups, 1, D].

veloxquant_mlx.quantizers.xquant.dequant_with_params_batched(codes, params, out_dtype=mlx.core.float16)[source]¶

Batched dequant_with_params() over an arbitrary leading batch shape.

Parameters:

codes (array) – [..., n_groups, group_size, D] fp32 integer codes.
params (GroupParams) – GroupParams (scale, zero, n_rows), broadcastable to codes.
out_dtype (mlx.core.Dtype)

Return type:

array

Returns:

Reconstructed [..., n_rows, D], cast to out_dtype (default fp16).

veloxquant_mlx.quantizers.xquant.quantize_residual_batched(x, recon, bits, group_size=32, out_dtype=mlx.core.float16)[source]¶

Batched quantize_residual() over an arbitrary leading batch shape.

Parameters:

x (array) – [..., N, D] this layer’s true keys/values.
recon (array) – [..., N, D] the shared-code reconstruction before correction.
bits (int) – Residual bit-width (>= 1).
group_size (int) – Tokens per group.
out_dtype (mlx.core.Dtype)

Return type:

array

Returns:

[..., N, D] quantized residual to add back to recon, cast to out_dtype (default fp16).

veloxquant_mlx.quantizers.xquant.cross_layer_similarity(a, b)[source]¶

Diagnostic: MSE and mean cosine similarity between two layers’ tensors.

Parameters:

a (array) – Same-shape [B, H, S, D] (or [S, D]) tensors.
b (array) – Same-shape [B, H, S, D] (or [S, D]) tensors.

Return type:

dict

Returns:

dict with mse and cosine (mean over all but the last axis).