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

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.Source: §1 · p. 1
• 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.Source: §1 · p. 1
• 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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2
• 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.Source: §1 · p. 2

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.Source: §1 · p. 2, Figure 1 · p. 2

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.Source: §2 · pp. 2–3

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.Source: §2 · pp. 2–3

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.Source: §2 · p. 3

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.Source: §2 · p. 3

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.Source: §2 · p. 3

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.Source: §3.1 · p. 3
• 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.Source: §3.1 · p. 3
• 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.Source: §3.1 · p. 3
• 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.Source: §3.1 · pp. 3–4

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.Source: §3.2 · pp. 4–5
• 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.Source: §3.2 · pp. 4–5
• 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.Source: §3.2 · p. 5

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.Source: Figure 2 · p. 4, §3.3 · p. 5

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.Source: §4 · p. 5, §4.1 · p. 5

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.Source: §4.2 · p. 6

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.Source: §4.2 · p. 6

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.Source: §4.2 · p. 7

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.Source: §4.3 · p. 7

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.Source: §4.3 · pp. 7–8

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

Calibration across depth

• SeqCalib replaces joint allocation with execution-order conditional decisions under the realized lower-layer policy.Source: §A.1 · p. 10

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.Source: §A.1 · p. 10
• The output difference is bounded by the maximum cross-set value distance ∆V, with the zero-omitted-mass case exact.Source: §A.1 · p. 10
• The omitted attention mass mw and value spread ∆V vary across queries and physical KV heads, producing heterogeneous truncation deviations.Source: §A.1 · p. 10
• SeqCalib averages query-head behavior within each GQA group and assigns the shared physical KV head its own window.Source: §A.1 · p. 10

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.Source: §A.1 · p. 10

Calibration score

• SeqCalib uses a stabilized cosine score to compare candidate and full-history outputs.Source: §A.2 · p. 10
• The full candidate scores one and supplies the replay reference for trace compatibility and finite-window comparisons.Source: §A.2 · p. 10
• Every feasible finite window satisfies an in-trace squared-error bound, while the full window has zero error.Source: §A.2 · p. 11

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.Source: §A.2 · p. 11
• 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.Source: §A.2 · p. 11

Operational distortion

• The budgeted objective trades predictive information absent from the policy-visible state and residual conditional mismatch against nominal configurable-KV storage.Source: §A.3 · p. 12

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.Source: §B.1 · p. 12

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.Source: §B.1 · p. 12
• 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.Source: §B.1 · p. 12

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.Source: §B.1 · p. 12

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.Source: §B.2 · p. 12

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.Source: §B.2 · p. 13

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.Source: §B.2 · pp. 12–13

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.Source: §B.2 · p. 13

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.Source: §B.2 · p. 13

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.Source: §B.3 · p. 13, Table 5 · p. 14

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.Source: §B.3 · p. 13

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%.Source: §B.3 · p. 13

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.Source: §B.4 · p. 14
• 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.Source: §B.4 · p. 14

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.Source: §B.4 · p. 14