Понимание и реализация KV-кэша в LLM с нуля: подробное руководство с кодом

KV-кэш (кэш ключей и значений) — одна из ключевых техник для эффективного инференса больших языковых моделей (LLM) в продакшене. При генерации текста модель на каждом шаге вычисляет ключи (K) и значения (V) для всех токенов, но для уже обработанных токенов эти вычисления повторяются, что ведёт к избыточности. KV-кэш сохраняет вычисленные K и V для повторного использования, что значительно ускоряет генерацию, особенно на длинных текстах. Однако он требует дополнительной памяти и усложняет код. В статье представлена реализация с нуля на PyTorch: добавляются буферы cache_k и cache_v в класс MultiHeadAttention, в forward-метод передаётся флаг use_cache, при включении которого новые ключи и значения конкатенируются с кэшированными. Также добавлен метод reset_cache для очистки между сессиями генерации. В GPTModel вводится счётчик current_pos для отслеживания позиции, чтобы новые запросы (queries) выравнивались с уже сохранёнными K и V. Приведён полный код функции generate_text_simple_cached, которая при use_cache=True передаёт модели только новый токен, а не весь контекст. Простое сравнение производительности показывает ускорение примерно в два раза на тестовом примере.