WUSH-KV: KV Cache Quantization with Data-Adaptive Transforms
Jiale Chen, Vage Egiazarian, Eldar Kurtić, Torsten Hoefler, Dan Alistarh
cs.LG
2026-09-30
WUSH-KV learns separate invertible key/value transforms for KV caches and folds the value side into weights; at 2-bit, Qwen3-8B perplexity hits 10.51 vs OSCAR's 13.74.
The KV cache is a hard cost in long-context inference: every attention layer stores the keys and values of past tokens, the cache grows linearly with sequence length and batch size, and each decoding step reads it in full. Quantizing it to 2 bits cuts storage per entry to one eighth of BF16, but at that precision the error can wreck the model. On Qwen3-8B, a Hadamard rotation with 2-bit quantization pushes WikiText-2 perplexity from 9.72 (BF16) to 258.29.
Applying an invertible linear transform before quantizing is the standard remedy. Fixed transforms such as Hadamard ignore the data; learned rotations (SpinQuant, FlatQuant, RotateKV) calibrate on data, but the concurrent OSCAR still confines its transforms to orthogonal matrices. The premise here: a cached key only enters the query-key dot product, and a cached value only enters the output projection, so a good transform must balance cache statistics against the sensitivity of those two paths. An orthogonal transform can rotate but not rescale directions anisotropically, so it cannot do that balancing in general. The ISTA and ETH Zürich team (including Dan Alistarh and Torsten Hoefler) adapts WUSH, a transform originally built for weight-activation quantization, to the KV cache, allows general invertible transforms, and proves a near-optimality guarantee.
WUSH is built in closed form: take the Gram matrix of the tensor being quantized (its second-order statistics) and the Hessian of the loss with respect to that tensor's quantization perturbation, and combine the two into a transform T. Quantize T X instead of X, and undo with T^{-1}. The original construction is provably optimal for floating-point formats and asymptotically optimal for integer ones, and the transform is block diagonal with block width equal to the quantization group size.
WUSH-KV learns one pair of transforms per KV head:
Placement decides the engineering. The value transform folds offline into the value and output projection weights, so it costs nothing at inference. The key transform cannot fold: folding requires it to commute with RoPE and the headwise normalization, and those joint constraints leave only sign flips, which redistribute nothing. Placing it before RoPE keeps the full freedom, but the compensation then depends on the cached position, forcing per-position work over the whole cache at every decode step. So the key transform stays after RoPE on the online path: one d×d product per new key and per query.
The transform is decoupled from the quantizer and pairs with per-token clipped quantizers. For the QuEST integer quantizer, under an additive rounding-noise model plus tail and clipping-alignment conditions, the ideal WUSH transform is near-optimal among all sensitivity-balanced invertible transforms, with an expected loss ratio of at most 1+O(αb^{-2}). Cache management follows convention: in the end-to-end runs the first 64 positions (the sink window) and the most recent 256 stay in BF16, the middle is quantized in pages, and the windows cover under 1% of the maximum sequence length. Calibration is forward-pass only: 128 FineWeb-Edu sequences, about 12 minutes once for Qwen3-8B on a single L40S with a 39 GiB peak, and the transforms are reused across bitwidths and tasks.
Three levels of evaluation. First, a controlled ablation that fixes the QuEST quantizer and varies only the transform, over all 36 attention modules of Qwen3-8B with keys and values both at 2 bits (geometric-mean module-output reconstruction error):
| Transform | Error |
| Hadamard | 0.325 |
| OSCAR | 0.309 |
| WUSH-KV | 0.208 |
| WUSH-A (attention-aware Hessian) | 0.195 |
Second, WikiText-2 perplexity on Qwen3-8B (BF16 baseline 9.72):
| Transform | 4-bit | 3-bit | 2-bit |
| Hadamard | 9.77 | 10.06 | 258.29 |
| OSCAR | 9.77 | 10.19 | 13.74 |
| WUSH-KV | 9.73 | 9.82 | 10.51 |
The gap widens as bits drop: at 2 bits WUSH-KV is 3.23 points below OSCAR and 0.79 above BF16. The Hadamard collapse has a clean mechanism: the first attention module's keys have spikes concentrated in essentially one coordinate, and the Hadamard transform spreads them into near-equal coefficients that land close to the midpoints between adjacent 2-bit reconstruction levels.
Third, end-to-end: integrated into SGLang with OSCAR's percentile-clipped affine quantizer, reusing OSCAR's tuned clipping ratios without retuning for WUSH, across Qwen3-4B-Thinking, 8B, and 32B on AIME 2025, MATH-500, GPQA Diamond, and LiveCodeBench v6 over three seeds. At 8B, WUSH-KV outscores OSCAR on all four tasks; where OSCAR degrades sharply the margin grows, notably on LiveCodeBench v6 at both 8B and 32B; at 4B and 32B the two are close overall. Fixed transforms do not hold up at 2 bits: Hadamard frequently exhausts the generation budget, and the TurboQuant-style variant scores zero on some tasks. On long-context RULER NIAH and MRCR, WUSH-KV holds higher accuracy than OSCAR at the longest tested lengths. Outputs do get longer, at most about 1.5x BF16 on average, still far below the 8x per-entry storage reduction.
This is direct evidence that a 2-bit KV cache can work for reasoning models. The 8x storage cut converts into larger batches or longer contexts, and the recipe is practical: a one-time 12-minute calibration, transforms reused across bitwidths, zero cost on the value side, and a working SGLang integration. Against the learned-rotation line, the closed form needs only forward-pass statistics, no backpropagation. The increment over OSCAR is easy to state: relax orthogonal to general invertible, and add product-aware Hessians. Honestly, it is an incremental improvement; the gains concentrate at 2 bits and at 8B, and at 4 bits everyone is near lossless.
The authors' own list: the dense key transform must run online after RoPE, adding a matrix product per token, and throughput or system-level evaluation is left to future work. Further points: the near-optimality proof covers only the QuEST quantizer, while the end-to-end runs use the OSCAR-style affine quantizer, which the authors say does not admit the same proof, so the guarantee and the deployed configuration do not match; the proof is for the undamped ideal transform and the experiments use damping. Clipping ratios are borrowed from OSCAR without retuning for the WUSH-transformed distributions, which leaves headroom. The OSCAR baseline is a re-run under the authors' pipeline because the concurrent paper's scores could not be fully reproduced, so the head-to-head carries a caveat. The attention-aware Hessian variant (WUSH-A) only trims reconstruction error slightly, so the marginal return on finer sensitivity modeling is uncertain. And 2-bit perplexity of 10.51 against 9.72 is cheap, not free.