Principales nouveautés

PyTorch 2.15 introduit le support natif du type torch.float8_e4m3fn sur le backend MPS. Chaque valeur occupe un octet, soit la moitié de la mémoire requise par le BF16 ou le FP16. La fonction F.scaled_mm permet de multiplier directement ces tenseurs en appliquant un facteur d’échelle tensor‑wise ou row‑wise. La version requiert macOS 15 ou ultérieur et Python 3.11+, ce qui rend le backend disponible uniquement sur les Macs Apple Silicon récents.

import torch, torch.nn.functional as F

def quantize_fp8(x):
    limit = torch.finfo(torch.float8_e4m3fn).max
    scale = x.abs().amax().clamp_min(1e-12) / limit
    return (x / scale).to(torch.float8_e4m3fn), scale

x = torch.randn(1, 1024, device="mps", dtype=torch.float32)
weight = torch.randn(1024, 1024, device="mps", dtype=torch.float32)

x_fp8, x_scale = quantize_fp8(x)
w_fp8, w_scale = quantize_fp8(weight)

y = F.scaled_mm(
    x_fp8,
    w_fp8.t(),
    scale_a=x_scale,
    scale_recipe_a=F.ScalingType.TensorWise,
    scale_b=w_scale,
    scale_recipe_b=F.ScalingType.TensorWise,
    output_dtype=torch.bfloat16,
)
print(y.shape, y.dtype)

Optimisations de décodage et de softmax

Les modèles mixture‑of‑experts (MoE) bénéficient d’une accélération de 5 à 6,5 × par rapport à la version 2.14. Sur un M3 Ultra, le modèle Qwen3.6‑35B‑A3B passe de 9,4 tokens/s à 61,3 tokens/s avec torch.compile. Les couches linéaires à une seule ligne utilisent désormais des kernels matrix‑vector (GEMV) intégrés, ce qui donne un gain moyen de 1,23 × sur 96 formes de couche BF16. Le softmax sur le vocabulaire est environ 8 × plus rapide, réduisant le temps de génération de chaque token pour tous les LLM.

Extensions algébriques et kernels spécialisés

PyTorch 2.15 ajoute le support complet des opérations complexes sur MPS : factorisation LU, résolution LU et résolution triangulaire sont exécutées sur le GPU, ouvrant la voie aux gradients de déterminants complexes. Exemple :

import torch
a = torch.randn(8, 8, device="mps", dtype=torch.complex64)
a = a + 8 * torch.eye(8, device="mps", dtype=torch.complex64)
b = torch.randn(8, 2, device="mps", dtype=torch.complex64)
lu, pivots = torch.linalg.lu_factor(a)
x = torch.linalg.lu_solve(lu, pivots, b)
print((a @ x - b).abs().max().item())

Les petites matrices (≤ 4 × 4) bénéficient d’un kernel unique qui combine factorisation et résolution, tandis que les décompositions QR et les résolutions triangulaires utilisent des mises à jour bloquées pour les tailles plus importantes. Les convolutions, le padding, le dropout et les interpolations backward sont également optimisés, ce qui se traduit par une consommation mémoire de 8 GiB au lieu de 270 GiB lors du décodage vidéo du VAE de Wan2.2.

Conséquences pour les utilisateurs Mac

Les flux MPS sont désormais exposés en Python, offrant un contrôle plus fin et une synchronisation d’événements plus fiable. Les modèles de génération d’images, comme Qwen‑Image‑2.1, corrigent un problème de padding qui faisait chuter le PSNR de 10,06 dB à 29,95 dB, alignant les résultats sur le CPU. Les exigences système (macOS 15+, Python 3.11+) sont clairement indiquées dans la matrice de compatibilité, et le mécanisme de vérification d’accès MPS signale explicitement l’absence de GPU Metal ou d’OS trop ancien. En pratique, les développeurs peuvent installer la version candidate avec uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu et profiter immédiatement d’une réduction de moitié de la mémoire FP8 et d’une décodage LLM nettement plus rapide.