Дата публикации
science

Как KV‑кэш делает генерацию текста в нейросетях в пять раз быстрее

Что открыли

Исследователи из сообщества Hugging Face разобрали на практическом примере, как работает KV caching (кэш ключей и значений) в трансформерах и что он даёт в реальной скорости.

Главный результат:

  • при генерации текста с включённым KV‑кэшем на GPU NVIDIA T4 модель HuggingFaceTB/SmolLM2-1.7B сгенерировала до 300 новых токенов за 11,7 секунды;
  • без KV‑кэша тот же сценарий занял 1 минуту 1 секунду;
  • ускорение составило примерно 5,21 раза.

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

Как исследовали

Авторы опираются на стандартную архитектуру трансформеров с автогенерацией текста (autoregressive modeling), такую же, как у GPT‑подобных моделей.

Они рассматривают работу механизма внимания, где для каждого токена модель вычисляет три матрицы:

  • Q (query) — запрос,
  • K (key) — ключ,
  • V (value) — значение.

Без KV‑кэша при генерации каждого нового токена модель заново вычисляет K и V для всей последовательности. KV caching предлагает другой подход: один раз посчитать K и V для уже сгенерированных токенов, сохранить их в кэше и дальше только дописывать новые K и V для появляющихся токенов.

Авторы показывают по шагам:

  • первый токен → в кэше [K1, V1];
  • второй токен → кэш расширяется до [K1, K2] и [V1, V2];
  • n‑й токен → кэш хранит [K1, K2, …, Kn] и [V1, V2, …, Vn].

Для демонстрации они используют модель HuggingFaceTB/SmolLM2-1.7B из библиотеки transformers и запускают генерацию текста на GPU NVIDIA T4. Сравнивают два режима: с включённым KV‑кэшем и без него.

Параллельно они показывают упрощённую реализацию кэша в PyTorch и базовый пример генерации с параметром use_cache=True.

Что это меняет на практике

KV caching меняет баланс между вычислениями и памятью при инференсе трансформеров.

Что происходит при стандартной генерации без кэша:

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

Что даёт KV‑кэш:

  • модель повторно использует уже посчитанные K и V;
  • на каждом новом шаге она считает только новые K и V и новую Q;
  • внимание вычисляется по кэшу + свежему токену;
  • скорость генерации почти не «проседает» при росте длины контекста.

Сравнение из работы:

  • Объём вычислений на слово
    — стандартный инференс: повторяет одни и те же вычисления для каждого токена;
    — KV‑кэш: переиспользует результаты, сокращая работу процессора или GPU.

  • Память
    — стандартный инференс: меньше памяти на шаг, но общие затраты растут с длиной текста;
    — KV‑кэш: требует дополнительную память под кэш, но делает вычисления эффективнее.

  • Скорость
    — стандартный инференс: становится медленнее по мере роста длины текста;
    — KV‑кэш: остаётся быстрым даже на длинных последовательностях.

  • Длинные тексты
    — стандартный инференс: заметные задержки;
    — KV‑кэш: лучше подходит для длинных ответов и чатов.

На практике это значит:

  • более быстрые чат‑боты и ассистенты;
  • комфортная работа с длинными промптами и документами;
  • меньше простаивающих GPU в продакшн‑серверах за счёт более эффективной загрузки.

Минус очевиден: KV caching съедает больше видеопамяти или оперативной памяти, потому что хранит все прошлые ключи и значения. Для очень длинных диалогов это может упереться в лимит памяти, особенно на компактных GPU.

Что это значит для вас

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

Разработчики и MLOps

Если вы:

  • разворачиваете LLM‑сервисы на своих серверах,
  • оптимизируете стоимость инференса,
  • строите продукты поверх моделей вроде SmolLM2,

то KV caching — обязательный элемент настройки:

  • включённый use_cache=True даёт реальное ускорение (в примере — ~5,21×);
  • вы осознанно обмениваете память на скорость;
  • при проектировании инфраструктуры нужно учитывать рост потребления памяти с длиной диалога.

Инженеры, пишущие промпты и прототипы

Даже если вы не лезете в низкоуровневый код, KV‑кэш влияет на ваш опыт:

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

Обычные пользователи

Если вы общаетесь с чат‑ботами, пишете тексты с помощью LLM‑ассистентов или генерируете код, KV‑кэш помогает вам косвенно:

  • ответы приходят быстрее;
  • сервисы выдерживают больше одновременных пользователей на тех же мощностях;
  • диалоги могут быть длиннее без драматического падения скорости.

При этом вы платите за это только тем, что провайдеру нужны чуть более мощные сервера. Для вас это может отразиться на цене подписки, но напрямую вы этого не увидите.

Как запустить

Авторы показывают два уровня работы с KV caching: ручная реализация кэша в PyTorch и использование встроенных возможностей библиотеки transformers.

Упрощённый KV‑кэш на PyTorch

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

class KVCache :
    def __init__ ( self ):
        self.cache = { "key" : None , "value" : None }

    def update ( self, key, value ):
        if self.cache[ "key" ] is None :
            self.cache[ "key" ] = key
            self.cache[ "value" ] = value
        else :
            self.cache[ "key" ] = torch.cat([self.cache[ "key" ], key], dim= 1 )
            self.cache[ "value" ] = torch.cat([self.cache[ "value" ], value], dim= 1 )

    def get_cache ( self ):
        return self.cache

Логика простая:

  • при первом вызове кэш пуст, туда просто записываются первые K и V;
  • при следующих шагах к существующим тензорам ключей и значений «пришиваются» новые по размерности последовательности (dim=1).

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

Генерация текста с KV‑кэшем в transformers

В transformers KV‑кэш включается параметром use_cache=True. По умолчанию он уже активирован, но авторы явно показывают это в коде.

Полный пример из статьи:

from transformers import AutoModelForCausalLM, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained( 'HuggingFaceTB/SmolLM2-1.7B' )
model = AutoModelForCausalLM.from_pretrained( 'HuggingFaceTB/SmolLM2-1.7B' ).cuda()


tokens = tokenizer.encode( "The red cat was" , return_tensors= "pt" ).cuda()


output = model.generate( tokens, max_new_tokens= 300 , use_cache = True )
output_text = tokenizer.batch_decode(output, skip_special_tokens= True )[ 0 ]

Ключевые моменты:

  • используется модель HuggingFaceTB/SmolLM2-1.7B;
  • генерация запускается на GPU (.cuda());
  • генерируется до 300 новых токенов;
  • use_cache=True включает KV‑кэш для всех слоёв внимания.

Авторы отдельно указывают, что через параметр cache_implementation в transformers можно выбирать разные реализации кэша, но в примере они не меняют дефолтное поведение.

Результаты бенчмарка

На GPU NVIDIA T4 получились такие цифры:

  • с KV‑кэшем: 11,7 секунды;
  • без KV‑кэша (стандартный инференс): 1 минута 1 секунда;
  • ускорение: примерно 5,21×.

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

Итог

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

На примере SmolLM2‑1.7B и GPU T4 видно, что грамотная работа с кэшем даёт ускорение более чем в пять раз. Для разработчиков LLM‑продуктов это прямой путь к более отзывчивым чат‑ботам и более дешёвому инференсу. Для пользователей — к ассистентам, которые не «задумываются» по минуте над каждым длинным ответом.


Читайте также