VeloxQuant-MLX
website·Indexed 29 days ago
https://veloxquant-mlx.netlify.app/A quantization library for LLM inference on Apple Silicon, utilizing Metal GPU kernels for compressed KV cache management. It implements various algorithms including VecInfer, RaBitQ, Ada-KV, SVDq, and KVQuant to reduce memory overhead. The framework provides a core API for quantizers and preconditioners, a CLI benchmarking tool, and integration patterns for mlx-lm.
What's inside VeloxQuant-MLX
- CommVQ is a residual Vector Quantization (VQ) structure designed for models using Rotary Position Embeddings (RoPE). Unlike standard quantization, which loses positional information because RoPE is applied after keys are cached, CommVQ commutes with RoPE. This allows position embeddings to be applied to quantized codes during the decode phase without needing to dequantize first. It is particularly useful for GQA models, long contexts (16k+), and scenarios requiring minimal cache format complexity.
Overview of Quantized Johnson-Lindenstrauss (QJL) algorithm
QJL is a simple algorithm in VeloxQuant-MLX that reduces key vectors to 1-bit sign sketches using a random Johnson-Lindenstrauss projection. It requires no calibration, no codebook, and no hyperparameters other than the sketch dimension. It approximates the inner product $\langle q, k \rangle$ as $(d/m) \cdot \langle Aq, b \rangle$, where $A$ is a random Gaussian matrix and $b$ is the sign sketch. The approximation error is bounded by $O(1/\sqrt{m})$ per the JL lemma.Overview of VeloxQuant-MLX KV Cache Compression
VeloxQuant-MLX is a native MLX library designed to compress the KV cache of LLMs on Apple Silicon to prevent memory exhaustion at high context lengths. It implements a pipeline based on TurboQuant research to compress Key (K) and Value (V) vectors. The process involves three main stages: Random Rotation to distribute variance, Lloyd-Max Scalar Quantization for efficient representation, and QJL Residual Correction to recover lost information.Overview of Q-Filters KV cache compression
Q-Filters is a KV cache compression method that identifies important tokens for retention by projecting cached keys onto a single fixed direction (uʰ) per attention head. This direction represents the shared drift of Query and Key distributions away from the origin (joint anisotropy). The expected attention logit for a cached key is proportional to its projection onto this direction, allowing for efficient eviction (project, rank, and drop the bottom) without needing the attention matrix or the current query at eviction time.Overview of Cross-Model KV Cache Transfer
The
veloxquant_mlx.transfermodule allows for transferring a KV cache from a source model to a target model within the same family. This is used to avoid re-running the full prefill on a receiving model when swapping models mid-session (e.g., for cost-quality cascading or routing).Unlike KV cache compression, this does not reduce the cache size; it reduces the receiver's prefill compute. It requires an offline fitting process for each specific model pair to create a mapper that linearly transforms the source cache into the target cache.
Overview of VecInfer vector-quantization
VecInfer is a vector-quantization method that optimizes key-value (KV) cache handling. It applies a per-(head, channel) smooth scale and a Walsh-Hadamard rotation to keys to suppress outliers before encoding them against a trained codebook. The keys are immediately dequantized back to fp16 before attention runs, meaning the live tensor remains in fp16 throughout the process.Overview of SKVQ Quantization
SKVQ (Sliding-window Key and Value Cache Quantization) is a method for reducing the memory footprint of KV caches in LLMs by pushing asymmetric group quantization to very low bit-widths. It employs three primary mechanisms:
- Channel Reordering: Permutes head-dimension channels so that those with similar dynamic ranges share a quantization group, preventing a single wide channel from stretching the scale for an entire group.
- Clipped Dynamic Quantization: Shrinks the quantization window by a factor α ∈ (0, 1] centered on the group midpoint. This saturates extreme values to provide finer resolution for the majority of the data. α is determined per group via grid search against reconstruction error.
- Sliding Window + Sink Filter: Maintains the most recent
skvq_windowtokens in fp16. Once tokens age out of this window, they are quantized once and frozen. The firstskvq_n_sinktokens (attention sinks) are kept in fp16 permanently.
Overview of NestedKV-adapted KV cache compression
NestedKV-adapted is a token-eviction method that compresses the KV cache by maintaining three parallel key-only continuum-memory statistics: stable (global), episodic (block-local), and current (recent-window). Unlike other methods that rely on a single importance signal, NestedKV scores tokens against all three independently and combines them using a training-free head-adaptive blend and a per-token surprise-gated route.
Key characteristics:
- Key-only: Does not require query or attention access.
- Prefill-only compression: Eviction occurs once at the end of the prefill phase. The cache is NOT bounded during the decode phase; newly decoded tokens are appended normally, meaning the cache can grow unboundedly during long generation runs.
- Implementation: This is a VeloxQuant-MLX adaptation of the NestedKV paper, not a faithful port, though it adheres to the paper's default constants.
Overview of AgeTieredKV quantization
AgeTieredKV is a KV-cache quantization method that assigns bit-widths based on a token's age (current_position - token_position). Unlike eviction methods, it retains all tokens but reduces their precision as they age. It uses a three-tier system with two boundaries to determine the bit-width of a token. The method uses an asymmetric min/max group quantizer and is fully deterministic.Overview of VeloxQuant-MLX KV Cache Compression
VeloxQuant-MLX is a library built on MLX for Apple Silicon that compresses the KV cache during inference. Unlike weight quantization (which is a one-time operation), KV cache quantization happens at runtime for every token. This prevents the 'memory wall' on M-series chips where unified memory is shared between the OS, applications, model weights, and the growing KV cache, which often leads to process kills or severe slowdowns during long-context tasks.Overview of VecInfer Vector Quantization
VecInfer is a vector-quantization method for KV caches that applies a per-(head, channel) smooth scale and a Walsh-Hadamard rotation to keys to suppress outliers before quantization. It encodes transformed keys against a trained codebook and immediately dequantizes them back to fp16 before attention runs. Because it uses a 'quantize then dequantize immediately' pattern, the live tensor remains in fp16, meaning reported compression ratios refer to hypothetical byte-accounting rather than actual resident memory savings.Overview of Zero-Calibration KV Cache Compression Algorithms
VeloxQuant-MLX implements a wide array of zero-calibration algorithms for KV cache compression. These methods do not require a calibration dataset and are categorized by their technical approach: quantization, inter-layer redundancy, latent storage, and eviction.