veloxquant_mlx.weight.quantized_linear module¶
TurboQuant-compressed drop-in replacement for mlx.nn.Linear.
Defines QuantizedLinear, which compresses a weight matrix row-by-row
by L2-normalizing each row, rotating it (Hadamard or QR), and quantizing
the rotated coordinates against a Lloyd-Max codebook fit for unit-norm
post-rotation vectors — the same rotate-then-quantize recipe as the
TurboQuant KV-cache quantizers, applied here to model weights instead of
keys/values. The forward pass dequantizes, unrotates, and rescales by the
stored per-row norm before the matrix multiply.
- class veloxquant_mlx.weight.quantized_linear.QuantizedLinear(in_features, out_features, bits=4, use_hadamard=True, bias=True, seed=42)[source]¶
Bases:
ModuleDrop-in replacement for nn.Linear with TurboQuant weight compression.
- Weight matrix W (out, in) is quantized row-by-row offline:
Normalize each row: w_norm = w / ||w||_2 (store norm)
Rotate: w_rot = preconditioner.apply(w_norm)
Quantize to b-bit Lloyd-Max indices (codebook fits N(0, 1/sqrt(in)))
Store indices + per-row norms
- Forward pass:
w_hat = norms * unrotate(centroids[indices]) out = x @ w_hat.T + bias
Per-row normalization is critical: weight rows have varying magnitudes but the Lloyd-Max codebook is calibrated for unit-norm post-rotation vectors. Without normalization the codebook range is completely wrong.
- Parameters:
in_features (
int) – Input dimension.out_features (
int) – Output dimension.bits (
int) – Bit-width for weight quantization (2, 3, or 4).use_hadamard (
bool) – Use randomized Hadamard (O(d log d)) instead of QR (O(d²)). Auto-falls back to QR if in_features is not Hadamard-compatible.bias (
bool) – Whether the layer has a bias term.seed (
int) – Random seed for the rotation matrix.