Оптимизация генерации токенов с помощью KV кэширования

Проблема медленной генерации токенов

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

Что такое KV кэширование и как оно ускоряет вывод?

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

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

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

Чтобы проверить KV кэширование, мы тестируем один и тот же запрос с включенным и выключенным кэшированием, измеряя среднее время генерации при неизменных параметрах модели. Результаты показывают, что повторное использование закэшированных ключей и значений значительно снижает вычислительную избыточность, что ускоряет вывод.

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.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
 
prompt = "Explain KV caching in transformers."
 
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,
            use_cache=use_cache,
            max_new_tokens=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 кэширования может сократить время генерации 1000 токенов с более чем 107 секунд без кэширования до около 21,7 секунд с кэшированием — почти в 5 раз быстрее. Это происходит потому, что модель избегает повторного вычисления внимания для всех токенов на каждом шаге, предотвращая экспоненциальный рост вычислений.

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