HARD-KV: Head-Adaptive Regularization for Decoding-time KV Compression

2026-06-27

AI Summary

This work addresses the tension between dynamic KV cache compression and the memory constraints of static inference engines in long-context large language model reasoning. It proposes HARD-KV, a framework that manages token lifecycles through a hierarchical cache structure, unifies importance metrics across heterogeneous attention heads via a logits calibration mechanism, and enables contiguous physical memory layout through dynamic index reordering. This design maintains compatibility with modern inference optimizations such as CUDA Graphs and PagedAttention. As the first complete solution to adapt dynamic Top-p KV compression to static inference engines, HARD-KV achieves up to 2× throughput improvement on mathematical reasoning benchmarks like AIME and U-Math in contexts exceeding 10k tokens, while preserving high-fidelity generation quality.

Abstract

Long-context LLM inference faces a fundamental conflict: head-adaptive compression algorithms (e.g., Top-$p$ nucleus sampling) offer superior accuracy by dynamically fluctuating memory budgets, yet modern inference engines (e.g., vLLM) demand rigid, static memory patterns to leverage CUDA Graphs and PagedAttention. We resolve this ``Static-Dynamic'' mismatch with HARD-KV, a unified framework that that bridges dynamic selection with rigid system constraints. HARD-KV introduces a Cascade Cache hierarchy, managing the token lifecycle across dense, sparse, and condensed tiers. Crucially, we propose a Logits Calibration mechanism that normalizes diverse importance metrics into a unified probability space, enabling consistent Top-$p$ budgeting across heterogeneous heads. To bridge the efficiency gap, we offer a system-level solution, which rewrites fragmented, dynamic indices into contiguous physical layouts compatible with high-performance inference engine. Extensive experiments on math-reasoning benchmarks (AIME, U-Math) verify that HARD-KV achieves up to 2$\times$ throughput improvement over static baselines while maintaining high-fidelity generation in 10k+ token scenarios. Code is available at https://github.com/SuDIS-ZJU/HARDInfer.

Problem

• KV compression
• long-context LLM inference
• head-adaptive
• static-dynamic mismatch
• memory budget

Innovation

• KV compression
• head-adaptive regularization
• CUDA Graphs
• PagedAttention
• logits calibration