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·가드·컴파일 지연 없이 어떤 형태에도 재컴파일 없이 동작해 강건하다
0단 자동
AI가 규칙대로 쓰고 그대로 게시했습니다. 사람이 따로 보지 않았습니다.
- 규칙 판
- 규칙 판 도입 이전 기사입니다.
- 남기는 것
- 규칙 판 · 모델 · 시각
- 판 기록
- 아직 없습니다.
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)보다 근소하게 느려도 형상 변화 시 재컴파일이 불필요해, 프로파일 트레이스로 커널 선택을 판단하는 실전 기준을 보여준다.
파이토치 모델의 성능 최적화는 한국 AI 개발 현장에서 GPU 자원 활용 효율성을 높이는 데 핵심적인 요소입니다. 특히 'nn.Linear'와 같은 기본 모듈을 최적화하여 퓨즈드 MLP를 만드는 기법은 실제 서비스 배포 시 속도와 비용 절감에 직접적인 영향을 미칠 수 있습니다. 국내 AI 스타트업 및 대기업 연구팀들이 대규모 모델을 운용하며 겪는 성능 저하 문제를 해결하는 데 중요한 시사점을 제공할 것입니다.
전체 내용이 궁금하다면?
원문을 직접 읽어보세요
이 글이 만들어진 과정
- 21:45AI 초안
