PyTorch 2.15 ускорил MoE-модели на Mac в 5–6,5 раза и принёс FP8

Релиз PyTorch 2.15 превращает Mac из терпимой площадки для экспериментов в вполне рабочую машину для локальных LLM и генерации видео. Команда PyTorch из Hugging Face, которая занимается бэкендом MPS, добавила поддержку FP8 и заменила часть операций на MPSGraph выделенными Metal-ядрами. На M3 Ultra с torch.compile модель Qwen3.6-35B-A3B вместо 9,4 токена в секунду теперь выдаёт 61,3, а декодирование видео в Wan2.2 больше не раздувает память драйвера до 270 GiB.

FP8 и MoE-декодинг: откуда взялись 61 токен в секунду

На MPS появился тип torch.float8_e4m3fn. Один байт на значение, то есть вдвое меньше, чем у FP16 или BF16. Перемножать такие тензоры можно через F.scaled_mm, причём поддерживаются и tensorwise-, и rowwise-масштабирование. Работают и остальные операции, доступные для этого типа на CUDA и CPU.

Основной прирост скорости связан с фазой декодирования. Там каждый новый токен идёт через модель одной строкой, и большие матричные умножения превращаются в матрично-векторные. В 2.14 GEMV-ядра были только у torch.mm, а F.linear шёл через MPSGraph. Теперь типичный вход [1, 1, hidden_size] использует их напрямую, а подходящие сложения с bias сливаются с ядром. На 96 формах BF16-слоёв из реальных моделей это даёт среднегеометрическое ускорение 1,23×. Сильнее всего выигрывают небольшие проекции, например гейты в MoE. Большие слои и LM-head и раньше упирались в пропускную способность памяти, так что там прироста около 1×.

Для MoE сделали нативный grouped matmul. Раньше смещения групп копировались на CPU, и операция крутилась по группам в цикле. Новые ядра читают смещения прямо на GPU. Они используют Metal Performance Primitives там, где те доступны, а в остальных случаях отдельный Metal-путь. Для маленьких групп, типичных при декодировании, есть собственное ядро. Дополнительно сократили синхронизации GPU и CPU и ускорили top-k-маршрутизацию.

Мелочи, которые съедали память и время

Самый наглядный результат дала правка constant padding. VAE-декодер Wan-AI/Wan2.2-TI2V-5B-Diffusers при полном декодировании из 28 тайлов делает 8 764 вызова паддинга. Через MPSGraph пиковая память драйвера доходила до 270,135 GiB, а с Metal-паддингом она составляет 8,108 GiB. Пик выделенных тензоров при этом почти не изменился: 4,422 против 4,095 GiB. Раньше разрасталась именно память, которую драйвер удерживал под графы.

Та же правка починила Qwen-Image-2.1 VAE encoder. В 2.14 он молча портил изображения на MPS, и это ломало сценарии редактирования в ComfyUI. Теперь round-trip «кодирование и декодирование» совпадает с CPU на уровне 29,95 дБ PSNR, раньше было 10,06 дБ.

Список остальных изменений 2.15:

• softmax и log_softmax: отдельные Metal-ядра, FP32 softmax по строке [1, 128256] ускорился с 57 до 7 мкс, примерно в 8 раз.
• conv1d forward: специализированные ядра вместо MPSGraph и путь im2col плюс matmul для коротких выходов, полезно для аудио и стриминговых моделей.
• Reflection и replication padding: новые Metal-реализации, быстрее и с исправленным backward.
• Dropout: nn.Dropout и F.dropout используют уже существующее fused-ядро.
• Upsampling backward: gather-ядро на Metal вместо графа градиента resize.
• Channels-last conv2d backward: остаётся на эффективном channels-last пути, без NCHW-графа и дорогого strided-доступа.
• Комплексная линейная алгебра: LU, LU solve и triangular solve для complex64, а вместе с ними комплексные определители и градиенты.
• Малые матрицы: батчи FP32 до 4×4 обращаются одним Metal-ядром, которое совмещает факторизацию и решение.
• QR и triangular solve: блочные обновления, значительная часть работы сводится к матричному умножению. Для Gated DeltaNet есть быстрый путь для систем 16×16, 32×32 и 64×64 в FP32 при восьми и более правых частях.
• SVD: корректно обрабатывает очень малые сингулярные значения, экстремальные масштабы входа и full_matrices=True, что улучшает псевдообратные и наименьшие квадраты.
• MPS streams: публичные torch.mps.Stream, контекст torch.mps.stream и Tensor.record_stream, плюс более надёжные замеры времени через события.
• Новые операции: torch.frexp, torch.logspace (пока без целочисленного и булева вывода для нескольких элементов), torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l и сравнения беззнаковых целых. В математической реализации SDPA исправлен broadcasting по батчу и головам.

Что в 2.16 и как попробовать

В ночных сборках уже лежит часть планов на 2.16. Command buffer будет коммититься каждые 16 ядер, чтобы GPU начинал работу, пока CPU кодирует остальное.

В сжатом BF16-декодировании при контексте 2048 токенов LFM2.5-8B-A1B вырастает с 108,8 до 126,0 токенов/с (+16%), Qwen3.6-35B-A3B с 52,9 до 62,8 (+19%), Qwen3.

8-27B с 12,1 до 12,6 (+4%). Внимание целиком уйдёт на Metal. На M5 Pro декодирование с GQA (32 query-головы, 8 KV-голов, head dim 80, контекст 4K) сократится с 465 до 65 мкс, а префилл 4K×4K с head dim 40 с 42 до 6 мс.

Из других изменений: bincount на 4 млн целых при почти одинаковых значениях вместо 6,8 с будет считаться за полмиллисекунды, а libtorch_cpu.dylib похудеет с 384 до 287 MiB благодаря сборке шейдеров под Metal 3.1 по умолчанию.

Практический порядок действий такой. Нужен Apple silicon Mac с macOS 15 или новее и Python 3.11+, иначе проверка доступности MPS теперь объяснит причину отказа. Релиз-кандидат ставится в отдельное окружение командой uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu. Тем, кто гоняет MoE-модели через transformers, имеет смысл включить torch.compile и сравнить токены в секунду с 2.14. Пользователям ComfyUI с Qwen-Image-2.1 стоит обновиться из-за исправления VAE, а владельцам Mac с большим объёмом памяти, генерирующим видео в Wan2.2, обновление избавит от неожиданного раздувания драйверной памяти. Если что-то по-прежнему медленно или не работает, команда просит открывать issue на GitHub.