Hugging Face Transformers가 Hub에서 다운로드 가능한 커스텀 커널과 MXFP4 양자화를 통합해 GPT-OSS 모델의 로딩·추론·파인튜닝 성능을 2~10배 향상

Context

커뮤니티에서 개발된 Flash Attention, Liger RMSNorm, MegaBlocks MoE 등의 커스텀 커널들이 서로 다른 라이브러리에 산재되어 있어 의존성 증가와 CUDA/C++ 컴파일 요구사항이 발생했다. 각 모델 통합 시마다 새로운 커널 라이브러리를 추가해야 하는 구조로 인해 라이브러리 복잡도가 증가하고 있었다.

Technical Solution

• Zero-build Kernels 패키지 도입: Hub에서 사전 컴파일된 커널 바이너리를 다운로드하고 @use_kernel_forward_from_hub() 데코레이터로 자동 선택하는 구조로 변경
• Liger RMSNorm 커널 통합: @use_kernel_forward_from_hub("RMSNorm") 데코레이터로 정규화 연산 최적화
• MegaBlocks MoE 커널 통합: @use_kernel_forward_from_hub("MegaBlocksMoeMLP") 데코레이터로 Mixture of Experts 연산 가속
• Flash Attention 3 통합: Attention Sinks를 지원하는 Flash Attention 3 커널을 Hopper 아키텍처 대상으로 추가
• MXFP4 양자화 커널 추가: Triton 기반 MXFP4 양자화 연산을 커스텀 커널로 제공
• 디바이스 자동 로딩 최적화: device_map="auto" 또는 Tensor Parallel 실행 시 멀티 GPU 로딩 속도 개선
• 커뮤니티 커널 자동 선택: CUDA/ROCm 여부 및 훈련/추론 모드에 따라 호환 커널을 자동 선택

Impact

• PyTorch 2.0의 torch.compile과 TorchInductor 백엔드는 2~10배 성능 향상 제공
• 커스텀 커널 사용 시 더 큰 배치 크기에서 최적 성능 달성 (Figure 1 벤치마크 결과)

Key Takeaway

커스텀 커널을 중앙 리포지토리(Hub)에서 사전 컴파일 바이너리로 배포하고 데코레이터 패턴으로 추상화하면, 의존성 증가와 컴파일 오버헤드를 제거하면서도 여러 모델에서 재사용 가능한 최적화 기법을 확산할 수 있다. 이는 커뮤니티 기여 커널을 참조 구현으로 제공함으로써 MLX, llama.cpp, vLLM 같은 다른 프레임워크의 학습 자료로도 활용된다.