«До 10% пропускной способности и 9% к TTFT»: AWS раскрыла оптимизации ядер MoE на SageMaker AI

Зачем оптимизировать ядра MoE

Публикация описывает MoE (mixture-of-experts) как разреженно-активируемую архитектуру, которая наращивает ёмкость больших языковых моделей, сохраняя низкие вычисления на токен. Поскольку операции MoE занимают большую долю общего времени вычислений, их неэффективность напрямую повышает задержку запросов и стоимость, и именно поэтому оптимизация ядер даёт наибольшую возможность для прироста производительности.

Fused MoE назван ключевой техникой оптимизации: он объединяет маршрутизацию токенов, перетасовку данных и вычисления экспертных сетей в одну пакетную операцию на GPU вместо последовательных циклов. Это даёт выигрыш по задержке, но усложняет реализацию и отладку ядер, а высокая загрузка GPU требует оптимизаций под конкретное железо. Поэтому в посте разобраны реализованные оптимизации ядер и показано, как из них извлечь пользу.

Swap-AB: развернуть матрицы под WGMMA

Первая оптимизация — Swap-AB — связана с инструкцией WGMMA (Warp Group Matrix Multiply Accumulate), на которую опирается архитектура NVIDIA Hopper для максимальной пропускной способности тензорных вычислений; инструкция выполняется коллективно по группе из 128 потоков. Асимметрия операндов заложена на уровне ISA: аппаратно фиксируется размер строки M = 64, и операнд A строго отображается на это фиксированное измерение, тогда как операнд B ложится на переменное измерение N.

На фазе декодирования MoE система обрабатывает по одному токену на последовательность, и размер батча падает до M = 1, из-за чего тензорные ядра простаивают. Решение — воспользоваться алгебраическим свойством транспонирования и переставить умножение: BT[N,K] x AT[K,M] -> CT[N,M]. Так большой экспертный размер N (например, 8192) попадает в первый операнд, легко заполняет обязательный порог в 64 строки, а заодно активируется асинхронная предвыборка через аппаратные копировальные движки Hopper.

Triton вместо CUTLASS для FP8-квантования

Вторая оптимизация — выбор бэкенда. Оракул MoE в vLLM маршрутизирует схемы квантования на Triton, FlashInfer+CUTLASS или другие бэкенды, и профилирование показало, что для тензорно-квантованных FP8-моделей на Hopper бэкенд Triton стабильно обходит путь FLASHINFER_CUTLASS. Причина в том, что Triton заменяет жёсткие матричные конфигурации уровня фреймворка на регистровое слияние уровня оператора и собирает весь конвейер в один проход по железу, вытягивая данные из глобальной памяти один раз.

Смена маршрута — чисто конфигурационное изменение без правок кода ядер, но оно дало 10-процентный прирост сквозной пропускной способности для FP8-квантованных fused MoE-моделей на H200, и этот выигрыш доступен в контейнере AWS Large Model Inference (LMI) v28.

Убрать лишние JIT-перекомпиляции Triton

Третья оптимизация борется с незаметной патологией: Triton как компилятор специализирует целочисленные аргументы, и каждое новое значение специализируемого параметра запускает свежую JIT-компиляцию. В обслуживающих нагрузках длины последовательностей и размеры батчей меняются непрерывно, поэтому компиляция стабильно надувает задержку, особенно время до первого токена.

В MoE-обслуживании ядро _fused_inv_rope_fp8_quant_per_head, отвечающее за объединённые rotary position embedding и FP8-квантование, перекомпилировалось на каждом новом значении num_tokens. Аннотации do_not_specialize для параметров переменной длины устранили эти перекомпиляции полностью и улучшили TTFT на 9% на H200 при 8-кратном тензорном параллелизме (PR #42849). В много-GPU кластере это убирает и простои: если одна из восьми GPU останавливается на перекомпиляцию, остальные семь ждут её, поэтому устранение перекомпиляций сглаживает координацию кластера. Оптимизацию выложили в upstream vLLM, и она доступна начиная с версии 0.22.0.

Что это даёт на практике

В заключение публикация сообщает, что в сотрудничестве с сообществом vLLM инженеры реализовали и открыли исходный код нескольких оптимизаций ядер для fused MoE-моделей и применяют их в образе LMI Deep Learning Container на Amazon SageMaker AI. Вместе эти оптимизации дали до 10% прироста пропускной способности и до 9% ускорения TTFT.

Материал подписан Моникой Моной (Mona Mona), специалистом-архитектором решений AWS по инфраструктуре машинного обучения и сервисам AI/ML, Садаф Фардин (Sadaf Fardeen), руководителем направления оптимизации вывода для SageMaker, HyperPod и Amazon Bedrock, и Синь Ян (Xin Yang), старшим инженером по разработке ПО, специализирующимся на крупномасштабном выводе больших языковых моделей и оптимизации производительности на GPU. Публикация относится к категории Announcements и разделу Amazon SageMaker AI с уровнем Advanced (300).