HeadWiseKV: Budgeted Per-Head Cache Residency for Hybrid Long-Context Language Models

2026-09-01

AI Summary

为解决长上下文推理中GPU内存消耗大、生成吞吐量低的问题，提出HeadWiseKV框架，通过压缩全局KV缓存并分配静态多级历史窗口来优化缓存需求。

Abstract

Long-context inference retains a growing key--value (KV) cache during decoding, which consumes substantial GPU memory and can reduce generation throughput. This bottleneck remains in hybrid language models because their residual global-attention layers can dominate context-dependent cache demand. We study how to allocate this state under an aggregate KV-residency budget. We introduce HeadWiseKV, a training-free framework that compresses the residual global KV caches of hybrid language models while preserving their native local, recurrent, and linear paths. It assigns each physical KV head a static, multilevel history window, making cache demand predictable before serving. We formulate this allocation as a restricted operational rate--distortion problem and propose SeqCalib as the core policy-generation algorithm in HeadWiseKV. SeqCalib processes layers in execution order and conditions each decision on the lower-layer policy used at deployment, thereby accounting for interactions across depth. A grouped-cache runtime materializes the selected policy as actual per-head KV residency rather than a mask over a full cache. We evaluate downstream quality across four hybrid long-context models and study physical residency and serving behavior on Qwen3.6-27B. HeadWiseKV retains near-Full-KV RULER and LoCoMo quality across the evaluated models. In the fixed-model systems study, it reduces sampled peak device memory by 8.59\% at a 112K context length and extends the largest verified successful context from 114K to 161K.

Problem

• long-context inference
• key-value cache
• GPU memory
• hybrid language models
• KV-residency budget

Innovation

• HeadWiseKV
• SeqCalib
• KV cache compression
• budgeted per-head cache residency
• hybrid language models