Fused MLP 설계를 통한 HBM 트래픽 제거 및 커널 최적화

Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP

2026년 6월 11일18분advanced

AI 요약

Context

PyTorch의 nn.Linear와 MLP 구조에서 발생하는 CPU Dispatch 오버헤드 및 GPU 메모리 병목 현상을 분석함. 특히 개별 커널 호출로 인한 HBM(High Bandwidth Memory)의 빈번한 데이터 읽기/쓰기가 전체 성능의 병목 지점으로 작용하는 한계를 식별함.

Technical Solution

GEMM Epilogue 활용을 통한 Bias-add 연산의 Matmul 커널 내 통합으로 메모리 액세스 횟수 최소화
aten::t(Transpose) 연산을 CPU 상의 Metadata 수정으로 처리하여 GPU 커널 런칭 오버헤드 제거
torch.compile 도입을 통해 GeLU, Mul, Reshape 연산을 단일 Triton 커널로 융합하여 중간 결과값을 레지스터 수준에서 유지
Liger Kernel 채택으로 Dynamo 및 Guard로 인한 Compile Latency를 제거하고 하드웨어 튜닝된 Launch Parameter 적용
Generic Kernel의 유연성과 Shape-specialized Kernel의 성능 사이의 Trade-off를 고려한 최적 커널 선택

실천 포인트

1. 단일 nn.Linear 최적화 시 torch.compile보다 내장된 addmm 커널의 Epilogue 활용 여부를 먼저 확인

2. MLP 구조 설계 시 활성화 함수와 요소별 연산이 분리되어 HBM 라운드트립이 발생하는지 Profiler로 검증

3. 런타임 Shape 변경이 잦은 환경에서는 Compile-time 특화 커널보다 견고한 Generic Fused Kernel 도입 검토

태그

#Triton #PyTorch Profiler #GEMM #HBM #Kernel Fusion

원문 읽기