PaperlayerView exact revision on arXiv

Source-linked AI summary

PuzzleKV: Page-Wise Low-Rank Decomposition for KV Cache Compression

Zizhong Wang, Jieying Wang, Zhao Zhang, Jiajia Li

arXiv:2608.23843v1cs.LG

TL;DR

Long-context LLM inference is constrained by the memory cost of the KV cache, motivating compression that preserves useful information at lower storage. PuzzleKV addresses this with training- and calibration-free, page-wise low-rank decomposition and direct mixed attention. At approximately 60% of original storage, it retains more than 96% of Full KV performance across evaluated models and benchmarks, while quantization enables 18.7% storage with more than 93% performance.

  • Problem

    Long-context inference is limited by the KV-cache memory footprint, which grows with context length and serving batch size.

  • Method

    PuzzleKV independently factorizes completed KV pages within each layer and KV head, directly attends over dense and factorized pages, and incrementally compresses newly eligible pages.

  • Results

    At approximately 60% of original storage, PuzzleKV achieves more than 96% of Full KV performance across both evaluated models and all benchmark settings.

  • Takeaways & Limitations

    Page-wise decomposition delivers strong accuracy under constrained KV-cache budgets and can be combined with quantization for more aggressive compression.

  • Takeaways & Limitations

    The current implementation is a batch-one standalone prototype, so integration into a batched paged serving engine remains future work.

Abstract

from arXiv · show

Long-context inference in large language models (LLMs) is increasingly limited by the memory required for the key-value (KV) cache. KV cache compression addresses this problem by reducing the storage cost of previous tokens. Among existing approaches, low-rank compression is particularly attractive because it represents every token in reduced dimensions. Previous low-rank methods typically derive fixed projection spaces from model weights, construct fixed spaces from calibration activations, or construct a shared basis over a broad cache region. Such representations may not capture detailed but important information. We partition each per-head KV cache into fixed-length logical pages and observe substantial low-rank structure within individual pages. Based on this observation, we propose PuzzleKV, a training- and calibration-free method that treats each completed page as an independent compression unit. PuzzleKV decomposes pages within each layer and KV head, computes attention directly over dense and factorized pages, and incrementally compresses newly eligible pages during autoregressive decoding. Experiments across models, context lengths, and benchmarks demonstrate the effectiveness of PuzzleKV under matched storage budgets. At approximately 60% of the original KV cache storage, PuzzleKV achieves more than 96% of Full KV performance across both evaluated models and all benchmark settings, with substantial gains over Global SVD on RULER and competitive performance on LongBench. To achieve a more aggressive compression ratio, PuzzleKV can be further combined with quantization while retaining more than 93% of Full KV performance using only 18.7% of the original storage.

1 Introduction

Longer contexts increase KV-cache memory pressure, motivating compression methods that preserve information while reducing storage. PuzzleKV uses independent page-wise low-rank decomposition and achieves strong accuracy under substantially reduced storage budgets.

  • Motivation: A 128K-token Qwen3-32B sequence requires about 32 GiB of BF16 KV cache, whose footprint grows linearly with context length and serving batch size.The cache contains 2LHkvTd KV elements before batching.
  • Motivation: Existing low-rank methods derive compression spaces from model weights, calibration activations, or broad cache regions, potentially missing detailed local information.These approaches use fixed or shared bases rather than independently adapting to each page.
  • Method: PuzzleKV performs direct attention over dense and factorized pages without reconstruction and incrementally converts newly eligible pages during decoding.Its implementation includes batched page decomposition and mixed attention processing.
  • Method: PuzzleKV independently factorizes completed KV pages within each layer and KV head while retaining every token through page-local low-rank representations.The method is training- and calibration-free and keeps sink and recent-window regions dense.
  • Results: At approximately 60% of Full KV storage, PuzzleKV achieves more than 96% of Full KV performance across all evaluated models and benchmark settings.It also substantially gains over Global SVD on RULER and remains competitive on LongBench.

2 Related Work

KV-cache compression methods primarily use low-rank decomposition, quantization, or eviction, with different trade-offs in representation and token retention. PuzzleKV distinguishes itself by independently factorizing each completed page within each layer and KV head while retaining every token.

  • Low-Rank KV Cache Compression: Low-rank decomposition stores keys and values in lower-dimensional spaces, reducing KV-cache storage without discarding tokens.Existing methods differ in whether their compression space comes from weights, calibration activations, online cache updates, or cross-layer redundancy.
  • Low-Rank KV Cache Compression: PuzzleKV independently factorizes each completed page within each layer and KV head rather than using a shared sequence-level basis.This page-wise granularity is the method’s central distinction from the other low-rank approaches described.
  • KV Cache Quantization: Quantization reduces memory by storing keys and values at lower numerical precision, using schemes including asymmetric low-bit and distortion-aware vector quantization.PuzzleKV reduces dimensionality instead and can apply quantization to its low-rank factors.
  • KV Cache Eviction: Eviction reduces storage by retaining selected useful tokens, but evicted tokens are permanently unavailable for future inference.PuzzleKV instead retains every token in compressed form.

3 Methodology

PuzzleKV uses KV pages as independent low-rank compression units, preserving locally relevant directions while supporting hybrid attention and incremental updates during decoding. Its page-wise design is motivated by consistently observed low-rank structure within pages across models, layers, heads, and page sizes.

  • Design principle: Each KV page receives its own compact basis, allowing local regions to use relevant directions while collectively covering richer sequence-level structure.This design depends on individual pages being strongly low-rank.
  • Design principle: Page-level low-rank structure appears across both evaluated models and all tested page sizes, with required ranks below full page rank at P = 32.The required ranks vary across layers and KV heads but remain below the full page rank across evaluated layers and heads.
  • Prefill construction: During prefill, PuzzleKV partitions each layer and KV head into fixed-size pages, factorizes completed pages, and keeps the attention sink and recent window dense.Historical pages are stored as low-rank factors in a hybrid cache.
  • Mixed attention: PuzzleKV computes attention jointly over dense and factorized pages using page-local states and online-softmax merging without reconstructing historical dense pages.The parenthesization contracts queries with small factors first, avoiding materialization of reconstructed keys and values.
  • Incremental updates: During decoding, newly filled pages leaving the local window are factorized once, moved into low-rank storage, and never revisited.This keeps the hybrid cache current while retaining every token in compressed form throughout generation.

4 Experiments

PuzzleKV is evaluated across two models and long-context benchmarks under matched storage budgets, with analyses of accuracy, quantization, design choices, runtime, and memory. It preserves strong task performance while reducing KV storage, though decomposition adds prefill latency.

  • Experimental Setup: PuzzleKV is evaluated on Qwen3-8B and Llama-3.1-8B-Instruct using RULER and LongBench under a matched KV storage budget of approximately 60%.RULER uses 16K and 32K inputs, while LongBench measures real-world long-context tasks.
  • Main Results: PuzzleKV achieves the highest compressed-method average on RULER in every model–context setting, retaining 96.2%–99.3% of Full KV performance.On Llama3.1, scores are 88.82 and 84.39 at 16K and 32K; on Qwen3, scores are 90.73 and 87.53.
  • Main Results: On LongBench, PuzzleKV retains 98.9% and 98.7% of Full KV performance on Llama3.1 and Qwen3, respectively, while outperforming Global SVD and Palu.It ranks first among compressed methods on Qwen3 and second on Llama3.1, within 0.28 points of H2O.
  • Quantization Compatibility: Combining PuzzleKV with INT4 reduces the KV storage ratio to 0.187 while remaining within 2.58 points on Llama3.1 and 0.53 points on Qwen3 of unquantized PuzzleKV.Factor quantization therefore extends page-wise low-rank compression to a more aggressive storage regime with limited accuracy loss.
  • Ablation Study: The default rank allocation is (rK, rV) = (16, 14), because increasing rank improves performance with diminishing gains beyond 16 and insufficient key rank cannot be offset by more value rank.The allocation keeps the total rank budget at 30 while assigning slightly more rank to keys.
  • Ablation Study: Page size P = 32 is selected because larger pages improve accuracy from 16 to 32 but provide smaller gains at 64 while noticeably increasing decomposition latency.PuzzleKV remains effective across all evaluated page sizes at the same compression ratio.
  • System Efficiency: At 32K context, PuzzleKV adds 20.9% TTFT over Raw KV, but increases TPOT by only 0.075 ms/token, concentrating runtime cost in prefill decomposition.The added decoding overhead is described as negligible.
  • System Efficiency: At 16K context with 256 generated tokens, PuzzleKV uses 58.76% of the Raw KV footprint, while incremental page conversion adds only 12.45 MiB at peak.The measured footprint is 0.17 percentage points above the theoretical page-level ratio.

5 Discussion

PuzzleKV’s page-based organization aligns compression with fixed-size cache blocks and supports regular batched GPU decomposition, but the current implementation remains a batch-one standalone prototype.

  • The current implementation is a batch-one standalone prototype.Its page-based abstraction has potential for integration into a batched paged serving engine such as vLLM.

6 Conclusion

PuzzleKV combines page-wise low-rank decomposition with direct mixed-format attention and incremental conversion during decoding. Across evaluated models, contexts, and benchmarks, it maintains strong accuracy under constrained storage, outperforms sequence-level low-rank compression, and supports more aggressive budgets with quantization.

  • PuzzleKV combines batched page decomposition, direct attention over dense and factorized pages, and incremental conversion during decoding.
  • Across models, context lengths, and benchmarks, page-wise decomposition achieves strong accuracy under constrained KV cache budgets and substantially outperforms sequence-level low-rank compression.
  • PuzzleKV can be combined with low-bit quantization to reach more aggressive storage budgets.
  • Compression granularity jointly shapes representation quality, cache management, and GPU execution.

A.1 Additional Evaluation Details

The evaluation uses fixed seeds and specified datasets, prompts, software environments, and baseline procedures. Baselines are configured to match storage or execution patterns where stated, while the environments differ between accuracy and efficiency studies.

  • All benchmark evaluations use fixed random seed 42, with RULER covering 13 tasks and LongBench using its test split.Qwen3 is evaluated without a chat template and with thinking disabled; Llama-3.1-8B-Instruct uses task-dependent prompt formatting on LongBench.
  • The accuracy and system-efficiency studies use different CUDA and PyTorch environments because the accuracy environment has an SVD performance issue.
  • Global SVD matches PuzzleKV’s dense sink and local-page policy while selecting input-dependent rank under the target storage budget.
  • Palu uses WikiText-2 calibration with 32 sequences for Fisher estimation and 256 for whitening, and its retained parameter ratio rounds to 0.60 for both models.
  • H2O selects entries independently per layer and KV head using accumulated post-softmax attention mass before evicting prompt-cache entries.For GQA models, attention mass from query heads mapped to one KV head is summed before selection.
  • OjaKV-PF is evaluated at a nominal retained KV cache ratio of 0.60 using rank r = 77, with online basis updates during evaluation.The basis is initialized by SVD on the WikiText-2 training split with calibration seed 2, and each online update samples at most 256 pooled rows.

B.1 Complete Quantization Compatibility Results

The complete quantization-compatibility results report all 13 RULER tasks at 16K, with storage expressed relative to Full KV storage. PuzzleKV uses page size P = 32 and ranks (rK, rV) = (16, 14).

  • Table 7 reports all 13 RULER tasks at 16K, with storage reported as a percentage of Full KV storage.
  • PuzzleKV uses page size P = 32 and ranks (rK, rV) = (16, 14) in the complete quantization-compatibility evaluation.

B.2 Complete LongBench Results

The complete LongBench results cover the English-language, code, and Chinese task subsets, with compressed methods evaluated at a shared storage ratio of ρ = 0.60.

  • Task coverage: Tables 8 and 9 provide complete per-task LongBench results for the English-language, code, and Chinese task subsets.The English-language and code results cover 14 English tasks and two code tasks; the Chinese results add five tasks.
  • Evaluation setting: All compressed methods use the same KV cache storage ratio ρ = 0.60 in the reported LongBench comparisons.

C OjaKV-PF Reproduction Results

The OjaKV-PF reproduction evaluates performance at ρ = 0.60 and reports substantial variation across models and tasks, including particularly large degradation on Qwen3-8B.

  • Evaluation setting: OjaKV-PF is evaluated at the same nominal retained KV cache ratio, ρ = 0.60, used by the compressed methods in the main evaluation.
  • Reported results: The reproduction reports complete OjaKV-PF results for RULER at 16K and LongBench across its English-language, code, and Chinese task panels.The reported tables include 16 RULER task results and 21 LongBench tasks.
  • Observed variation: OjaKV-PF performance varies substantially across models and tasks, with particularly large degradation and several zero scores on Qwen3-8B.
  • Sensitivity analyses: Page-size and rank-selection experiments report matched factorized-page storage configurations and RULER scores under a fixed dense-page policy.The page-size configurations use symmetric ranks, while the rank grid holds page size and dense-page policy fixed.

D.3 RoPE Key Analysis

The RoPE key analysis compares pre- and post-RoPE key factorization on 13 RULER tasks at 16K, finding a higher average score for pre-RoPE keys.

  • Average performance: Pre-RoPE key factorization raises the average RULER score from 90.40 to 90.88.The comparison uses Llama-3.1-8B-Instruct at 16K with P = 32 and (rK, rV) = (16, 16).
  • Evaluation scope: Table 14 reports percentages over 13 RULER tasks, with QA-1 and QA-2 denoting SQuAD and HotpotQA.
Loading 2608.23843v1…