Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Where Matters More Than What: Decoding-aligned KV Cache Compression via Position-aware Pseudo Queries

Zhenxu Tian, Yi Su, Juntao Li, Min Zhang

Abstract

The Key-Value (KV) cache is crucial for efficient Large Language Models (LLMs) inference, but excessively long contexts drastically increase KV cache memory footprint. Existing KV cache compression methods typically rely on input-side attention patterns within a prompt observation window to estimate token importance during the prefill stage. They fail to preserve critical tokens for future generation since these assessments are not derived from the decoding process. Intuitively, an effective observation window should mirror the decoding-stage queries to accurately reflect which tokens the generation process will attend to. However, ground-truth decoding queries are inherently unavailable during inference. For constructing pseudo queries to approximate them, we find that positional information plays a more critical role than semantic content. Motivated by this insight, we propose decoding-aligned KV cache compression via position-aware pseudo queries (DapQ), a novel and lightweight eviction framework that leverages position-aware pseudo queries to simulate the output tokens, thereby establishing an effective observation window for importance assessment. It aligns closely with the actual generation context and enables precise token eviction. Extensive evaluations across multiple benchmarks and LLMs demonstrate that DapQ achieves superior performance, particularly under strict memory constraints (e.g., up to nearly lossless performance 99.5% on NIAH with 3% KV cache budgets).

Where Matters More Than What: Decoding-aligned KV Cache Compression via Position-aware Pseudo Queries

Abstract

The Key-Value (KV) cache is crucial for efficient Large Language Models (LLMs) inference, but excessively long contexts drastically increase KV cache memory footprint. Existing KV cache compression methods typically rely on input-side attention patterns within a prompt observation window to estimate token importance during the prefill stage. They fail to preserve critical tokens for future generation since these assessments are not derived from the decoding process. Intuitively, an effective observation window should mirror the decoding-stage queries to accurately reflect which tokens the generation process will attend to. However, ground-truth decoding queries are inherently unavailable during inference. For constructing pseudo queries to approximate them, we find that positional information plays a more critical role than semantic content. Motivated by this insight, we propose decoding-aligned KV cache compression via position-aware pseudo queries (DapQ), a novel and lightweight eviction framework that leverages position-aware pseudo queries to simulate the output tokens, thereby establishing an effective observation window for importance assessment. It aligns closely with the actual generation context and enables precise token eviction. Extensive evaluations across multiple benchmarks and LLMs demonstrate that DapQ achieves superior performance, particularly under strict memory constraints (e.g., up to nearly lossless performance 99.5% on NIAH with 3% KV cache budgets).
Paper Structure (39 sections, 16 equations, 5 figures, 15 tables)

This paper contains 39 sections, 16 equations, 5 figures, 15 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Long-Context LLMs.
  4. KV Cache Compression.
  5. Observation
  6. Position Drives Query Representation
  7. Precise Positional Alignment is Necessary
  8. Position-Based Queries Enable More Accurate Token Eviction
  9. Method
  10. Construct Decoding-Aligned Q_pseudo
  11. Importance Assessment and Eviction
  12. Experiments
  13. Settings
  14. Improvement on Accuracy
  15. Analysis on Efficiency
  16. ...and 24 more sections

Figures (5)

  • Figure 1: An overview of DapQ. A synthetic pseudo context (length $N$) is appended to the original context (length $L_p$), forming an extended sequence of length $L_p{+}N$. The model processes this sequence during the prefill phase and then obtains pseudo queries for the synthetic tokens, which are endowed with the correct positional encodings of the first $N$ decoding steps. These pseudo queries compute attention scores with all keys from the original prompt, establishing the token importance distribution. The $topK$ tokens are retained in the compressed KV cache, while the others, along with all the synthetic tokens are evicted. Autoregressive decoding then begins from position $L_p$.
  • Figure 2: Analysis of Positional Dominance and Offset Sensitivity in Query Similarity. We set the pseudo queries of fixed length 32. (a) Boxplot of query similarity distributions for a 4k context under different content and position conditions, each aggregated from 100 independent trials. DC: pseudo-queries content is constructed by randomly sampling 32 tokens from the model's vocabulary; DP: pseudo-queries positions are assigned by randomly sampling a consecutive span of 32 index positions from the context length range [0, m] (e.g., [0, 4000]). (b) Query similarity curves over offset positions for contexts of lengths 2k, 4k, 6k, and 8k. The x-axis denotes the starting position assigned to pseudo queries (e.g., an x-axis value of 3500 corresponds to position IDs $3500\sim3531$).
  • Figure 3: Recall Performance of different methods across various Window Sizes.
  • Figure 4: Ablation analysis of pseudo queries with respect to quality, length and insert positions. (a) Performance under a fixed Q_pseudo length $32$, with varying semantic content of Q_pseudo: Fm_Ln is the concatenation of the first m and the last n tokens from the input context; RS_C is constructed by concatenating 32 randomly sampled individual tokens from the context; RC_C is a randomly sampled consecutive span of 32 tokens from the context; and Fix_C is a fixed, repetitive nonsensical sequence (e.g., "Sorry, I don't know. Sorry, I don't know…"). (b) Performance under varying observation window sizes $N$. (c) Performance under varying insert positions of Q_pseudo. Left Parts: we divide the original context interval $[0, L_p)$ equally into three segments: part1 (beginning), part2 (middle), and part3 (end). For each segment, we randomly select a continuous sequence of 32 positions and map the position IDs of Q_pseudo to it. Right Parts: we only move backwards the position IDs of Q_pseudo to continuous 32 positions at different offsets after the context $L_p$ (e.g., $[L_p+0, L_p+32)$), $[L_p+32, L_p+64)$).
  • Figure 5: Visualization of Needle-in-a-Haystack results. We take LLaMA-3-8B-Instruct (8k context, 256 KV size) as a representative example to demonstrate the performance differences. The vertical axis represents the depth percentage, and the horizontal axis represents the token length.