KAIST убрал калибровку из MoE-декодинга и ускорил блоки до 2,69× к vLLM

Декодирование Mixture-of-Experts на малых батчах упирается не в вычисления, а в память: на каждый токен приходится подгружать веса маршрутизируемых экспертов. Исследователи из KAIST AI Геунхён Ким и Сун Джу Хван показали, что популярный приём ускорения, отбор важных каналов внутри эксперта, можно сделать без калибровочных данных. Их ридер, построенный только на весах модели, сохраняет 98,1% точности полной модели при чтении 17,8% весов экспертов. Блоки MoE на Triton-ядрах работают в 1,60–1,90 раза быстрее оптимизированного vLLM на batch one.

Порог обещает бюджет, но не держит его

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

Авторы показывают, что у такой схемы две слабости. Во-первых, порог не фиксирует объём чтения: число каналов, прошедших отсечку, меняется от входа к входу. Prox с порогами, откалиброванными под бюджет 18%, на самом деле читал в среднем 12,9% на их бенчмарках. Инженер, который планирует задержку и пропускную способность, получает не ту цифру, на которую рассчитывал.

Во-вторых, выигрыш в точности при заданном бюджете тоже не гарантирован.

Top-k вместо порога

Альтернатива не требует данных вообще: для каждого токена оставляют k каналов с самым высоким скором. Бюджет чтения известен заранее и совпадает с заявленным. При сравнении на одних и тех же скорах top-k-версия селектора совпала с пороговой или превзошла её в 40 из 43 случаев при бюджетах 13% и 18%. На задачах генерации выиграла во всех 23 сравнениях.

Есть и качественный эффект. Откалиброванные пороги приводили Gemma 4 к коллапсу в петли повторов, а top-k-версии этого избегают. Вывод авторов: калибровка порога не даёт ни одной из двух вещей, ради которых её делают, то есть ни известного бюджета чтения, ни стабильного прироста точности при этом бюджете.

Ридер из одних весов

Вторая часть работы — собственный ридер, которому не нужны ни калибровочные данные, ни дообучение. Он состоит из двух шагов:

• каналы грубо отсеиваются по тернарному эскизу (ternary sketch) весов;
• пропущенные каналы компенсируются низкоранговым отображением, которое тоже строится из весов модели.

Ключевые цифры по трём пост-тренированным MoE-моделям:

• точность: 98,1% от полной модели;
• объём чтения весов маршрутизируемых экспертов: 17,8%;
• скорость MoE-блоков на Triton-ядрах на batch one: в 1,60–1,90 раза выше оптимизированного vLLM;
• на больших батчах ускорение достигает 2,69 раза на тех же трёх моделях;
• Prox при калибровке под 18% читал фактически 12,9%;
• top-k против порога: 40 из 43 сравнений при бюджетах 13% и 18%, 23 из 23 на генерации.

Кому это пригодится

Главный выигрыш получают те, кто крутит большие MoE-модели на железе с узкой памятью. Один из авторов в обсуждении на Hugging Face прямо назвал проблему активационного прунинга острой для GPU с медленной памятью вроде DGX Spark. Там пропускная способность памяти, а не FLOPS, определяет скорость генерации токенов. Следующий режим разреженности, по его словам, придёт именно из активационной разреженности при обслуживании на устройствах с жёстким лимитом полосы.

Практический плюс отсутствия калибровки в том, что метод применим к любой свежей open-weights-модели сразу после выхода. Не нужно искать репрезентативную выборку, и нет риска, что порог, подобранный на одном домене, развалит генерацию на другом, как это случилось с Gemma 4. Фиксированный бюджет чтения упрощает расчёты: заложенные 17,8% превращаются в предсказуемую экономию трафика памяти и задержки.

Код опубликован в репозитории calibration-free-moe на GitHub. Для команд, которые сами пишут ядра на Triton или оптимизируют локальный инференс, это готовая отправная точка. Стоит проверить её на своих моделях, особенно в интерактивных сценариях с batch one, где память упирается сильнее всего.