veloxquant_mlx.quantizers.composite module

Mixed-precision quantizer that splits a vector’s channels by magnitude.

CompositeQuantizer implements outlier-aware quantization: a small set of high-magnitude “outlier” channels is routed to a high-bit-width child quantizer while the remaining “inlier” channels go to a low-bit-width child, and the two child ``EncodedVector``s are nested under one composite encoding. This lets any two existing per-channel quantizers be combined to protect the few dimensions that dominate reconstruction error without paying full bit-width across the whole vector.

class veloxquant_mlx.quantizers.composite.CompositeQuantizer(outlier_quantizer, inlier_quantizer, outlier_idx, total_dim)[source]

Bases: Quantizer

Split-channel quantizer that routes outlier and inlier channels separately.

Outlier channels are quantized with a high-bit-width quantizer. Inlier channels are quantized with a low-bit-width quantizer.

The outlier_encoded and inlier_encoded fields of the returned EncodedVector nest the children’s outputs.

Parameters:

  • outlier_quantizer (Quantizer) – Quantizer for high-magnitude (outlier) channels.

  • inlier_quantizer (Quantizer) – Quantizer for low-magnitude (inlier) channels.

  • outlier_idx (ndarray) – Integer indices of outlier channels in the original vector.

  • total_dim (int) – Total dimension d of the original vector.

encode(x)[source]

Encode by routing channels to respective child quantizers.

Parameters:

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

Return type:

EncodedVector

Returns:

CompositeEncodedVector with nested outlier and inlier encodings.

decode(ev)[source]

Decode and reassemble channels.

Parameters:

ev (EncodedVector) – CompositeEncodedVector.

Return type:

Any

Returns:

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

estimate_inner_product(q, ev)[source]

Estimate ⟨q, k⟩ by summing contributions from both children.

Parameters:

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

  • ev (EncodedVector) – CompositeEncodedVector.

Return type:

Any

Returns:

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