Инженерия ядер Mixture-of-Experts на Amazon SageMaker AI
aws ·
В данной статье подробно описываются оптимизации инженерии ядер для моделей Mixture-of-Experts (MoE) в среде Amazon SageMaker AI, с акцентом на повышение производительности при инференсе больших языковых моделей. Авторы описывают три ключевые техники: транспонирование Swap-AB для согласования размерностей матриц с аппаратными требованиями NVIDIA Hopper WGMMA, корректировки бэкенд-маршрутизации для использования Triton вместо CUTLASS при квантовании FP8 и устранение ненужных перекомпиляций JIT в Triton для снижения задержек. Эти оптимизации решают специфические узкие места в архитектурах MoE, такие как малые размеры батчей во время декодирования и переменное количество токенов. Внедрение этих изменений привело к увеличению пропускной способности до 10% и улучшению показателя Time to First Token (TTFT) на 9% на GPU H200. Несколько из этих улучшений были интегрированы в открытое сообщество vLLM. Работа демонстрирует, как глубокая настройка под конкретное оборудование может значительно повысить эффективность моделей с разреженной активацией.
aws