PyTorch 2.14에서 NVGEMM을 max_autotune GEMM 후보에 넣으면 달라지는 점 | DAKER 커뮤니티
PyTorch 2.14에서 행렬곱 뒤 후처리까지 한 번에 묶을 수 있는지 보려면, 먼저 Inductor가 어떤 GEMM 백엔드를 후보로 올리는지부터 확인하는 것이 좋습니다. NVGEMM은 기본적으로 이름만 알려진 기능이 아니라, max_autotune 환경에서 실제 비교 대상에 들어가야 의미가 생깁니다.
특히 mm, addmm, scaled_mm 뒤에 bias나 activation 같은 후처리가 붙는 경우라면, 커널이 나뉘는지 한 번에 처리되는지가 성능에 직접 연결될 수 있습니다. 이 글은 NVGEMM이 어떤 조건에서 후보가 되고, epilogue fusion이 어디까지 가능한지 원문 기준으로 정리한 내용입니다.
NVGEMM이 하는 일
PyTorch 2.14의 NVGEMM은 Inductor가 mm, addmm, scaled_mm 후보를 고를 때 Triton, ATen과 함께 경쟁하는 CuTeDSL 기반 GEMM 백엔드입니다.
NVGEMM은 max_autotune 아래에서 max_autotune_gemm_backends에 넣어야 실제 후보로 비교됩니다.
즉, 후보 목록에 NVGEMM이 없으면 이 경로는 아예 선택 과정에 들어가지 않습니다. 행렬곱 뒤 activation을 따로 돌리고 있다면, 먼저 커널이 몇 번 뜨는지부터 확인해 볼 만합니다.
재현 조건
원문에서 제시한 조건은 단순합니다. PyTorch 2.14와 호환 CUTLASS DSL 패키지 환경을 준비하고, torch.compile에서 max autotune을 켠 뒤, GEMM 백엔드 목록에 NVGEMM을 포함하면 됩니다.
이 조건에서 Inductor는 NVGEMM 후보를 Triton, ATen과 함께 비교합니다. 그리고 표현 가능한 epilogue는 커널 안으로 합칩니다. 반대로 표현하지 못하는 epilogue는 Triton 쪽으로 떨어져 기존 fusion을 유지합니다.
bias, pointwise, reduction epilogue는 같은 커널에 붙일 수 있습니다.
epilogue fusion이 의미하는 것
핵심은 행렬곱만 빠르게 끝내는 것이 아니라, 그 뒤에 이어지는 후처리까지 한 커널 안에서 처리할 수 있느냐입니다. 원문 기준으로 NVGEMM은 bias, pointwise, reduction epilogue를 같은 커널에 붙일 수 있습니다.
다만 모든 후처리가 항상 NVGEMM 안으로 들어가는 것은 아닙니다. Inductor가 표현하지 못하는 epilogue는 Triton 경로로 내려가며, 이 경우에는 기존 fusion 방식이 유지됩니다.
성능을 볼 때 함께 확인할 점
저정밀, 스케일 GEMM 경로에서도 pointwise 후처리가 epilogue로 들어갈 수 있습니다. 다만 일부 초저정밀 경로는 특정 GPU 세대 제약이 있습니다.
또 하나는 autotune 비용입니다. 후보가 늘어나면 autotune 시간도 늘어날 수 있으니, 벤치를 볼 때는 warm compile과 steady-state를 구분하는 편이 좋습니다. 디스크 캐시가 있다면 재컴파일 비용은 줄어들 수 있지만, 클린 환경과 캐시 환경을 섞어서 보고하지 않는 것이 좋습니다.
관련 자료
함께 보면 좋은 DAKER 학습 자료는 아래와 같습니다.
참고 자료
https://daker.ai/public/learning/materials/pytorch-advanced-custom-loop-amp-debugging-1
https://daker.ai/public/learning/materials/pytorch-advanced-deployment-torchscript-onnx-perfo
실제로 NVGEMM이 후보에 들어간 뒤 커널 수나 후처리 방식이 어떻게 달라졌는지, 여러분 환경에서는 어떤 차이가 있었는지 궁금합니다.