MPS in PyTorch 2.15

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

PyTorch 2.15는 Mac의 MPS에 FP8 지원과 Metal 커널 최적화를 도입해 LLM 디코딩, 영상 메모리 사용량, 이미지 처리 정확도와 연산 지원을 개선한다.

📌 핵심 요약

• PyTorch 2.15의 MPS는 Apple silicon Mac과 macOS 15 이상, Python 3.11 이상을 요구한다. torch.float8_e4m3fn은 값당 1바이트로 FP16·BF16의 절반 공간을 사용하며, F.scaled_mm으로 텐서별·행별 스케일링을 지원하는 GPU 행렬 곱셈을 수행한다.
• M3 Ultra와 torch.compile 환경에서 MoE 디코딩은 2.14 대비 약 5~6.5배 빨라져 Qwen3.6-35B-A3B는 9.4→61.3 tokens/s, LFM2.5-8B-A1B는 23.6→123.6 tokens/s를 기록했다. F.linear의 BF16 실제 모델 레이어 96개 형상에서는 기하평균 1.23배 향상이 있었고, [1, 128256] FP32 softmax는 57→7µs로 약 8배 빨라졌다.
• Metal 상수 패딩 적용으로 Wan2.2의 28타일 VAE 디코딩에서 최대 드라이버 메모리가 270.135→8.108 GiB로 감소했다. 같은 패딩 변경은 Qwen-Image-2.1 VAE 인코더의 이미지 손상을 수정해 CPU 기준 인코딩·디코딩 왕복 PSNR을 10.06→29.95 dB로 개선했다.
• complex64 LU 분해·풀이와 삼각 풀이, 작은 행렬 역행렬, QR 및 SVD가 개선됐다. 합성곱·패딩·드롭아웃·업샘플링 역전파에도 최적화가 적용됐으며, Python에서 MPS 스트림을 제어할 수 있고 추가 연산과 attention 브로드캐스팅 정확도 수정이 제공된다.
• PyTorch 2.16에서는 커널 16개마다 명령 버퍼를 제출하는 변경이 소개됐다. 2048토큰 문맥의 컴파일된 BF16 디코딩에서 LFM2.5-8B-A1B는 108.8→126.0 tokens/s, Qwen3.6-35B-A3B는 52.9→62.8 tokens/s로 개선됐으며, 소개된 변경 중 일부는 main·nightly에 반영됐고 일부는 검토 중이어서 출시 전 달라질 수 있다.

🧩 주요 포인트

1. MPSGraph 경로를 전용 Metal 커널로 교체 → 토큰 생성 지연, 패딩 메모리 급증과 이미지 정확도 문제를 함께 개선.
2. 디코딩 입력과 행렬 형상에 맞춘 최적화 → MoE와 작은 투영층의 이득이 크며, 이미 메모리 대역폭 한계에 가까운 큰 레이어는 개선 폭이 제한적.
3. FP8·선형대수·스트림·연산 지원 확대 → Mac에서 실행 가능한 작업 범위와 제어 수단이 늘지만, 운영체제·Python 요구사항과 2.16의 미확정 상태를 고려할 필요.

🧠 상세 정리

1. MPS 개선의 배경과 실행 요구사항

2026년 10월 20일 공개된 이 글은 Hugging Face의 PyTorch 팀이 PyTorch 유지관리자와 커뮤니티 기여자들과 함께 진행한 MPS 개선을 설명한다. 팀은 2026년 초만 해도 작은 모델에서 미지원 연산, 급격한 성능 저하, 예상하지 못한 메모리 증가를 우회해야 하는 경우가 많았다고 전한다. 이에 Hugging Face Hub의 언어·이미지·영상 모델을 실제로 실행해 문제를 찾고, MPSGraph를 전용 Metal 커널로 대체하는 등의 수정을 PyTorch에 반영해 왔다. PyTorch 2.15의 MPS 실행에는 Apple silicon Mac과 macOS 15 이상, Python 3.11 이상이 필요하다. 운영체제가 오래됐거나 Metal GPU가 없으면 MPS 가용성 검사가 그 이유를 알려주며, 글은 별도 환경에서 릴리스 후보를 설치해 시험하는 방법도 제시한다.

2. FP8 저장과 스케일 행렬 곱셈

PyTorch 2.15는 MPS에 torch.float8_e4m3fn 텐서를 추가해 Mac에서 모델 데이터를 더 적은 공간에 저장할 수 있게 한다. 이 형식은 값 하나에 1바이트를 사용하므로 FP16이나 BF16에 필요한 저장 공간의 절반을 차지한다. F.scaled_mm은 FP8 텐서를 GPU에서 곱하며, 텐서 전체에 하나의 스케일을 적용하는 방식과 행마다 스케일을 적용하는 방식을 모두 지원한다. 원문의 예제는 입력의 최대 절댓값과 FP8 표현 한계로 스케일을 계산한 뒤 입력과 가중치를 변환하고, BF16 결과를 출력하는 행렬 곱셈을 보여준다. 글이 여기서 강조하는 효과는 저장 공간 절감과 연산 지원이며, FP8 사용에 따른 구체적인 처리량 향상이나 모델 정확도 측정치는 제시하지 않는다.

3. LLM 디코딩의 선형층과 MoE 최적화

LLM 디코딩에서는 새 토큰이 한 행의 입력으로 모델을 통과하므로 큰 행렬 곱셈이 행렬·벡터 곱셈으로 바뀌며, 2.15는 기존 torch.mm의 GEMV 커널을 F.linear에도 연결하고 가능한 경우 편향 덧셈을 융합한다. 실제 모델의 BF16 레이어 96개 형상에서 기하평균 1.23배 빨라졌지만, 이득은 작은 투영층에서 크고 이미 대역폭 한계에 가까운 큰 레이어와 LM 헤드는 약 1배 수준에 머문다. MoE용 네이티브 grouped matmul은 2.14의 그룹 오프셋 CPU 복사와 그룹별 반복을 GPU에서 직접 오프셋을 읽는 Metal 경로로 대체하고, 디코딩에 흔한 작은 그룹에는 별도 커널을 사용한다. M3 Ultra와 torch.compile 환경의 전체 개선 결과로 Qwen3.6-35B-A3B는 9.4에서 61.3 tokens/s, LFM2.5-8B-A1B는 23.6에서 123.6 tokens/s로 빨라졌다. 원문은 이 전체 개선을 약 5~6.5배로 소개하는 한편, 별도 디코딩 커널 설명에서는 약 1.3배의 토큰 처리량 향상도 언급하므로 두 수치를 같은 비교로 단정해서는 안 된다.

4. 복소수와 다양한 크기의 선형대수 개선

complex64의 LU 분해, LU 풀이와 삼각 풀이가 MPS에서 실행되면서 복소수 행렬식처럼 이 연산에 의존하는 기능과 기울기 계산도 지원 범위가 넓어진다. 카메라 기하와 이미지 변환에 쓰이는 최대 4×4 FP32 행렬 배치는 분해와 풀이를 하나의 Metal 커널에 결합해 역행렬을 계산한다. 더 큰 문제에서는 QR 분해에 블록 단위 갱신을 적용하고, 삼각 풀이의 많은 작업을 효율적인 행렬 곱셈으로 바꾼다. Gated DeltaNet에는 우변이 최소 8개인 16×16, 32×32, 64×64 FP32 삼각 시스템 배치를 위한 전용 고속 경로가 마련됐다. SVD도 매우 작은 특이값, 매우 크거나 작은 입력 크기, full_matrices=True에서의 기저 완성을 올바르게 처리하도록 수정돼 의사역행렬과 최소제곱 결과의 신뢰성이 개선된다.

5. 오디오·영상·학습 연산의 Metal 경로

오디오 모델, 영상 디코더와 학습 루프에서 비용이 큰 연산에도 전용 Metal 경로가 추가됐다. conv1d 순전파는 MPSGraph 대신 특화 커널을 사용하고, 출력이 짧은 경우에는 im2col과 행렬 곱셈을 결합한 경로를 제공해 오디오와 스트리밍 모델의 형상에 대응한다. 상수 패딩은 Metal로 이동해 영상 VAE의 메모리 급증을 해결하며, 반사·복제 패딩의 새 구현은 속도와 역전파 정확도를 개선한다. 별도 출력을 만드는 nn.Dropout과 F.dropout은 기존 융합 커널을 사용하고, 업샘플링 역전파는 MPSGraph의 크기 조정 기울기 경로를 Metal gather 커널로 대체한다. channels-last conv2d 역전파도 접근 비용이 큰 NCHW 그래프를 거치던 방식에서 효율적인 channels-last 경로를 유지하는 방식으로 바뀐다.

6. 어휘 softmax와 토큰 샘플링

PyTorch 2.15는 토큰 샘플링에 쓰이는 softmax와 log_softmax를 전용 Metal 커널로 옮긴다. 원문은 큰 어휘에서 이득이 특히 크다고 설명하며, [1, 128256] 형상의 FP32 로짓 한 행에 대한 softmax 실행 시간이 MPSGraph의 57µs에서 7µs로 줄었다고 제시한다. 이는 해당 연산 기준 약 8배의 속도 향상이며, 글은 다른 일반적인 어휘 크기에서도 개선이 이어진다고 설명한다. 이 결과는 디코딩 선형층과 MoE 커널 개선에 더해 토큰 생성 과정의 다른 비용도 줄였다는 의미가 있지만, softmax의 배수를 모델 전체 처리량 향상으로 그대로 해석할 근거는 없다. 특히 원문에서 제시한 57µs와 7µs는 특정 형상과 FP32 자료형에 대한 측정이므로 그 조건을 함께 보존할 필요가 있다.

7. 영상 메모리 급증과 이미지 손상 수정

상수 패딩의 Metal 전환은 Wan-AI/Wan2.2-TI2V-5B-Diffusers의 VAE 디코딩에서 두드러진 결과를 보였다. 전체 28타일 디코딩 동안 패딩이 8,764번 호출되는 이 작업에서 최대 드라이버 메모리는 270.135 GiB에서 8.108 GiB로 감소했다. 반면 텐서 할당 최대치는 4.422 GiB에서 4.095 GiB로 소폭 줄었고 할당자가 예약한 메모리의 최대치는 8.070 GiB로 같아, 표에 제시된 큰 변화는 드라이버 메모리 지표에 집중된다. 같은 패딩 변경은 2.14에서 Qwen-Image-2.1 VAE 인코더가 MPS 이미지를 조용히 손상시키던 문제도 해결한다. ComfyUI의 이미지 편집 흐름에서 보고된 이 문제를 수정한 뒤 CPU를 기준으로 한 인코딩·디코딩 왕복 PSNR은 10.06 dB에서 29.95 dB로 개선됐다.

8. Python 스트림 제어와 연산 지원 확대

이제 Python에서 torch.mps.Stream과 torch.mps.stream 문맥 관리자를 사용해 작업을 실행할 MPS 스트림을 선택할 수 있다. Tensor.record_stream은 다른 스트림에서 텐서가 사용되는 사실을 할당자가 추적하게 하며, 예제는 별도 스트림에 연산을 보내고 결과를 읽기 전에 해당 스트림을 동기화한다. 글은 이벤트 시간 측정의 신뢰성 개선도 주요 변화로 소개한다. 추가 지원 연산에는 torch.frexp, torch.logspace, torch.take, Tensor.put_, torch.heaviside, torch.special.laguerre_polynomial_l과 부호 없는 정수 비교가 포함된다. 다만 logspace는 원소가 둘 이상인 정수형·불리언 출력을 아직 지원하지 않으며, scaled dot-product attention의 MPS 수학 구현은 배치와 헤드 브로드캐스팅을 올바르게 처리하도록 수정됐다.

9. PyTorch 2.16의 디코딩 개선 예고와 불확실성

원문은 2.16에서 attention의 Metal 전환과 추가 MPSGraph 대체, LLM 디코딩 개선을 예고하며, 일부 변경은 main과 nightly에 있지만 일부는 검토 중이어서 출시 전 세부 사항이 달라질 수 있다고 명시한다. 구체적으로 소개한 문제는 디코딩 경로에 MPSGraph 연산이 없어지면 CPU가 모든 커널 인코딩을 끝낼 때까지 GPU가 해당 토큰의 작업을 시작하지 못한다는 점이다. 이를 개선하기 위해 2.16에서는 커널 16개마다 명령 버퍼를 제출해 CPU가 나머지를 인코딩하는 동안 GPU가 작업을 시작하게 한다. 2048토큰 문맥에서 컴파일된 BF16 디코딩을 측정한 결과 LFM2.5-8B-A1B는 108.8에서 126.0 tokens/s로 16%, Qwen3.6-35B-A3B는 52.9에서 62.8 tokens/s로 19% 개선됐다. 이 비교는 앞서 제시한 2.15 대 2.14 수치와 별도로 읽어야 하며, 제공된 원문은 이 표 이후에서 끊겨 후속 내용까지 확인할 수는 없다.

🧾 핵심 주장 / 시사점

• 전용 Metal 커널의 효과는 속도에만 한정되지 않는다. 상수 패딩 변경처럼 드라이버 메모리 급증과 이미지 정확도 문제를 함께 해결하는 경우도 있다.
• 최적화의 이득은 작업 형상과 기존 병목에 따라 다르다. 작은 투영층과 MoE 디코딩의 큰 개선을 이미 대역폭 한계에 가까운 모든 레이어에 일반화하기는 어렵다.
• Wan2.2에서는 텐서 할당량과 예약 메모리 변화가 작아도 드라이버 메모리는 크게 감소했다. 서로 다른 메모리 지표를 구분해야 개선의 실제 범위를 파악할 수 있다.

✅ 액션 아이템

• PyTorch 2.15 적용 환경의 Apple silicon Mac, macOS 15 이상, Python 3.11 이상 충족 여부 확인.
• Qwen3.6-35B-A3B와 LFM2.5-8B-A1B의 M3 Ultra·torch.compile 측정 조건을 고려해 사용 환경의 디코딩 개선 폭 검증.
• Wan2.2의 최대 드라이버 메모리와 Qwen-Image-2.1의 CPU 기준 PSNR을 비교해 패딩 변경 효과 확인.

❓ 열린 질문

• FP8의 FP16·BF16 대비 절반 저장 공간 이점은 사용하려는 모델의 전체 메모리 사용량에 얼마나 반영되는가?
• MoE와 작은 투영층의 개선이 큰 반면 큰 레이어의 이득은 제한적인 상황에서, 실제 사용 모델의 디코딩 처리량은 얼마나 개선되는가?
• PyTorch 2.16에서 소개된 커널 16개마다 명령 버퍼를 제출하는 변경은 정식 출시에서 어떤 상태와 성능으로 제공되는가?