Hugging Face의 PyTorch 프로파일링 시리즈 2편으로, nn.Linear부터 GeGLU MLP까지 eager·torch.compile·Liger 수제 커널을 A100에서 비교 분석한다. nn.Linear의 bias 덧셈은 이미 GEMM 에필로그(addmm)로 융합돼 단일 cuBLAS 커널로 실행되므로 단일 Linear에서는 torch.compile이 융합할 것이 없고, CPU 디스패치 오버헤드(aten::t 뷰 처리)만 줄어든다. MLP(GEMM 3개 + GeLU + mul)에서는 eager가 5개 GPU 커널을 띄우고 [8192,3072] 중간 텐서가 HBM을 왕복하는데, torch.compile은 GeLU·mul·reshape를 단일 Triton 커널(triton_poi_fused__unsafe_view_gelu_mul_0)로 융합해 중간값을 레지스터에 유지한다. Liger 수제 커널은 같은 융합을 내장하되 Dynamo·가드·컴파일 지연이 없어, 89.4µs인 Inductor보다 92.8µs로 약간 느리지만 입력 형태가 바뀌어도 재컴파일이 필요 없다는 점에서 강건하다.
- •nn.Linear의 bias 덧셈은 GEMM 에필로그(addmm)로 융합돼 별도 add 커널 없이 단일 cuBLAS 커널로 실행되며, aten::t 전치는 GPU 커널 없이 스트라이드 메타데이터만 재작성한다
- •단일 Linear에서는 torch.compile이 융합할 GPU 작업이 없어 동일한 cuBLAS GEMM 커널을 쓰고 CPU 디스패치 오버헤드만 제거한다
- •eager MLP는 GEMM 3개·GeLU·mul로 총 5개 GPU 커널을 실행하며 약 50MB의 [8192,3072] 중간 텐서가 HBM을 왕복한다
- •torch.compile은 GeLU·mul·reshape를 단일 Triton 커널로 융합해 중간값을 레지스터에 유지하고 HBM 왕복을 제거하지만 정적 형태에 특화돼 형태 변경 시 재컴파일 비용이 든다
- •Liger 수제 커널(92.8µs)은 Inductor 융합 커널(89.4µs)보다 약간 느리나 Dynamo·가드·컴파일 지연 없이 어떤 형태에도 재컴파일 없이 동작해 강건하다
Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP
- 1.nn.Linear의 bias 덧셈은 addmm 에필로그로 융합돼 별도 add 커널 없이 단일 cuBLAS GEMM으로 실행
- 2.aten::t 전치는 텐서 stride 메타데이터만 재작성할 뿐 GPU 커널을 띄우지 않음(CUDA 0µs)
- 3.GeGLU MLP는 eager에서 GEMM 3개+GeLU+mul 등 5개 커널 실행, 중간 텐서가 HBM을 왕복
- 4.torch.compile은 GeLU·mul·reshape를 단일 Triton 커널로 융합해 HBM 왕복을 제거
왜 중요한가?
단일 Linear에는 torch.compile이 융합할 여지가 거의 없지만, MLP처럼 연산이 쌓이면 pointwise 융합으로 중간 텐서의 HBM 왕복을 없애 성능이 개선된다. Liger의 수작업 Triton 커널(92.8µs)은 컴파일 커널(89.4µs)보다 근소하게 느려도 형상 변화 시 재컴파일이 불필요해, 프로파일 트레이스로 커널 선택을 판단하는 실전 기준을 보여준다.
🏷️ 언급 프로젝트
전체 내용이 궁금하다면?
원문을 직접 읽어보세요
