CONF-KV: Вытеснение KV-кэша с учетом уверенности и смешанной точностью хранения для длинногоризонтных LLM

Аннотация

Долгосрочный логический вывод LLM превращает кэш ключ-значение (KV) в основного потребителя памяти GPU и делает внимание на каждый токен всё более затратным. Многие распространённые политики вытеснения используют статические окна недавности или историческое внимание, оставляя неиспользованным сигнал, вычисляемый на каждом шаге декодирования: текущую неопределённость модели. Мы представляем CONF-KV — менеджер KV-кэша, который преобразует распределение следующего токена в скалярный показатель уверенности и использует его для выбора бюджета кэша на каждом шаге, сохраняя больше контекста, когда модель не уверена, и агрессивно сокращая его, когда она уверена. В рамках каждого бюджета токены ранжируются по композиту накопленной массы внимания и недавности, при этом защищённое окно недавности сохраняет локальную связность. Мы объединяем эту политику с блочным онлайн-softmax вниманием, смешанным хранением FP16/INT8 и пирамидальным вариантом бюджета на слой. На четырёх семействах моделей и длинах генерации до 4K CONF-KV остаётся близким по объёму памяти к фиксированному скользящему окну из 512 токенов, при этом отклоняясь от полного KV на 1,5–2,1 пункта перплексии. В задаче «Игла в стоге сена» с длиной до 32K токенов CONF-KV достигает точности извлечения 91,4% против 53,8% для скользящих окон и 80,6% для H2O; на 75 задачах VisualWebArena он сохраняет 95,3% успеха полного KV при в 2,8 раза меньшем пиковом объёме памяти.

English

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.