mlx_turboquant

mlx_turboquant is a powerful standalone adapter designed for TurboQuant integration with mlx-lm on Apple silicon. It features a custom Metal kernel that facilitates non-uniform quantization, enabling improved performance and efficiency.

Overview of TurboQuant

TurboQuant, as presented in the work by Zandieh et al. (2025) (paper link), is a data-oblivious vector quantizer that operates without the need for calibration. Its innovative approach leverages a random rotation technique known as the Randomized Hadamard Transform. This method disperses outliers across dimensions, resulting in a near-Gaussian distribution that allows low-bit scalar quantizers to perform optimally. Importantly, the orthogonal rotation maintains the inner product, ensuring that attention scores derived from rotated queries and keys remain unaffected—making it exceptionally effective for KV-cache quantization.

Key Functionalities

Weight Quantization (MSE regime)

• The adapter rotates each weight matrix and subsequently quantizes it. This rotation is a robust, consistent advantage. Users can opt for MLX's efficient affine quantized_matmul (default) or utilize a specialized non-uniform Lloyd–Max LUT Metal kernel via the --mode lut option.

KV Cache Management (Inner-product regime)

• A drop-in TurboQuantKVCache is implemented, enabling the storage of keys in the rotated frame. Queries are also rotated accordingly, highlighting TurboQuant's superior capabilities.

Example Use Cases

Quantizing a Model (Weights)

To convert a model, execute:

turboquant convert --model mlx-community/Qwen3-0.6B-bf16 --out ./qwen3-tq4 --bits 4

This command generates a standardized mlx-lm model directory comprising a safetensors file and a config.json file with a specified quantization_config.

Running the Model

In Python, the model can be loaded and utilized as follows:

import mlx_turboquant as tq

tq.register()  # Register TurboQuant with mlx-lm
from mlx_lm import load, generate
model, tok = load("./qwen3-tq4")
print(generate(model, tok, prompt="Why is the sky blue?", max_tokens=128, verbose=True))

Alternatively, users can run the model via the command line interface:

turboquant generate --model ./qwen3-tq4 --prompt "Why is the sky blue?"

Serving Models via OpenAI-Compatible HTTP API

The command turboquant serve integrates TurboQuant hooks with mlx_lm.server, allowing it to support TurboQuant-quantized model directories seamlessly. The command provides a drop-in OpenAI-compatible endpoint, offering flexibility and ease of use in deploying AI models.

# To serve a quantized TurboQuant model with default settings:
turboquant serve --model ./qwen3-tq4 --port 8080

# To serve with a rotated 4-bit KV cache:
turboquant serve --model ./qwen3-tq4 --port 8080 --kv-bits 4 --qjl

Key TurboQuant-specific flags include:

Memory Efficiency and Performance

TurboQuant demonstrates significant improvements in memory usage and processing speed. This performance is particularly relevant when working with long context sequences. In particular, at 4-bit KV, TurboQuant maintains near-fp16 quality, providing enhanced throughput and reduced memory footprint compared to traditional methods. The innovative 1-bit QJL residual correction further boosts the quality of low-bit KV caches, showcasing TurboQuant's commitment to optimization.

Conclusion

mlx_turboquant significantly enhances the functionality and efficiency of ML models on Apple silicon by integrating TurboQuant's advanced quantization techniques. It is designed to improve performance metrics such as throughput, latency, and memory consumption, making it an effective tool for AI practitioners aiming to leverage data-oblivious quantization methods.