Hugging Face와 AMD가 MI300X GPU용 커스텀 커널 3개를 개발해 Llama 3.1 405B FP8 추론의 디코딩 레이턴시를 최대 141.88% 단축

Context

새로운 AI 모델 최적화는 주로 NVIDIA GPU를 먼저 지원하는 "Day 0 커널"로 구현되어, AMD GPU 같은 다른 하드웨어 플랫폼이 배제되는 문제가 있었다. ChatGPT 수준의 대규모 모델 추론에서는 커널 레벨의 최적화가 1% 지연시간 또는 전력 감소도 상당한 비용 절감으로 이어지므로 중요하다.

Technical Solution

• Fused Residual Connection + RMS Norm + FP8 변환 커널 개발: 잔차 연결, 정규화, 데이터 타입 변환을 하나의 연산으로 통합해 메모리 접근 횟수 감소
• Fused SwiGLU 활성화 함수 + FP8 변환 커널 개발: 게이트 선형 유닛 활성화와 양자화를 단일 커널로 결합
• Skinny GEMM 커널 개발: M(행) 크기가 작은 행렬 곱셈(QKV 프로젝션, Gate/Up 프로젝션, Down 프로젝션)을 MI300X MFMA 명령어에 맞춰 최적화
• hf-rocm-kernels 리포지토리 공개: 세 커널의 소스 코드, Python 바인딩, 벤치마크 스크립트, 테스트 스위트 제공으로 커뮤니티가 커널을 독립적으로 사용하거나 커스텀 커널을 개발할 수 있도록 지원
• VLLM AMD 포크 통합: 최적화된 커널을 VLLM에 통합하여 추론 엔진이 자동으로 활용

Impact

• Skinny GEMM 커널: M=1일 때 QKV 프로젝션(N=2304, K=16384)에서 127.84% 속도 개선, Gate/Up 프로젝션(N=13312, K=16384)에서 141.88% 속도 개선
• M=8일 때 QKV 프로젝션에서 132.07% 속도 개선, Gate/Up 프로젝션에서 137.15% 속도 개선
• M=16일 때 QKV 프로젝션에서 120.02% 속도 개선, Gate/Up 프로젝션에서 125.45% 속도 개선
• M=32일 때 성능 향상이 감소하여 QKV 프로젝션에서 98.80%, Gate/Up 프로젝션에서 118.30% 개선에 그침

Key Takeaway

GPU 커널 최적화는 하드웨어 메모리 계층(레지스터, 캐시, 로컬 메모리)과 명령어 집합(MFMA)을 깊이 있게 이해한 후 구체적인 행렬 크기와 연산 특성에 맞춰 튜닝해야 하며, 오픈소스 공개를 통해 다른 플랫폼에 최적화 기법을 확산할 수 있다.