Как работает KV cache и от чего зависит его память
Что такое KV cache, почему его можно переиспользовать при decode и от каких факторов зависит его объем?
Короткий ответ
KV cache хранит key/value vectors прошлых токенов по слоям. При decode прошлые K/V не меняются, поэтому их не надо пересчитывать. Память зависит от batch, context length, layers, KV heads, head dim и dtype.
Полный разбор
В causal self-attention новый токен attends ко всем предыдущим токенам. Key и value для уже обработанных токенов зависят от их hidden states и весов слоя, но не меняются при генерации следующих токенов. Поэтому после prefill их можно сохранить и дальше доставать из cache.
На каждом decode step модель считает Q/K/V только для нового токена, добавляет новый K/V в cache и делает attention текущего Q к сохраненным K/V.
Оценка памяти примерно:
`batch * seq_len * num_layers * num_kv_heads * head_dim * 2(K,V) * bytes_per_value`.
Для MHA `num_kv_heads` обычно равно числу attention heads, для MQA/GQA меньше. Поэтому MQA/GQA сильно уменьшают KV cache.
Теория
KV cache ускоряет autoregressive decode ценой памяти. Для long-context serving именно память под cache часто ограничивает batch size.
Типичные ошибки
- Сказать, что cache хранит logits или все hidden states.
- Не учитывать layers и batch size.
- Не отличать attention heads от KV heads в MQA/GQA.
Как отвечать на собеседовании
- Объясни неизменность K/V прошлых токенов в causal decode.
- Дай формулу памяти и отдельно упомяни dtype.