PyTorch 2.14 Inductor NVGEMM, CUTLASS 에필로그 융합과 NVFP4 후보 선택 정리 | DAKER 커뮤니티
PyTorch 2.14에서 Inductor의 GEMM 경로를 볼 때, 단순히 어떤 커널이 빠른지만 보는 것으로는 부족할 때가 있습니다. mm, addmm, scaled_mm 같은 연산이 어떤 후보로 만들어지고, bias나 활성화 같은 후처리가 커널 바깥이 아니라 안으로 들어가는지가 실제 성능에 직접 연결되기 때문입니다.
이번 글은 Inductor의 NVGEMM이 CuTeDSL·CUTLASS 경로를 통해 어떤 후보를 만들고, 어떤 조건에서 에필로그 융합을 처리하는지 짧게 정리합니다. 특히 max_autotune 아래에서 Triton·ATen과 함께 경쟁한다는 점, 그리고 NVFP4 전역 스케일이 에필로그에서 처리된다는 점이 핵심입니다.
PyTorch 2.14 Inductor의 NVGEMM은 CuTeDSL·CUTLASS 경로로mm·addmm·scaled_mm후보를 만들고, bias·활성화·축소 같은 에필로그를 커널 안에 붙입니다.
여기에 더해 NVFP4 전역 스케일도 에필로그에서 처리됩니다. 이 후보들은 max_autotune 아래에서 Triton·ATen과 같이 겨룹니다.
NVGEMM에서 먼저 볼 핵심 포인트
이번 주제는 이전에 다룬 combo kernel·reorder_for_locality와는 다른 축입니다. 여기서는 GEMM 자체보다, GEMM 뒤에 붙는 에필로그가 얼마나 커널 내부로 흡수되는지를 중심으로 보면 됩니다.
실제로는 NVGEMM이 mm, addmm, scaled_mm에 대한 후보를 만들고, bias·pointwise 활성화·축소 같은 후처리를 에필로그 형태로 붙입니다. 이 구성이 성립하면 별도 커널이 줄어들 수 있고, 반대로 표현하지 못하는 에필로그는 Triton 경로로 돌아갈 수 있습니다.
표현 못 하는 에필로그는 Triton으로 돌아가니, 실패를 성능 회귀로만 볼 필요는 없습니다.
재현해 볼 실험 조건
가볍게 확인하려면 작은 GEMM 또는 addmm 벤치를 하나 고른 뒤, NVGEMM을 max_autotune_gemm_backends에 넣은 경우와 뺀 경우를 같은 형태로 비교하면 됩니다.
nvidia-cutlass-dsl4.6+가 있는 환경에서 작은mm/addmm루프를 만듭니다.max_autotune과NVGEMM백엔드를 켠 채 컴파일하고, 끈 채와 스텝 시간·커널 이름을 비교합니다.- 에필로그( bias·pointwise )가 따로 남지 않았는지 프로파일만 확인합니다. NVFP4는 Blackwell에서만 시험합니다.
실험 노트에는 PyTorch·CUTLASS DSL 버전, 도형, 선택된 백엔드 이름을 함께 남기는 것이 좋습니다. 이 경로는 API·DSL 버전에 민감할 수 있어, 재현성 확보에 버전 기록이 중요합니다.
실무에서 바로 볼 지점
학습·추론 모두에서 GEMM이 핫패스라면 autotune 후보에 NVGEMM을 올려 보는 것이 좋습니다. 특히 작은 GEMM이나 addmm에서 에필로그 융합 여부가 체감 차이로 이어질 수 있습니다.
다만 모든 경우에 NVGEMM이 선택되는 것은 아닙니다. max_autotune 아래에서 Triton·ATen과 함께 경쟁하므로, 실제 선택 결과는 도형과 환경에 따라 달라질 수 있습니다. 그래서 on/off 비교와 커널 이름 확인을 같이 보는 편이 안전합니다.
작은 GEMM·addmm 벤치를 하나 고른 뒤,NVGEMM을max_autotune_gemm_backends에 넣고 끈 상태를 같은 형태로 비교하면 됩니다.
관련 DAKER 학습
참고 자료
https://daker.ai/public/learning/tracks/pytorch-practice-advanced-track
https://daker.ai/community
여러분은 GEMM autotune을 볼 때 커널 자체와 에필로그 융합 중 어느 쪽이 더 큰 차이를 만든다고 느끼는지 궁금합니다.