NVIDIA открыла исходный код KVzap: метод обрезки кэша KV, обеспечивающий сжатие данных почти без потерь в 2–4 раза

Когда длина контекста достигает десятков и сотен тысяч токенов, кэш ключевых значений в декодерах-трансформерах становится основным узким местом при развёртывании. Кэш хранит ключи и значения для каждого слоя и заголовка с формой (2, L, H, T, D). Для простого трансформера, такого как Llama1-65B, размер кэша составляет около 335 ГБ при 128 тыс. токенов в bfloat16, что напрямую ограничивает размер пакета и увеличивает время до первого токена.

Архитектурное сжатие не затрагивает ось последовательности

Производственные модели уже сжимают кэш по нескольким осям. Grouped Query Attention разделяет ключи и значения между несколькими запросами и обеспечивает коэффициент сжатия 4 в Llama3, 12 в GLM 4.5 и до 16 в Qwen3-235B-A22B — всё это по оси заголовков. DeepSeek V2 сжимает измерения ключей и значений через Multi head Latent Attention. Гибридные модели смешивают внимание со скользящим оконным вниманием или слоями пространства состояний, чтобы уменьшить количество слоёв, поддерживающих полный кэш.

Эти изменения не сжимают ось последовательности. Разреженное внимание и внимание в стиле поиска извлекают только подмножество кэша на каждом шаге декодирования, но все токены по-прежнему занимают память. Поэтому для практического использования длинных контекстов нужны методы, которые удаляют записи из кэша, которые незначительно повлияют на будущие токены.

Проект KVpress от NVIDIA

Проект KVpress от NVIDIA собирает более двадцати таких методов обрезки в одной кодовой базе и представляет их через общедоступную таблицу лидеров на Hugging Face. Такие методы, как H2O, Expected Attention, DuoAttention, Compactor и KVzip, оцениваются согласованным образом.

KVzip и KVzip+

KVzip в настоящее время является наиболее эффективным базовым методом обрезки кэша в списке лидеров KVpress. Он определяет показатель важности для каждой записи в кэше, используя задачу копирования и вставки претекста. Модель работает с расширенным запросом, где её просят повторить исходный контекст точно. Для каждой позиции токена в исходном запросе оценка равна максимальному весу внимания, который любая позиция в повторяющемся сегменте назначает обратно этому токену, по головам в одной группе, когда используется групповое запросное внимание. Записи с низкими оценками удаляются до тех пор, пока не будет достигнут глобальный бюджет.

KVzip+ уточняет эту оценку. Он умножает вес внимания на норму вклада значения в остаточный поток и нормализует по норме получающегося скрытого состояния. Это лучше соответствует фактическому изменению, которое токен вызывает в остаточном потоке, и улучшает корреляцию с точностью на выходе по сравнению с исходной оценкой.

Эти оценки оракула эффективны, но дороги. KVzip требует предварительного заполнения расширенного запроса, что удваивает длину контекста и делает его слишком медленным для производства. Он также не может работать во время декодирования, поскольку процедура оценки предполагает фиксированный запрос.

KVzap, суррогатная модель на скрытых состояниях

KVzap заменяет оценку оракула небольшой суррогатной моделью, которая работает непосредственно со скрытыми состояниями. Для каждого слоя-трансформера и каждой позиции последовательности t модуль получает скрытый вектор hₜ и выводит прогнозируемые логарифмические оценки для каждой ключевой головной ценности. Рассматриваются две архитектуры: одиночный линейный слой (KVzap Linear) и двухуровневый MLP с GELU и скрытой шириной, равной одной восьмой размера скрытого слоя модели (KVzap MLP).

Обучение использует подсказки из образца набора данных Nemotron Pretraining. Исследовательская группа фильтрует 27 тыс. подсказок длиной от 750 до 1250 токенов, выбирает до 500 подсказок для каждого подмножества, а затем выбирает 500 позиций токенов для каждой подсказки. Для каждой ключевой головной ценности они получают около 1,2 миллиона обучающих пар и набор для проверки из 23 тыс. пар. Суррогат учится переходить от скрытого состояния к логарифмической оценке KVzip+.

Во всех моделях квадратичная корреляция Пирсона между прогнозами и оценками оракула достигает примерно от 0,63 до 0,77, причём вариант MLP последовательно превосходит линейный вариант.

Пороговая обработка, скользящее окно и незначительные накладные расходы

Во время логического вывода модель KVzap обрабатывает скрытые состояния и выдаёт оценки для каждой записи в кэше. Записи со значениями ниже фиксированного порога удаляются, а скользящее окно из 128 последних токенов всегда сохраняется. Исследовательская группа предоставляет краткую функцию в стиле PyTorch, которая применяет модель, устанавливает значения локального окна в бесконечность и возвращает сжатые тензоры ключей и значений.

В ходе всех экспериментов обрезка применяется после операции внимания. KVzap использует пороговую обработку оценок, а не фиксированный выбор топ-k. Один порог даёт разные эффективные коэффициенты сжатия на разных тестах и даже в рамках одних и тех же тестов. Исследовательская группа сообщает о колебаниях коэффициента сжатия до 20 процентов при фиксированном пороге, что отражает различия в плотности информации.

Вычислительные накладные расходы невелики. Анализ на уровне слоёв показывает, что дополнительные затраты на KVzap MLP составляют не более 1,1 процента от FLOPs линейной проекции, а линейный вариант добавляет около 0,02 процента. Относительные накладные расходы памяти соответствуют тем же значениям. В режимах с длинным контекстом квадратичная стоимость внимания доминирует, поэтому дополнительные FLOPs практически незначительны.

Результаты на RULER, LongBench и AIME25

KVzap оценивается на тестах длинного контекста и рассуждениях с использованием Qwen3-8B, Llama-3.1-8B Instruct и Qwen3-32B. Поведение в длинном контексте измеряется на RULER и LongBench. RULER использует синтетические задачи на длинах последовательностей от 4 тыс. до 128 тыс. токенов, в то время как LongBench использует документы из реального мира из нескольких категорий задач. AIME25 предоставляет рабочую нагрузку для математических рассуждений с 30 задачами уровня Олимпиады, оцениваемыми по критериям «проход на 1» и «проход на 4».

На RULER KVzap соответствует базовому уровню полного кэша с небольшим запасом точности, удаляя при этом значительную часть кэша. Для Qwen3-8B лучшая конфигурация KVzap достигает удаляемой доли выше 0,7 на RULER 4k и 16k, сохраняя средний балл в пределах нескольких десятых долей балла от полного кэша. Аналогичное поведение наблюдается для Llama-3.1-8B Instruct и Qwen3-32B.

На LongBench те же пороговые значения приводят к более низким коэффициентам сжатия, поскольку документы менее повторяются. KVzap остаётся близким к базовому уровню полного кэша вплоть до сжатия примерно в 2–3 раза, в то время как методы с фиксированным бюджетом, такие как Expected Attention, ухудшаются на нескольких подмножествах по мере увеличения сжатия.

На AIME25 KVzap MLP поддерживает или немного улучшает точность прохода на 4 при сжатии примерно в 2 раза и остаётся пригодным для использования даже при удалении более половины кэша. Чрезвычайно агрессивные настройки, например, линейные варианты при высоких порогах, которые удаляют более 90 процентов записей, как и ожидалось, снижают производительность.

В целом, таблица выше показывает, что лучшая конфигурация KVzap для каждой модели обеспечивает среднее сжатие кэша примерно в 2,7–3,5 раза, сохраняя при этом результаты задач очень близкими к базовому уровню полного кэша на RULER, LongBench и AIME25.

Ключевые выводы

KVzap — это адаптивная по входным данным аппроксимация KVzip+, которая учится предсказывать оценки важности KV оракула из скрытых состояний, используя небольшие суррогатные модели для каждого слоя, либо линейный слой, либо неглубокий MLP, а затем удаляет пары KV с низкими оценками.

Обучение использует подсказки Nemotron pretraining, где KVzip+ обеспечивает контроль, создавая около 1,2 миллиона примеров на голову и достигая квадратичной корреляции в диапазоне от 0,6 до 0,8 между прогнозируемыми и оракуловыми оценками, что достаточно для достоверного ранжирования важности кэша.

KVzap применяет глобальный порог оценки с фиксированным скользящим окном недавних токенов, поэтому сжатие автоматически адаптируется к плотности информации в подсказке, и исследовательская группа сообщает о колебаниях до 20 процентов в достигнутом сжатии по подсказкам при том же пороге.

На Qwen3-8B, Llama-3.1-8B Instruct и Qwen3-32B на RULER, LongBench и AIME25 KVzap достигает сжатия кэша KV примерно в 2–4 раза, сохраняя точность очень близкой к полному кэшу, и достигает современных компромиссов в рейтинге NVIDIA KVpress Leaderboard. Дополнительные вычисления невелики — не более 1,1 процента дополнительных FLOPs для варианта MLP, а KVzap реализован в открытой среде kvpress с готовыми к использованию контрольными точками на Hugging Face, что делает его практичным для интеграции в существующие стеки обслуживания длинных контекстов LLM.

1. Какие проблемы решает метод KVzap, представленный в статье?

Метод KVzap решает проблему обрезки кэша в декодерах-трансформерах при работе с длинными контекстами. Он позволяет удалять записи из кэша, которые незначительно повлияют на будущие токены, что уменьшает размер кэша и ускоряет время до первого токена.

2. Какие методы сжатия кэша уже существуют и как они работают?

В статье упоминаются несколько методов сжатия кэша, таких как Grouped Query Attention, DeepSeek V2 и гибридные модели. Они работают по-разному: некоторые разделяют ключи и значения между несколькими запросами, другие сжимают измерения ключей и значений через Multi head Latent Attention. Однако все эти методы не затрагивают ось последовательности.

3. Как работает метод KVzap и чем он отличается от других методов?

KVzap — это суррогатная модель на скрытых состояниях, которая заменяет оценку оракула небольшой моделью. Для каждого слоя-трансформера и каждой позиции последовательности модуль получает скрытый вектор и выводит прогнозируемые логарифмические оценки для каждой ключевой головной ценности. В отличие от других методов, KVzap использует пороговую обработку оценок, а не фиксированный выбор топ-k.

4. Какие результаты показал метод KVzap на тестах длинного контекста и рассуждениях?

На тестах длинного контекста и рассуждениях метод KVzap показал хорошие результаты. Он соответствует базовому уровню полного кэша с небольшим запасом точности, удаляя при этом значительную часть кэша. На RULER KVzap достигает удаляемой доли выше 0,7 на RULER 4k и 16k, сохраняя средний балл в пределах нескольких десятых долей балла от полного кэша. На LongBench KVzap остаётся близким к базовому уровню полного кэша вплоть до сжатия примерно в 2–3 раза.

5. Какие выводы можно сделать из статьи о методе KVzap?

Из статьи можно сделать вывод, что метод KVzap является эффективным инструментом для обрезки кэша в декодерах-трансформерах. Он позволяет сохранять точность задач очень близкой к полному кэшу, достигая сжатия кэша KV примерно в 2–4 раза. Кроме того, дополнительные вычисления невелики, что делает метод практичным для интеграции в существующие стеки обслуживания длинных контекстов LLM.

Источник