Technical Summary: ProxyKV

Problem Statement

Efficient long-context inference in Large Language Models (LLMs) is currently bottlenecked by the memory requirements of the Key-Value (KV) cache. As sequence length ($N$) increases, the linear growth of the KV cache often exceeds GPU memory capacity, leading to prohibitive I/O overhead and inference latency.

Existing KV cache pruning methods face a fundamental trade-off between scoring efficiency and pruning precision:

• Heuristic Methods (e.g., H2O, SnapKV): These are fast and low-latency but rely on local observation windows, often missing global semantic dependencies required for complex reasoning.
• Reconstruction-Based Methods (e.g., KVZip): These achieve high precision by performing a secondary, teacher-forced prefilling pass on the target model to reconstruct global context. However, this incurs a prohibitive prefilling overhead, negating the latency benefits of pruning.

The core challenge is bridging the gap between low-latency heuristics and high-precision reconstruction without incurring the heavy compute cost of a secondary prefill on the target model.

Methodology

The authors propose ProxyKV, a cross-model proxy pruning framework designed to offload the importance scoring task from the target model to a lightweight, asynchronous "Small-Model Proxy."

System Architecture

The system operates via a three-component pipeline:

1. Large-Model Target ($M_l$): The primary model responsible for the prefill critical path and the entire decode phase.
2. Small-Model Proxy ($M_s$): An intra-family smaller model (e.g., 1.5B or 4B parameters) that runs asynchronously on a separate execution stream (or CUDA stream). It extracts cross-head attention features ($X$) from the input context.
3. HybridAxialMapper ($\Phi_\theta$): A learnable module that maps the proxy's features ($X$) to target-aligned importance scores ($\hat{Y}$).

Workflow:
During inference, the input context is dispatched to both models concurrently. The proxy generates attention features, which are processed by the HybridAxialMapper to produce a binary pruning mask ($M$) via Top-K thresholding. The target model then performs decoding using only the pruned KV cache. Crucially, the proxy's KV cache is released immediately after the prefill phase, making the memory overhead transient.

Key Components

1. HybridAxialMapper Architecture

To bridge the architectural gap between heterogeneous models (different head counts and layer depths), the mapper employs a three-stage pipeline that disentangles temporal feature extraction from cross-head alignment:

• Stage 1: Temporal Feature Extraction: A 1D Convolutional Stem projects raw proxy features into a model-agnostic latent space, normalizing the head configuration.
• Stage 2: Time-Axis Context Encoder: A Transformer encoder processes the sequence with sinusoidal positional encodings to capture global context, utilizing a sliding window approach for long sequences (up to 170k tokens).
• Stage 3: Head-Axis Cross-Attention: Learnable target queries aggregate synthetic head tokens from the proxy into the target's specific head configuration. This stage allows the mapper to scale to arbitrary target head counts without retraining the entire network.

2. Multi-Granularity Hybrid Loss

Recognizing that the deployment metric is Top-K retention quality rather than pointwise score reconstruction, the authors design a composite loss function ($L_{total}$) consisting of five terms:

• $L_{bin}$ (Multi-Ratio Binary): Supervises the binary decision boundary across multiple retention ratios, emphasizing aggressive compression.
• $L_{mse}$ (Value-Weighted MSE): Ensures numerical magnitude alignment for salient tokens.
• $L_{fine}$ (Intra-Top-K Rank): Enforces relative ordering within the retained Top-K set.
• $L_{global}$ (Mass-Aware Global Rank): Optimizes the separation between Top-K and non-Top-K tokens (signal-to-noise boundary).
• $L_{cos}$ (Cosine Alignment): Aligns the global orientation of the predicted score vector with the target.

This loss design shifts the objective from rigid regression to relative ranking consistency, which is more directly aligned with pruning quality.

Key Contributions

1. Asynchronous Proxy Framework: A novel pruning framework that removes the scoring bottleneck from the target's critical path by offloading it to an asynchronous small model, achieving up to 3.21× prefilling speedup on Llama-3.1-8B (dual-GPU) and sustaining speedups up to 170k tokens.
2. HybridAxialMapper & Hybrid Loss: A learnable architecture and loss function that successfully bridge structural gaps between heterogeneous models (e.g., 8B target vs. 1.5B proxy), overcoming the rigidity of static head-to-head alignment used in prior work.
3. Empirical Validation: Extensive experiments across Llama-3.1, Qwen-2.5, and Qwen-3 families (7B–32B targets) demonstrating that ProxyKV recovers ~98.7% of the KVZip oracle's mean accuracy while significantly reducing latency.

Results

• Accuracy: ProxyKV matches the performance of the KVZip oracle (the high-precision baseline) across LongBench, SCBench, and RULER benchmarks. It recovers ~98.7% of KVZip's mean accuracy and tracks within 1–2 points on most tasks, even at aggressive retention ratios ($\rho \le 0.2$).
• Efficiency:
  • On Llama-3.1-8B, ProxyKV delivers a 3.21× speedup in prefilling latency compared to KVZip (dual-GPU) and ~1.5× on a single shared GPU.
  • The HybridAxialMapper itself contributes only 4.7%–6.8% of the prefill wall time.
  • The method scales effectively to 32B targets (Qwen-3) with an ~8× target/proxy size ratio.
• Robustness: The framework demonstrates robust zero-shot generalization to held-out datasets (e.g., SCBench RepoQA) and maintains performance on complex multi-hop reasoning tasks where heuristic methods (SnapKV) fail.

Significance and Claims

The paper claims that ProxyKV successfully bridges the "scoring-cost–accuracy gap" that has constrained long-context LLM inference. By leveraging intra-family attention correlations and a learnable cross-axial mapping, it achieves the high precision of reconstruction-based methods without the prohibitive prefilling overhead.

The authors highlight that while static layer-to-layer alignment (as used in IAM or KVZip+IAM) becomes a binding constraint as target models scale (e.g., to 32B), the learnable HybridAxialMapper remains effective. This suggests that ProxyKV offers a scalable path for efficient long-context inference, recovering near-oracle accuracy while significantly reducing latency and memory pressure during the prefill phase. The method is presented as a deploy-time solution that requires only a lightweight, intra-family proxy and a small, trainable mapper, avoiding the need for full model retraining or complex hardware modifications.