Как и зачем ускоряют LLM

Зачем вообще нужен KV cache?

Сегодня вы узнаете об одном из самых важных понятий в современном ИИ — о KV cache. С помощью него происходит оптимизация трансформера.

Мы знаем, что сегодняшние большие языковые модели опираются на архитектуру transformer, а самое важное понятие в трансформере — механизм attention(он же механизм внимания). Кратко напомним: attention помогает уловить смысл слова, учитывая все окружающие слова. Это делается через attention-вычисления.

Модели генерации текста вроде Chat GPT тоже основаны на архитектуре трансформера, но используют только decoder-часть. Они генерируют текст авторегрессивно — проще говоря, по одному токену за раз.

Чтобы сгенерировать даже один токен, модель внутри выполняет очень много вычислений. Если заглянуть, как именно рождается токен, то становится видно, что внутри модели — масса расчётов происходит только ради одного токена.

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

Почему без кэша это неэффективно

Внутри модели самая важная часть это — вычисление attention. Если присмотреться внимательнее, видно: вычислительная сложность растёт по мере роста размера контекста — то есть по мере того, как текст становится всё длиннее и длиннее.

Если подумать, это действительно неэффективно. При генерации тысяч токенов такие вычисления становятся очень медленными. Кроме того, если посмотреть на attention-расчёты, которые мы делаем при генерации текста, заметно, что в вычислениях много повторений.

Допустим, для первого токена мы посчитали attention score. Но когда генерируем следующий токен, мы снова считаем attention scores, которые уже были посчитаны раньше. И чем больше токенов генерируем, тем больше раз мы пересчитываем эти scores снова и снова. Отсюда видно: есть повторение вычислений, которое на самом деле не нужно.

Это особенно важно, потому что attention имеет квадратичную сложность по размеру контекста. Именно здесь нам помогает KV cache.

Итак, если понятно, зачем нужен KV cache, пойдём дальше и разберём, как он решает эту неэффективность.

Как работает KV cache

Чтобы понять, как работает KV cache, используем его во время inference и посмотрим, как он себя ведёт. KV cache используется только на инференсе, а не на обучении, и помогает снизить латентность генерации токенов.

Пример: возьмём пустой контекст и начнём генерировать текст. Первый шаг — посчитать embedding для токена. Назовём его embedding для токена 1.

Затем в блоке attention считаем соответствующие query, key и value для этого токена. По query и key считаем attention score, а с помощью этого score и value получаем новый контекстный embedding. В финальном слое получаем предсказание следующего токена.

Теперь сохраним key и value в кэш-памяти и попробуем использовать их для следующего токена. Если у вас несколько attention heads и несколько decoder-блоков, вы делаете это для каждого из них и кладёте посчитанные keys и values в кэш. Но здесь для простоты предположим, что у нас одна attention head и один decoder-блок.

Во время инференса мы дописываем токен в контекст и генерируем следующий. Снова начинаем с расчёта embedding для нового токена. Затем в блоке attention считаем query, key и value для последнего сгенерированного токена.

Для key и value у нас теперь есть значения, посчитанные для токена 2. А в кэше уже лежат значения для токена 1. Используя и кэшированные, и только что посчитанные значения, мы собираем key- и value-матрицы для всего контекста.

Дальше в attention считаем score через query-вектор токена 2 и key-матрицу. Так получается attention score для токена 2. Для финального embedding умножаем attention score на value-матрицу — получаем новый контекстный embedding для токена 2 и генерируем следующий токен. Key и value для токена 2 тоже кладём в KV cache.

Могут возникнуть вопросы: почему мы берём только последний токен и почему не кэшируем query, а считаем attention только для токена 2? Чтобы это понять, сравним инференс с KV cache и без.

Разница подходов с KV Chache и без него

В контексте сейчас два сгенерированных токена. Посмотрим дальнейшие расчёты.

Сначала считаем embeddings. Без KV cache считаем embeddings для обоих токенов. С KV cache используем embedding только последнего токена.

В блоке attention без KV cache, раз мы используем embeddings всех токенов, считаем обычные матрицы query, key и value — те, что вам знакомы. В случае KV cache считаем только query последнего токена, а для key и value берём данные из кэша плюс значения, посчитанные для последнего токена.

В attention без KV cache мы считаем attention для всех токенов. С KV cache — только для последнего.

В этой архитектуре у нас один attention-блок и один decoder-слой, поэтому пропустим остальные компоненты вроде layer normalization и feed-forward сети: единственное отличие двух случаев — размер считаемых матриц.

Перейдём к финальному линейному слою: он даёт выходные logits, а через softmax — вероятности для генерации следующего токена. Если вспомнить: когда KV cache не используем, мы не берём вероятности предыдущих токенов и учитываем только вероятности последнего. Следующий токен сэмплируем только по ним. В случае KV cache у нас и так один набор вероятностей — по нему и сэмплируем.

Это делается только на инференсе. На обучении мы используем все значения, чтобы считать ошибку и делать backpropagation. На инференсе же берём только logits последнего токена.

Почему кэшируют K и V, а не Q

Если это понятно, сделаем шаг назад и посмотрим, что мы сделали. Раз мы не используем logits предыдущих токенов, то и считать эти logits не нужно. Значит, в модели можно безопасно пропустить эти вычисления и считать только необходимое. В нашем случае нужен attention-расчёт только для последнего токена.

Если понимать, как считается attention, становится ясно, почему нам нужен только query последнего токена, но при этом всё равно нужны полные key- и value-матрицы всех токенов. В attention, чтобы получить контекстную информацию от всех токенов, мы умножаем query последнего токена на key-матрицу всех токенов. В результате embedding содержит всю контекстную информацию — это и есть суть attention.

Поэтому в KV cache мы берём только последний токен и используем закэшированные keys и values для предыдущих. Так генерируем новые токены. По мере генерации продолжаем сохранять и доставать пары key–value из KV cache.

Эти данные лежат в памяти GPU. Требование к памяти зависит от нескольких факторов: длины последовательности, размерности модели, числа слоёв и batch size. Чем больше модель, тем больше нужно GPU-памяти.

Плюсы, минусы и цена в памяти

Посмотрим на плюсы и минусы KV cache. На графике compute против числа сгенерированных токенов видно: инференс без KV cache требует значительно больше вычислений. Причина — attention масштабируется квадратично с длиной последовательности, когда keys и values пересчитываются на каждом шаге.

С KV cache модель избегает этого повторного вычисления. На каждом шаге декодирования attention считается только для последнего токена, поэтому per-token attention растёт линейно с длиной контекста — инференс становится гораздо эффективнее.

Минус KV cache — требования к памяти на хранение keys и values. Объём кэша зависит от числа decoder-блоков, batch size, числа attention heads, размерности на голову и длины контекста. Поскольку храним и keys, и values, памяти нужно вдвое больше. При точности 16 bit каждое значение занимает 2 байта.

По расчёту маленькая модель GPT-2 требует примерно 36 MB памяти — это терпимо, но по сегодняшним меркам модель очень маленькая.

Если взять чуть большую модель — примерно втрое относительно предыдущей архитектуры — требование к памяти растёт примерно до 500 MB (0.5 GB).

Для ещё более крупных моделей — около 60 слоёв и контекст порядка 100 000 токенов — требование к памяти может дойти почти до 400 GB, что крайне много только на кэш.

Итак: KV cache сильно ускоряет инференс, но требует много памяти — и это становится вызовом для современных моделей с миллиардами параметров. Тем не менее инференс всё равно становится намного быстрее. Если сравнить генерацию токенов с KV cache и без, без кэша инференс замедляется по мере роста контекста. Любой современный LLM продукт от Deepseek до Claude не обходится без кеширования.

Надеюсь, теперь понятно вам стало понятнее что такое KV cache, зачем мы его используем и какие сложности с ним связаны.

Спасибо за прочтение!

Источник: habr.com