KV-кешування: чому генерація швидка і куди йде пам'ять

Attention-head article розрізняє вивчені матриці проекцій від запитів, ключів і значень, обчислених для введення. Блокова стаття слідувала за їхніми сумішами до залишкових оновлень. Тут ми додаємо часовий вимір: що може повторно використовувати причинний декодер, коли він додає інший маркер?

Підказка обробляється в паралельному прямому проході. Generation потім додає маркери по одному. Попередні причинно-наслідкові представлення не залежать від щойно доданих токенів, тому їхні ключі та значення можна зберігати та повторно використовувати. У цій статті описано цей кеш і вартість його пам’яті.

Чому покоління повторюється

Згадайте один структурний факт, який робить декодер декодером: каузальна маска. У блоці декодера самоувага маскується так, що маркер t може звертатися лише до маркерів 1..t — себе та всього до нього, ніколи нічого після. (Це єдиний біт, який відокремлює декодер GPT від кодера BERT; усе інше всередині блоку ідентичне.) Ця маска дозволяє моделі генерувати зліва направо: вихід кожної позиції залежить лише від уже записаних позицій.

Отже, генерація — це цикл. Надішліть підказку, прочитайте розповсюдження в останній позиції, зразок маркера, додайте його до послідовності та подайте все назад, щоб отримати наступний. Повторюйте, поки не закінчите. Вирішальною деталлю є те, що модель повинна обчислювати щоразу. Оскільки Q, K і V створюються свіжо під час кожного прямого проходу, проста реалізація, яка просто повторно запускає модель на зростаючій послідовності, повторно обчислює ключі та значення кожного токена, на кожному рівні, на кожному окремому кроці — навіть якщо всі токени, крім найновішого, ідентичні попередньому кроку. Ключовий вектор маркера 5 є однаковим незалежно від того, чи обчислює його модель, коли послідовність має довжину 5 маркерів або 500 маркерів. Наївний цикл все одно обчислює це 495 разів.

Порахуйте роботу. Щоб випустити токен у позиції t, повний прохід вперед над токенами t створює матрицю оцінки уваги t × t на голову — робота O(t²). Зробіть це для кожного зі згенерованих n токенів, і загальна сума буде O(n³): кубічні витрати на створення лінійної кількості тексту. Термін «квадратичне значення на крок» — чисте марнотратство, оскільки єдиною справді новою інформацією на кроці t є єдиний останній маркер. Все інше безпричинно перераховується.

Кеш: обчисліть новий маркер, запам’ятайте решту

Виправлення випливає безпосередньо зі спостереження. Якщо ключі та значення минулих токенів не змінюються, збережіть їх. Підтримуйте робочий буфер — кеш KV — зберігаючи K і V для кожного маркера, який ви бачили на даний момент, на кожному рівні. Тоді кожен крок генерації виконує лише невідновну роботу:

1. Обчисліть Q, K, V для одного нового маркера (один рядок, а не вся послідовність).
2. Додайте цей маркер K і V до кешу.
3. Зверніть увагу на один запит нового маркера щодо всього кешованого K і V.
4. Прочитайте результат, зразок наступного токена, перейдіть до 1.

Вартість кроку зменшується з O(t²) до O(t) — один запит до ключів t замість повної матриці t × t — і загальна сума з O(n³) до O(n²). Це вся причина, чому інтерактивна генерація досить швидка, щоб передавати відповідь слово за словом замість того, щоб довше зупинятися на кожному маркері.

Зверніть увагу на асиметрію, яка дає назву кешу. Ми зберігаємо K і V, але викидаємо всі попередні запити. Це не помилка — запит минулого маркера справді більше ніколи не потрібен. Єдине завдання запиту полягає в тому, щоб вивести той маркер, який стався на кроці, що маркер був найновішим і ніколи не буде переглянуто. Ключі та значення, навпаки, зчитуються кожним запитом маркера future. Отже, те, що варто кешувати, — це саме K і V, а те, що можна безпечно викинути, — це Q. Недаремно це кеш KV, а не QKV.

Два режими: попереднє заповнення та декодування

Наведений вище цикл приховує поділ, який варто назвати, оскільки він пояснює, чому «час до першого маркера» та «час для наступного маркера» регулюються різними апаратними обмеженнями. Обробка запиту є одним паралельним проходом по всіх його маркерах одночасно — точно обчислення повної послідовності зі статті блоку — і вона заповнює кеш для кожної позиції запиту за один знімок. Ця фаза попереднього заповнення пов’язана з обчисленнями: великі матриці множення, які насичують арифметичні блоки GPU. Декодування — наступний цикл «один жетон за раз» — навпаки. Кожен крок виконує крихітні matmuls, але повинен витягувати весь кеш із пам’яті, щоб обслуговувати його, тому він обмежений пропускною здатністю пам’яті і залишає більшість обчислювальних процесів графічного процесора неактивними. (Така сама відмінність у статті Модель виконання графічного процесора між ядрами, пов’язаними з обчисленням і пропускною здатністю.) Висновок: шкала вартості попереднього заповнення з швидкою довжиною та обчисленням; декодуйте масштаби витрат із розміром кешу, який вам доведеться перечитувати кожен крок. Що вказує прямо на кеш як на те, що має значення.

Книга витрат: куди йде пам’ять

Кешування обмінює обчислення на пам’ять, а рахунок приходить у гігабайтах. Для стандартної уваги з кількома головками кожен рівень зберігає один ключ і одне значення на маркер для всіх головок. Із загальною шириною ключа та значення, яка дорівнює ширині моделі d, N шарів, і фіксованою точністю зберігання:

KV bytes per token ≈ 2 · N · d · (bytes per value)

2 є парою (K і V); N — кількість шарів, кожен з яких має власний кеш; d — це ширина моделі, підсумована за всіма головами. Для GPT-3 на fp16 це 2 × 96 × 12288 × 2 bytes ≈ 4.7 MB per token. Контекст 2048 маркерів містить близько 9,6 ГБ кешу; гіпотетично розширити ту саму конфігурацію до контексту 200K-токенів, а лише кеш-пам’ять становить ~940 ГБ — це набагато більше, ніж будь-який один GPU.

Тепер поставте це число поруч із вагою моделі, тому що вся суть у порівнянні. Ваги GPT-3 фіксовані: 175B × 2 bytes = 350 GB при fp16, однакові незалежно від того, підключений один користувач чи тисяча. Кеш — це те, що росте — лінійно залежно від довжини контексту та знову лінійно залежно від кількості групованих запитів, оскільки кожен одночасний запит містить власний кеш. На обслуговуючому вузлі ваги є одноразовим відрахуванням із пам’яті; все, що залишилося, розподіляється між кешами живих запитів. Таким чином, кількість користувачів, яких ви можете обслуговувати одночасно, і найдовший контекст, який ви можете запропонувати, встановлюються не розміром моделі, а скільки кеш-пам’ять KV поміщається в пам’яті, ваги ще не заявлені. У тривалому контексті кеш-пам’ять зменшує вагу — 940 ГБ проти 350 ГБ — і модель, на навчання якої ви витратили мільйони, стає маленьким об’єктом у бюджеті пам’яті. (Це сторона висновку зображення куди йде пам’ять викладається повністю.)

Ось чому цілий рівень системної інженерії існує навколо кешу, а не навколо моделі. PagedAttention керує нею як віртуальною пам’яттю операційної системи, розподіляючи її на фіксовані сторінки, щоб усунути фрагментацію. Кешування префіксів повторно використовує кеш для загальної системної підказки для багатьох запитів замість того, щоб щоразу перебудовувати його. Квантування кеш-пам’яті KV зберігає K і V в int8 або fp8, щоб зменшити площу вдвічі або вчетверть. Кожен є прямим нападом на байти 2 · N · d · context · batch вище; стаття про оптимізацію висновків пояснює, як вони це роблять.

Назва вузького місця

Зніміть гросбух до важеля. Розмір кешу 2 · N · d · context · batch · bytes. Контекст і партія — це вимоги користувачів — ви не можете зменшити їх, не пропонуючи гірший продукт. Байти на значення – це те, що квантування атакує. Це залишає термін 2 · N · d, вбудований в архітектуру, де d насправді є n_heads · d_head — і його частина, яку розробник моделі може змінити, не торкаючись того, що вивчає модель, — це кількість головок ключа/значення.

Це єдине спостереження є гачком для всього, що буде далі. Якщо більша частина кешу становить «один K і один V на головку уваги на шар», тоді спосіб зменшити її полягає в тому, щоб мати менше ключів і значень, ніж головок, які виконують запити. Є два способи зробити це:

• Поділіться ними — дозвольте кільком головкам запиту читати одну спільну головку K/V. Це увага групового запиту, предмет наступної статті, і зараз це типове значення майже в кожному новому відкритому декодері.
• Стисніть їх — проектуйте K і V у невеликий спільний прихований вектор і кешуйте його замість цього. Це багатоголова латентна увага, яку в блозі розглядають як місце призначення прихованого вузького місця.

Обидва починаються з номера на цій сторінці. KV-кешування — це те, що робить генерацію доступною; кеш, який він створює, робить генерацію дорогою; а решта сучасного дизайну уваги — це постійний аргумент щодо того, як зробити цей кеш меншим, не погіршивши модель.