LLM пожирают много памяти, а PagedAttention решает эту проблему

LLM, такие как GPT и PaLM, меняют методы работы и взаимодействия, обеспечивая работу всего: от помощников-программистов до универсальных чат-ботов. Однако, запуск этих невероятно мощных моделей обходится очень дорого, зачастую в 10 раз дороже, чем традиционный поиск по ключевым словам.

Скрытый пожиратель памяти: кэш KV
LLM имеет в основе модель Transformer, которая генерирует текст по одному слову за раз, пишет xrust. Для эффективной работы необходимо запоминать «контекст» предыдущих токенов. Эта память хранится в так называемом кэше «ключ-значение» (KV). Можно представить его как кратковременную память LLM для разговора.

Проблема в том, что этот кэш KV огромен, и его размер динамически увеличивается и уменьшается для каждого запроса. Существующие системы сталкиваются с этой проблемой, поскольку обычно хранят кэш KV в одном непрерывном блоке памяти. Такой подход приводит к двум серьёзным проблемам. Первая - Фрагментация памяти:

• Внутренняя фрагментация. Системы заранее выделяют большой объём памяти для каждого запроса, исходя из максимально возможной длины выходных данных (например, 2048 токенов). Однако, если запрос генерирует только короткие выходные данные, большая часть зарезервированной памяти остаётся неиспользованной, что приводит к значительным потерям;
• Внешняя фрагментация. Поскольку разные запросы резервируют фрагменты разного размера, память графического процессора оказывается разбросанной, заполняясь небольшими неиспользуемыми промежутками, что затрудняет обработку новых запросов даже при наличии всей свободной памяти. В существующих системах только 20,4–38,2% кэш-памяти KV фактически используется для хранения состояний токенов, а остальная часть не используется.

Вторая проблема – Нет разделения памяти.
Продвинутые методы декодирования, такие как параллельная выборка или лучевой поиск, часто генерируют несколько выходных данных из одной последовательности, что позволяет им совместно использовать части кэша KV. Однако существующие системы не могут легко совместно использовать эту память, поскольку кэш KV каждой последовательности находится в отдельном непрерывном блоке.

Эта неэффективность серьезно ограничивает количество запросов, которые могут быть обработаны одновременно («размер пакета»), что напрямую влияет на пропускную способность системы (сколько токенов/запросов она может обработать в секунду).

На сцену выходит PagedAttention. Исследователи разработали программу, которая решает эти проблемы. В основе – виртуальная память и подкачка страниц.

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

Гибкое разделение памяти
PagedAttention позволяет совместно использовать блоки KV между различными последовательностями, даже между разными запросами. Например, при параллельной выборке или поиске по лучу несколько выходов могут совместно использовать кэш KV исходного запроса, что значительно экономит память. Кроме того, используется механизм копирования при записи для блоков, которые необходимо модифицировать разными последовательностями, что обеспечивает эффективное совместное использование без ненужного дублирования.