ExFold ускоряет инференс MoE-моделей до 2,45 раза без дообучения

В новой научной работе, опубликованной на arXiv, представлен фреймворк ExFold для ускорения инференса моделей на основе смеси экспертов (MoE). Разработка не требует дополнительного обучения и реализована в виде подключаемого плагина к популярной системе vLLM.

MoE-модели масштабируют число параметров, сохраняя вычислительную нагрузку на токен за счет разреженной активации экспертов. Однако обслуживание таких моделей с низкой задержкой остается сложной задачей, поскольку инференс включает две фазы с разными узкими местами: prefill, где доминируют экспертные вычисления на токен, и decode, где ключевое значение имеет трафик памяти из-за активируемого набора экспертов для всего батча.

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

ExFold объединяет обе фазы в единую задачу приближенного вычисления с бюджетом. На этапе инференса выполняется только допустимый набор экспертов, а вклад экспертов, не попавших в бюджет, проецируется на удержанных с помощью калиброванных скалярных проекторов. Авторы заметили, что выходы многих экспертов совпадают по направлению, но отличаются по величине, что позволяет использовать такую проекцию.

Матрица парных скалярных проекторов калибруется на неразмеченных данных и применяется на этапе инференса. В этой модели ускорение prefill соответствует сворачиванию по верхним K токенам, а ускорение decode — сворачиванию пула экспертов на уровне батча. Фазы различаются только способом выбора удержанных экспертов, а механизм восстановления исключенных вкладов у них общий.

Авторы реализовали ExFold как плагин для vLLM с легковесным CUDA-ядром для сворачивания экспертов. В экспериментах фреймворк обеспечил ускорение до 1,41 раза для времени до первого токена (TTFT) и до 2,45 раза для времени генерации одного токена (TPOT) при сохранении примерно 99% исходного среднего качества.

Предложенный подход позволяет снижать задержки при обслуживании MoE-моделей без затрат на дообучение и может быть полезен для production-систем, где важна высокая скорость отклика. Детали метода доступны в архиве препринтов arXiv.