PyTorch 2.15 accélère les modèles MoE sur Mac de 5 à 6,5 fois et ajoute le FP8

La version 2.15 de PyTorch fait du Mac une machine réellement exploitable pour les LLM locaux et la génération vidéo, alors qu'il n'était jusqu'ici qu'un terrain d'expérimentation tout juste acceptable. L'équipe PyTorch de Hugging Face, chargée du backend MPS, a ajouté la prise en charge du FP8 et remplacé une partie des opérations MPSGraph par des noyaux Metal dédiés. Sur M3 Ultra avec torch.compile, le modèle Qwen3.6-35B-A3B produit désormais 61,3 tokens par seconde au lieu de 9,4, et le décodage vidéo de Wan2.2 ne fait plus gonfler la mémoire du pilote jusqu'à 270 GiB.

FP8 et décodage

MoE : d'où viennent les 61 tokens par seconde

Le type torch.float8_e4m3fn fait son apparition sur MPS. Il occupe un octet par valeur, soit deux fois moins que FP16 ou BF16. Ces tenseurs se multiplient via F.scaled_mm, qui gère à la fois la mise à l'échelle tensorwise et rowwise. Les autres opérations disponibles pour ce type sur CUDA et CPU fonctionnent également.

L'essentiel du gain de vitesse concerne la phase de décodage. Là, chaque nouveau token traverse le modèle sous la forme d'une seule ligne, et les grosses multiplications matricielles deviennent des produits matrice-vecteur. En 2.14, seul torch.mm disposait de noyaux GEMV, tandis que F.linear passait par MPSGraph. Désormais, l'entrée typique [1, 1, hidden_size] les utilise directement, et les additions de bias compatibles sont fusionnées dans le noyau. Sur 96 formes de couches BF16 issues de modèles réels, cela donne une accélération moyenne géométrique de 1,23×. Les petites projections en profitent le plus, par exemple les gates des MoE. Les grandes couches et le LM-head étaient déjà limités par la bande passante mémoire : le gain y reste d'environ 1×.

Pour les MoE, l'équipe a développé un grouped matmul natif. Auparavant, les décalages de groupes étaient copiés sur le CPU et l'opération bouclait groupe par groupe. Les nouveaux noyaux lisent ces décalages directement sur le GPU. Ils s'appuient sur Metal Performance Primitives lorsque ceux-ci sont disponibles, et sur un chemin Metal distinct dans les autres cas. Un noyau spécifique couvre les petits groupes, typiques du décodage. L'équipe a aussi réduit les synchronisations entre GPU et CPU et accéléré le routage top-k.

Les détails qui grignotaient mémoire et temps

La correction du constant padding a donné le résultat le plus parlant. Lors d'un décodage complet en 28 tuiles, le décodeur VAE Wan-AI/Wan2.2-TI2V-5B-Diffusers effectue 8 764 appels de padding. Avec MPSGraph, le pic de mémoire du pilote atteignait 270,135 GiB ; avec le padding Metal, il tombe à 8,108 GiB. Le pic de tenseurs alloués a à peine bougé : 4,422 contre 4,095 GiB. Ce qui gonflait auparavant, c'était la mémoire que le pilote réservait pour les graphes.

Le même correctif répare l'encodeur VAE de Qwen-Image-2.1. En 2.14, il corrompait silencieusement les images sur MPS, ce qui cassait les scénarios d'édition dans ComfyUI. Désormais, l'aller-retour « encodage puis décodage » coïncide avec le résultat CPU à 29,95 dB de PSNR, contre 10,06 dB auparavant.

Voici la liste des autres changements de la 2.15 :

• softmax et log_softmax : noyaux Metal dédiés ; le softmax FP32 sur la ligne [1, 128256] passe de 57 à 7 µs, soit un gain d'environ 8 fois.
• conv1d forward : noyaux spécialisés à la place de MPSGraph, et chemin im2col plus matmul pour les sorties courtes, utile pour l'audio et les modèles de streaming.
• Reflection et replication padding : nouvelles implémentations Metal, plus rapides et avec un backward corrigé.
• Dropout : nn.Dropout et F.dropout utilisent le noyau fusionné déjà existant.
• Upsampling backward : noyau gather sur Metal à la place du graphe de gradient de resize.
• Channels-last conv2d backward : reste sur le chemin channels-last efficace, sans graphe NCHW ni accès strided coûteux.
• Algèbre linéaire complexe : LU, LU solve et triangular solve pour complex64, ainsi que les déterminants et gradients complexes associés.
• Petites matrices : les batchs FP32 jusqu'à 4×4 sont inversés par un seul noyau Metal qui combine factorisation et résolution.
• QR et triangular solve : mises à jour par blocs, la majeure partie du travail se ramenant à des multiplications matricielles. Pour Gated DeltaNet, un chemin rapide existe pour les systèmes 16×16, 32×32 et 64×64 en FP32 avec huit seconds membres ou plus.
• SVD : gère correctement les très petites valeurs singulières, les échelles d'entrée extrêmes et full_matrices=True, ce qui améliore les pseudo-inverses et les moindres carrés.
• MPS streams : torch.mps.Stream devient public, avec le contexte torch.mps.stream et Tensor.record_stream, ainsi que des mesures de temps plus fiables via les événements.
• Nouvelles opérations : torch.frexp, torch.logspace (pour l'instant sans sortie entière ni booléenne pour plusieurs éléments), torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l et les comparaisons d'entiers non signés. Dans l'implémentation mathématique de SDPA, le broadcasting sur le batch et les têtes est corrigé.

Ce qui attend la 2.16 et comment essayer

Une partie des plans pour la 2.16 figure déjà dans les nightly builds. Le command buffer sera validé tous les 16 noyaux, afin que le GPU démarre son travail pendant que le CPU encode le reste.

En décodage BF16 compressé avec un contexte de 2048 tokens, LFM2.5-8B-A1B passe de 108,8 à 126,0 tokens/s (+16 %), Qwen3.6-35B-A3B de 52,9 à 62,8 (+19 %) et Qwen3.

8-27B de 12,1 à 12,6 (+4 %). L'attention passera entièrement sur Metal. Sur M5 Pro, le décodage avec GQA (32 têtes de requête, 8 têtes KV, head dim 80, contexte 4K) tombera de 465 à 65 µs, et le prefill 4K×4K avec head dim 40 de 42 à 6 ms.

Parmi les autres changements : bincount sur 4 millions d'entiers aux valeurs presque identiques passera de 6,8 s à une demi-milliseconde, et libtorch_cpu.dylib passera de 384 à 287 MiB grâce à la compilation des shaders pour Metal 3.1 par défaut.

Voici la marche à suivre. Il faut un Mac Apple silicon avec macOS 15 ou plus récent et Python 3.11+ ; sinon, la vérification de disponibilité de MPS indiquera désormais la raison du refus.

La release candidate s'installe dans un environnement séparé avec la commande uv pip install torch==2.15.0 --index-url https://download.

pytorch.org/whl/test/cpu. Pour ceux qui exécutent des modèles MoE via transformers, il est logique d'activer torch.

compile et de comparer les tokens par seconde avec la 2.14. Les utilisateurs de ComfyUI avec Qwen-Image-2.

1 ont intérêt à mettre à jour à cause du correctif VAE, et les propriétaires de Mac dotés de beaucoup de mémoire qui génèrent de la vidéo avec Wan2.2 éviteront grâce à cette version le gonflement inattendu de la mémoire du pilote. Si quelque chose reste lent ou ne fonctionne pas, l'équipe demande d'ouvrir une issue sur GitHub.