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