Серия интервью с AI #4: объясните, что такое KV-кэширование

Серия интервью с AI #4: объясните, что такое KV-кэширование

Вопрос:

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

Если вычисления не являются основным узким местом, то какая неэффективность вызывает такое замедление, и как можно перепроектировать процесс логического вывода, чтобы генерация токенов стала значительно быстрее?

Что такое KV-кэширование и как оно ускоряет генерацию токенов?

KV-кэширование — это метод оптимизации, используемый при генерации текста в больших языковых моделях для избежания повторных вычислений. В авторегрессивной генерации модель создаёт текст по одному токену за раз, и на каждом шаге обычно заново вычисляет внимание ко всем предыдущим токенам. Однако ключи (K) и значения (V), вычисленные для более ранних токенов, никогда не меняются.

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

Такое повторное использование прошлых вычислений значительно сокращает избыточную работу, делая вывод быстрее и эффективнее — особенно для длинных последовательностей — за счёт дополнительной памяти для хранения кэша.

Оценка влияния KV-кэширования на скорость вывода

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

```python

import numpy as np

import time

import torch

from transformers import AutoModelForCausalLM, AutoTokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"

model_name = "gpt2-medium"

tokenizer = AutoTokenizer.frompretrained(modelname)

model = AutoModelForCausalLM.frompretrained(modelname).to(device)

prompt = "Объясните, что такое KV-кэширование в трансформерах."

inputs = tokenizer(prompt, return_tensors="pt").to(device)

for use_cache in (True, False):

times = []

for _ in range(5):

start = time.time()

model.generate(

inputs,

usecache=usecache,

maxnewtokens=1000

)

times.append(time.time() - start)

print(

f"{'с' if use_cache else 'без'} KV-кэширования: "

f"{round(np.mean(times), 3)} ± {round(np.std(times), 3)} секунд"

)

```

Результаты наглядно демонстрируют влияние KV-кэширования на скорость в...

Читать далее