🌟 MetaShuffling от PyTorch: ускоряем вывод Llama 4 MoE без лишних вычислений и задержек.

PyTorch представил MetaShuffling — решение для ускорения вывода в Llama 4 MoE, которое решает проблемы динамизма и разреженности моделей, связанных с маршрутизацией токенов. Вместо традиционных методов padding или slicing, MetaShuffling использует переупорядочивание токенов по экспертам, избавляясь от ненужных операций и синхронизации между CPU и GPU. Это снижает использование памяти и устраняет задержки, связанные с обработкой «пустышек» или множественными запусками ядер.

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

Благодаря этому MetaShuffling совместим с механизмами графов (CUDAGraph, torch.compile), которые ускоряют исполнение, избегая повторных синхронизаций. Решение особенно эффективно для Llama 4, где каждый MoE-слой активирует лишь часть экспертов, что на практике создает динамические нагрузки.

▶️Ключевыми инновациями стали оптимизация ядер GroupedGEMM и IndexShuffling:

🟢GroupedGEMM, написанный на Triton, обрабатывает несколько матриц в одном вызове, используя статические и динамические разбиения размеров, позволяя пропускать неактивных экспертов и «лишние» токены без дополнительных затрат.

🟢IndexShuffling, в свою очередь, выполняет сортировку токенов и подсчет их количества на каждом эксперте за один проход, что по тестам оказалось в 5–13 раз быстрее, чем стандартные реализации PyTorch.

Результаты экспериментов на H100 80GB выглядят многообещающими.

Prefill Llama 4 Maverick с FP8 GroupedGEMM достигает 1,197 TFlops при 286 мкс, что близко к теоретическому пределу GPU.

В задачах декодирования метрики также демонстрируют высокую эффективность: 44,88 TFlops за 59 мкс. Даже при малом количестве токенов (128) MetaShuffling показывает 80% использования пропускной способности памяти.

Для multi-host сценариев MetaShuffling предлагает гибкость между «динамическими» и «статичными» формами данных. В режиме eager (без графов) используется минимальное заполнение без синхронизации CPU-GPU.

В graph mode — статичные формы с минимальным паддингом, что сокращает сетевой трафик и память. Также реализована дедупликация коммуникаций, которая распределяет нагрузку между узлами, снижая задержки в распределенных вычислениях.

MetaShuffling поддерживает интеграцию с FBGEMM Generative AI Kernel Library, позволяя применять оптимизации для vLLM и SGLang.


📌 Полная статья в блоге Pytorch


@ai_machinelearning_big_data

#AI #ML #MetaShuffling #Pytorch