PyTorch 2.14에서 bf16 그룹 GEMM에 cuBLASLt 경로를 쓰는 방법과 확인 포인트 | DAKER 커뮤니티
PyTorch 2.14에서 그룹 GEMM 경로를 살펴보고 있다면, 이번 변화는 한 번쯤 짚어볼 만합니다. Hopper·Blackwell 환경에서는 cuBLASLt 그룹 GEMM 백엔드가 제공되고, float16은 기본으로 이 경로를 탈 수 있지만 bfloat16은 별도 옵트인이 필요하기 때문입니다.
특히 같은 bf16 그룹 GEMM 모양에서도 플래그 설정에 따라 선택되는 백엔드가 달라질 수 있어, 성능 비교를 할 때는 다른 변수보다 이 차이를 먼저 분리해 보는 것이 좋습니다. 이번 글은 플래그 off/on에 따른 차이와, 정렬 조건 때문에 생길 수 있는 패딩·슬라이스 이슈를 함께 정리합니다.
PyTorch 2.14는 Hopper·Blackwell에서 cuBLASLt 그룹 GEMM 백엔드를 제공하며, bfloat16은 torch.backends.cuda.matmul.prefer_cublaslt_grouped_gemm = True로 옵트인합니다.
bf16 그룹 GEMM에서 달라지는 점
float16 적격 부하는 기본으로 cuBLASLt 그룹 GEMM 경로를 탈 수 있습니다. 반면 bfloat16은 torch.backends.cuda.matmul.prefer_cublaslt_grouped_gemm = True를 켜야 같은 경로를 사용합니다. 이 설정은 torch.compile과 CUDA Graph와도 함께 사용할 수 있습니다.
다만 행렬과 leading dimension이 16바이트 정렬이 아니면 그대로는 원하는 경로를 타지 못할 수 있습니다. 이 경우 패딩이나 슬라이스가 필요할 수 있어, 커널 자체의 속도만이 아니라 복사와 정렬 보정 비용까지 같이 봐야 합니다.
이번 실험에서 비교할 범위
이번 비교는 범위를 좁게 잡는 것이 핵심입니다. 같은 bf16 그룹 GEMM 모양을 두고 플래그를 끈 경우와 켠 경우만 비교하고, 스텝 시간과 선택된 백엔드의 흔적을 기록합니다.
즉, 어제 다룬 TunableOp 후보 수, 핀드 호스트 스냅샷, 워밍업 연장 같은 요소와는 분리해서 봅니다. 이번에는 그룹 GEMM 옵트인 효과만 확인하면 됩니다.
재현해 볼 실험 조건
- CUDA 13.3 이상·Hopper 또는 Blackwell에서 같은 bf16 그룹 GEMM을 플래그를 끈 상태로 한 번, 켠 상태로 한 번 측정합니다.
- 정렬이 어긋난 모양은 16바이트에 맞게 패딩한 뒤 다시 측정해, 실패·폴백과 성공 경로를 구분합니다.
- 가능하면
torch.compile또는 CUDA Graph 캡처 안에서도 같은 플래그 전후를 한 번씩만 비교합니다.
실무에서 먼저 볼 포인트
전역 기본값을 바로 바꾸기보다, 먼저 핫 패스에 있는 그룹 GEMM 몇 개를 골라 bf16 옵트인 효과를 확인하는 편이 안정적입니다. 같은 모델 안에서도 모든 연산이 같은 이득을 내는 것은 아니기 때문입니다.
또한 정렬과 패딩 비용을 빼고 보면 커널은 빨라졌는데 end-to-end 시간은 그대로일 수 있습니다. 그래서 측정할 때는 커널 시간뿐 아니라 패딩이나 복사에 들어간 시간도 함께 적어두는 것이 좋습니다.
TunableOp·NVGEMM 실험과 겹친다면, 먼저 이 플래그만 분리해 기록해야 비교가 쉬워집니다.
이미 TunableOp나 NVGEMM 관련 실험을 함께 돌리고 있다면, 변수들이 섞이지 않도록 이번 플래그만 따로 분리해 기록하는 편이 해석에 도움이 됩니다.
관련 자료
TunableOp cuBLASLt 후보
Inductor NVGEMM 에필로그 융합
PyTorch 실전 고급 트랙
같은 bf16 그룹 GEMM 모양에서 플래그 전후를 비교해 보셨다면, 어떤 조건에서 차이가 가장 분명했는지 함께 나눠주셔도 좋겠습니다.