PaperlayerView exact revision on arXiv

Source-linked AI summary

A JoLT for the KV cache: Near-Lossless KV Cache Compression via Joint Rank-bit Allocation

Rahul Krishnan, Volker Schulz

arXiv:2607.12550v4cs.LGcs.CLmath.OC

TL;DR

Long-context inference stores a growing KV cache, while existing compression methods do not jointly allocate low-rank structure and precision under a shared budget. JoLT combines partial Tucker decomposition with a rotated quantized residual and a joint allocator. Across five models, it reports 2–3× compression with less than 0.2% perplexity drift without retraining.

  • Problem

    The KV cache can outgrow model weights during long-context, large-batch inference, and existing approaches do not jointly allocate rank and precision under a fixed byte budget.

    Source: §1 · pp. 1–2

  • Method

    JoLT applies partial Tucker decomposition to grouped prefill caches, quantizes the rotated truncation residual, and jointly allocates ranks and residual bit-widths under a global byte budget.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Across five models from four architecture families, JoLT achieves 2–3× cache compression with less than 0.2% perplexity drift and no retraining.

    Source: §1 · p. 2

  • Takeaways & Limitations

    In the reported free zone, partial Tucker decomposition and rotated low-bit residuals make low-rank compression competitive with quantization without retraining.

    Source: §5 · p. 8

    The fused decode kernel is slower than uncompressed attention at every measured point, with latency 7.5× to 65.9× the dense baseline in the reported microbenchmark.

    Source: §4.5 · p. 8

Abstract

from arXiv · show

The key-value (KV) cache is the dominant memory bottleneck in long-context language model inference. Existing compression methods apply low-rank factorization or quantization independently, without jointly allocating rank and precision under a shared storage budget. We introduce JoLT, a training-free compressor that treats grouped prefill caches as fourth-order tensors and applies partial Tucker decomposition along the token and feature modes, the two axes that carry low-rank structure, while leaving the head and layer modes intact. A rotated low-bit quantizer captures the truncation residual, and a single Lagrangian dual allocates per-group Tucker ranks and residual bit-widths under a global byte constraint. FlashJoLT replaces the exact token-mode SVD with a randomized approximation that matches JoLT within the free zone at a fraction of the compression cost, and a fused Triton decode kernel evaluates attention directly over the stored factors without materializing dense KV tensors. Across five models from four architecture families, covering multi-head attention, grouped-query attention, and mixture-of-experts architecture, JoLT achieves 2 - 3x compression with less than 0.2% perplexity degradation, without retraining. On RULER at 64K context with LLaMA-3.1-8B, retrieval accuracy remains near-lossless through 3x and declines by only 0.90 and 2.40pp at 4x and 5x, respectively. JoLT demonstrates that tensor-aware low-rank decomposition and quantized residuals, unified under a single storage budget, achieve near-lossless KV-cache compression across diverse model architectures without retraining.

1 Introduction

Long-context inference is constrained by a KV cache that can exceed model-weight memory, while existing low-rank and quantization approaches do not jointly allocate rank and precision. JoLT combines tensor-aware truncation, a quantized residual, and a global byte-budget allocator, with near-lossless results across models and downstream tasks.

1 Introduction

  • The KV cache grows with batch size, context length, layers, and KV-head dimensions, and can outgrow model weights at large batch sizes and long contexts.Reducing cache memory is presented as a direct way to serve longer contexts and larger batches on a device.Source: §1 · p. 1
  • Prior methods factor a single cache unfolding, quantize every entry, or combine both with a preset split rather than jointly allocating rank and precision.The allocation gap motivates a shared byte budget for the two compression choices.Source: §1 · p. 1
  • JoLT truncates token and feature modes, stores the truncation residual after a random orthogonal rotation, and jointly allocates ranks and residual bit-widths under one byte budget.The head and layer modes remain intact, and keys and values receive separate allocations.Source: §1 · p. 2
  • FlashJoLT approximates the costly token-mode SVD, while a fused Triton kernel computes attention over stored factors and packed residuals without dense reconstruction.The overview describes the shared Lagrangian allocator distributing bytes across Tucker ranks and residual bit-widths for all cells.Source: §1 · p. 2, Figure 1 · p. 2
  • At 2–3× compression, JoLT keeps perplexity drift below 0.2% across five models and retrieval near-lossless through 3× on RULER at 64K context.For LLaMA-3.1-8B, retrieval accuracy drops by 0.90 and 2.40 pp at 4× and 5×, respectively.Source: §1 · p. 2

2 Related Work

Existing KV-cache compression methods use low-rank factorization, quantization, or hybrids, but generally set rank and precision independently. JoLT uses the tensor structure of the cache and jointly allocates those resources under a shared budget.

KV-cache quantization

  • KV quantizers lower the bit-width of every entry, but operate at fixed rates; JoLT instead covers 2–10× continuously by jointly optimizing rank and precision.The passage describes quantizing model weights and Tucker factors as complementary routes to further compression.Source: §2 · pp. 2–3

Low-rank and structural KV compression

  • Matrix-based structural methods factor a single two-dimensional cache slice, while hybrids typically fix rank, precision, or sparsity rather than trading rank and precision under one budget.Token-eviction methods are described as orthogonal and potentially combinable with JoLT.Source: §2 · p. 3

Tensor decompositions and rate allocation

  • JoLT uses partial Tucker decomposition to leave index-like head and layer modes intact, and applies rotated low-bit quantization to the truncation residual.Its Lagrangian allocation jointly chooses Tucker ranks and residual bit-widths under the byte budget, a combination not previously applied to KV-cache compression in the passage’s account.Source: §2 · p. 3

3 Method: Joint Lagrangian Tucker (JoLT)

JoLT compresses grouped prefill caches by truncating token and feature modes, quantizing the residual, and allocating ranks and bits under a shared storage budget. FlashJoLT approximates the costly token-mode decomposition while retaining near-identical reconstruction error in the reported free zone.

3 Method: Joint Lagrangian Tucker (JoLT)

  • JoLT groups consecutive layers whose token-mode singular vectors align, and decomposes keys before rotary position embedding because post-RoPE keys have higher reconstruction error at matched storage.The grouping is computed from the cache being compressed without separate calibration data; the passage gives 11 groups for Mistral-7B-v0.3 at T=1024.Source: §3 · p. 3
  • Each key or value cache in a layer group is a fourth-order tensor; JoLT truncates its token and feature modes while keeping layer and head modes at full rank.The token and feature bases are shared across heads and layers within the group, and the token truncation precedes feature truncation.Source: §3 · p. 3
  • Because Tucker truncation leaves a large spectral tail, especially for values, JoLT stores the residual after a feature-axis rotation and quantizes it at a selected low bit-width.The seeded rotation is shared across cells and regenerated at decode, and b=0 means no residual is stored.Source: §3 · p. 4
  • The allocator minimizes summed reconstruction error subject to a global byte budget, jointly choosing ranks and residual bits because their best split varies by cell.Its surrogate multiplies the discarded-mass estimate by the fraction unrecovered by the rotated quantizer; the coordinate-wise rank search is a heuristic with two truncated modes.Source: §3 · pp. 4–5
  • The surrogate treats discarded Frobenius mass uniformly and misses asymmetric key-error effects; beyond the free zone on LLaMA-2, it can make extra key rank appear cheaper than residual bits.A token-rank cap is introduced to guard against this failure mode.Source: §3 · p. 5

3.1 Fast variant: FlashJoLT

  • FlashJoLT replaces the costly exact token-mode SVD with a randomized approximation that computes leading singular vectors, while retaining the exact feature-mode SVD and the same allocator.The token-mode unfolding’s exact SVD is costly in the typical long-context regime where T exceeds |g| nh dh.Source: §3.1 · p. 5
  • FlashJoLT estimates uncomputed singular-value energy from the tensor norm and appends it as a synthetic tail entry so the allocator can make rank-bit decisions on the truncated spectrum.Its token-rank cap also bounds rank spending to guard against the surrogate limitation.Source: §3.1 · p. 5
  • Across 2–3× compression on Mistral at T=1024, FlashJoLT’s reconstruction-error gap from JoLT is |∆K| ≤0.0093 and |∆V| ≤0.0011, with end-to-end compression speedups of 1.04× to 3.97×.Timing includes the full compression pipeline, including layer grouping.Source: §3.1 · p. 5

4 Experiments

Experiments test JoLT across five models and downstream tasks, finding near-lossless perplexity in the 2–3× free zone and examining behavior beyond it. Ablations and decode measurements expose trade-offs in residuals, allocation, and latency.

4.1 Setup

  • The evaluation covers five models from four architecture families; prompts are split in half, with the first half compressed and the second half scored for perplexity.Perplexity uses a deterministic WikiText-2 and C4 mixture and reports the mean over three seeds unless otherwise noted.Source: §4.1 · p. 5

4.2 Perplexity: the near-lossless free zone

  • Every 2–3× free-zone cell across five models stays below 0.2% perplexity drift, with the largest at +0.17%; at 4×, LLaMA-2 reaches +27.28%.Degradation beyond the free zone varies by model rather than attention type.Source: §4.2 · p. 6, Table 1 · p. 6

4.3 Downstream evaluation

  • At 64K on LLaMA-3.1-8B, RULER retrieval stays near-lossless through 3×, then drops by −0.90 pp at 4× and −2.40 pp at 5×.At 5×, RULER drops only 2.40 pp despite +51.27% perplexity degradation, because the metrics assess different slices of quality.Source: §4.3 · p. 6, Figure 3 · p. 7
  • On GSM8K, 2×–3× results are near-lossless on Mistral and LLaMA-2; LongBench on Mistral stays within 0.52 pp of baseline through 4×.GSM8K free-zone confidence intervals include zero; LongBench pooled mean shifts are −0.09, +0.05, and −0.52 pp at 2×, 3×, and 4×.Source: §4.3 · pp. 6–7

4.4 Ablations

  • Dropping the residual raises 2× perplexity from 6.54 to 7.27, while joint allocation beats K/V-decoupled allocation at 8× and per-group ranks beat uniform ranks across all 15 grid cells.The 8× comparison is 8.14 versus 8.42; per-group allocation's peak advantage is 0.97 PPL at n=512, 4×.Source: §4.4 · p. 8
  • Residual bits saturate at four: moving from 4 to 8 bits improves perplexity by less than 0.001 at 2×.The ablation compares residual bit-widths at the 2× compression ratio.Source: §4.4 · p. 8

4.5 Fused decode kernel

  • At batch 1, the fused kernel uses 1.43×–4.73× less peak KV memory at tested targets but is 7.5×–65.9× slower than dense attention; keys require full tile reconstruction because RoPE does not commute with factorization.The kernel avoids writing dense KV buffers, but the reported RULER task latency is 20.7s and 18.2s at 2× and 3× versus 1.6s uncompressed.Source: §4.5 · p. 8

5 Conclusion

JoLT combines partial Tucker compression with rotated low-bit residuals under a single byte budget, achieving a 2–3× near-lossless range without per-model tuning. The conclusion also identifies decode latency and allocation and streaming limits beyond that range.

5 Conclusion

  • JoLT combines partial Tucker truncation and rotated low-bit residuals under one Lagrangian budget, sustaining 2–3× compression with less than 0.2% perplexity drift.The allocation works at any target byte budget without per-model tuning.Source: §5 · p. 8
  • The fused decode kernel reduces peak KV memory by up to 4.73×, but incurs substantial latency because RoPE requires key reconstruction per tile.The position-dependent rotation does not commute with token-mode factorization, preventing rank-space scoring.Source: §5 · p. 8
  • Folding RoPE into factor scoring could remove per-tile key reconstruction and is the primary open kernel problem.Source: §5 · p. 9
  • Past the free zone, the Lagrangian surrogate can over-allocate key rank because it prices Frobenius mass equally without modeling key-error effects through softmax.A tighter surrogate or improved cap policy could extend the near-lossless range beyond 3×.Source: §5 · p. 9
  • The current global prefill pass suits prompt caching but does not support streaming generation without efficient online rank updates that avoid full SVD recomputation.Source: §5 · p. 9

Reproducibility Statement

The paper reports public code, detailed method and evaluation specifications, and the hardware and software used for experiments.

Reproducibility Statement

  • JoLT, FlashJoLT, and baseline-harness source code is available, and the method is specified in Section 3 and Section 3.1 and Appendix A.Specifications cover layer grouping, partial Tucker decomposition, rotated residuals, allocation, and the randomized variant’s cap policy.Source: Reproducibility Statement · p. 9
  • Evaluation protocols, seeds, and evaluated-token counts are stated alongside experiments, and all five models and listed datasets are public.Source: Reproducibility Statement · p. 9
  • Experiments used NVIDIA A100 GPUs and, for two 4K and 8K LLaMA-3.1-8B RULER rows, one RTX 3090, with PyTorch 2.5.1 and CUDA 12.1.Models were in bf16 and decomposition numerics in float32; fused-kernel memory measurements reproduced bit-identically across the two A100 PCIe SKUs.Source: Reproducibility Statement · p. 9

AI Use Statement

The authors report using generative AI for implementation assistance, manuscript editing, and related-work search, while retaining responsibility for the final content.

AI Use Statement

  • Generative AI assisted with software implementation, experiment scaffolding, manuscript drafting and editing, and related-work search and citation checks.Source: AI Use Statement · p. 9
  • The authors state that AI was not used to formulate research questions, design methods or protocols, analyze results, or formulate scientific claims.Source: AI Use Statement · p. 9
  • The authors reviewed and verified AI-assisted code and text and take responsibility for the final content.Source: AI Use Statement · p. 9

A Fast variant: speedup and cap calibration

FlashJoLT improves isolated token-mode SVD timing, but measured end-to-end speedups vary by cache length. The section also evaluates cap calibration and its effect on allocation and perplexity beyond the free zone.

A Fast variant: speedup and cap calibration

  • 5.24×–9.97× speedups are reported for LLaMA-2 on the isolated token-mode SVD stage, not the complete production pipeline.The Mistral isolated timing uses T=1024 and qcap=32 on pre-grouped tensors.Source: A Fast variant: speedup and cap… · p. 13, Table 3 · p. 13
  • End-to-end FlashJoLT speedups were 1.04× at cache length 821 (C=3.7) and 3.97× at cache length 1024 (C=4) on an A100 SXM4 40 GB.These timings include layer grouping; isolated-stage speedups do not transfer directly to the production path.Source: A Fast variant: speedup and cap… · p. 13
  • For C ≤5, the Mistral calibration sweep passed the ∆K ≤0.025 reconstruction-error gate at every context length from 512 to 8192.The cap policy is a no-op for T ≤1024 relative to the production default it replaced.Source: A Fast variant: speedup and cap… · p. 13
  • Against matched uncompressed perplexity of 5.56, FlashJoLT, JoLT, and capped JoLT reached +14.21%, +26.47%, and +14.46%, respectively.The capped exact solver was within 0.25 pp of FlashJoLT and 12 pp from uncapped JoLT.Source: A Fast variant: speedup and cap… · p. 13
  • Without the cap, six of fourteen groups assigned keys token rank 55–81 and spent 18.8 MB on factors; the cap reduced key-factor storage to 7.1 MB.The capped groups spent 34.8 MB on key residuals instead of 23.7 MB; value allocation was unchanged.Source: A Fast variant: speedup and cap… · p. 13

B Full perplexity grids

Beyond the free zone, perplexity degradation varies by model rather than attention type, and operating points above 3× require per-model validation. OLMoE degrades gradually through 4×, while LLaMA-2 shows a sharp 4×–5× step and non-monotone behavior at higher ratios.

B Full perplexity grids

  • On LLaMA-2, the 4×–5× perplexity step is sharp, with peaks at 7× for n=512 and n=1024 and at 8× for n=2048 before values fall through 10×.The allocation log shows monotone changes from 6× to 10×, which the authors say do not explain the non-monotone perplexity; they attribute it to a mismatch between the Frobenius surrogate and downstream perplexity.Source: B Full perplexity grids · p. 13, Table 5 · p. 14
  • On Mistral, perplexity degradation past the free zone is gradual at every context length.The full grid reports perplexity relative to the uncompressed baseline across context lengths.Source: B Full perplexity grids · p. 13, Table 4 · p. 14
  • At 8×/3×, the perplexity multiplier is 1.24× for Mistral, 1.70× for Qwen, 2.01× for LLaMA-2, and 2.31× for LLaMA-3.1.The degradation ordering does not track attention type or KV-head count; above 3×, the authors advise per-model validation rather than an architecture rule.Source: B Full perplexity grids · pp. 13–14
  • OLMoE’s 3-seed perplexity drift is +0.02% at 2×, +0.11% to +0.14% at 3×, and +3.0% to +3.2% at 4×, with no cliff.Its attention projections are standard dense layers, so its KV cache is a standard tensor to which JoLT applies without modification.Source: B Full perplexity grids · p. 14, Table 6 · p. 14

C Downstream details

Downstream results are near baseline at moderate compression, while performance at higher ratios varies by task, model, and context length. The section also reports a matched-memory GSM8K comparison and a 16K LongBench evaluation setup.

C Downstream details

  • On GSM8K, 2× and 3× are within noise of baseline on both models; at 4×, Mistral drops 8.49 pp while LLaMA-2 degrades by −2.20 pp.The LLaMA-2 confidence interval still touches zero.Source: C Downstream details · p. 14, Table 7 · p. 15
  • At roughly 4× matched memory on Mistral GSM8K, KIVI 2-bit scores 40.18% (+3.11 pp over int4’s 37.07%), TurboQuant scores 38.89% (+1.82 pp), and JoLT scores 35.94% (−1.14 pp versus int4).KIVI and TurboQuant use a slightly looser budget; the int4-native byte rate pushes JoLT beyond its free zone at short context.Source: C Downstream details · p. 15, Table 8 · p. 15
  • On Mistral-7B RULER, 2× and 3× remain near-lossless through 32K; the single-seed 4× penalty shrinks from −6.10 pp at 4K to −0.20 pp at 32K.At 8K, 2× and 3× partial-credit deltas are −0.20 pp and −0.70 pp on a 97.37 baseline; pooled 16K and 32K deltas are −0.45 pp at 2× and −0.225 pp at 3×.Source: C Downstream details · p. 15, Table 9 · p. 16
  • Table 10 evaluates FlashJoLT against the uncompressed baseline on six LongBench tasks at 16K, with 200 examples per task and seed 0.The reported task metrics are F1, ROUGE-L, and classification accuracy, depending on the task.Source: C Downstream details · p. 16, Table 10 · p. 16

D Full ablation grids

The ablations show that allocation choices matter most under tighter budgets: joint allocation can pull ahead at 8×, and per-group ranks outperform uniform ranks across the tested cells. The reported ablation grids use a 100-chunk subset and are not directly comparable to full-corpus values.

D Full ablation grids

  • The ablation perplexities use a 100-chunk subset with 51,100 evaluated tokens per cell, so they are comparable within the appendix but not with full-corpus Table 1 values.The joint and per-group headline configurations are separate runs that agree within 0.001 PPL at n=1024 and 8×.Source: D Full ablation grids · p. 16, Table 11 · p. 17
  • At 8×, the joint solver pulls ahead of K/V-decoupled allocation except at n=2048, where they remain within 0.17 PPL; at 2× they track closely across contexts.The comparison is reported by context length in Table 12.Source: D Full ablation grids · p. 16, Table 12 · p. 17
  • Per-group rank allocation beats uniform rank at every tested cell, with the largest gap, 0.97 PPL, at n=512 and 4×.Table 13 resolves this comparison by context length.Source: D Full ablation grids · p. 16, Table 13 · p. 17

E Cross-method reconstruction and perplexity benchmark

JoLT and FlashJoLT rank strongly among low-rank methods at 2×, but the best reconstruction or perplexity result near 4× depends on the cache component and baseline. The comparisons also note implementation deviations for several reproduced baselines.

E Cross-method reconstruction and perplexity benchmark

  • Palu, KQ-SVD, and TurboQuant are reimplementations rather than reproductions of originally reported results, with method-specific deviations from their papers.The deviations include prompt-cache SVD or uniform rank choices for some low-rank arms, and uniform 4-bit quantization without an outlier-channel split for TurboQuant.Source: E Cross-method reconstruction and… · pp. 17–18
  • At 2×, JoLT and FlashJoLT rank first and second among low-rank methods on keys and values on both architectures; on Mistral values, JoLT error is 0.00571 versus 0.15546 for Tucker without residual.The Mistral value error is roughly 27× lower; the fixed int8 baseline has lower key error (0.0047 on Mistral) but higher value error.Source: E Cross-method reconstruction and… · p. 18, Table 14 · p. 18
  • At 4×, int4 per-channel has lower key error than JoLT on Mistral (0.080 against 0.118) and LLaMA-2 (0.077 against 0.120), while JoLT has lower value error on both.For values, the respective comparisons are 0.090 against 0.132 on Mistral and 0.087 against 0.124 on LLaMA-2.Source: E Cross-method reconstruction and… · p. 18, Table 14 · p. 18, Table 15 · p. 18
  • On Mistral perplexity, Tucker-only reaches +10.81% at 2× and +21.23% at 4×, while Per-head SVD reaches +8.39% and +15.74%.The LLaMA-2 grid reports FlashJoLT at +14.16% at 4×; near 4×, TurboQuant and KIVI outperform it on perplexity, while the JoLT family’s advantage is in the 2–3× band where fixed-rate quantizers do not operate.Source: E Cross-method reconstruction and… · p. 19, Table 16 · p. 19
  • At 2×, fixed-rate 8-bit baselines match free-zone perplexity quality on both models, but neither int8 per-channel nor fp8 e4m3 has an operating point above 2×.Their perplexity stays within 0.02% for int8 and 0.04% for fp8 of the uncompressed baseline; int8’s ratio is 2T/(T+4) < 2.Source: E Cross-method reconstruction and… · p. 19
Loading 2607.12550v4…