turboquant-pytorch: a from-scratch TurboQuant port that drops the paper's QJL stage

The problem it targets: KV cache growth in long-context inference

Every token a transformer generates adds a key vector and a value vector to the KV cache, and those vectors stay resident for the whole sequence. At long context the cache becomes a large share of GPU memory, which caps batch size and context length on a fixed card. TurboQuant, the paper this repository implements, attacks that growth with vector quantization: store the cache in fewer bits per element and reconstruct on the fly.

This repository is not a serving integration. It is a research implementation with test harnesses, aimed at someone who wants to read the algorithm as PyTorch code and run the validation scripts themselves. The README states it was tested on Windows 11 with an RTX 3060 at 12GB, and that it also works on Linux. The audience is narrow: people who want to reproduce or extend a quantization method, not people who want a drop-in cache backend for vLLM.

The mechanism: random rotation, then scalar quantization per coordinate

The core idea is described in the README as multiplying each vector by a random orthogonal matrix. After rotation, every coordinate follows a predictable bell-curve distribution, so a single scalar quantizer can be applied to each coordinate independently. The repository uses a Lloyd-Max solver (lloyd_max.py) to find optimal centroids for that distribution. Quantizing means normalize, rotate, round each coordinate to the nearest centroid, and store the indices plus the norm. Dequantizing reverses the order: look up centroids, apply the inverse rotation, restore the norm.

The paper adds a second stage, QJL, which stores one bit of sign information to make inner-product estimates unbiased. The repository implements that as V2. The README's argument against it is specific: attention pushes scores through softmax, softmax amplifies variance exponentially, and QJL's unbiasedness comes with noise that gets magnified. MSE-only quantization has biased inner products but lower variance. The README cites a measurement from another project, scos-lab, of +300% error with QJL against +7.6% without on GPT-2, and reports its own V2 result of 0/27 generation tests passing against V3's 18/18. It also notes QJL does work for vector search, where there is no softmax, which is the paper's other use case.

V3 changes: asymmetric bits, bit-packed storage, protected layers

Four changes define V3. First, QJL is removed and all bits go to reconstruction quality through MSECompressor in compressors_v3.py. Second, keys and values get different bit widths, exposed as TurboQuantV3(key_bits=4, value_bits=2). The README's reasoning is that keys decide which tokens get attended to and need precision, while values are averaged together and their errors cancel. K4/V2 at 5.1x compression reaches 0.9996 cosine similarity and 94% top-1 attention match in the 8K attention test.

Third, storage is bit-packed rather than storing tensors. The README states V2 stored tensors 38% larger than uncompressed, which means the theoretical ratio was not the real one. MSECompressor.compress() uses bit-shifting to get actual packed bytes. Fourth, protected_layers=4 keeps the first and last layers at higher precision, which drops compression from 5.1x to 3.6x but lifts top-1 match from 94% to 99% and top-5 from 97% to 100% in the same attention test.

Getting it running: install paths and the four test entry points

The README gives pip install -r requirements.txt as the base install, pip install -e . for editable local development, and pip install torch --index-url https://download.pytorch.org/whl/cu128 for CUDA PyTorch. Requirements listed are Python 3.10+ and a CUDA-capable NVIDIA GPU.

There are four runnable modules. python -m turboquant.generation_test is the recommended one: it downloads Qwen2.5-3B-Instruct (roughly 2GB on first run) and tests multiple configurations across context lengths for actual text output. python -m turboquant.validate_v3 compares V3 and V2 attention score accuracy side by side. python -m turboquant.test_turboquant runs synthetic algorithm tests against theoretical bounds from the paper and needs no model. python -m turboquant.validate runs the original V2 attention comparison. The repository has no published releases, so the install path is the git tree, not a package index.

The correction that matters more than the headline

The README opens with a correction dated 2026-03-30. An earlier version claimed 18/18 perfect generation at 5x compression. That result came from a test where residual_window=0 caused no compression to happen at all, so the model was running on an uncompressed cache. The issue is linked as #14 and credited to a community member. The corrected generation table tells a different story. K6/V4 with a 128-token fp16 residual window gives EXACT output at both 2K and 4K context, but only about 2x compression. K8/V4 rw=128 is EXACT at roughly 1.6x. K4/V4 rw=128 is PARTIAL at 2K (the model returns AURORA7749 without the hyphen) and MISS at 4K. K4/V4 rw=0 and K4/V2 rw=0 both miss at both lengths, at roughly 3.4x and 5x compression respectively.

The README states the conclusion plainly: 3-4 bit compression without a residual window produces garbage, and high attention score similarity does not guarantee working generation. That second point is the one to carry away. The 0.9996 cosine similarity at 5.1x compression in the attention table is measured on captured KV tensors, not through the cache path, and the README flags it as valid for what it measures while still not predicting text quality. Two tables in the same document point in different directions, and the generation table is the one that reflects deployment.

Where it is the wrong tool

If you need 5x compression in production, this repository does not currently give it to you with correct output. The configuration that produces EXACT generation lands near 2x. The configuration that reaches 5x produces MISS on the needle test. That gap is the central limitation, and it is documented by the author rather than buried.

The residual window is the second constraint. Keeping 128 recent tokens in fp16 caps the compression ratio, because a fixed slice of the cache is never compressed. On very long contexts that slice becomes a smaller fraction and the ratio improves, but the README does not present measurements beyond 4K for generation, so the behavior at 32K or 128K is not established here. The attention tests go to 8K, and those are a different measurement.

There is also a maintenance signal. The last push is 2026-04-23 and no releases were retrieved. There is no published package. Anyone adopting this is tracking a branch. The README itself shows the cost of that: a correctness bug in a test harness invalidated a headline claim for a period, and it was found by an outside contributor rather than by the author's own runs.

Alternatives and the actual difference in approach

The obvious comparison is the reference TurboQuant implementation the paper implies, which includes QJL. The difference is not a matter of tuning. QJL is a residual correction stage that stores one bit per coordinate to make inner-product estimates unbiased. This repository removes that stage entirely and spends the freed bits on reconstruction, on the argument that softmax converts unbiased-but-noisy estimates into worse attention than biased-but-quieter ones. If you believe the softmax variance argument, V3 is the better design for attention. If you are using quantized vectors for retrieval rather than attention, the README concedes QJL works there, and the removal would be the wrong call.

Against a general-purpose quantization library, the difference is scope. This repository implements one algorithm and its variants, with test scripts written for that algorithm. It does not offer calibration pipelines, kernel fusion, or serving hooks. A library would give you those and less visibility into the quantizer itself. The trade is legibility against integration work you would have to do yourself.

Licence and the cost of following a research branch

The repository is MIT licensed, which permits commercial use and modification with the licence text retained. That is the permissive end of the spectrum, and it removes the licence as a reason not to read or fork the code. It does not remove the engineering cost. There are no releases, so upgrades mean pulling master and re-running the test modules to see whether the numbers moved. The README's correction history suggests that is not a formality: a change in a test parameter silently disabled compression once already.

There is no changelog in the supplied material, and no versioning scheme is described. If you build on this, your upgrade procedure is git pull followed by python -m turboquant.generation_test and python -m turboquant.validate_v3, and your acceptance criterion is that the generation table still shows EXACT for your configuration with compressed token counts logged. That is a concrete, project-specific check, and it is the only one the material supports.

Editorial conclusion

Adopt this if you want to study or extend a KV cache quantizer in plain PyTorch and you can accept roughly 2x real compression with a 128-token fp16 residual window. Do not adopt it if you need 3-bit uniform compression at 5x for production serving, or if you need a maintained library with releases. Verify first that python -m turboquant.generation_test reproduces EXACT output on your own model and context lengths, and confirm the compressed token counts are logged, because the README's own correction shows a residual_window=0 setting can silently disable compression.

Official sources

1. Issues
2. License: MIT
3. README
4. tonbistudio/turboquant-pytorch on GitHub