Как се реализира KV кешът при инференса на авторегресивен трансформър?

В авторегресивните трансформатори по време на inference е важно да се използват изчисленията ефективно, за да не се пресмята цялата последователност отново. За това се използва KV кеш, който съхранява ключовете (Key) и стойностите (Value) от слоевете на вниманието (self-attention).

Идеята е, че при генериране на следващия токен вече сме изчислили ключовете и стойностите за всички предишни токени. Вместо да ги пресмятаме отново, ги запазваме в паметта (кеш) и при изчисляване на вниманието за новия токен използваме този кеш.

Това работи приблизително така:

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

Това значително ускорява inference, като намалява сложността от квадратична по дължина на последователността до линейна.

Пример за псевдокод за един слой:

# 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 кешът е механизъм за съхраняване на междинните ключове и стойности за ускоряване на последователната генерация в трансформаторите.