Source-linked AI summary
WUSH-KV: KV Cache Quantization with Data-Adaptive Transforms
Jiale Chen, Vage Egiazarian, Eldar Kurtić, Torsten Hoefler, Dan Alistarh
TL;DR
Long-context inference faces growing KV-cache memory and bandwidth costs, while aggressive quantization can introduce errors in attention. WUSH-KV calibrates separate key and value transforms for low-bit cache quantization. It reduces reconstruction error and achieves the lowest perplexity among tested transforms, with competitive 2-bit downstream results against OSCAR.
Problem
KV-cache memory and bandwidth costs grow with context length and batch size, and aggressive quantization can introduce large errors into attention scores and outputs.
Source: §1 · p. 1Method
WUSH-KV uses calibration data to build separate key and value transforms, pairing them with clipped quantization and folding the value transform into model weights.
Source: §1 · p. 2Results
At 2-bit, WUSH-KV scores higher than OSCAR on all four Qwen3-8B tasks and is competitive at 4B and 32B; its perplexity is 10.51 versus OSCAR’s 13.74.
Source: §1 · p. 2, §5.2 · p. 10Takeaways & Limitations
The results show that WUSH transforms can improve the quality of low-bit KV-cache quantization.
Source: §6 · p. 13WUSH-KV reuses OSCAR’s tuned clipping ratios without retuning them for WUSH-transformed distributions, so those ratios may be suboptimal.
Source: Setup · p. 11
Abstract
from arXiv · showhide
KV cache memory and bandwidth costs grow with context length and batch size, which limits efficient long-context inference. To address this bottleneck, we introduce WUSH-KV for low-bit KV-cache quantization. It adapts WUSH, which constructs a data-aware transform from the second-order statistics of both factors in a matrix product to reduce quantization error. WUSH-KV uses calibration data to construct separate key and value transforms, with the value transform folded into the model weights and the key transform applied after RoPE. The transforms can be paired with clipped quantizers. For one such quantizer, QuEST INT, we show that, under mild assumptions, the WUSH transform is near-optimal. With this quantizer, WUSH-KV reduces layerwise reconstruction error and achieves the lowest end-to-end perplexity among other tested transforms. For end-to-end evaluation, we integrate WUSH-KV into SGLang using OSCAR-style percentile-clipped affine quantization. At 2-bit, WUSH-KV performs comparably to or outperforms the OSCAR transform across all evaluated models and downstream tasks.
1 Introduction
Long-context inference is constrained by KV-cache memory and bandwidth costs, while aggressive quantization can damage attention scores and outputs. WUSH-KV adapts WUSH to build separate, calibration-based key and value transforms, and the paper reports improved reconstruction and end-to-end quality.
1 Introduction
- The KV cache grows with sequence length and batch size and is repeatedly read during generation, making it a major long-context inference bottleneck.Aggressive low-bit quantization can introduce large errors in attention scores and outputs.Source: §1 · p. 1
- WUSH builds an invertible transform from second-order statistics of both factors in a matrix product, extending its product-aware approach to the KV cache.Key errors affect query-key products, while value errors affect projected attention outputs, motivating separate transforms that account for cache statistics and downstream sensitivity.Source: §1 · pp. 1–2
- WUSH-KV adapts WUSH to KV-cache quantization, constructing separate key and value transforms for each KV head from calibration-time statistics.The value transform can be folded into model weights; the key transform and query-side compensation are applied after normalization, when present, and RoPE.Source: §1 · p. 2
- With QuEST, WUSH-KV reduces attention reconstruction error and achieves the lowest perplexity among tested methods; with OSCAR-style quantization, it is higher on all four Qwen3-8B tasks and competitive at 4B and 32B.The reported downstream comparison is at 2-bit.Source: §1 · p. 2
2 Related Work
Prior KV-cache quantization work addresses cache structure, outliers, and selective precision, while transform-based approaches alter vector representations before quantization. WUSH-KV is positioned against methods with fixed or restricted transforms by allowing general invertible transforms.
2 Related Work
- KIVI uses per-channel key and per-token value quantization, while KVQuant exploits key channel structure, outlier handling, and attention-sink retention; AQUA-KV predicts keys and values across layers and quantizes residuals.These approaches focus on quantizer design, outliers, and selective high-precision retention.Source: §2 · p. 2
- QuaRot redistributes outliers with randomized Hadamard transforms, SpinQuant learns calibration-based orthogonal rotations, and FlatQuant learns Kronecker-structured transforms.These methods change vector representations before quantization.Source: §2 · p. 2
- RotateKV targets 2-bit compression, while TurboQuant uses random rotations and scalar quantizers; concurrent OSCAR learns separate orthogonal key and value transforms and clipping parameters.Unlike OSCAR’s orthogonal transforms, WUSH-KV allows general invertible transforms, permitting anisotropic rescaling to balance cache statistics and downstream sensitivity.Source: §2 · p. 3
3 Background on the WUSH Transform
The WUSH transform is motivated by quantization error in a matrix product: the error’s effect depends on both the quantized tensor and its product partner. WUSH constructs a transform from their second-order statistics.
3 Background on the WUSH Transform
- WUSH builds an invertible transform from second-order statistics of both factors in a matrix product, rather than choosing a transform from the quantized tensor alone.The paper follows Chen et al. (2026) and writes the construction as a function of a Gram matrix and a Hessian.Source: §3 · p. 3
- For a linear layer, quantizing transformed activations and undoing the transform leaves a perturbation whose effect is measured as quadratic output loss.The loss Hessian is formed from the exact weight factor, connecting quantization error to the layer output.Source: §3 · pp. 3–4
- The construction uses the tensor’s Gram matrix and the loss Hessian; damping makes the matrices positive definite, and a scalar factor approximately preserves the transformed tensor’s Frobenius norm.The scalar preserves the norm exactly when γ = 0.Source: §3 · p. 4
- WUSH is invariant to independently rescaling its Gram-matrix and Hessian inputs, so their normalization need not be tracked.The paper expresses the transform using these matrices to simplify notation for the KV-cache derivation.Source: §3 · pp. 3–4
4 WUSH-KV
WUSH-KV adapts product-aware transforms to grouped-query attention by calibrating separate transforms for cached keys and values. It combines clipped quantization with a QuEST guarantee under stated assumptions, retains selected cache positions in full precision, and incurs online computation for key-side transforms.
4.1 KV Cache in Grouped-Query Attention
- In grouped-query attention, multiple query heads share each key/value head, and autoregressive generation caches keys and values because later queries reuse them.Stored keys already include headwise normalization, when present, and rotary positional encoding.Source: §4.1 · pp. 4–5
4.2 WUSH-KV Transforms
- WUSH-KV calibrates one key and one value transform per KV head from the cached tensors’ Gram matrices and local surrogate Hessians.Value transforms are folded into projection weights offline, whereas dense key transforms remain after RoPE and require online application; moving them before RoPE would require position-dependent compensation.Source: §4.2 · pp. 5–6
4.3 Quantizer and Near-Optimality
- The QuEST setup quantizes each token’s d channels as one group, using RMS-based clipped levels; its near-optimality guarantee assumes Gaussian-tail and clipping-alignment conditions.The guarantee applies to the ideal undamped, sensitivity-balanced WUSH transform; experiments use damping.Source: §4.3 · pp. 6–7
4.4 Full-Precision Cache Window
- WUSH-KV keeps the first S_sink positions and newest S_keep positions in full precision, quantizing eligible middle entries in batches of S_flush.New chunks are appended and used for attention in full precision before eligible entries are quantized.Source: §4.4 · p. 7
4.5 Storage and Computation Costs
- Full-precision windows typically cover less than 1% of maximum task sequence length, while each new key and query incurs a d × d product for its key head.Value transforms add no online cost, and key transforms are stored once per key/value head.Source: §4.5 · p. 8
5 Experiments
The experiments compare WUSH-KV with other transforms using attention-stage reconstruction error, WikiText-2 perplexity, and downstream reasoning benchmarks. WUSH performs strongly across these evaluations, with its clearest reported advantages at 2-bit quantization.
5 Experiments
- WUSH-KV is evaluated through controlled attention-stage reconstruction, WikiText-2 perplexity, and downstream reasoning accuracy.For perplexity, the Qwen3-8B transforms are calibrated on 128 FineWeb-Edu sequences of length 32,768 and reused across bitwidths; the WikiText-2 test split is evaluated in 2,048-token sequences.Source: §5 · p. 8, §5.2 · p. 10
5.1 Attention-Stage Quantization Error
- The attention-stage comparison fixes the QuEST quantizer to isolate transform effects, using Qwen3-8B modules calibrated on 128 FineWeb-Edu sequences and evaluated on 32 disjoint sequences.It measures reconstruction at 2, 3, and 4 bits without full-precision cache windows, against activations from the clean BF16 trajectory.Source: §5.1 · pp. 8–9
- At 2-bit, WUSH and WUSH-A have geometric-mean module-output errors of 0.208 and 0.195, respectively, versus 0.325 for H and 0.309 for OSCAR.Across bitwidths, WUSH and WUSH-A have lower module-output error than H and OSCAR for most attention modules; WUSH-A’s gain over WUSH is modest, so the main method uses simpler Hessians.Source: Figure 1 · p. 9, §5.1 · pp. 9–10
5.2 Perplexity Evaluations
- 10.51 versus 13.74 perplexity: at 2-bit, WUSH achieves lower Qwen3-8B WikiText-2 perplexity than OSCAR, and WUSH is lowest among quantized transforms at every tested bitwidth.The full-precision baseline is 9.72; Hadamard’s extreme 2-bit perplexity is attributed to concentrated key-vector spikes spreading into coefficients near reconstruction-level midpoints.Source: §5.2 · p. 10, Table 1 · p. 10
5.3 Downstream Benchmarks
- Downstream tests use OSCAR’s percentile-clipped affine quantizer with WUSH-KV in SGLang, enabling direct comparison with OSCAR but not the exact QuEST proof setting.The evaluation covers three Qwen3 models and four benchmarks; OSCAR’s clipping ratios are reused for WUSH-KV without retuning and may be suboptimal for its transformed distributions.Source: §5.3 · p. 11, Setup · p. 11
- At 2-bit, WUSH-KV scores higher than OSCAR on all four benchmarks for Qwen3-8B, and retains higher accuracy on long-context RULER tasks at the longest tested lengths.Across model sizes, WUSH-KV’s degradation relative to BF16 is comparatively consistent; its generated outputs are at most about 1.5× as long as BF16 outputs, below the nominal 8× reduction in storage per quantized KV entry.Source: Table 2 · p. 12, Setup · pp. 12–13
6 Conclusion
WUSH-KV uses calibration-derived transforms for low-bit KV-cache quantization, paired with clipped quantizers; the conclusion reports a near-optimality guarantee for QuEST under stated assumptions. The appendices describe constraints on key-transform placement around RoPE and the clipping-aware proof strategy.
6 Conclusion
- WUSH-KV constructs calibration-based transforms for cached keys and values, and the QuEST analysis gives a clipping-aware near-optimality guarantee under its stated noise model and assumptions.The conclusion notes that the dense key-side transform remains online after RoPE, adding inference computation, while the value-side transform can be folded into model weights.Source: §6 · p. 13
Appendix A. Alternative Key-Transform Placements
- The alternative key-transform placements appendix contrasts RoPE-commuting transforms with unrestricted pre-RoPE transforms and frames the tradeoff between placement constraints and compensation work.Under the fixed-normalization setting, RoPE commutation and normalization restrict exact folding.Source: Appendix A · p. 16, §A.1 · p. 17, §A.2 · p. 17
A.1 RoPE-commuting transforms
- For fixed normalization gains, exact transforms that commute with RoPE are tightly constrained; generic learned gains leave only paired sign flips, which do not redistribute magnitudes for a symmetric quantizer.The appendix restricts this discussion to fixed-normalization reparameterizations; allowing gains to change can admit additional transforms but adds design choices.Source: §A.1 · p. 17
A.2 Full pre-RoPE transforms
- An unrestricted transform before RoPE preserves full transform freedom but requires position-dependent key compensation; the post-RoPE placement instead uses fixed online transforms for each new key and query.Decoding a pre-RoPE transform would require restoring and rotating cached keys or applying position-specific compensation.Source: §A.2 · p. 17
Appendix B. Clipping-Aware Near-Optimality for the QuEST Quantizer
- The appendix proves clipping-aware near-optimality by starting from the unclipped rounding loss minimized by WUSH, then accounting for clipping error and removed rounding noise.Both clipping effects are included in the comparison.Source: Appendix B · p. 17
B.1 Rounding and clipping errors
This analysis separates clipped quantization loss into in-range rounding noise and exact clipping error, then establishes WUSH as the minimizer of the unclipped trace-product loss. The clipped comparison builds on that optimum.
B.1 Rounding and clipping errors
- The clipped loss combines in-range rounding noise with exact clipping error, and clipping can reduce the total by removing rounding noise from clipped coordinates.The rounding model assumes fixed calibration data and Hessian, with independent uniform noise approximating rounding within the clipping range while clipping remains exact.Source: §B.1 · p. 18
B.2 The unclipped optimum
- WUSH minimizes the transform-dependent trace product over all invertible transforms, attaining the unclipped lower bound.The proof sets γ = 0 in Eq. (1) and substitutes the WUSH transform to attain the bound.Source: §B.2 · p. 19
- An orthogonal transform leaves both traces unchanged, reaching the minimum only when M and H are proportional.Source: §B.2 · p. 19
B.3 The effect of clipping
The clipping analysis bounds WUSH’s clipped loss relative to sensitivity-balanced competitors under stated Gaussian-tail and clipping-alignment conditions. The appendix derives position-dependent query, key, and value Hessians for grouped-query attention and aggregates them to construct WUSH-A.
B.3 The effect of clipping
- The scalar Gaussian clipping analysis shows the optimal endpoint α_b grows without bound as bitwidth increases.The proof rules out a bounded subsequence because its Gaussian-tail clipping error would contradict the vanishing total quantization error.Source: §B.3 · p. 22
- With fixed assumption bounds, the clipping comparison factor is 1+O(α_b^-2); the comparison is restricted to sensitivity-balanced transforms.The bounds must remain fixed across bitwidths for the asymptotic statement, and balancing an arbitrary transform can change its clipped loss.Source: §B.3 · pp. 22–23
Appendix C. Exact Hessians for Grouped-Query Attention
- Appendix C derives exact grouped-query-attention Hessians that depend on the position of the perturbed vector; its final subsection aggregates key and value Hessians for WUSH-A.Source: Appendix C · p. 23
C.1 Setup
- The setup indexes query heads that read each key/value head; causal masking excludes cache positions later than the query position.Queries and keys are taken after normalization and RoPE, and output-projection blocks encode coupling among query heads sharing a key/value head.Source: §C.1 · pp. 23–24
C.2 Hessian from Attention-Output Derivatives
- The appendix differentiates attention outputs with respect to a perturbed query, key, or value vector to obtain local Hessians.Source: §C.2 · p. 24
C.3 Query Hessian
- For a query perturbation, the logit derivative uses the keys, yielding a query Hessian for the affected head and position.Source: §C.3 · pp. 24–25
C.4 Key Hessian
- A key perturbation affects every query head that reads it at causal positions t ≥ s, with derivatives involving the residual v_h,s − o_h,g,t.The resulting key Hessian is d_k × d_k.Source: §C.4 · p. 25
C.5 Value Hessian
- Value perturbations leave attention probabilities fixed because values do not enter the softmax, and the appendix derives the corresponding attention-output derivative.Source: §C.5 · p. 26
D.1 Additional Attention-Stage Errors
The additional attention-stage figures report relative squared L2 errors for several transforms at different bitwidths. Figure 2 isolates keys in attention-stage and single-sided module-output diagnostics, while Figure 3 covers module-output errors at 3 and 4 bits.
D.1 Additional Attention-Stage Errors
- Figure 2’s 2-bit relative squared L2 panels compare I, R, H, OSCAR, WUSH, and WUSH-A; WUSH reduces query-key dot-product, softmax, and single-sided module-output errors.The top row measures dot-product and softmax errors with keys quantized; the bottom row measures module-output error with only keys or only values quantized.Source: §D.1 · p. 26, Figure 2 · p. 27
- Figure 3 compares relative squared L2 module-output errors for I, R, H, OSCAR, WUSH, and WUSH-A at 3 and 4 bits.Source: Figure 3 · p. 27
D.2 Quantizer and Calibration Choices
With WUSH-KV held fixed, OSCAR-style quantization scores higher than QuEST on all four tasks. Calibration comparisons also show that more calibration data did not improve OSCAR’s results in the tested setup.
Quantizer choice
- OSCAR-style quantization yields higher end-to-end scores than QuEST on all four tasks with WUSH-KV and FineWeb-Edu calibration held fixed.Table 3 covers the main Qwen3-8B configurations and quantizer or calibration alternatives.Source: §D.2 · p. 27, Table 3 · p. 28
- QuEST has lower local L2 quantization error for both transforms, although that error need not predict downstream autoregressive quality.A separate single-run AIME 2025 diagnostic scored 0.0% for GPQA Diamond-calibrated OSCAR with QuEST.Source: §D.2 · p. 28
Calibration choice
- Size-matched FineWeb-Edu calibration scored lower than released GPQA Diamond-calibrated OSCAR on all four tasks.Two preliminary AIME runs with GPQA-calibrated WUSH-KV and OSCAR-style quantization scored 16.7% and 10.0%, suggesting that setup was suboptimal for WUSH-KV in these tests.Source: §D.2 · p. 28
D.3 Long-Context Evaluations
Long-context evaluations compare 2-bit WUSH-KV with OSCAR and BF16 on Qwen3 models. WUSH-KV’s relative performance varies by task and model, with stronger results against OSCAR at longer contexts in several comparisons.
Evaluation setup
- The evaluations use 2-bit KV caches and BF16 baselines; MRCR is reported across 0k–128k, with scores averaged over multiple seeds.Qwen3-8B results cover RULER NIAH and MRCR, while the separate Qwen3-32B results are MRCR.Source: §D.3 · p. 28
RULER NIAH
- WUSH-KV scores above OSCAR at every tested RULER NIAH context length, with the largest gap at 128k.It remains close to BF16 at 4k and 8k, then degrades as context length increases.Source: §D.3 · p. 29, Table 4 · p. 29
MRCR
- On Qwen3-8B MRCR, WUSH-KV generally degrades more gradually and retains higher scores at long contexts than OSCAR.MRCR has substantial losses for all quantized methods relative to BF16.Source: §D.3 · p. 29, Table 5 · p. 29, Table 6 · p. 30
- On Qwen3-32B MRCR, OSCAR has slightly higher aggregate scores, while WUSH-KV scores higher at 64k–128k for all three needle counts.The evaluation reports scores by needle count, context-length buckets, overall scores, and AUC.Source: §D.3 · p. 29, Table 7 · p. 30