PyTorch 2.15, MPS에서 FP8 지원 및 MoE 모델 디코딩 5~6.5배 가속
MPS in PyTorch 2.15
Hugging Face의 PyTorch 팀은 PyTorch 2.15의 MPS 백엔드에서 FP8을 지원하고, LLM 디코딩을 가속하며, Metal 커널로의 전환을 진행했다고 발표했다. M3 Ultra에서는 Qwen3.6-35B-A3B가 2.14의 9.4토큰/초에서 61.3토큰/초로 향상됐다. Wan2.2의 비디오 VAE 디코딩에서는 드라이버 메모리 피크가 270 GiB에서 8 GiB로 떨어졌다.
FP8에서는 torch.float8_e4m3fn 텐서를 MPS에서 사용할 수 있게 되어 BF16의 절반 용량으로 보관할 수 있다. F.scaled_mm은 tensorwise 및 rowwise 스케일링을 지원한다. LLM 디코딩에서는 단일 행 linear 레이어에 GEMV 커널을 연결해 96종류의 BF16 레이어 형태에서 기하 평균 1.23배가 됐다. MoE용으로는 grouped matmul의 네이티브 구현을 추가하고 디코딩 전용 커널도 마련했다. 어휘 크기가 큰 softmax는 Metal 커널이 되어 [1, 128256] FP32 행에서 57 µs에서 7 µs로 단축됐다.
선형대수에서는 complex64의 LU 분해, LU solve, 삼각 solve가 MPS에서 동작한다. 4×4 이하 FP32 행렬의 역행렬은 단일 Metal 커널로 처리되며, QR은 블록화된 업데이트를 사용한다. SVD도 극단적으로 작은 특이값이나 크기가 극단적인 입력을 올바르게 처리하도록 개선됐다. 음성, 동영상, 학습용으로는 conv1d forward, constant padding, reflection·replication padding, dropout, upsampling backward 등이 Metal을 통해 처리된다. Wan2.2의 타일형 VAE 디코딩은 8,764회의 padding 호출을 포함하는데, Metal화로 메모리 폭주가 해소됐다. 2.14에서 Qwen-Image-2.1의 VAE 인코더가 이미지를 깨뜨리던 문제도 같은 padding 변경으로 해결되어 왕복 PSNR이 10.06 dB에서 29.95 dB가 됐다.
Python에서는 torch.mps.Stream과 torch.mps.stream으로 MPS 스트림을 지정할 수 있으며 Tensor.record_stream도 사용할 수 있다. frexp, logspace 등 여러 연산도 MPS에서 동작한다. 실행 환경은 macOS 15 이상과 Python 3.11 이상이 필요하게 됐다. 릴리스 후보는 uv pip install torch==2.15.0 --index-url https://download.pytorch.org/whl/test/cpu 명령으로 시험해 볼 수 있다.
PyTorch 2.16에서는 10여 개 커널마다 커맨드 버퍼를 커밋하는 방식으로 디코딩을 더 빠르게 할 계획이며, Qwen3.6-35B-A3B는 52.9에서 62.8토큰/초가 된다. 어텐션은 전면적으로 Metal로 옮겨지며, MPSGraph에서의 이전은 batch norm 등을 중심으로 계속된다. Metal 3.1을 기본값으로 설정하면 libtorch_cpu.dylib 크기가 384 MiB에서 287 MiB로 줄어들 것으로 예상된다.
MPS 백엔드의 변경 사항을 MoE 모델의 토큰/초, 메모리 사용량, 필요 환경 수치로 구체적으로 확인할 수 있어 Mac에서의 활용 여부를 판단하는 데 도움이 된다.
출처: Hugging Face Blog · huggingface.co