Source-linked AI summary
CateKV: On Sequential Consistency for Long-Context LLM Inference Acceleration
Haoyun Jiang, Haolin Li, Jianwei Zhang, Fei Huang, Qiang Hu, Minmin Sun, Shuai Xiao, Yong Li, Junyang Lin, Jiangchao Yao
TL;DR
Long-context LLM inference is constrained by rising memory use and latency from storing and retrieving KV caches. CateKV identifies sequentially consistent attention heads with a coefficient-of-variation-based method and applies hybrid KV retention policies. Across long-context benchmarks, it maintains performance comparable to full attention while reducing memory and latency for single inputs and increasing batch throughput.
Problem
Long-context LLM inference requires storing and retrieving all KV caches, increasing computational cost, memory consumption, and runtime.
Method
CateKV uses pre-filling observations and coefficient-of-variation-based scores to classify heads as consistent or adaptive, retaining few critical KV pairs for consistent heads and most for adaptive heads.
Results
Across long-context benchmarks, CateKV maintains performance comparable to full attention while reducing memory usage by up to 2.72×, decoding latency by 2.18×, and increasing batch throughput by 3.96×.
Takeaways & Limitations
Sequential consistency enables a plug-and-play hybrid KV cache that reduces cache and decoding overhead while remaining combinable with existing acceleration approaches.
Takeaways & Limitations
The theory analysis assumes probabilities for correct and incorrect classification of consistent and adaptive heads, along with a probability λ for the query-attention ordering hypothesis.
Abstract
from arXiv · showhide
Large language models (LLMs) have demonstrated strong capabilities in handling long-context tasks, but processing such long contexts remains challenging due to the substantial memory requirements and inference latency. In this work, we discover that certain attention heads exhibit sequential consistency in their attention patterns, which can be persistently identified using a coefficient-of-variation-based algorithm. Inspired by this observation, we propose CateKV, a hybrid KV cache method that retains only critical token information for consistent heads, thereby reducing KV cache size and computational overhead, while preserving the majority of KV pairs in adaptive heads to ensure high accuracy. We show the unique characteristics of our algorithm and its extension with existing acceleration methods. Comprehensive evaluations on long-context benchmarks show that, while maintaining accuracy comparable to full attention, CateKV reduces memory usage by up to $2.72\times$ and accelerates decoding by $2.18\times$ in single-sample inputs, and boosts throughput by $3.96\times$ in batch scenarios.
1. Introduction
Long-context LLM inference becomes increasingly costly because all KV caches must be stored and retrieved. CateKV exploits sequentially consistent attention heads with a hybrid cache, reducing resource use while preserving performance comparable to full attention.
- Method: CateKV is a hybrid KV cache method that reduces memory usage while maintaining high accuracy and performance close to full attention.The method is plug-and-play, orthogonal to, and combinable with many existing acceleration approaches.
- Motivation: Context expansion increases inference cost because autoregressive LLMs must store and retrieve all key-value caches.For Llama-3-8B-Instruct-262k with FlashAttention, extending context from 1K to 1M tokens increases inference latency by over 3000 times.
- Observation: Certain attention heads show sequentially consistent patterns across pre-filling and decoding, while others exhibit rich activation dynamics.The phenomenon appears across different LLMs and layers, motivating head-specific cache treatment.
- Method: CateKV classifies attention heads with a coefficient-of-variation-based score and routes them to different KV retention policies.It retains a small proportion of critical KV pairs for consistent heads and most KV pairs for adaptive heads.
- Results: 2.72× lower memory usage and 2.18× lower decoding latency are achieved for single inputs, while batch throughput increases by 3.96×.These results come from extensive long-context benchmark evaluations while maintaining performance comparable to full attention.
2. Related Work
Long-context acceleration methods primarily evict KV pairs or retrieve relevant pairs during decoding. Eviction reduces cache size but can lose essential information, whereas retrieval preserves accuracy at the cost of high memory usage.
- KV Cache Eviction: KV cache eviction methods discard KV pairs according to predefined policies to reduce cache size.Examples include retaining attention sinks and recent tokens or maintaining heavy-hitter tokens in a fixed-size cache.
- KV Cache Retrieval: KV cache retrieval methods retain the complete cache and dynamically retrieve important KV pairs to reduce inference latency.Quest retrieves relevant pages using approximate attention scores, while InfLLM offloads much of the cache to CPU and ShadowKV stores low-rank key caches.
- Head Classification: Attention-head classification methods assign distinct sparse patterns or split heads into retrieval and streaming types.MInference uses A-shape, Vertical-Slash, and Block-Sparse patterns, while RazorAttention and DuoAttention determine between full attention and StreamingLLM-like mechanisms.
3. Method
CateKV identifies attention heads with stable or dynamic patterns from pre-filling observations, then applies different KV-retention strategies to them. Its analysis links retention accuracy to head-classification accuracy, retention budgets, and the quality of token-score ordering.
- Sequential Consistency of Attention Heads: Consistent heads focus on a limited set of tokens across decoding steps, whereas adaptive heads have variable distributions requiring a larger attention space.
- Observation Matrix: CateKV forms an observation matrix from the last pre-filling queries while temporarily excluding initial and recent tokens during head and critical-token identification.The observation window is intended to reflect attention patterns while reducing identification cost.
- Coefficient of Variation (CV) Score: A coefficient-of-variation score measures attention concentration and similarity after percentile-threshold binarization, producing token frequencies that identify critical tokens.The binarized matrix yields a frequency vector counting how often each token is marked critical.
- CateKV Cache Policy: Consistent heads retain top-k critical KV pairs, while adaptive heads retain nη selected KV pairs to preserve flexible token interactions.The algorithm assigns separate cache-selection branches to the two head types during pre-filling and retrieves cached keys and values during decoding.
- Reference-Based Static Identification: CateKV uses a reference dataset to derive fixed head types from cross-sample identification frequencies, avoiding expensive dynamic classification during inference.The adaptive-head ratio r determines the percentile threshold used for model-wise classification.
- Theory Analysis: Theoretical analysis identifies head-classification accuracy, token-score ordering quality, and budget control as factors influencing token-retention accuracy.The analysis models correct and incorrect classification cases using retention ratios for consistent and adaptive heads.
4. Experiments
The experiments evaluate CateKV across five long-context LLMs and three challenging benchmarks, comparing it with eviction-, retrieval-, and head-wise classification methods under exact prefilling and sparse decoding.
- Setup: Five long-context LLMs are evaluated: LLaMA-3-8B-Instruct-1048K, Phi-3-Mini-128K, Llama-3.1-8B, Yi-9B-200K, and Qwen2.5-7B.
- Setup: The evaluation uses RULER, LongBench, and Needle in a Haystack, with LongBench averaging performance across 21 tasks.
- Setup: Experiments run on a single NVIDIA A100-80G GPU, using a reference set distinct from the test set for head identification.
- Baselines: CateKV is compared with StreamingLLM, SnapKV, PyramidKV, Quest, ShadowKV, and Duoattention.
- Baselines: All approaches use an exact pre-filling stage and sparse attention during decoding.
Appendix B.3 for detailed data.
Across long-context benchmarks and efficiency tests, CateKV preserves accuracy close to full attention while reducing cache and decoding costs. Its hybrid head treatment also remains compatible with other acceleration methods.
- Benchmark results: CateKV outperforms baselines and remains comparable to full attention on 12 synthetic 128K-context tasks while retaining approximately 41% of the KV cache.The configuration uses adaptive head ratio r = 0.4, retention ratio η = 1.0, and a 2048 sparse budget for consistent heads.
- Benchmark results: 42% KV-cache size with no more than a 0.3% average accuracy decline is achieved across five LLMs on LongBench.Combining CateKV with Quest or ShadowKV reduces memory with an accuracy decrease not exceeding 0.2% under the same computational budget.
- Benchmark results: From 20K to 1M tokens, CateKV accurately retrieves relevant information on Needle In A Haystack while reducing memory and computational cost by more than half.Additional tests report overall comparable performance across models, context windows, and needle depths.
- Efficiency: 2.72× lower memory usage and 2.18× lower decoding latency are achieved on Llama-3 single-sample inference, with accuracy decline below 0.25%.On Phi-3, the corresponding reductions are 2.11× in memory and 1.79× in latency.
- Efficiency: 3.96× throughput and 4.33× batch-size increases are achieved for Llama-3 after further KV-cache compression while maintaining accuracy.Under generic settings, gains reach 2.38× in batch size and 2.25× in throughput.
- Method: CateKV classifies heads using coefficient-of-variation scores, retaining few critical KV pairs in consistent heads and most pairs in adaptive heads.The method is presented as plug-and-play and can integrate with other acceleration approaches.
A. Further Observations
Attention-weight heatmaps across four models show both consistent and adaptive heads across various layers, supporting the generality of sequential consistency.
- Further observations: Llama-3-8B-1M, Llama-3.1-8B, Phi-3-128K, and Yi-9B-200K all exhibit consistent and adaptive heads across various layers.The observation is based on visualized attention-weight heatmaps.
B.1. Implementation Details of Experiments
The experiments use a separate Variable Tracking reference set for head identification and evaluate CateKV across context lengths, while baseline configurations are specified for comparison.
- Implementation: CateKV identifies heads with a 64-token observation window and a separate 100-sample, 128K-token Variable Tracking reference set.The sparse budget during identification is 2048.
- Implementation: SnapKV and PyramidKV use observation windows of 32, PyramidKV uses β = 20, and StreamingLLM uses 128 initial tokens.DuoAttention experiments use attention patterns provided by its code.
- Context-length evaluation: Across RULER context lengths from 8K to 256K, CateKV remains comparable to full attention and slightly improves performance at certain lengths.These results use r = 0.4 and η = 1.0.
B.2.2. COMBINE WITH KV RETRIEVAL METHODS
CateKV can be combined with KV retrieval methods to reduce memory usage under matched computational budgets while preserving or improving accuracy relative to KV-eviction baselines.
- KV retrieval integration: CateKV combined with Quest and ShadowKV reduces memory usage under the same computational budget.The evaluation is conducted at 128K context length.
- KV retrieval integration: 41% memory usage is achievable for KV retrieval methods while maintaining accuracy.The table defines budget as the computational budget for sparse attention.
- KV retrieval integration: CateKV significantly outperforms KV-eviction methods in accuracy at the same computational load while reducing Quest and ShadowKV memory usage.This comparison is reported for the combined methods.
B.3. Full Longbench Results
CateKV integrates with Full Attention and KV retrieval methods across all 21 Longbench tasks without significant accuracy loss, despite retaining only 42% of the KV cache.
- 42% KV cache retention produced no significant per-task accuracy drop across all 21 Longbench tasks.Average accuracy even exceeded the original methods.
- CateKV integration improved average accuracy over the original Full Attention and KV retrieval methods.The integrated methods were Quest and ShadowKV.
- Around half of Longbench tasks showed slight performance improvements after combining CateKV with the original methods.
B.4. Additional Results in Needle In A Hystack
On Needle In A Haystack, CateKV maintains overall performance comparable to full attention across models, context windows, and needle depths. This indicates preserved long-context semantic retrieval capacity.
- CateKV maintains overall comparable Needle In A Haystack performance to full attention across three evaluated models.The models are Llama3.1-8B, Phi-3-Mini-128K, and Yi-9B-200K.
- CateKV performance varies across context windows and needle depths relative to full attention.
- The results suggest CateKV does not significantly affect access to and retrieval of long-context semantic information.
B.5. Additional Results on Larger Models
CateKV scales to larger models evaluated at their supported context lengths, achieving near full-attention accuracy on 30B and 14B models and outperforming SnapKV and PyramidKV.
- CateKV achieves near full-attention accuracy on the evaluated 30B and 14B models.The models are Qwen2.5-32B, Yi-34B-200K, and Phi-4-14B.
- CateKV outperforms baseline methods including SnapKV and PyramidKV on the larger-model evaluation.
- The evaluation uses each model’s maximum supported context length: 128k for Qwen2.5-32B and Yi-34B-200K, and 16k for Phi-4-14B.