Парадокс прост - при том же железе и той же модели генерация токенов замедляется по мере роста последовательности. Причина не в вычислительной мощности, а в избыточных вычислениях механизма внимания: модель на каждом шаге повторно пересчитывает проекции прошлых токенов.
Эксперимент на gpt2-medium демонстрирует эффект и предлагает простое решение. При генерации 1000 новых токенов автор прогонял один и тот же промпт дважды - с KV caching и без него. Результат: с кешем время генерации около 21.7 секунд, без кеша - более 107 секунд, то есть почти пятикратное замедление. Такой разрыв объясняется тем, что без кеша вычисления внимания растут квадратично с длиной последовательности, а с кешем - почти линейно.
Что такое KV caching и почему это работает
KV caching - это оптимизация инференса для автогрессивной генерации. Идея в том, что ключи (K) и значения (V), вычисленные для ранее сгенерированных токенов, не меняются после своего первого вычисления. KV caching сохраняет эти K и V при первом проходе и на последующих шагах повторно использует их.
Вместо повторного расчета проекций K и V для всех прошлых токенов модель при каждом новом токене вычисляет только запросы (Q) и K, V для нового токена. Затем внимание считается по уже закешированным прошлым K и V плюс проекции нового токена. Это устраняет большую часть лишней работы в механизме внимания.
Эксперимент и его методика
В тесте использовалась стандартная схема: фиксированная модель, один и тот же промпт и одинаковая длина генерируемой последовательности. В коде - AutoModelForCausalLM с параметром use_cache включенным или выключенным. Автор прогонял по 5 запусков в каждом режиме и брал среднее и стандартное отклонение времени.
- Модель: gpt2-medium.
- Промпт: "Explain KV caching in transformers."
- Длина генерации: 1000 новых токенов.
- Результат: с KV caching ~21.7 s, без KV caching >107 s - почти 5x разница.
Практические выводы - как переделать инференс
Из эксперимента следует простая инженерная формула: если генерация замедляется с ростом последовательности, сначала проверьте, не пересчитывает ли ваша реализация K и V на каждом шаге.
- Храните вычисленные K и V после первого шага и переиспользуйте их на следующих шагах.
- На каждом шаге вычисляйте только Q и K,V для нового токена.
- Готовьтесь выделить дополнительную память под кеш - это явный компромисс, но на практике стоимость памяти обычно значительно ниже стоимости повторных вычислений.
Куда это ведет индустрию
Эксперимент подтверждает тезис: KV caching необходим для эффективного production-инференса автогрессивных LLM. Это не фича ради оптимизации - это базовая инженерная практика для масштабируемой генерации длинных текстов.
Следующий шаг индустрии - интеграция кеширования в стандартные инференс-стеки, где оптимизация памяти и управление кешем станут ключевыми задачами инженеров инференса. В ближайшие месяцы внедрение KV caching будет обязательным для проектов с длинным контекстом, а команды, игнорирующие эту простую оптимизацию, потеряют в латентности и стоимости.
Для тех, кто хочет воспроизвести результаты и посмотреть код, автор прикрепил практический ноутбук с тестами и примерами использования параметра use_cache в библиотеке transformers.