Идея от 1962 г. направи LLM моделите 2-4 пъти по-бързи за "serving"

Когато LLM генерира текст, той пази KV cache: ключовете и стойностите от attention за всеки токен досега, във всеки слой. Кешът расте с всеки нов токен и не знаеш предварително колко дълъг ще е отговорът.

До 2023 г. повечето системи за сървиране заделяха едно непрекъснато парче GPU памет за всяка заявка, оразмерено за максималната дължина. Заявка, която спира след 200 токена, пак държи памет за 2048. Авторите на vLLM измерват, че само 20-38% от паметта за KV cache съдържа реални данни. Останалото е заделено, но неизползвано място и фрагментация между парчетата.

Същият проблем, 60 години по-рано

Операционните системи са имали същия проблем с RAM в началото на 60-те. Компютърът Atlas в Манчестър го решава през 1962 г. с paging. Паметта се дели на страници с фиксиран размер, всяка програма има таблица, която свързва логическите ѝ страници със свободни физически рамки, и страница се заделя само когато потрябва.

PagedAttention

PagedAttention прави същото с KV cache:

• Кешът се дели на блокове с фиксиран брой токени.
• Всяка последователност има таблица, която свързва логическите ѝ блокове с физически блокове в GPU паметта.
• Нов блок се заделя едва когато предишният се напълни.

Загубата пада до незапълнената част от последния блок на всяка последователност, под 4% според статията. По-малко загуба означава по-големи batch-ове на същия GPU, и оттам идва 2-4 пъти по-високият throughput спрямо FasterTransformer и Orca.

Таблицата дава и още един трик от ОС. Две последователности могат да сочат към един и същ физически блок, така че паралелни отговори на един prompt споделят кеша му, с copy-on-write, когато се разделят. Същата идея като fork().

Новата инженерна работа е в attention kernel-а, който трябва да чете ключове и стойности от несъседни блокове, а не от един плосък буфер.

Статията: Efficient Memory Management for Large Language Model Serving with PagedAttention

Защо пиша това

Последна година съм в ТУЕС и започвам Systems Society Bulgaria: безплатни месечни срещи в София. Всяка среща взима една такава идея от системното програмиране и я проследява до съвременните изследвания в AI и квантовите изчисления, обяснено от нулата. Каним инженери и изследователи да изнасят лекции.

• Записване: forms.gle/ru6TKrAv4tCBDLn29

Кои други стари идеи от системите според вас се крият в днешната ML инфраструктура?