A JoLT for the KV Cache: Near-Lossless KV Cache Compression via Joint Tucker and JL-Residual Allocation for LLMs

2026-07-14

AI Summary

This work addresses the memory bottleneck of key-value (KV) cache in large language models during long-context inference by introducing FlashJoLT, the first framework to model the KV cache as a third-order tensor. FlashJoLT integrates partial Tucker decomposition to compress token and feature dimensions, Johnson–Lindenstrauss random projections combined with low-bit quantization to preserve residual information, and a Lagrangian dual formulation to jointly optimize compression ratio and reconstruction accuracy. Randomized SVD further accelerates tensor decomposition. Evaluated on Mistral-7B and LLaMA-2-13B, FlashJoLT achieves 2–3× KV cache compression with relative Frobenius errors as low as 0.006–0.009, while maintaining perplexity, GSM8K accuracy, and RULER retrieval performance on par with baselines, and accelerating decomposition by 5–13×.

Abstract

The key-value (KV) cache has become the dominant memory cost of transformer inference. It grows with batch size, context length, and depth, and at long context it, rather than the model weights, sets the ceiling on throughput. Two families of methods reduce it. Low-rank methods factor two-dimensional slices of the cache, either per-head matrices or cross-layer feature blocks, and quantization methods lower the bit-width of every entry. Neither family exploits the fact that the cache at a layer is naturally a third-order tensor whose three axes, the heads, the tokens, and the features, carry very different amounts of redundancy. We take this tensor view directly. Our method, JoLT, applies a partial Tucker decomposition that compresses only the token and feature axes while leaving the head and layer axes intact, and then restores the energy that truncation discards with a Johnson-Lindenstrauss (JL) rotated low-bit residual. A single Lagrangian dual allocates the Tucker ranks and the residual bit-widths together, per layer group and separately for keys and values, under one byte budget. The result is a near-lossless 2-3x compression: perplexity, GSM8K accuracy, and RULER needle-in-a-haystack retrieval all stay at or within statistical noise of the uncompressed baseline on both a grouped-query-attention model (Mistral-7B-v0.3) and a multi-head-attention model (LLaMA-2-13B). At 2x, JoLT reconstructs the cache to relative Frobenius error 0.009 (K) and 0.006 (V) on both architectures, roughly an order of magnitude below cross-layer SVD and 4-bit quantization. A randomized-SVD variant, FlashJoLT, delivers a 5-13x compression-time speedup at matched quality.

Problem

• KV cache
• memory compression
• tensor decomposition
• large language models
• inference efficiency

Innovation

• KV cache compression
• Tucker decomposition
• Johnson-Lindenstrauss residual
• tensor-based compression
• near-lossless compression