mlx_turboquant is a standalone, pip-installable adapter for integrating TurboQuant with MLX LM on Apple silicon. It features a custom Metal kernel for non-uniform quantization and enables efficient vector quantization through a Randomized Hadamard Transform, enhancing performance in deep learning applications.
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 lutoption.
KV Cache Management (Inner-product regime)
- A drop-in
TurboQuantKVCacheis 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:
| Flag | Default | Effect |
|---|---|---|
--kv-bits N | off | Allows enabling of the rotated TurboQuant KV cache at N bits. |
--kv-group-size N | 64 | Determines KV quantization group size. |
--qjl | off | Integrates an unbiased 1-bit QJL residual estimator, ensuring near-fp16 scores at low KV bits. |
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.
No comments yet.
Sign in to be the first to comment.