veloxquant_mlx.quantizers.turboquant_prod module

TurboQuant Prod — two-stage unbiased-inner-product quantizer.

The flagship TurboQuant variant: stage 1 rotates and MSE-quantizes the vector at b-1 bits (as in TurboQuantMSE); stage 2 applies QJL (quantizers/qjl.py) to the leftover residual r = x - x_hat_mse. Combining the two stages’ contributions yields an inner-product estimator that is unbiased rather than merely low-error, at roughly one extra bit of storage over single-pass MSE quantization. Also defines TurboQuantProdAdaptive, a subclass that defaults use_adaptive_codebook=True so the stage-1 codebook is refit from observed post-rotation data instead of using the fixed Gaussian/Beta prior.

class veloxquant_mlx.quantizers.turboquant_prod.TurboQuantProd(d, b=3, m=None, seed=42, store=None, use_hadamard=False, use_adaptive_codebook=False, n_calib=64, **kwargs)

Bases: Quantizer

TurboQuant two-stage unbiased inner-product quantizer.

Stage 1: MSE quantize at bit-width (b-1) using rotation + codebook. Stage 2: QJL on the residual r = x - x̂_mse.

Stored representation:

indices: (batch, d) uint8 — MSE codebook indices signs: (batch, m) int8 — QJL signs of residual residual_norm: (batch,) fp16 — ‖x - x̂_mse‖₂

Inner product estimation (unbiased):

IP(q, k) ≈ ⟨q, x̂_mse⟩ + √(π/2)/m · ‖r‖ · ⟨Sq, sign(Sr)⟩

Parameters:

• d (int) – Input dimension.

• b (int) – Total effective bit-width. MSE stage uses (b-1) bits.

• m (int | None) – JL projection dimension.

• seed (int) – Random seed.

• store (ArtifactStore | None) – Optional ArtifactStore.

• use_hadamard (bool) – If True, use randomized Hadamard (O(d log d), Metal-accelerated) instead of QR rotation (O(d²), CPU matmul).

• use_adaptive_codebook (bool)

• n_calib (int)

• kwargs (Any)

static m_default(d, b)

Default JL sketch dimension as a function of (d, b).

At b<=2 the MSE residual is large, so allocate the full d sketch dimensions for the QJL correction. At b>=3 the residual is small and m=min(d, 64) is enough.

Return type:

int

Parameters:

• d (int)

• b (int)

encode(x)

Two-stage encode: MSE + QJL residual.

Parameters:

x (Any) – Array of shape (batch, d), fp16.

Return type:

EncodedVector

Returns:

EncodedVector with indices, signs, and residual_norm.

decode(ev)

Reconstruct approximate vector from two-stage encoding.

x̃ = x̂_mse + ‖r‖ · (√(π/2)/m) · S^T · sign(S·r)

Parameters:

ev (EncodedVector) – EncodedVector with indices, signs, residual_norm.

Return type:

Any

Returns:

Reconstructed array of shape (batch, d), fp16.

estimate_inner_product(q, ev)

Estimate ⟨q, k⟩ unbiasedly using both MSE and QJL stages.

IP(q,k) ≈ ⟨q, x̂_mse⟩ + √(π/2)/m · ‖r‖ · ⟨Sq, sign(Sr)⟩

MSE term is evaluated as ŷ @ q_rot where q_rot = Π·q, so the costly per-key unrotation (x_hat = ŷ @ Π) is replaced by a single query rotation: O(d²) once + O(n·d) dots vs O(n·d²) previously.

Parameters:

• q (Any) – Query vector, shape (d,) or (1, d).

• ev (EncodedVector) – Encoded keys.

Return type:

Any

Returns:

Estimated inner products, shape (batch,), fp16.

class veloxquant_mlx.quantizers.turboquant_prod.TurboQuantProdAdaptive(*args, **kwargs)

Bases: TurboQuantProd

TurboQuantProd with use_adaptive_codebook=True forced on by default.

Parameters:

• args (Any)

• kwargs (Any)