Скрытая цена LLM: как KV-cache увеличивает стоимость инференса и как эту проблему решает Google TurboQuant

В сфере искусственного интеллекта произошло заметное событие. konstantin_kozhin 13 минут назад Скрытая цена LLM: как KV-cache увеличивает стоимость инференса и как эту проблему решает Google TurboQuant Уровень сложности Средний Время на прочтение 5 мин Охват и читатели 442 Блог компании МТС Искусственный интеллект Машинное обучение * Natural Language Processing * Сжатие данных * Обзор При инференсе LLM общее потребление памяти определяется не только размером самой модели, но и промежуточными данными, накапливаемыми в процессе ее работы. С ростом контекста объем этих данных растет почти линейно и может стать сопоставимым или даже превышать размер самой модели. В основе этой проблемы лежит KV-cache.

Пример : у LLaMA 2 7B веса занимают около 14 ГБ, но при контексте 8K токенов KV-cache весит уже примерно 4 ГБ. Всего при четырех параллельных запросах это около 16 ГБ. Это и есть скрытая цена инференса, которая не так очевидна на первый взгляд.

Технические детали

Содержание: Что такое KV-cache Почему KV-cache сложнее сжать, чем веса модели Как TurboQuant сжимает KV-cache Почему это важно Что в итоге Что такое KV-cache Чтобы понять, что такое KV-cache, посмотрим на главный механизм LLM, который называется вниманием (attention). Модель генерирует текст по одному токену за раз, и каждый новый токен зависит от всех предыдущих. Например, если уже есть «Я люблю», то, чтобы сгенерировать следующее слово, модели нужно учесть и «Я», и «люблю».

Схематично вычисление внимания выглядит так: Для каждого токена заранее вычисляются векторы Q, K и V. Далее они комбинируются следующим образом: Q каждого токена сравнивается с K всех токенов (через скалярное произведение); результат масштабируется (деление на ); полученные значения нормализуются с помощью softmax; затем они умножаются на V. В результате каждый токен учитывает все остальные токены (включая токены из предыдущего контекста).

Это и есть механизм внимания. Обратите внимание на размерность seq_len на схеме — это длина всей последовательности, то есть количество токенов в контексте. Очевидно, что это долго и дорого.

Отраслевые последствия

Но K и V нет необходимости пересчитывать каждый раз (А вот Q нужно вычислять заново, потому что он относится к текущему токену и меняется на каждом шаге генерации), потому что они зависят только от уже обработанных токенов и не меняются на следующих шагах. Поэтому их можно сохранить и использовать снова на следующем токене — это и есть KV-cache. Дополним предыдущую схему: На первом шаге вычисляются Q, K и V для всей последовательности, и K с V сохраняются в кеш.

На следующих шагах считается только новый Q, а K и V берутся из кеша и дополняются значениями для нового токена. Если снова воспользоваться аналогией, то без KV-cache модель каждый раз перечитывала бы книгу с начала, а с ним — просто открывает нужные страницы по закладкам и продолжает чтение. Но, как уже говорилось выше, решение одной проблемы привело к появлению другой: экономия вычислений обернулась ростом потребления памяти.

Этот прогресс даёт важные сигналы о будущем отрасли, и технологический мир внимательно наблюдает.