AMS KV Compression Framework: A New Solution to KV Cache Bottlenecks in Long-Context Inference

Long-context inference is a key requirement for large language model applications, but the linear growth of KV cache limits efficiency. Existing global Top-k compression methods cause the "region erasure" problem (important continuous reasoning blocks are discarded entirely). The AMS (Adaptive Mass-Segmented) framework replaces global Top-k with region-aware quota allocation to solve region erasure, can be seamlessly integrated into inference frameworks like vLLM, and improves inference quality and memory efficiency.

Original author team: Paper author team (arXiv submission) Source: arXiv (May 22, 2026) Original link: http://arxiv.org/abs/2605.23200v1

KV Cache Dilemma in Long-Context Inference

In autoregressive generation, KV cache stores key-value vectors of previous tokens to avoid repeated computation, but increasing context length leads to linear cache growth: a model with 100K context may occupy tens of GB of memory, and frequent access increases latency, limiting batching and concurrency capabilities. Existing solutions use global Top-k to select important tokens, but have fatal flaws.

Region Erasure Problem

Global Top-k causes "region erasure": in tasks like mathematical reasoning and code generation, logically related tokens form continuous reasoning blocks (e.g., derivation steps, code functions). Global Top-k may discard these blocks entirely, breaking the reasoning chain, losing logical coherence, and reducing output quality. The reason is that global Top-k assumes equal competition among tokens, ignoring the importance of regional structure—individual token importance cannot be evaluated out of context.

Core Idea of the AMS Framework

Shift from token-level competition to region-aware quota allocation:

1. Attention quality distribution analysis: adaptively partition the KV cache;
2. Region-level quota guarantee: key reasoning segments get memory quotas;
3. EMA smoothing mechanism: prevent boundary jitter.

Technical Implementation

1. Adaptive Partitioning: Divide regions at "valleys" of attention quality score distribution to retain high-attention-quality regions;
2. Region Quota Allocation: Each region gets a minimum guaranteed quota, with the remaining allocated by importance—tokens within a region share a quota pool;
3. EMA Smoothing: Apply time-domain smoothing to region boundaries to stabilize the decoding process.

Universality and Compatibility

• Orthogonal to existing scorers (TOVA, Expected Attention, etc.), responsible for region division and quota allocation;
• Compatible with paged KV service frameworks like vLLM, supports efficient gather-and-compact execution, no additional steady-state attention overhead, and no need to reconstruct infrastructure.

Experimental Task Validation

Validated effectiveness across multiple task sets:

• Mathematical Reasoning: Reduced inference failures caused by region erasure and improved accuracy in MATH500, AIME, and GSM8K tasks;
• Code Generation: Preserved key variable definitions, function calls, and logic flow, resulting in more coherent and executable code;
• Open-Domain QA: Retained key document information blocks and improved answer accuracy;
• Sparse Retrieval: Enhanced precision for retrieving specific information from long documents.

Advantages Over Baselines

Compared to global Top-k methods:

• Structural Integrity: Continuous reasoning blocks are better preserved;
• Reasoning Coherence: Higher success rate for multi-step reasoning;
• Compression Efficiency: Better performance at the same compression ratio;
• Stability: More stable and consistent decoding outputs.

Practical Significance and Deployment Advice

When to Use AMS:

1. Long-context inference (>32K);
2. Structured generation tasks (mathematical reasoning, code generation, etc.);
3. Scenarios requiring high reliability;
4. Existing KV compression infrastructure needs performance improvement.

Deployment Notes:

1. Low computational overhead for region division;
2. Need to tune the number of regions and quota ratio;
3. Can be combined with KV quantization for further memory compression.

Technical Insights

System design must consider data structure characteristics: global optimization easily breaks structure, while structure-aware optimization can compress efficiently while maintaining performance—this is insightful for model pruning, knowledge distillation, etc.

Limitations and Future Directions

Limitations: Region division heuristics are not optimal, task-dependent, and boundaries need adjustment for dynamic attention patterns; Future: Learning-based region division, task-adaptive quota allocation, and combination with sparse attention.

Conclusion

AMS provides a practical and efficient solution for long-context inference. By solving the region erasure problem via region-aware quota allocation and being compatible with existing frameworks, it is an option for upgrading production environments. As models move toward longer contexts, structure-aware compression technology will become increasingly important.