Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

HeteroCache: A Dynamic Retrieval Approach to Heterogeneous KV Cache Compression for Long-Context LLM Inference

Zhiyuan Shi, Qibo Qiu, Feng Xue, Zhonglin Jiang, Li Yu, Jian Jiang, Xiaofei He, Wenxiao Wang

TL;DR

HeteroCache is a training-free dynamic compression framework that applies a fine-grained weighting strategy that allocates larger cache budgets to heads with rapidly shifting attention to capture context changes, thereby addressing the inefficiency of coarse-grained strategies.

Abstract

The linear memory growth of the KV cache poses a significant bottleneck for LLM inference in long-context tasks. Existing static compression methods often fail to preserve globally important information, principally because they overlook the attention drift phenomenon where token significance evolves dynamically. Although recent dynamic retrieval approaches attempt to address this issue, they typically suffer from coarse-grained caching strategies and incur high I/O overhead due to frequent data transfers. To overcome these limitations, we propose HeteroCache, a training-free dynamic compression framework. Our method is built on two key insights: attention heads exhibit diverse temporal heterogeneity, and there is significant spatial redundancy among heads within the same layer. Guided by these insights, HeteroCache categorizes heads based on stability and redundancy. Consequently, we apply a fine-grained weighting strategy that allocates larger cache budgets to heads with rapidly shifting attention to capture context changes, thereby addressing the inefficiency of coarse-grained strategies. Furthermore, we employ a hierarchical storage mechanism in which a subset of representative heads monitors attention shift, and trigger an asynchronous, on-demand retrieval of contexts from the CPU, effectively hiding I/O latency. Finally, experiments demonstrate that HeteroCache achieves state-of-the-art performance on multiple long-context benchmarks and accelerates decoding by up to $3\times$ compared to the original model in the 224K context. Our code will be open-source.

HeteroCache: A Dynamic Retrieval Approach to Heterogeneous KV Cache Compression for Long-Context LLM Inference

TL;DR

HeteroCache is a training-free dynamic compression framework that applies a fine-grained weighting strategy that allocates larger cache budgets to heads with rapidly shifting attention to capture context changes, thereby addressing the inefficiency of coarse-grained strategies.

Abstract

The linear memory growth of the KV cache poses a significant bottleneck for LLM inference in long-context tasks. Existing static compression methods often fail to preserve globally important information, principally because they overlook the attention drift phenomenon where token significance evolves dynamically. Although recent dynamic retrieval approaches attempt to address this issue, they typically suffer from coarse-grained caching strategies and incur high I/O overhead due to frequent data transfers. To overcome these limitations, we propose HeteroCache, a training-free dynamic compression framework. Our method is built on two key insights: attention heads exhibit diverse temporal heterogeneity, and there is significant spatial redundancy among heads within the same layer. Guided by these insights, HeteroCache categorizes heads based on stability and redundancy. Consequently, we apply a fine-grained weighting strategy that allocates larger cache budgets to heads with rapidly shifting attention to capture context changes, thereby addressing the inefficiency of coarse-grained strategies. Furthermore, we employ a hierarchical storage mechanism in which a subset of representative heads monitors attention shift, and trigger an asynchronous, on-demand retrieval of contexts from the CPU, effectively hiding I/O latency. Finally, experiments demonstrate that HeteroCache achieves state-of-the-art performance on multiple long-context benchmarks and accelerates decoding by up to compared to the original model in the 224K context. Our code will be open-source.
Paper Structure (32 sections, 10 equations, 6 figures, 4 tables)

This paper contains 32 sections, 10 equations, 6 figures, 4 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. KV Cache Eviction.
  4. Observation
  5. Temporal Heterogeneity.
  6. Intralayer Redundancy.
  7. Method
  8. Head Profiling and Taxonomy
  9. Stability-Based Budget Allocation
  10. Hierarchical Storage & Retrieval
  11. Experiments
  12. Experiment Settings
  13. Backbone LLMs and Baselines.
  14. Implementation.
  15. Hyperparameter Configuration.
  16. ...and 17 more sections

Figures (6)

  • Figure 1: Analysis of attention heads heterogeneity and redundancy. (a) illustrates temporal heterogeneity: the blue line represents a stable head maintaining consistent focus, while the red line indicates a decaying head with rapid attention shift; (b) visualizes intralayer redundancy, where the darker diagonal blocks indicate that attention heads within the same layer share significantly higher similarity than those across different layers; (c) and (d) present the classification of heads based on their stability and similarity scores, respectively.
  • Figure 2: The workflow of HeteroCache. (A, B) Offline Calibration: Heads are categorized into functional roles (A) to determine stability-based budgets for compressed heads (B). (C, D) Online Inference: Guided by (B), (C) initializes hierarchical storage in prefill stage. In (D), Pivot heads monitor drift to trigger asynchronous CPU retrieval for updating satellite heads in decode stage.
  • Figure 3: Results for Llama-3.1-8B-Instruct on NIAH, evaluated on context lengths from 16K to 128K tokens.
  • Figure 4: End-to-end latency results for HeteroCache and the baselines. The left and right plots show the latency for the prefill and decode stage, respectively.
  • Figure 5: Ablation study of the threshold with Llama-3.1-8B-Instruct.
  • ...and 1 more figures