Как и зачем ускоряют LLM: ключ к скорости — KV cache

Современные большие языковые модели основаны на архитектуре трансформера, и важнейшим её элементом является механизм внимания, или attention. Именно благодаря нему модель определяет, какие слова в контексте связаны друг с другом, и выстраивает смысловые зависимости. Однако вычислительная стоимость attention огромна: каждая операция сводится к многократному перемножению матриц, а поскольку параметров в моделях миллиарды, любой шаг генерации требует колоссальных ресурсов. Тем не менее чат-боты и другие сервисы отвечают практически мгновенно. Секрет такой скорости кроется в наборе оптимизаций, среди которых особое место занимает механизм, известный как KV cache.

В процессе генерации модель для каждого токена вычисляет два специальных набора данных — ключи и значения, обозначаемые K и V. Именно они участвуют в матричных операциях, определяющих, насколько важны друг для друга слова в последовательности. Для уже обработанных токенов модель раньше уже вычислила эти данные, и потому было бы бессмысленно повторять ту же работу при создании каждого следующего слова. Поэтому трансформер сохраняет ключи и значения всех ранее встреченных токенов в отдельном буфере в памяти — этот буфер и называют KV cache. В результате на каждом следующем шаге генерации модели нужно пересчитать только ключи и значения нового токена, а для остальных берутся готовые значения из кэша. Такой подход кратно уменьшает объём вычислений и позволяет получать ответы без пауз и задержек.

Однако у быстродействия есть обратная сторона: чем длиннее контекст, тем больше памяти занимает этот кэш. При работе с большими объёмами текста KV cache становится узким местом, и исследователи ищут способы его сжимать и оптимизировать. Разрабатываются алгоритмы, которые позволяют удалять малоиспользуемые ключи или удерживать объём кэша в разумных пределах. Именно поэтому инженеры ведут постоянную борьбу между качеством генерации и эффективностью вычислений. Понимание того, как устроен KV cache, помогает увидеть, почему современные LLM работают так быстро, и какие решения позволят сделать их ещё быстрее в будущем.