Как сообщает AWS Machine Learning, архитектура Mixture-of-Experts (MoE) представляет собой систему с разреженной активацией, которая позволяет увеличить мощность больших языковых моделей (LLM), сохраняя при этом низкие затраты на вычисления для каждого токена. Поскольку операции MoE занимают значительную долю общего времени вычислений, их неэффективность напрямую увеличивает задержку запросов и стоимость. Поэтому оптимизация операций MoE с помощью проектирования ядер предлагает наибольшую возможность для повышения производительности. Одной из ключевых оптимизаций для архитектуры MoE является система Fused MoE, которая объединяет маршрутизацию токенов, перемешивание данных и специализированные сетевые вычисления в одну пакетную операцию на графическом процессоре (GPU). Это заменяет выполнение маршрутизации и вычислений экспертов как отдельных последовательных операций. Хотя это дает преимущества по задержке, это также приводит к сложности в реализации ядер и отладке. Высокая загрузка GPU для этой архитектуры требует оптимизаций, специфичных для аппаратного обеспечения. В этой статье мы описываем различные оптимизации ядер, которые мы реализовали для моделей Fused MoE, и показываем, как вы можете извлечь выгоду из этой работы. Мы использовали представительную модель MoE в качестве основного примера. Для максимизации пропускной способности обработки тензоров архитектура NVIDIA Hopper полагается на протокол инструкций Warp Group Matrix Multiply Accumulate (WGMMA), который выполняется совместно в рамках группы 128 потоков. Эта асимметрия операндов является фундаментальным свойством архитектуры наборов инструкций (ISA). В определении ассемблера wgmma.mma_async аппаратное обеспечение фиксирует размер матрицы по строкам (M) на постоянном уровне 64, чтобы определить свою топологию параллельного отображения потоков. Размер по столбцам (N) остается высоко переменным. Операнд A строго отображается по фиксированному размеру M (с использованием регистров или общей памяти), а операнд B отображается по переменному размеру N (через дескриптор общей памяти). Эти слоты операндов, таким образом, структурно не взаимозаменяемы. В результате, максимальная загрузка аппаратного тайла определяется этими жесткими кодировками форм, требуя, чтобы размер M равнялся или превышал 64, чтобы предотвратить простое выполнение вычислительных линий. Эта структурная предпочтительность возникает из того, как ядра Tensor Hopper разбивают блоки памяти. Аппаратное обеспечение навязывает фиксированную минимальную форму компоновки для ведущего измерения первого операнда, фиксируя высоту блока на уровне 64 (M=64). Когда математическая матрица имеет большой, непрерывный ведущий размер в этом основном слоте, внутренние планировщики аппаратного обеспечения Hopper могут четко сопоставить и распределить данные по физической вычислительной сетке. Эта структурная асимметрия становится серьезным узким местом производительности в фазе декодирования сетей MoE. В стандартном слое MoE линейная алгебра, управляющая вычислениями экспертов, формулируется с использованием традиционной матричной нотации: A[M,K] x B[K,N] —> C[M,N], где (M) представляет количество строк токенов (или размер пакета), (K) представляет скрытое размерное измерение, а (N) представляет расширенное размерное измерение выходов экспертов. В фазе декодирования система обрабатывает только один токен за последовательность пользователя за раз. Следовательно, пространство пакета или размер строки токенов обычно оказывается очень тонким, регулярно опускаясь до профиля, где (M < 64), и часто до одной строки (M = 1). При этом стандартном структурном макете матрица (A) превращается в тонкий, фрагментированный горизонтальный вектор (например, 1 x 4096). Матрица (B) (основные веса экспертов) остается в большом, плотном вертикальном блоке (например, 4096 x 8192). Поскольку количество строк (M) мало, матрица (A) физически не может заполнить фиксированный (64)-строчный макет для операнда A. Когда инструкция WGMMA выполняется по этому макету, она оставляет значительную часть физической емкости ядра Tensor полностью голодной. Кремний тратит тактовые циклы на обработку пустых пространственных слотов, что приводит к плохой загрузке аппаратного обеспечения. В стандартном плотном слое LLM весовые матрицы статичны, едины и идеально непрерывны. Однако в слое MoE рабочая нагрузка фрагментирована по флоту независимых подсетей (экспертов). Это создает уникальные узкие места памяти и компоновки, которые делают ее гиперчувствительной к ограничениям WGMMA графического процессора Hopper. В стандартном плотном слое умножение матриц умножает вектор токена на единую матрицу весов слоя. Размер выходного измерения (N) равен скрытому размеру модели (например, 4096). В ядре Fused MoE графический процессор не только умножает на одного эксперта. Он группирует и выполняет вычисления для всех активных экспертов одновременно, чтобы избежать запуска отдельных ядер для каждого эксперта. Для стандартного плотного слоя N — это скрытый размер (4096). Для слоя Fused MoE N — это количество выбранных экспертов, умноженное на промежуточный размер эксперта. Для модели, работающей по рутине topk=8, эффективное измерение N внутри объединенного ядра взрывается до десятков или сотен тысяч параметров. Поскольку размер N так непропорционально велик по сравнению с крошечным размером токена декодирования (M=1), асимметрия компоновки становится экстремальной. Транспонирование матриц перемещает это большое, объединенное измерение MoE в основной слот, разблокируя предвыборку аппаратного обеспечения в масштабе, который стандартные плотные слои не могут сопоставить. Чтобы разрешить это аппаратное несоответствие без изменения основной математической целостности модели, мы используем алгебраическое свойство транспонирования матриц: BT[N,K] x AT[K,M] —> CT[N,M]. Программно перевернув умножение матриц до того, как оно достигнет слоя ассемблера PTX, мы выровняли программное обеспечение, чтобы соответствовать точной топографии кремния Hopper. Это транспонирование вызывает два немедленных архитектурных преимущества на графическом процессоре. Во-первых, оно насыщает физическую компоновку. Широкое измерение эксперта (N = 8192) продвигается на первую позицию операнда (BT), и поскольку (N) велико, оно легко заполняет обязательный (64)-строчный порог аппаратного обеспечения. Ядра Tensor затем работают с максимальной локальной плотностью без потерь циклов выполнения. Во-вторых, оно активирует асинхронную предвыборку. Специальные аппаратные копирующие движки Hopper могут асинхронно предвыбирать адреса памяти из общей памяти прямо в регистры до того, как блок выполнения математики запустится. Размещение непрерывных, больших весовых блоков (BT) в ведущем слоте позволяет этим двигателям плавно отслеживать компоновку данных и поддерживать насыщенность вычислительного конвейера без вмешательства ЦП. Это обеспечило улучшение производительности на 2–4 процента по итогам работы с большими моделями MoE на H200. Мы внедрили эту оптимизацию в vLLM, и она доступна в версии 0.24.0 и выше. Выбор бэкенда для выполнения вычислений MoE также оказывает большое влияние на производительность. MoE-оракул vLLM направляет различные схемы квантования в Triton, FlashInfer+CUTLASS или другие бэкенды. Профилирование показало, что для моделей с квантованием тензоров FP8 на Hopper бэкенд Triton постоянно превосходит стандартный путь FLASHINFER_CUTLASS. Бэкенд Triton превосходит CUTLASS в условиях низкой конкуренции для рабочих нагрузок MoE в первую очередь потому, что он заменяет жесткие конфигурации матричной математики на уровне фреймворка на индивидуальную, операторную фузию регистров. CUTLASS полагается на стандартные шаблоны C++, которые рассматривают масштабирование тензоров FP8, умножение матриц и агрегацию экспертов как дискретные математические операции. Бэкенд Triton, наоборот, консолидирует весь этот конвейер в одну аппаратную операцию. Данные извлекаются из глобальной памяти с высокой пропускной способностью (HBM) ровно один раз и полностью обрабатываются внутри сверхбыстрых локальных регистров SRAM графического процессора и файлов выполнения вычислений. Затем они передаются обратно в конечное местоположение матрицы. Это устраняет медленные глобальные операции чтения/записи VRAM, которые компрометируют CUTLASS, когда ему приходится постоянно проходить через промежуточные этапы масштабирования тензоров. Кроме того, эта фузия на уровне регистров специально сочетается с динамическими планировками компоновки, разработанными для переменного распределения токенов, устраняя серьезный структурный недостаток традиционных шаблонов компиляции. CUTLASS изначально оптимизирован для плотных, больших и предсказуемых размерностей матриц. Когда он сталкивается с жидкими, ультратонкими профилями пакетов (M < 64 или M = 1), типичными для фазы генерации живого декодирования, его стандартные структуры блочной плитки и отображения потоков становятся глубоко неэффективными. В результате возникает значительное недоиспользование ворков или дорогостоящие накладные расходы на заполнение памяти. Напротив, бэкенд Triton отлично справляется с динамическим отображением этих небольших, фрагментированных массивов токенов по физическим границам активных весов экспертов. Это позволяет внутренним потокам предвыборки аппаратного обеспечения и конвейерам Tensor Core оставаться полностью насыщенными, даже при обработке сильно нестабильных рабочих нагрузок при очень низких размерах пакетов. Обновление политики маршрутизации является чисто конфигурационным изменением без изменения кода ядра. Тем не менее, это обеспечило 10-процентное улучшение производительности по всей длине для моделей Fused MoE с квантованием FP8 на H200. Это улучшение доступно в контейнере AWS Large Model Inference (LMI) версии 28. Одним из самых значительных путей повышения производительности в GPU-инференции является также один из наименее очевидных: ненужная повторная компиляция ядер. Triton, компилятор, лежащий в основе многих пользовательских ядер vLLM, по умолчанию специализирует целочисленные аргументы. Каждый раз, когда ядро вызывается с ранее не видимым значением для специализированного параметра, Triton запускает новую компиляцию в режиме реального времени (JIT). При обслуживании рабочих нагрузок, где длины последовательностей и размеры пакетов варьируются, это создает постоянный поток накладных расходов на компиляцию, который увеличивает задержку, особенно время до первого токена (TTFT). В обслуживании моделей MoE ядро _fused_inv_rope_fp8_quant_per_head (ответственное за объединенное позиционное встраивание и квантование FP8) повторно компилировалось при каждом новом значении num_tokens. Добавление аннотаций do_not_specialize к параметрам переменной длины полностью устранило эти повторные компиляции, улучшив TTFT на 9 процентов на H200 при 8-кратном тензорном параллелизме. В многопроцессорном кластере все 8 графических процессоров должны выполнять свои ядра в идеальной синхронизации на уровне микросекунд. Если хотя бы один графический процессор должен остановиться, чтобы повторно скомпилировать свое локальное ядро _fused_inv_rope_fp8_quant_per_head, потому что изменилось количество токенов в его локальном пакете, другие семь графических процессоров остаются без дела и ждут его. Устранение повторных компиляций по всему кластеру сгладило накладные расходы на координацию многопроцессорных систем, обеспечив 9-процентное увеличение скорости для времени начальной задержки пользователя. Принцип здесь заключается в том, чтобы рассматривать компилятор как поверхность оптимизации первого класса. Понимание того, как Triton специализирует семантику, и распознавание, когда параметры мертвы или переменные, открывает улучшения задержки с минимальными изменениями кода. Мы внедрили эту оптимизацию в vLLM, и она доступна в версии 0.22.0 и выше. В заключение, благодаря нашему сотрудничеству с сообществом vLLM мы реализовали и открыли несколько оптимизаций ядер для моделей Fused MoE. Мы применяем эти оптимизации в контейнере глубинного обучения LMI на Amazon SageMaker AI. В совокупности эти оптимизации обеспечили до 10 процентов более высокую пропускную способность и до 9 процентов более быстрое время до первого токена. Чтобы начать работу с пользовательским хостингом моделей, смотрите документацию по хостингу Amazon SageMaker AI. Материал подготовлен нейросетью на основе оригинальной публикации. ИИ может ошибаться в деталях (цифры, имена, даты, цитаты) — для важной информации сверяйтесь с первоисточником.