PyTorch 2.15 acelera modelos MoE no Mac em 5–6,5 vezes e traz FP8

O PyTorch 2.15 faz o Mac deixar de ser uma plataforma apenas tolerável para experimentos e passar a ser uma máquina perfeitamente utilizável para LLMs locais e geração de vídeo. A equipe do PyTorch na Hugging Face, responsável pelo backend MPS, adicionou suporte a FP8 e substituiu parte das operações do MPSGraph por kernels Metal dedicados. Em um M3 Ultra com torch.compile, o modelo Qwen3.6-35B-A3B, que gerava 9,4 tokens por segundo, agora entrega 61,3. Já a decodificação de vídeo no Wan2.2 não infla mais a memória do driver até 270 GiB.

FP8 e decodificação de

MoE: de onde vêm os 61 tokens por segundo

O MPS ganhou o tipo torch.float8_e4m3fn. Ele usa um byte por valor, ou seja, metade do espaço de FP16 ou BF16. Esses tensores podem ser multiplicados via F.scaled_mm, com suporte a escalonamento tensorwise e rowwise. As demais operações disponíveis para esse tipo em CUDA e CPU também funcionam.

O ganho principal de velocidade está na fase de decodificação. Nela, cada novo token passa pelo modelo como uma única linha, e as grandes multiplicações de matrizes viram multiplicações matriz-vetor. No 2.14, os kernels GEMV existiam apenas para torch.mm, enquanto F.linear passava pelo MPSGraph. Agora a entrada típica [1, 1, hidden_size] usa esses kernels diretamente, e somas de bias compatíveis são fundidas ao kernel. Em 96 formatos de camadas BF16 extraídos de modelos reais, isso resulta em uma aceleração média geométrica de 1,23×. Quem mais ganha são as projeções pequenas, como os gates de MoE. Camadas grandes e o LM-head já esbarravam na largura de banda da memória, então o ganho ali fica em torno de 1×.

Para MoE, a equipe criou um grouped matmul nativo. Antes, os offsets dos grupos eram copiados para a CPU, e a operação rodava em um loop sobre os grupos. Os novos kernels leem os offsets diretamente na GPU. Eles usam o Metal Performance Primitives onde ele está disponível e, nos demais casos, um caminho Metal separado. Para grupos pequenos, típicos da decodificação, há um kernel próprio. Além disso, as sincronizações entre GPU e CPU foram reduzidas, e o roteamento top-k ficou mais rápido.

Detalhes que consumiam memória e tempo

O resultado mais visível veio da correção do constant padding. O decodificador VAE do Wan-AI/Wan2.2-TI2V-5B-Diffusers, na decodificação completa de 28 tiles, faz 8 764 chamadas de padding. Via MPSGraph, o pico de memória do driver chegava a 270,135 GiB; com o padding em Metal, fica em 8,108 GiB. O pico de tensores alocados mal mudou: 4,422 contra 4,095 GiB. O que crescia era a memória que o driver retinha para os grafos.

A mesma correção consertou o encoder VAE do Qwen-Image-2.1. No 2.14, ele corrompia imagens silenciosamente no MPS, o que quebrava fluxos de edição no ComfyUI. Agora o round-trip de "codificação e decodificação" coincide com o da CPU em 29,95 dB de PSNR; antes, eram 10,06 dB.

Lista das demais mudanças do 2.15:

• softmax e log_softmax: kernels Metal dedicados; o softmax FP32 sobre a linha [1, 128256] caiu de 57 para 7 µs, cerca de 8 vezes mais rápido.
• conv1d forward: kernels especializados no lugar do MPSGraph e um caminho im2col mais matmul para saídas curtas, útil para áudio e modelos de streaming.
• Reflection e replication padding: novas implementações em Metal, mais rápidas e com o backward corrigido.
• Dropout: nn.Dropout e F.dropout usam o kernel fused que já existia.
• Upsampling backward: kernel gather em Metal no lugar do grafo de gradiente do resize.
• Channels-last conv2d backward: permanece no caminho eficiente channels-last, sem grafo NCHW e sem acesso strided custoso.
• Álgebra linear complexa: LU, LU solve e triangular solve para complex64, além de determinantes e gradientes complexos.
• Matrizes pequenas: lotes FP32 de até 4×4 são invertidos por um único kernel Metal, que combina fatoração e resolução.
• QR e triangular solve: atualizações em blocos, com boa parte do trabalho reduzida a multiplicação de matrizes. Para o Gated DeltaNet, há um caminho rápido para sistemas 16×16, 32×32 e 64×64 em FP32 com oito ou mais lados direitos.
• SVD: trata corretamente valores singulares muito pequenos, escalas extremas de entrada e full_matrices=True, o que melhora pseudoinversas e mínimos quadrados.
• MPS streams: torch.mps.Stream público, contexto torch.mps.stream e Tensor.record_stream, além de medições de tempo mais confiáveis via eventos.
• Novas operações: torch.frexp, torch.logspace (por enquanto sem saída inteira e booleana para vários elementos), torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l e comparações de inteiros sem sinal. Na implementação matemática do SDPA, foi corrigido o broadcasting por batch e por cabeças.

O que vem no 2.16 e como testar

Parte dos planos para o 2.16 já está nas builds noturnas. O command buffer será commitado a cada 16 kernels, para que a GPU comece a trabalhar enquanto a CPU codifica o restante.

Na decodificação BF16 comprimida com contexto de 2048 tokens, o LFM2.5-8B-A1B sobe de 108,8 para 126,0 tokens/s (+16%), o Qwen3.6-35B-A3B de 52,9 para 62,8 (+19%) e o Qwen3.

8-27B de 12,1 para 12,6 (+4%). A atenção irá inteira para Metal. No M5 Pro, a decodificação com GQA (32 cabeças de query, 8 cabeças KV, head dim 80, contexto de 4K) cairá de 465 para 65 µs, e o prefill 4K×4K com head dim 40, de 42 para 6 ms.

Entre as outras mudanças, o bincount sobre 4 milhões de inteiros com valores quase idênticos, que leva 6,8 s, passará a levar meio milissegundo, e o libtorch_cpu.dylib emagrecerá de 384 para 287 MiB, porque os shaders passam a ser compilados para Metal 3.1 por padrão.

Na prática, o roteiro é este. É preciso um Mac com Apple silicon, macOS 15 ou mais recente e Python 3.11+; caso contrário, a verificação de disponibilidade do MPS agora explicará o motivo da recusa.

O release candidate é instalado em um ambiente separado com o comando uv pip install torch==2.15.0 --index-url https://download.

pytorch.org/whl/test/cpu. Quem roda modelos MoE via transformers deve ativar o torch.

compile e comparar os tokens por segundo com os do 2.14. Usuários do ComfyUI com Qwen-Image-2.

1 devem atualizar por causa da correção do VAE. Para donos de Mac com bastante memória que geram vídeo no Wan2.2, a atualização elimina o inchaço inesperado da memória do driver.

Se algo continuar lento ou não funcionar, a equipe pede que abram uma issue no GitHub.