PaperlayerView exact revision on arXiv

Source-linked AI summary

ClusterKV: Manipulating LLM KV Cache in Semantic Space for Recallable Compression

Guangda Liu, Chengwei Li, Jieru Zhao, Chenqi Zhang, Minyi Guo

arXiv:2412.03213v2cs.LGcs.AIcs.PF

TL;DR

Long-context inference strains memory and latency, while permanent eviction can miss tokens that become important during decoding. ClusterKV recalls tokens through semantic clusters, reporting negligible accuracy loss with small cache budgets and gains in latency and decoding throughput.

  • Problem

    Long contexts increase KV-cache memory use and decoding latency, while permanently evicted tokens that later become important cannot be recalled.

    Source: §I · p. 1

  • Method

    ClusterKV groups key vectors into semantic clusters and selects tokens by scoring cluster centroids rather than individual tokens.

    Source: §I · p. 1, B. Clustering in the Semantic Space · p. 3, C. Selection at the Granularity of… · p. 3

  • Results

    Across tasks at context lengths up to 32k, ClusterKV reports negligible accuracy loss with a 1k–2k KV cache budget, up to 2× latency speedup, and up to 2.5× higher decoding throughput.

    Source: §I · p. 2

  • Takeaways & Limitations

    ClusterKV preserves model accuracy while significantly enhancing inference efficiency through recall at semantic-cluster granularity.

    Source: §VI · p. 6

    Clustering costs O(niCLd), higher than obtaining page representations such as Quest’s per-channel maximal keys at O(Ld).

    Source: D. Efficiency Concerns · p. 3

Abstract

from arXiv · show

Large Language Models (LLMs) have been widely deployed in a variety of applications, and the context length is rapidly increasing to handle tasks such as long-document QA and complex logical reasoning. However, long context poses significant challenges for inference efficiency, including high memory costs of key-value (KV) cache and increased latency due to extensive memory accesses. Recent works have proposed compressing KV cache to approximate computation, but these methods either evict tokens permanently, never recalling them for later inference, or recall previous tokens at the granularity of pages divided by textual positions. Both approaches degrade the model accuracy and output quality. To achieve efficient and accurate recallable KV cache compression, we introduce ClusterKV, which recalls tokens at the granularity of semantic clusters. We design and implement efficient algorithms and systems for clustering, selection, indexing and caching. Experiment results show that ClusterKV attains negligible accuracy loss across various tasks with 32k context lengths, using only a 1k to 2k KV cache budget, and achieves up to a 2$\times$ speedup in latency and a 2.5$\times$ improvement in decoding throughput. Compared to SoTA recallable KV compression methods, ClusterKV demonstrates higher model accuracy and output quality, while maintaining or exceeding inference efficiency. Our code is available at https://github.com/sjtu-zhao-lab/ClusterKV.

I. INTRODUCTION

Long contexts strain inference efficiency through growing KV-cache memory use and latency, while permanent eviction can miss tokens whose importance changes. ClusterKV recalls tokens by semantic cluster, and the paper reports negligible accuracy loss and improved efficiency at small cache budgets.

I. INTRODUCTION

  • Long contexts can push KV-cache memory beyond GPU capacity and increase decoding latency through extensive memory access.The paper notes that this can cause inference failure or extremely high latency.Source: §I · p. 1
  • Permanent eviction cannot recover tokens that become important later in decoding, which can degrade accuracy and output quality.The comparison figure marks tokens selected for attention computation with green boxes.Source: §I · p. 1, Figure 1 · p. 1
  • Page-based recall can waste cache budget because position-based pages may include unimportant tokens alongside important ones.The paper proposes recalling semantic clusters instead, noting that nearby tokens in semantic space exhibit similar attention weights.Source: §I · p. 1, Figure 2 · p. 1
  • With 32k contexts and a 1k–2k KV cache budget, ClusterKV reports negligible accuracy loss, up to 2× latency speedup, and up to 2.5× higher decoding throughput.It also reports higher accuracy and output quality than state-of-the-art recallable compression methods while maintaining or exceeding their inference efficiency.Source: §I · p. 2

A. LLM Inference and KV Cache

The section introduces the attention and KV-cache operations used in autoregressive inference, then explains why long contexts motivate cache compression. Because only a subset of tokens contributes substantially to attention, a fixed cache budget can bound cache size and decoding cost.

A. LLM Inference and KV Cache

  • Multi-head attention projects inputs into query, key, and value tensors, then combines attention outputs across heads for later feedforward and normalization operations.The passage gives the attention computation as softmax of query–key scores applied to values.Source: A. LLM Inference and KV Cache · p. 2
  • During autoregressive decoding, the KV cache stores prior tokens’ keys and values for reuse, while the latest query attends to those cached tensors.Inference consists of prefill, which processes the input and computes the cache, followed by decoding.Source: A. LLM Inference and KV Cache · p. 2

B. Long Context Inference and KV Cache Compression

  • Long-context decoding raises memory and computation costs because KV-cache size and attention complexity grow linearly with context length.These costs can reduce inference efficiency or cause inference failures.Source: B. Long Context Inference and KV Cache… · p. 2
  • Attention sparsity motivates selecting a subset of tokens to approximate attention, with selected keys and values stored under a fixed cache budget.The passage states that this keeps cache size and decoding cost stable regardless of context length; tokens are typically selected by attention weights.Source: B. Long Context Inference and KV Cache… · p. 2
  • Figure 3 covers changes in token importance over decoding steps and fragmentation of important tokens across pages of 16 consecutive tokens.Source: Figure 3 · p. 2

C. Related Work and Motivation

Existing compression approaches face a trade-off between recalling dynamically important tokens and limiting selection overhead. Permanent eviction can miss tokens that later matter, while page-based recall reduces selection cost but can waste cache budget on unimportant tokens.

C. Related Work and Motivation

  • Permanent eviction misses tokens whose attention importance changes during decoding, potentially reducing model accuracy and output quality.The cited example describes a token becoming crucial after 20 decoding steps, while other tokens’ importance fluctuates.Source: C. Related Work and Motivation · p. 2
  • InfiniGen reduces query and key dimensions with offline SVD-generated partial weights, but still has selection cost linear in context length and stores partial keys alongside original keys.Source: C. Related Work and Motivation · pp. 2–3
  • Quest lowers selection cost to O(Ld/page size) by scoring pages, but 16-token pages can contain only one or two important tokens and waste cache budget.In some cases, two pages are needed to include two consecutive important tokens.Source: C. Related Work and Motivation · p. 3

III. ALGORITHM DESIGN

The algorithm targets token selection that approximates full attention by grouping semantically similar keys and scoring cluster representations rather than individual tokens. It uses cosine-based clustering and applies distinct clustering procedures to prompt and generated tokens.

A. Problem Formulation and Design Rationale

  • The selection objective is to choose the B tokens with the largest attention weights so that selected-token attention approximates the original computation.Source: A. Problem Formulation and Design… · p. 3
  • Because nearby keys in semantic space tend to have similar attention weights, scoring cluster representations can reduce recall overhead; the number of clusters is typically an order of magnitude smaller than the number of tokens.Source: A. Problem Formulation and Design… · p. 3

B. Clustering in the Semantic Space

  • Semantic distance is defined using cosine similarity between key vectors, which the authors prefer to L2 or inner-product distance because outlier channels can distort the latter measures.The distance is smaller for vectors with larger cosine similarity.Source: B. Clustering in the Semantic Space · p. 3
  • K-means forms semantic clusters by assigning each key to its nearest centroid by cosine similarity and updating each centroid to the mean of assigned keys until assignments stop changing.The figure depicts key vectors and cluster centroids.Source: B. Clustering in the Semantic Space · p. 3, Figure 4 · p. 3
  • After prefill, clustering retains the first 16 tokens as attention sinks and clusters the subsequent prompt tokens into C0 centroids.Source: B. Clustering in the Semantic Space · p. 3
  • For generated tokens, clustering is performed every 320 decoding steps on generated tokens alone, creating four new centroids to amortize clustering cost.Source: B. Clustering in the Semantic Space · p. 3

C. Selection at the Granularity of Semantic Clusters

ClusterKV ranks semantic-cluster centroids by their attention weights to a query, then retrieves tokens from the highest-ranked clusters until the top-B token budget is met.

C. Selection at the Granularity of Semantic Clusters

D. Efficiency Concerns

Cluster-based selection avoids page-based internal fragmentation and can improve accuracy, but introduces clustering and indexing costs that the system design addresses.

D. Efficiency Concerns

  • Cluster selection avoids page-based internal fragmentation of important tokens and can achieve higher accuracy, but incurs efficiency concerns.Source: D. Efficiency Concerns · p. 3
  • Clustering costs O(niCLd), higher than obtaining page representations such as per-channel maximal key vectors, which costs O(Ld).Source: D. Efficiency Concerns · p. 3
  • Unlike fixed-size pages, variable-size clusters with discontinuous token positions make the number of clusters needed and selected-token indices harder to determine.Source: D. Efficiency Concerns · p. 4

A. System Overview

  • The system clusters keys on GPU during prefill, offloads generated KV to CPU, and uses decode-time cluster scores and metadata to select and load KV.A GPU cache retains selected KV; every m decoding steps, clustering and KV offloading are performed for the m generated tokens.Source: A. System Overview · p. 4

B. Semantic Clustering

ClusterKV optimizes clustering at the system and kernel levels, overlapping work with inference operations and using a custom parallel centroid-update kernel.

B. Semantic Clustering

  • Clustering is launched asynchronously after key computation and overlaps current-layer attention and FFN, plus next-layer QKV projection and RoPE.This overlap minimizes clustering overhead.Source: B. Semantic Clustering · p. 4
  • For batched clustering across heads, ClusterKV uses available Torch kernels for assignment and a custom CUDA kernel for centroid updates.Individual ThreadBlocks process different heads in parallel.Source: B. Semantic Clustering · p. 4
  • The centroid-update kernel accumulates keys assigned to each cluster and counts them in shared memory, then computes new centroid means.The example uses an 8-thread ThreadBlock with the channel dimension split into two partitions.Source: B. Semantic Clustering · p. 4, Figure 7 · p. 4
  • For BlockSize = 512 and 128 channels, offline profiling found optimal performance with P = 16 or P = 32.Larger P reduces keys processed in parallel, while smaller P increases per-thread channel iterations and may increase write conflicts.Source: B. Semantic Clustering · p. 4

C. Selection and Indexing

ClusterKV stores cluster metadata for decoding-time selection and indexing, and caches selected KV to reduce CPU-to-GPU transfers.

C. Selection and Indexing

  • After clustering, ClusterKV stores centroids and metadata including cluster sizes, prefix sums, and sorted token indices.Source: C. Selection and Indexing · p. 4
  • During decoding, clusters are ranked by attention weights to q; their sizes and prefix sums determine which clusters fit the KV budget.If a selected cluster would take the total above budget, tokens are trimmed from the last selected cluster.Source: C. Selection and Indexing · pp. 4–5, Figure 8 · p. 5
  • CUDA indexing kernels process KV heads in parallel, with frequently accessed cluster-size metadata stored in shared memory.Source: C. Selection and Indexing · p. 5

D. Caching KV of Selected Tokens

  • The GPU cache retains selected KV and cluster labels from the last R decoding steps, loading only clusters absent from the cache.The authors find R = 1 strikes a good balance between memory usage and caching effectiveness.Source: D. Caching KV of Selected Tokens · p. 5

A. Experimental Setup

The evaluation covers long-context accuracy and output quality on LongBench and PG19, then compares token recall across methods and ClusterKV configurations. ClusterKV approaches full-cache accuracy at low budgets and improves on competing compression methods across most LongBench settings.

A. Experimental Setup

  • Accuracy and output quality were tested on eight LongBench datasets and PG19 with GLM4-9B-Chat; inference comparisons used Llama-3.1-8B.Quest and InfiniGen support efficient inference only for Llama-architecture models, and the first two layers used full KV cache for aligned comparisons.Source: A. Experimental Setup · p. 5

B. Model Accuracy

  • ClusterKV outperforms Quest and InfiniGen in most LongBench settings and approaches full-KV accuracy with budgets as low as 1k–2k tokens.Table I summarizes average scores across eight tasks.Source: Figure 9 · p. 5, B. Model Accuracy · p. 5, Table I · p. 5
  • Perplexity deviation from full KV is up to 0.5 for ClusterKV, versus approximately 4 for Quest and 2 for InfiniGen.The comparison uses a 1024-token KV cache budget and PG19 prompts up to 32,000 tokens.Source: B. Model Accuracy · pp. 5–6, Figure 10 · p. 6
  • On a 32k NarrativeQA sample, ClusterKV achieves higher important-token recall than other methods across all tested budgets.Recall rates are averaged across layers, heads, and decoding steps; budgets range from 256 to 2048 tokens.Source: Figure 11 · p. 6, B. Model Accuracy · p. 6
  • Cosine similarity outperforms L2 and inner-product distance, while recall gains from increasing C0 become less significant above 400.The authors choose C0 = 400 as a balance between accuracy and efficiency.Source: B. Model Accuracy · p. 6

C. Inference Efficiency

ClusterKV reduces inference latency versus full KV cache and maintains similar latency to recallable compression baselines. Its caching mechanism also raises decoding throughput compared with loading clusters directly from CPU memory.

C. Inference Efficiency

  • For P = 32k and D = 1024, latency speedup is 2× with a 1024-token budget, and decoding throughput increases by up to 2.5× versus full KV cache.Clustering accounts for 6% to 8% of prefill and less than 2% of total inference time.Source: Figure 12 · p. 6, C. Inference Efficiency · p. 6
  • Against InfiniGen on OPT-6.7B, ClusterKV achieves an average speedup of 2.3×; its selection overhead is about 5% of total decoding latency.InfiniGen latency is comparable to full-KV inference in this comparison.Source: Figure 13 · p. 6, C. Inference Efficiency · p. 6
  • Against Quest on Llama-3.1-8B with a 1k-token budget, ClusterKV’s latency deviations are up to 5%, with significantly higher model accuracy.The comparison is shown in Fig. 13b.Source: Figure 13 · p. 6, C. Inference Efficiency · p. 6
  • For R = 1 and R = 2, average cluster-cache hit rates are 63% and 74%, and throughput improves 2.3× and 3× over direct CPU-memory loading, respectively.Source: C. Inference Efficiency · p. 6

VI. CONCLUSION

ClusterKV compresses KV cache by recalling tokens at the granularity of semantic clusters, preserving model accuracy while improving inference efficiency.

VI. CONCLUSION

  • ClusterKV achieves efficient and accurate KV-cache compression by recalling tokens at the granularity of semantic clusters.The paper concludes that it preserves model accuracy while significantly enhancing inference efficiency.Source: §VI · p. 6
Loading 2412.03213v2…