KVzap: Nvidia научились в 3–4 раза эффективнее использовать память на инференсе

KV-кэш сегодня – главная ахиллесова пята трансформеров при масштабировании контекста. Он растет линейно по длине последовательности и хранится для каждого слоя и каждой головы.

Например, для LLaMA-подобной модели на 65B параметров KV-cache при 128k токенов занимает ~335 ГБ памяти. И по времени это тоже больно бьет.

При этом большинство оптимизаций сокращают KV-cache по слоям или по головам. Хотя главный потенциал – именно по оси токенов: далеко не все из них реально нужны модели.

Первый рабочий способ сокращения KV по токенам придумали авторы KVzip: до 4× сжатия при нулевых потерях качества. Но на практике метод оказался слишком медленным.

Nvidia взяли эту идею, немного ее переделали и получили почти то же самое, но практически бесплатно.

Они просто обучают маленькую модель, которая по hidden state токена предсказывает, насколько важен его KV. Для каждого слоя она своя, но это либо линейная модель, либо двухслойный MLP — максимум 1–2 матричных умножения.

И все, никаких дорогих операций и пересчетов (для сравнения: в KVzip промпт по сути приходилось гонять дважды). Дальше просто выкидываются KV-пары, у которых важность ниже заданного порога.

Compute overhead — около 0.02% FLOPs для линейных моделей. На длинном контексте это шум на фоне квадратичного внимания.

Деградация на бенчмарках – около нуля, сжатие – 3–4×. Прямо сказка какая-то (хотя, конечно, многое все еще зависит от движка).

Снимаем шляпу перед Nvidia за отличную работу. Все в опенсорсе, кстати.

https://arxiv.org/abs/2601.07891