Как реализуется KV-кэш в инференсе авторегрессивного трансформера?

В авторегрессивных трансформерах при инференсе для каждого нового токена важно эффективно использовать вычисления, чтобы не пересчитывать всю последовательность заново. Для этого применяется KV-кэш — кэш ключей (Key) и значений (Value) из слоёв внимания (self-attention).

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

Примерно это работает так:

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

Это значительно ускоряет инференс, снижая вычислительную сложность с квадратичной по длине последовательности до линейной.

Пример псевдокода для одного слоя:

# kv_cache хранит ключи и значения для предыдущих токенов
new_key, new_value = compute_kv(new_token)
kv_cache.keys = concatenate(kv_cache.keys, new_key)
kv_cache.values = concatenate(kv_cache.values, new_value)
output = attention(query=new_key, keys=kv_cache.keys, values=kv_cache.values)

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