OasisKV ускорил инференс LLM в 1,69, 2,1 раза, вынеся KV-кэш из HBM
Обслуживание запросов к большим языковым моделям всё чаще упирается не в вычисления, а в память. При длинном контексте и рассуждающих (reasoning) моделях KV-кэш, хранилище пар «ключ-значение» для уже обработанных токенов, занимает основную часть памяти и трафика на этапе генерации токенов (decode). HBM (память с высокой пропускной способностью, установленная прямо на GPU), дефицитный и дорогой ресурс, и именно он ограничивает размер батча запросов и итоговую пропускную способность системы.
Исследователи предложили OasisKV, систему инференса, ориентированную на память, которая разрывает связь между полным KV-кэшем и HBM во время decode. Идея опирается на наблюдение: внимание модели на этапе генерации токенов по своей природе разреженное, то есть реально нужна лишь малая часть накопленного кэша. Поэтому OasisKV держит в HBM только записи KV для самых релевантных токенов. Какие токены понадобятся дальше, система предсказывает заранее, по «заглядывающим вперёд» токенам, которые генерирует спекулятивное декодирование (speculative decoding). Фоновый конвейер вычисления внимания определяет нужные блоки KV, они подгружаются из более ёмких уровней памяти (памяти хоста или удалённой памяти) и заранее оказываются в HBM, до того, как понадобятся на следующем шаге генерации.
OasisKV реализован поверх vLLM. Точность предсказания оказалась достаточно высокой: при бюджете KV в 2048 токенов разрыв с полным вниманием (без урезания кэша) не превышает 0,7 пункта точности. За счёт этого разреженность превращается в прирост скорости: на рабочей нагрузке с рассуждениями пропускная способность выросла в 1,69 раза относительно плотного (dense) vLLM при потере точности всего 0,1 пункта, а при обслуживании длинного контекста на нескольких GPU прирост достигает 2,1 раза. При разнесении стадий обработки запроса (prefill) и генерации (decode) по разным узлам OasisKV показывает примерно двукратную пропускную способность относительно dense-варианта, принимая на обработку каждого запроса в 6,5, 9,7 раза меньше данных KV и потребляя в 2,2, 2,6 раза меньше памяти на узле decode, чем при полной передаче KV-кэша.
Ключевые факты
- OasisKV выносит основную часть KV-кэша из HBM, оставляя там только блоки самых важных токенов
- Нужные токены предсказываются заранее по данным спекулятивного декодирования; фоновый конвейер внимания находит нужные блоки KV и подгружает их из памяти хоста/удалённой памяти в HBM до следующего шага генерации
- При бюджете KV в 2048 токенов точность отличается от полного внимания не более чем на 0,7 пункта
- Пропускная способность растёт в 1,69 раза на задачах с рассуждениями (при потере точности 0,1 пункта) и до 2,1 раза при обслуживании длинного контекста на нескольких GPU
- При разнесении prefill и decode по разным узлам система выходит на почти двукратную пропускную способность, потребляя в 6,5, 9,7 раза меньше KV на запрос и в 2,2, 2,6 раза меньше памяти узла decode
Почему это важно
Инференс LLM всё чаще ограничен не вычислениями, а памятью: KV-кэш растёт вместе с длиной контекста и занимает основную часть дорогой HBM-памяти на GPU, из-за чего падает размер батча и итоговая пропускная способность сервера. OasisKV снимает это ограничение, не трогая архитектуру самой модели: она использует то, что внимание на этапе генерации токенов реально разреженное, и держит в быстрой памяти лишь нужный минимум KV-кэша, подгружая остальное заранее и почти без потери точности.
Кому это важно
Прежде всего инженерным командам, которые эксплуатируют инференс-серверы LLM, провайдерам API и компаниям с собственной инфраструктурой обслуживания моделей, особенно там, где востребован длинный контекст или рассуждающие модели с большим KV-кэшем. Также релевантно разработчикам vLLM и смежных систем инференса, поскольку OasisKV реализован как надстройка именно над vLLM.
Как это применить
OasisKV встраивается в существующий стек инференса поверх vLLM, а не требует переобучения или замены модели: механизм работает на уровне управления памятью и предсказания нужных блоков KV через спекулятивное декодирование. В источнике не указаны ни лицензия, ни доступность кода, ни условия использования, судить об этом преждевременно.
Можно ли доверять
Материал, исследовательская работа, опубликованная на HuggingFace Papers, с пока скромным откликом (8 отметок, 1 комментарий на момент публикации). В тексте не названы ни организация или институт авторов, ни конкретная модель/её размер, использованные в тестах, ни название бенчмарка с «рассуждающей» нагрузкой, ни дата/место публикации, это ограничивает возможность независимо проверить и воспроизвести заявленные цифры.
Риски и подводные камни
Заявленный прирост скорости идёт рука об руку с небольшой, но ненулевой потерей точности (0,1, 0,7 пункта в зависимости от сценария), и насколько стабильно предсказание «важных» токенов работает на разных типах нагрузок, из текста не следует. Схема требует управления несколькими уровнями памяти и фонового конвейера предсказания/подгрузки, что добавляет инженерную сложность по сравнению с плотным (dense) инференсом. Ключевые детали экспериментальной установки (модель, оборудование, бенчмарк) в источнике не раскрыты.