DOCSEARCHbyAlgolia
Submit Docs

TurboQuant+ Documentation

repository·main·Indexed 27 days ago

https://github.com/thetom/turboquant_plus

A research-driven implementation of TurboQuant (ICLR 2026) for compressing transformer KV caches by 3.8-6.4x using PolarQuant and Walsh-Hadamard rotation. It includes integrations for vLLM, llama.cpp, mlx-swift-lm, and a Python reference implementation. The repository also features REFRACT (v0.3.2.3), a multi-axis KV-cache fidelity scoring tool for LLMs supporting llama.cpp, MLX, vLLM, and SGLang.

Tokens
67.9K
Snippets
98
Records
382
Agent score
92%

What's inside TurboQuant+

  1. Overview of longctx retrieval and rescue modes

    main

    longctx Overview

    longctx is an open-source retrieval companion service designed for million-token inference. It is implemented as a single-purpose FastAPI service that provides two primary capabilities:

    1. Indexing: Chunks and embeds text spans into a FAISS index, which can be scoped per session or per repository.
    2. Retrieval: Returns top-K spans for a given query, with support for optional reranking.

    Operating Modes

    longctx can be used in two distinct architectural patterns:

    • Standalone Retrieval (RAG Layer): Acts as a code-aware RAG layer wired in front of any OpenAI-compatible inference engine. It handles the indexing and retrieval of context spans independently of the model's KV cache.
    • Rescue Layer (TriAttention V3 Integration): Acts as a safety net for engines using query-aware KV-cache eviction (specifically TriAttention V3). When the eviction policy removes necessary information from the KV cache, longctx retrieves that information and serves it back to the next prefill to prevent recall failure.

  2. Compare TurboQuant performance across hardware configurations

    main

    The hardware comparison matrix tracks performance metrics (Prefill, Decode, and Perplexity) for different hardware setups using q8_0 and turbo3 modes.

    Key Hardware Performance Characteristics:

    • M5 Max (Apple10, Tensor API: true): Shows gradual decode degradation (~4% per context doubling). Stable through 48K+ context.
    • M1 Max/Ultra (Apple7, Tensor API: false): Experiences catastrophic performance drops at long context due to constant cache thrashing (e.g., M1 Max Decode 42K drops to ~0.09x of baseline).
    • NVIDIA (via unixsysdev implementation): Performance is lower than M1 Metal because it lacks TurboQuant optimizations like graph-side WHT, block-32, fp16 LUT, and split LUT.

    Performance Metrics Definitions:

    • Prefill: Speed of processing initial tokens.
    • Decode: Speed of generating subsequent tokens.
    • turbo3/q8_0: The ratio of turbo3 performance relative to the q8_0 baseline.
    • PPL: Perplexity (lower is better; turbo3 typically shows a slight increase in PPL compared to q8_0).

  3. Understand the TurboQuant+ Architecture

    main

    TurboQuant+ compresses KV cache vectors using a pipeline of norm extraction, random rotation (Walsh-Hadamard Transform + random sign flips), and optimal scalar quantization (Lloyd-Max).

    Compression Ratios:

    • turbo4: 16 centroids (4-bit) $\approx$ 3.8x compression
    • turbo3: 8 centroids (3-bit) $\approx$ 5.1x compression
    • turbo2: 4 centroids (2-bit) $\approx$ 7.5x compression

    Important Implementation Note: For production use, avoid the QJL (1-bit error-correction) stage as it amplifies variance and introduces attention noise. Instead, use TurboQuantMSE for V cache, or implement 4 to 8-bit PolarQuant on K.

  4. Understand the TurboQuant4 (turbo4) 4-bit format

    main
    The optimized turbo4 format is a 4-bit quantization scheme for KV cache vectors. Unlike the original broken version that used 3-bit PolarQuant plus a 1-bit QJL (Quantized Johnson-Lindenstrauss) residual correction, the improved version uses the full 4 bits for PolarQuant with 16 optimal centroids. This approach provides better quality than turbo3 and q4_0 while maintaining high decode speeds and compression ratios.

  5. Understand the distinction between TurboQuant and TurboQuant+

    main

    When evaluating quantization performance, it is critical to distinguish between the original Google Research algorithm and the TurboQuant+ implementation.

    • TurboQuant (Google paper, 2026): An abstract algorithm using QJL on V residuals (_prod) and a deterministic post-rotation step (_mse).
    • TurboQuant+ (TheTom): An implementation using WHT rotation, matched-norm scaling, and Lloyd-Max centroids. It skips QJL on V cache and uses asymmetric K/V quantization (e.g., K=q8_0 and V=turbo) by default.
    • TurboQuant+ in production: Implementations in llama.cpp and MLX running asymmetric K=q8_0 / V=turbo4 on GQA ≥6:1 models.

  6. Evaluate MLX Quality Suite for Qwen2.5-7B-Instruct-8bit

    main
    The MLX Quality Suite provides benchmarking results for the mlx-community/Qwen2.5-7B-Instruct-8bit model using the turbo4 configuration. It evaluates model performance across three primary dimensions: NIAH (Needle In A Haystack) for retrieval accuracy, KL Divergence for distribution fidelity, and Context Size Scaling for throughput and memory efficiency.

  7. Understand why MSE is an unreliable metric for KV cache quantization

    main

    Mean Squared Error (MSE) is a poor proxy for the actual quality of KV cache compression in transformer models. While MSE optimizes for uniform error reduction across all coordinates, the attention mechanism is non-linear and concentrates sensitivity on a sparse subset of coordinates.

    Key failure modes:

    • Attention Sensitivity: Small L2 errors in $K$ can lead to large errors in the direction of $Q$ that matters for the next token.
    • Softmax Discontinuity: Perturbations can flip the dominant attention weight (the 'max-pre-softmax bucket'), causing discrete output changes that MSE cannot predict.
    • Sign Inversion: Improvements in MSE can systematically correspond to worse model behavior (higher KL divergence) because MSE-improving changes can trigger catastrophic attention weight flips.

    When MSE is acceptable: MSE is only a valid proxy when the downstream operator is approximately linear (e.g., distributed-mean-estimation, top-k inner-product search, or weight-quantization in small-perturbation regimes). Attention is non-linear and falls outside this regime.

  8. Summary of TurboQuant+ Production Capabilities

    main

    As of 2026-03-26, the following features are available in the TurboQuant+ ecosystem:

    • Compression: 4.6x KV cache compression at 99% prefill speed.
    • Quality: PPL +1.1% vs q8_0 (utilizing norm correction).
    • Retrieval: 80% single needle, 100% multi-key distractors through 32K context.
    • Layer-Adaptive Mode 2: q8_0 quality at 3.5x effective compression.
    • Temporal Decay: 30-34% additional memory savings at zero decode cost.
    • Testing: 511 total tests across all modules with 100% coverage.

  9. Why MSE is an insufficient metric for KV quantization

    main

    Mean Squared Error (MSE) is an unreliable proxy for the actual quality of KV cache quantization in deployment. Key reasons include:

    • MSE-KL Disconnect: Improvements in K-MSE do not necessarily translate to improvements in KL divergence. Some methods show 13% MSE improvements while suffering 90% mean KL regressions.
    • Operator Non-linearity: The non-linearity of the attention softmax breaks the link between MSE and KL.
    • Prefill-Decode Distribution Shift (PDS): Standard calibration often uses short prefill-only sequences. The distribution shift during decoding (PDS) is undetectable at fit time but causes catastrophic failures during generation.
    • Centroid Fragility: Centroid optimization based on MSE can be 'calibration-positive' but 'deployment-negative'.

  10. Understand TurboQuant performance bottlenecks

    main

    TurboQuant currently experiences a significant speed regression compared to standard quantization like q8_0. While it meets compression targets (e.g., 4.9×), it generates at ~2.4 tok/s compared to ~85.0 tok/s for q8_0 on certain hardware (M5 Max).

    Root Cause: The dequantize_turbo3_0 function is called too frequently within the flash attention kernel. Instead of dequantizing a full block once, it performs a full 128-element dequantization (including expensive Inverse WHT rotations and stack allocations) for every small chunk (4 or 16 elements). This results in 8 to 32 redundant calls per block, increasing compute requirements from ~2,000 ops to up to 64,000 ops per block.

  11. Understand the WHT inverse rotation optimization strategy

    main

    The project investigates moving the Walsh-Hadamard Transform (WHT) inverse rotation from per-block dequantization (which has $O(128)$ complexity per 4-element access) to a graph-level approach involving Q forward and V inverse rotation ($O(1)$ amortized per access).

    Current Status:

    • The graph-level rotation approach currently results in a Perplexity (PPL) of 23.5, whereas the baseline dequant inverse approach achieves a PPL of 6.194.
    • As a result, the Dequant inverse rotation has been restored to maintain quality, though it is slower (~10.7 tok/s).
    • The graph rotation implementation is preserved as TODO comments for future investigation.
    • Virtual method infrastructure (get_turbo_rot_forward/get_turbo_rot_inverse) and rotation tensor allocation/initialization in the KV cache remain in the codebase.

  12. Performance characteristics of turbo4 4-bit PolarQuant

    main

    The turbo4 4-bit quantization provides performance that matches or exceeds q8_0 across various context lengths. Key metrics include:

    • Perplexity (PPL): +0.23% vs q8_0 (highest quality tested).
    • Decode Speed: ~79.87 tok/s (93% of q8_0).
    • Prefill Speed: 101-105% of q8_0 across all tested context lengths (2K to 32K).
    • Needle In A Haystack (NIAH): 93.9% accuracy (31/33), beating q8_0 (90.9%).
    • KLD: 0.0096 (40% lower than turbo3).
    • Dense Model Decode: Matches q8_0 (17.25 vs 17.17 tok/s).