Search papers, labs, and topics across Lattice.
The paper introduces CONF-KV, a KV-cache management strategy for long-horizon LLM inference that leverages the model's next-token distribution to estimate confidence and dynamically adjust the cache budget. This method retains more context when the model is uncertain and prunes aggressively when confident, using a composite ranking of accumulated attention mass and recency within each budget. Combined with blockwise online-softmax attention, mixed-precision storage, and a pyramidal per-layer budget, CONF-KV achieves near-full KV performance with significantly reduced memory footprint.
LLMs can maintain near-perfect long-context recall (91.4% on Needle-in-a-Haystack) while slashing memory use by intelligently pruning the KV cache based on the model's own confidence.
Long-horizon LLM inference turns the key--value (KV) cache into the dominant GPU memory consumer and makes per-token attention increasingly expensive. Many common eviction policies use static recency windows or historical attention, leaving unused a signal computed on every decoding step: the model's current uncertainty. We introduce CONF-KV, a KV-cache manager that converts the next-token distribution into a scalar confidence score and uses it to choose the per-step cache budget, retaining more context when the model is uncertain and pruning aggressively when it is confident. Within each budget, tokens are ranked by a composite of accumulated attention mass and recency, while a protected recent window preserves local coherence. We combine the policy with blockwise online-softmax attention, mixed FP16/INT8 storage, and a pyramidal per-layer budget variant. Across four model families and generated lengths up to 4K, CONF-KV stays near the footprint of a fixed 512-token sliding window while remaining within 1.5--2.1 perplexity points of full KV. On Needle-in-a-Haystack up to 32K tokens, CONF-KV reaches 91.4% retrieval accuracy versus 53.8% for sliding windows and 80.6% for H2O; on 75 VisualWebArena tasks it retains 95.3% of full-KV success at 2.8 times lower peak memory.
