PyTorch 2.15 на Mac: MoE декодируется в разы быстрее

В PyTorch 2.15 бэкенд MPS на Apple Silicon получил пакет ускорений под декод языковых моделей: отдельные ядра Metal вместо обходов через MPSGraph, быстрее линейные слои и softmax при выборе следующего токена. Команда PyTorch в Hugging Face вместе с мейнтейнерами гоняет модели с Hub и переносит фиксы в основную ветку.

Замеры опираются на transformers, diffusers и ComfyUI — не на синтетические микробенчмарки.

Быстрее декод

На этапе декода каждый токен проходит сеть одной строкой, поэтому тяжёлые matmul превращаются в умножение матрицы на вектор. В 2.15 GEMV-ядра подключили к F.linear; softmax и log_softmax переехали на Metal — для строки logits из 128256 элементов время падает с 57 до 7 мкс, примерно в 8 раз.

Смешанные эксперты относительно 2.14 ускорились в 5–6,5 раза. На M3 Ultra с torch.compile Qwen3.6-35B-A3B идёт с 9,4 до 61,3 токена/с, LFM2.5-8B-A1B — с 23,6 до 123,6. F.grouped_mm больше не копирует смещения групп на CPU: grouped matmul читает offsets прямо на GPU.

Видео, картинки и FP8

Constant padding на Metal убрал раздувание памяти драйвера при tiled-декоде: у Wan2.2 пик падает с 270 до 8 ГиБ на проходе из 28 тайлов и 8764 вызовов padding.

На MPS появился torch.float8_e4m3fn и умножение через F.scaled_mm — веса в FP8 занимают вдвое меньше, чем в BF16. Encoder VAE для Qwen-Image-2.1 на 2.14 тихо портил картинки в ComfyUI; после фикса encode-decode даёт 29,95 dB PSNR вместо 10,06 dB относительно CPU.

Требования

MPS теперь требует macOS 15 и выше, PyTorch 2.15 — Python 3.11+. Из Python доступны torch.mps.Stream и понятная проверка, если GPU или OS не подходят. Релиз-кандидат ставят отдельным окружением через test-индекс PyTorch. В 2.16 уже в работе attention целиком на Metal и commit command buffer каждые 16 ядер, чтобы GPU меньше ждал CPU на длинном декоде.

Источник: MPS in PyTorch 2.15.