Внимание в LLM · часть 4

KV cache

При генерации новый токен смотрит на весь префикс. Ключи и значения старых токенов уже известны, но без кэша модель вычисляет их снова на каждом шаге. KV cache (кэш ключей и значений) сохраняет это прошлое и дописывает только новую пару kt,vt.

Сравнить два режима →

Прошлое вычисляется снова

Пусть префикс уже содержит несколько токенов. Чтобы добавить следующий, нужен новый запрос qt и все ключи и значения прошлого. Прямая реализация заново применяет WK и WV ко всему префиксу, хотя старые результаты не изменились.

Зафиксируйте длину префикса и сравните вычисление с кэшем и без него. Затем увеличьте префикс: какая работа повторяется, а какая сохраняется? Отдельно проверьте самый короткий префикс.

Повторный расчёт или KV cache

шаг 3
Режим вычисления
Длина префикса
3
Чтобы обработать «молоко», модель заново вычисляет K и V для всех 3 токенов префикса.
позиция 1K1 · V1вычислено повторно
позиция 2K2 · V2вычислено повторно
позиция 3K3 · V3вычислено сейчас
вычислено сейчас6
переиспользовано0
Всего проекций с начала: 12молоко → attention по 3 позициям → потому
Повторный расчёт против сохранённого прошлогоКэш сохраняет ранее вычисленные ключи и значения. На очередном шаге добавляется новая пара, но чтение доступного прошлого при вычислении внимания сохраняется.

Переключайте режимы на одном и том же шаге генерации. С KV cache старые блоки читаются, но не пересчитываются; вычисляется только пара для нового токена.

Дописывать вместо пересчитывать

После каждого шага новые ключ и значение добавляются к уже сохранённым:

Kcache←[Kcache;kt],Vcache←[Vcache;vt]

На следующем шаге attention читает весь cache, но вычисляет только новые kt,vt.

Сколько копий прошлого хранить

Cache растёт с длиной контекста, числом слоёв и числом KV-heads. При длинной генерации узкое место смещается: вместо повторного расчёта модель всё чаще ждёт, пока сохранённые K,V будут прочитаны из памяти.

Компромисс между вычислениями и постоянной памятью
  1. Без cacheмало постоянной памяти; много повторных проекций.
  2. KV cacheкаждый токен занимает память; старые проекции не повторяются.

В обычном MHA у каждой query-head собственная KV-head. MQA (multi-query attention) оставляет одну общую KV-head, а GQA (grouped-query attention) делит запросы на несколько групп. Число query-heads при этом не меняется.

Начните с MHA и сосчитайте группы ключей и значений. Перейдите к GQA, затем к MQA: что должно произойти с памятью, если число голов запросов сохраняется? Проверьте это по схеме связей.

Сколько запросов делят K и V

MHA
Схема attention
MHAсвои ключи и значения для каждого запроса
query 1Q1KV-группа 1
query 2Q2KV-группа 2
query 3Q3KV-группа 3
query 4Q4KV-группа 4
query 5Q5KV-группа 5
query 6Q6KV-группа 6
query 7Q7KV-группа 7
query 8Q8KV-группа 8
размер KV cache100%

16 блоков K/V на токен вместо 16 в MHA.

Как query-heads делят K и VЧисло голов запросов сохраняется, а несколько таких голов получают общие ключи и значения. Это уменьшает кэш, но меняет устройство слоя.