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

TL;DR

This paper introduces DapQ, a position-aware pseudo query method for KV cache compression in LLMs, improving memory efficiency while maintaining high accuracy during inference.

Contribution

It proposes a novel decoding-aligned KV cache compression framework that uses position-aware pseudo queries to better simulate decoding queries for importance assessment.

Findings

Achieves up to 99.5% performance retention with only 3% KV cache.

Outperforms existing compression methods across multiple benchmarks.

Effectively reduces memory footprint during LLM inference.

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).