HARD-KV: Head-Adaptive Regularization for Decoding-time KV Compression
Yuxuan Yang ⋅ Feiyang Ren ⋅ Bowen Zeng ⋅ Dalin Zhang ⋅ Jinpeng Chen ⋅ Gang Chen ⋅ Huan Li
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.
Lay Summary
Large language models are increasingly used for tasks that require reading or generating very long texts, such as solving complex math problems, analyzing long documents, or carrying out multi-step reasoning. However, these models must keep track of a growing amount of intermediate information as the text gets longer, which consumes a large amount of GPU memory and slows down inference. This paper studies how to make long-context language model inference faster and more memory-efficient without sacrificing answer quality. Existing methods can improve accuracy by dynamically deciding which past information is important to keep, but these dynamic decisions are difficult to run efficiently on modern GPU inference systems, which prefer regular and predictable memory layouts. We propose HARD-KV, a framework that makes dynamic memory-saving methods compatible with efficient inference systems. HARD-KV organizes stored information into multiple levels, calibrates different importance scores into a common scale, and reorganizes irregular memory access patterns into layouts that GPUs can handle efficiently. Experiments on challenging math-reasoning benchmarks show that HARD-KV can improve inference throughput by up to 2$\times$ while preserving high-quality generation for long outputs exceeding 10k tokens. Our code will be made publicly available.
Successful Page Load