PyTorch 2.15 acelera los modelos MoE en Mac entre 5 y 6,5 veces e incorpora FP8

La versión PyTorch 2.15 convierte el Mac, de una plataforma tolerable para experimentar, en una máquina perfectamente utilizable para LLM locales y generación de vídeo. El equipo de PyTorch de Hugging Face, responsable del backend MPS, ha añadido soporte para FP8 y ha sustituido varias operaciones de MPSGraph por kernels de Metal dedicados. En un M3 Ultra con torch.compile, el modelo Qwen3.6-35B-A3B pasa de 9,4 tokens por segundo a 61,3, y la decodificación de vídeo en Wan2.2 ya no infla la memoria del driver hasta los 270 GiB.

FP8 y decodificación

MoE: de dónde salen los 61 tokens por segundo

MPS incorpora el tipo torch.float8_e4m3fn. Ocupa un byte por valor, la mitad que FP16 o BF16. Estos tensores se pueden multiplicar con F.scaled_mm, que admite escalado tanto tensorwise como rowwise. También funcionan el resto de operaciones disponibles para este tipo en CUDA y CPU.

La mayor parte de la mejora de velocidad se concentra en la fase de decodificación. Ahí cada token nuevo atraviesa el modelo como una sola fila, y las grandes multiplicaciones de matrices se convierten en multiplicaciones matriz-vector. En la 2.

14, los kernels GEMV solo existían para torch.mm, mientras que F.linear pasaba por MPSGraph.

Ahora la entrada típica [1, 1, hidden_size] los usa directamente, y las sumas de bias que encajan se fusionan con el kernel. En 96 formas de capas BF16 de modelos reales, esto supone una aceleración media geométrica de 1,23×. Las más beneficiadas son las proyecciones pequeñas, como las puertas (gates) de los MoE.

Las capas grandes y el LM-head ya estaban limitados por el ancho de banda de memoria, así que ahí la mejora ronda 1×.

Para los MoE se ha creado un grouped matmul nativo. Antes, los desplazamientos de los grupos se copiaban a la CPU y la operación recorría los grupos en un bucle. Los nuevos kernels leen los desplazamientos directamente en la GPU. Usan Metal Performance Primitives donde están disponibles y, en el resto de los casos, una ruta Metal independiente. Para los grupos pequeños, típicos de la decodificación, hay un kernel propio. Además, se han reducido las sincronizaciones entre GPU y CPU y se ha acelerado el enrutamiento top-k.

Los pequeños detalles que consumían memoria y tiempo

El resultado más llamativo viene de corregir el constant padding. El decodificador VAE de Wan-AI/Wan2.2-TI2V-5B-Diffusers realiza 8 764 llamadas de padding al decodificar por completo 28 tiles. Con MPSGraph, el pico de memoria del driver llegaba a 270,135 GiB; con el padding en Metal queda en 8,108 GiB. El pico de tensores asignados apenas cambió: 4,422 frente a 4,095 GiB. Lo que crecía era la memoria que el driver mantenía reservada para los grafos.

Esa misma corrección arregló el encoder VAE de Qwen-Image-2.1. En la 2.14 estropeaba las imágenes en silencio en MPS, lo que rompía los flujos de edición en ComfyUI. Ahora el ciclo completo de codificación y decodificación coincide con la CPU en 29,95 dB de PSNR; antes era de 10,06 dB.

Lista del resto de cambios de la 2.15:

• softmax y log_softmax: kernels de Metal independientes; el softmax FP32 sobre la fila [1, 128256] pasó de 57 a 7 µs, unas 8 veces más rápido.
• conv1d forward: kernels especializados en lugar de MPSGraph y una ruta im2col más matmul para salidas cortas, útil para audio y modelos de streaming.
• Reflection y replication padding: nuevas implementaciones en Metal, más rápidas y con el backward corregido.
• Dropout: nn.Dropout y F.dropout usan el kernel fusionado que ya existía.
• Upsampling backward: kernel gather en Metal en lugar del grafo de gradiente de resize.
• Channels-last conv2d backward: se mantiene en la ruta eficiente channels-last, sin grafo NCHW ni accesos strided costosos.
• Álgebra lineal compleja: LU, LU solve y triangular solve para complex64, y con ellos determinantes y gradientes complejos.
• Matrices pequeñas: los lotes FP32 de hasta 4×4 se invierten con un único kernel de Metal que combina factorización y resolución.
• QR y triangular solve: actualizaciones por bloques, con buena parte del trabajo reducida a multiplicación de matrices. Para Gated DeltaNet hay una ruta rápida para sistemas de 16×16, 32×32 y 64×64 en FP32 con ocho o más lados derechos.
• SVD: maneja correctamente valores singulares muy pequeños, escalas de entrada extremas y full_matrices=True, lo que mejora las pseudoinversas y los mínimos cuadrados.
• MPS streams: torch.mps.Stream público, el contexto torch.mps.stream y Tensor.record_stream, además de mediciones de tiempo más fiables mediante eventos.
• Nuevas operaciones: torch.frexp, torch.logspace (por ahora sin salida entera ni booleana para varios elementos), torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l y comparaciones de enteros sin signo. En la implementación matemática de SDPA se ha corregido el broadcasting por lote y por cabezas.

Qué llega en la 2.16 y cómo probarlo

En las compilaciones nocturnas ya está parte de los planes para la 2.16. El command buffer se confirmará cada 16 kernels, para que la GPU empiece a trabajar mientras la CPU codifica el resto.

En la decodificación BF16 comprimida con un contexto de 2048 tokens, LFM2.5-8B-A1B sube de 108,8 a 126,0 tokens/s (+16%), Qwen3.6-35B-A3B de 52,9 a 62,8 (+19%) y Qwen3.

8-27B de 12,1 a 12,6 (+4%). La atención pasará por completo a Metal. En un M5 Pro, la decodificación con GQA (32 cabezas de query, 8 cabezas KV, head dim 80, contexto de 4K) bajará de 465 a 65 µs, y el prefill de 4K×4K con head dim 40, de 42 a 6 ms.

Entre otros cambios, bincount sobre 4 millones de enteros con valores casi idénticos, que ahora tarda 6,8 s, se calculará en medio milisegundo, y libtorch_cpu.dylib adelgazará de 384 a 287 MiB gracias a que los shaders se compilarán por defecto para Metal 3.1.

El orden práctico de pasos es este. Hace falta un Mac con Apple silicon, macOS 15 o posterior y Python 3.11+; de lo contrario, la comprobación de disponibilidad de MPS ahora explicará el motivo del rechazo.

La release candidate se instala en un entorno aparte con el comando uv pip install torch==2.15.0 --index-url https://download.

pytorch.org/whl/test/cpu. Quienes ejecuten modelos MoE con transformers deberían activar torch.

compile y comparar los tokens por segundo con los de la 2.14. Los usuarios de ComfyUI con Qwen-Image-2.

1 deberían actualizar por la corrección del VAE, y a los propietarios de Mac con mucha memoria que generen vídeo en Wan2.2 la actualización les evitará el inesperado inflado de la memoria del driver. Si algo sigue lento o no funciona, el equipo pide abrir un issue en GitHub.