При инференсе LLM общее потребление памяти определяется не только размером самой модели, но и промежуточными данными, накапливаемыми в процессе ее работы. С ростом контекста объем этих данных растет почти линейно и может стать сопоставимым или даже превышать размер самой модели.
В основе этой проблемы лежит KV-cache. Пример: у LLaMA 2 7B веса занимают около 14 ГБ но при контексте 8K токенов KV-cache весит уже примерно 4 ГБ. Всего при четырех параллельных запросах это около 16 ГБ.
Это и есть скрытая цена инференса, которая не так очевидна на первый взгляд.
Содержание:
Что такое 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 модель каждый раз перечитывала бы книгу с начала, а с ним — просто открывает нужные страницы по закладкам и продолжает чтение.
Но, как уже говорилось выше, решение одной проблемы привело к появлению другой: экономия вычислений обернулась ростом потребления памяти.
Почему KV-cache сложнее сжать, чем веса модели
Во-первых, KV-cache динамический и появляется прямо во время инференса. Обычно методы квантования весов модели применяются заранее и аккуратно, чтобы не потерять качество. Эти методы здесь не подходят.
Во-вторых, при сжатии неизбежно появляются ошибки, которые в случае KV-cache напрямую влияют на механизм внимания. Из-за этого искажения в KV-cache сильнее сказываются на качестве генерации.
В-третьих, для KV-cache важно сохранять взаимное расположение векторов. Здесь опять вычисления опираются на сравнение Q и K. Если при сжатии изменить эти расположения, меняется распределение внимания. В результате модель начинает иначе интерпретировать контекст, что может ухудшить качество генерации.
Именно поэтому сжатие KV-cache — это отдельная и более сложная задача. Здесь важно не просто уменьшить объем данных, но и сделать это быстро, точно, а также не разрушить структуру этих представлений.
Как TurboQuant сжимает KV-cache
И тут самое время рассказать о новом решении исследователей из Google — TurboQuant. Их метод как раз позволяет одновременно эффективно и точно сжимать KV-cache. Работает он в два этапа:
Сначала TurboQuant случайно поворачивает векторы. После такого поворота координаты в высоких размерностях становятся почти независимыми и получают удобное распределение. Благодаря этому каждую координату можно квантовать отдельно почти оптимально. В результате большая часть сжатия достигается уже на этом этапе, но точность представления все-таки немного уменьшается.
После этого применяется метод QJL (Quantized Johnson–Lindenstrauss), который работает с ошибкой после сжатия. Он запоминает ошибку в существенно упрощенном виде — не через сами значения ошибки, а через набор знаков +1 и −1 после случайного преобразования. Такой способ дешев по памяти и позволяет при вычислении скалярных произведений во внимании учитывать эту ошибку — не идеально, но без систематического перекоса.
TurboQuant решает проблему с ограничением по памяти не через уменьшение самой модели, а через оптимизацию KV-cache, то есть тех данных, которые модель хранит во время инференса.
И вот что это дает, по заявлению авторов:
представления сжимаются примерно с 16 бит до ~3–4 бит;
в экспериментах KV-cache сжимается минимум примерно в 4,5 раза, а при 2,5 битах теоретическая экономия по представлению может приближаться к 6-кратной;
снижается объем данных, которые нужно хранить и передавать, что потенциально может ускорять инференс.
При этом KV-cache сжимается на лету во время работы модели.
Почему это важно
После всего описанного логично задать вопрос: окей, мы научились сжимать KV-cache, и что дальше? На самом деле последствия тут довольно заметные.
Раньше модели в основном сравнивали по размеру: 7B, 13B, 70B и так далее. Но на практике быстро оказывается, что на стоимость влияет не только сама модель, но и KV-cache, который она держит в памяти во время работы.
Начинают появляться другие вопросы: сколько памяти уходит на один токен, сколько стоит один запрос, сколько пользователей можно посадить на один GPU? То есть фокус постепенно смещается с размера модели на стоимость ее работы в смысле потребления памяти.
Если KV-cache уменьшается в несколько раз, это напрямую влияет на инфраструктуру. На том же железе можно обслуживать больше пользователей и держать более длинные сессии.
И здесь важный момент, который часто неправильно интерпретируют. Кажется, что если модель стала потреблять меньше памяти, значит спрос на железо должен упасть. Но на практике обычно происходит обратное. Когда что-то становится дешевле, его начинают использовать чаще. Поэтому снижение стоимости инференса чаще приводит не к снижению нагрузки, а к ее росту.
Еще один эффект — это длинные контексты. Обычно кажется, что ограничение здесь в вычислениях. Но на практике основная проблема — это память. KV-cache растет с каждым токеном, и именно он ограничивает длину диалога. Если его удается сжать, то можно держать больше контекста в той же памяти, становится проще строить сложные сценарии и работать с агентными системами.
Что в итоге
При этом важно не переоценивать эффект TurboQuant. Пока что он показывает хорошие результаты в контролируемых условиях. В реальных условиях эффект может быть ниже. Кроме того, он не устраняет KV-cache, а только делает его компактнее. То есть сама проблема никуда не исчезает.
Также есть альтернативные подходы — например, вместо хранения KV-cache можно пытаться частично пересчитывать его заново. Это другой компромисс, при котором требуется уже больше вычислительных ресурсов, чем памяти. Поэтому TurboQuant — это скорее начало новой ветки оптимизаций, связанной с инференсом и управлением памятью.
Мы привыкли оценивать LLM по размеру модели. Но в реальности ключевую роль играет стоимость ее работы, где и значительную часть этой стоимости формирует KV-cache.
TurboQuant показывает довольно неплохие результаты по оптимизации памяти. И это, скорее всего, отдельное направление развития, которое в ближайшие годы будет только усиливаться.