PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Renjie Xie, Juncheng Yang, Aoting Hu, Mingxi Zhang, Liyao Wu, Zheheng Hong, Wei Xu

arXiv:2609.02029v1cs.AI

TL;DR

Long-context hybrid models still face growing KV-cache memory demands in residual global-attention layers. HeadWiseKV assigns static multilevel windows to physical KV heads, calibrates them sequentially under deployed lower-layer policies, and materializes them as physical residency. Across evaluated models it preserves near-Full-KV quality, while the fixed-model study reduces memory and extends verified context capacity.

  • Problem

    Residual global-attention layers retain histories that grow with context, so hybridization relocates rather than eliminates the KV-cache bottleneck.

  • Method

    HeadWiseKV is a training-free framework that assigns static multilevel per-head suffix windows, uses deployment-conditioned SeqCalib, and materializes the policy in grouped physical caches.

  • Results

    HeadWiseKV preserves near-Full-KV quality across evaluated hybrid models and reduces sampled peak device memory by 8.59% at 112K while extending verified context from 114K to 161K.

  • Takeaways & Limitations

    HeadWiseKV provides predictable physical KV residency while retaining long-context quality and maintaining decode throughput in the fixed-model systems study.

  • Takeaways & Limitations

    Quality and systems panels use separately recorded operating points, and separate control campaigns do not establish complete joint quality–systems operating points.

Abstract

from arXiv · show

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.

1 Introduction

Hybrid models reduce but do not eliminate long-context KV pressure because residual global-attention layers retain growing histories. HeadWiseKV addresses this bottleneck with static, multilevel per-head physical residency calibrated under deployment-relevant lower-layer policies.

  • Motivation: Residual global-attention layers can determine whether a hybrid model fits a long request in memory.Hybridization relocates rather than eliminates the KV bottleneck.
  • Motivation: Prompt-dependent compression can preserve isolated distant evidence, but its systems benefit depends on runtime cache realization.Logical sparse masks may reduce attention work without reducing allocated KV storage.
  • Motivation: Static common windows can waste memory on tolerant heads while truncating sensitive heads, and early-cache changes alter later-layer inputs.Deployment requires fine-grained capacities chosen under the lower-layer decisions that will actually be active.
  • HeadWiseKV: HeadWiseKV assigns each residual global KV head a fixed multilevel suffix window and stores only the selected histories before prefill.This replaces request-specific token selection with a predictable, prompt-independent memory plan.
  • HeadWiseKV: SeqCalib selects per-head windows layer by layer while conditioning each decision on the previously selected lower-layer policy.It compares candidate outputs with conditional full-history references and chooses the lowest-cost entry meeting a similarity threshold.
  • Results: 8.59% lower sampled peak device memory at 112K and a verified context increase from 114K to 161K accompany near-Full-KV RULER and LoCoMo quality.The systems and quality panels use separately recorded operating points.

2 Related Work

Prior KV-cache methods differ in when retention is chosen and whether it changes physical residency. HeadWiseKV combines prompt-independent, multilevel per-head capacities with deployment-conditioned calibration and physical allocation before prefill.

  • Comparison axes: KV-cache methods vary along policy timing and cache realization: decisions may be request-dependent or offline, and storage may be physical or merely logical.These axes separate prompt dependence from whether allocated memory changes.
  • Dynamic methods: Request-dependent methods adapt token or block access to prompts, queries, attention, positions, or auxiliary context.Examples include QUEST, TokenSelect, InfLLM, and PyramidInfer.
  • Static methods: Static methods fix retention before requests, but StreamingLLM and several head-role methods primarily use sink–recent or binary full-versus-compressed histories.Such layouts are regular to deploy but can be coarse for heterogeneous hybrid-model heads.
  • HeadWiseKV: HeadWiseKV fixes a multilevel suffix length for each physical KV head, calibrates higher layers under selected lower-layer policies, and allocates capacities before prefill.The comparison explicitly separates prompt-independent retention, physical residency, and deployment-conditioned calibration.
  • Scope: HeadWiseKV targets residual global-attention histories while leaving bounded-memory local or recurrent paths unchanged.Its retention policy is distinct from cache precision and serving-time memory management.

3 Method

HeadWiseKV formulates physical KV residency as a budgeted per-head allocation and uses SeqCalib to select windows sequentially under deployed prefixes. A grouped-cache runtime materializes the resulting policy as physically sized per-head caches.

  • Structured Residency Model: HeadWiseKV defines feasible history lengths and storage costs for physical KV heads in residual full-attention layers.Under GQA, one physical KV-head decision is shared by its mapped query heads.
  • Structured Residency Model: Each finite assignment retains a contiguous suffix, and the complete policy is a layer-by-KV-head matrix of history lengths.Finite windows are ordered below full history, with full representing the entire context.
  • Structured Residency Model: At 32K context, an 8K assignment stores one quarter as many KV elements as a full assignment for the same head.Retained storage grows linearly with history length.
  • Structured Residency Model: The normalized residency rate compares selected storage with all-full storage and excludes model weights, workspaces, temporary buffers, and allocator overhead.Whole-device VRAM is measured separately.
  • Sequential Prefix-Conditioned Calibration: SeqCalib evaluates candidate windows one layer at a time under the frozen lower-layer policy and chooses the shortest window meeting a similarity floor.It replays candidates from a shared trace, avoiding joint search over all heads.
  • Sequential Prefix-Conditioned Calibration: The calibration score compares each candidate attention output with the conditional full-history output over sampled calibration queries.Candidate windows restrict the available key/value positions to recent suffixes.
  • Sequential Prefix-Conditioned Calibration: SeqCalib is stage-wise exact for the realized lower-layer prefix but does not guarantee global optimality over all policy matrices.A threshold grid then records policies and selects the best budget-feasible record by the specified lexicographic rule.
  • Grouped-Cache Runtime: The runtime loads the selected matrix once and materializes grouped physical caches with different per-head lengths rather than masking an all-full cache.Heads sharing a window may share group configuration while retaining disjoint KV contents.

4 Evaluation

HeadWiseKV is evaluated for downstream quality, physical memory, serving efficiency, and cross-model transfer. It preserves near-Full-KV quality while reducing physical memory and extending verified context capacity, with transfer and cross-campaign comparisons requiring careful interpretation.

  • Experimental scope: The evaluation covers generation quality, serving efficiency, and transfer across four hybrid models, with controlled fixed-model comparisons and additional-model transfer tests.The fixed-model study uses Qwen3.6-27B; quality includes RULER, LoCoMo, and agent tasks.
  • Downstream quality: HeadWiseKV is the only compressed method that consistently preserves Full-KV behavior across retrieval, conversational memory, and both evaluated agent domains at 68.95% retention.Other baselines preserve some dimensions but show uneven retrieval, conversational-memory, or agent-task performance.
  • Memory and throughput: 8.59% lower sampled peak memory at 112K accompanies 1.092× prefill and 1.017× decode throughput relative to Full-KV.The matched Full-KV and HeadWiseKV cohort shows decode throughput remains close to Full-KV, while prefill is more context-dependent.
  • Capacity: HeadWiseKV succeeds at 161K while Full-KV succeeds at 114K, producing a 1.41× larger maximum verified successful context on the tested 1K grid.The adjacent-grid probes report OOM at 162K for HeadWiseKV and 115K for Full-KV.
  • Retention sensitivity: At 68.95% retention, RULER decreases by 0.35 points while LoCoMo remains unchanged; moderate compression preserves quality across the broader evaluated range.The strongest evaluated compression is the only point where the reported RULER degradation appears.
  • Model transfer: Across models, τ = .995 policies remain within 0.61 RULER points of Full-KV, but LoCoMo losses reach 2.08 points and induced retention varies by model.Capacity planning should use each model’s produced policy, and the target workload should be validated before deployment.

5 Conclusion

HeadWiseKV provides predictable physical KV residency for hybrid long-context models through training-free multilevel per-head windows and deployment-aware SeqCalib. Across evaluated models it preserves Full-KV quality more consistently while reducing memory, extending verified context, and maintaining decode throughput.

  • HeadWiseKV is a training-free framework for predictable physical KV residency in hybrid long-context models.
  • SeqCalib assigns multilevel per-head suffix windows under the lower-layer decisions used at deployment.
  • The grouped-cache runtime materializes the selected policy without retraining.
  • HeadWiseKV preserves Full-KV quality more consistently than matched-retention baselines across the evaluated models.
  • The fixed-model systems study reports reduced physical memory, an extended verified context range, and maintained decode throughput.

A Design Rationale and Calibration Properties

The design rationale treats static KV allocation as a depth-coupled, head-sensitive problem. HeadWiseKV addresses it with per-head windows and sequential calibration under the lower-layer policy active at deployment.

  • Calibration across depth: SeqCalib replaces joint allocation with execution-order conditional decisions under the realized lower-layer policy.
  • Head-dependent truncation: For one query, the full attention output is a mixture of retained and omitted value means, while suffix renormalization returns the retained mean.
  • Head-dependent truncation: The output difference is bounded by the maximum cross-set value distance ∆V, with the zero-omitted-mass case exact.
  • Head-dependent truncation: The omitted attention mass mw and value spread ∆V vary across queries and physical KV heads, producing heterogeneous truncation deviations.
  • Head-dependent truncation: SeqCalib averages query-head behavior within each GQA group and assigns the shared physical KV head its own window.
  • Calibration across depth: Shortening an earlier cache can change the state presented to every later block, so later calibration must use the deployed lower-layer prefix.

A.2 Calibration Score and Fixed-Prefix Selection

SeqCalib scores candidate windows against full-history outputs under fixed realized prefixes, then selects feasible low-cost policies and evaluates their operational distortion against storage.

  • Calibration score: SeqCalib uses a stabilized cosine score to compare candidate and full-history outputs.
  • Calibration score: The full candidate scores one and supplies the replay reference for trace compatibility and finite-window comparisons.
  • Calibration score: Every feasible finite window satisfies an in-trace squared-error bound, while the full window has zero error.
  • Fixed-prefix selection: Conditional on the realized lower-layer prefix, SeqCalib terminates finitely and selects minimum nominal layer cost among choices satisfying local similarity floors.
  • Fixed-prefix selection: A per-head policy costs no more than a shared feasible window, with strict improvement when at least one head admits a strictly cheaper feasible candidate.
  • Operational distortion: The budgeted objective trades predictive information absent from the policy-visible state and residual conditional mismatch against nominal configurable-KV storage.

B.1 Experimental Protocol

The experiments evaluate downstream quality across standardized long-context tasks and measure systems behavior with repeated, controlled Qwen3.6-27B runs. The protocol reports explicit timing, memory, capacity, and baseline controls.

  • Quality evaluation: RULER evaluation covers five retrieval tasks at context lengths from 4K through 128K, using official code, seed 42, and 40 examples per task–context cell.
  • Systems evaluation: Systems measurements use Qwen3.6-27B with F16 K/V, FlashAttention, full GPU offload, automatic batch fitting disabled, and three fresh-server repetitions per successful method–context pair.
  • Systems evaluation: Capacity scans vary only context length on a fixed setup and repeat the last success and first OOM three times to confirm the 1K-grid boundary.
  • Baseline controls: Baseline controls distinguish prefill compaction from preallocated compressed caches and reuse matched Full-KV measurements before independent adapted-control runs begin at 96K.

B.2 Runtime Realization and Deployment

HeadWiseKV deploys a static layer–KV-head window matrix through a grouped physical-cache runtime, allocating only the selected histories while preserving native local and recurrent cache paths.

  • Policy and calibration: Each runtime entry selects full history or a positive local window, and inference never recomputes the policy from prompts or request history.The deployed matrix is indexed by configurable full-attention layer and physical KV head.
  • Physical cache realization: The runtime partitions full-attention storage into active (window, KV head) groups, reusing group types while retaining disjoint layer-specific K/V entries.Unassigned layers and heads allocate no entries in that group.
  • Deployment assumptions: Native local and recurrent layers retain their original cache path while configurable full-attention layers use the selected physical residency layout.The runtime contract assumes known full-attention layers, GQA head mappings, and a cache-group-representable window codebook.
  • Physical cache realization: Each physical KV head reads and writes only its assigned K/V slice, while GQA query groups follow the mapped KV head.Per-head outputs are concatenated before unchanged gating and output projection, so the prompt is not globally truncated.
  • Accounting: Peak VRAM and confirmed context capacity are measured on complete runtime invocations, whereas normalized retention counts token–head entries only in configurable full-attention layers.Weights, workspaces, allocator effects, and non-attention state remain outside normalized retention accounting.

B.3 Qwen3.5-9B Quality and Systems Results

The Qwen3.5-9B evaluation separates audited LongMemEval-S quality from standardized physical-memory measurements. HeadWiseKV provides a lower-memory operating point with comparable decode throughput, while quality varies across calibration settings and baselines.

  • Evaluation protocol: The quality comparison uses 500 audited LongMemEval-S questions, while peak VRAM uses a separate standardized single-GPU workload with median fresh-server measurements.The reference context is 131,072 tokens and the workload generates 128 tokens.
  • Quality: 70.6% Reviewed ACC at τ = .999 retains 80.27% of configurable KV, compared with 70.2% for Full-KV.The 0.4-point difference corresponds to two questions and does not establish a ranking under this audit.
  • Baseline comparison: AdaKV and HeadKV-R2 compact only after full-context prefill, leaving whole-invocation peaks close to Full-KV despite decode-steady peaks of 8109 and 7529 MiB.HeadKV-R2’s 53.71% entry is its configured query-head budget equivalent, while unioned physical payload retention is 39.44%.

B.4 Limitations

HeadWiseKV’s conclusions are bounded by fixed prompt-independent profiles, calibration assumptions, reported workloads, and hardware-specific measurements.

  • Scope and assumptions: HeadWiseKV cannot recover evicted evidence or adapt to requests with atypical long-range dependencies because its window matrix is prompt-independent.SeqCalib profiles use in-sample reconstruction cosine and require revalidation when models, cache formats, contexts, or workloads change.
  • Scope and assumptions: Layerwise selection need not minimize a global allocation objective, and quality conclusions cover only the reported tasks and operating points.Aggregate scores do not support paired per-example uncertainty analysis.
  • Measurement boundary: Memory and throughput depend on tested hardware, backend, batch settings, quantization, and context length; token–head retention alone does not determine bytes or peak VRAM.Baseline absolute measurements characterize only the reported implementations because their cache lifecycles and adaptations differ.
Loading 2609.02029v1…