Hugging Face 로 ROCm 커널을 쉽게 구축하고 공유하기

맞춤형 커널 (Custom kernels) 은 고성능 딥러닝의 핵심으로, 작업에 맞춰 GPU 연산을 정밀하게 가능하게 합니다. 이미지 처리, 텐서 변환, 또는 다른 계산 집약적 작업이든 간에 해당됩니다. 하지만 이러한 커널을 올바른 아키텍처를 위해 컴파일하고, 모든 빌드 플래그를 연결하며, PyTorch 확장에 깔끔하게 통합하는 것은 CMake/Nix, 컴파일러 오류, ABI 문제의 혼란으로 인해 매우 즐겁지 않습니다. Hugging Face 의 kernels 라이브러리는 kernel-builder 를 사용하여 이러한 커널을 구축하고 kernels-community 와 공유하는 것을 쉽게 만들어줍니다. CUDA, ROCm, Metal, XPU 를 포함한 여러 GPU 및 가속기 백엔드를 지원하여 PyTorch 와 원활하게 통합되어 빠르고 포터블하며 사용자의 커널이 됩니다.

이번 가이드에서는 ROCm 호환 커널에 집중하고, kernels 를 사용하여 구축, 테스트, 공유하는 방법을 보여줍니다. AMD GPU 에서 효율적으로 실행되는 커널을 만드는 방법과 재현성, 패키징, 배포를 위한 모범 사례를 배우게 됩니다.

이 ROCm 특화 워크스루는 원래 kernel-builder 가이드의 간소화된 버전입니다. 더 넓은 CUDA 중심 버전을 찾으시는 경우, 여기에서 찾을 수 있습니다: A Guide to Building and Scaling Production-Ready CUDA Kernels.

우리는 RadeonFlow_Kernels 의 GEMM 커널을 예로 사용합니다. 가이드에 바로 가기 원하시다면, 여기 클릭하세요.

이 섹션은 RadeonFlow GEMM kernel 작성자들에 의해 작성되었습니다.

저작자: ColorsWind, Zesen Liu, Andy

RadeonFlow GEMM 커널은 AMD Instinct MI300X GPU 를 위해 최적화된 고성능, FP8 블록 기반 행렬 곱셈 구현입니다. GEMM (General Matrix Multiplication) 은 대부분의 딥러닝 워크로드 뒤에 있는 핵심 구성 요소입니다: 두 행렬 A 와 B 를 주면, 그들 곱 C = A × B 를 계산합니다. 여기서는 FP8 에서 구현됩니다. 저정밀도 부동소수점 형식으로 정확도를 조금 희생하여 훨씬 높은 처리량과 낮은 메모이 대역폭을 제공합니다. 이 커널은 AMD Developer Challenge 2025 에 개발되었습니다. 2025 년 6 월에 🏆 Grand Prize 를 수상하여 AMD 하드웨어에서의 성능과 혁신의 우수성을 인정받았습니다.

커널은 e4m3fnuz 부동소수점 형식을 사용하여 양자화된 입력을 작동시키고, 저정밀도 계산 동안 정확도를 보존하기 위해 블록별 스케일링을 적용합니다. e4m3fnuz 형식은 4 지수 비트와 3 만다나 비트를 가진 FP8 변형으로, 신경망 워크로드에 효율적으로 설계되었습니다. FP16/FP32 보다 훨씬 작은 동적 범위를 가진 FP8 이므로, 각 블록의 값이 계산 전후로 수치적으로 "편안한" 범위에서 재스케일되도록 블록별 스케일링 인자 (a_scale 와 b_scale) 를 적용합니다. 이는 저정밀도에도 불구하고 정확도를 보존하는 데 도움이 됩니다. 다음 인수들을 취합니다:

(a, b, a_scale, b_scale, c)

여기서 a 와 b 는 입력 행렬이며, a_scale 와 b_scale 는 각각 a 와 b 의 스케일링 인자이고, c 는 출력 행렬입니다:

a 는 K × M 에서 e4m3fnuz 형식
b 는 K × N 에서 e4m3fnuz 형식
a_scale 는 (K // 128) × M 에서 fp32
b_scale 는 (K // 128) × (N // 128) 에서 fp32
c 는 M × N 에서 bf16

커널은 특정 행렬 형식을 위해 미리 컴파일되어 있고, 경쟁에 의해 요구되는 전치 메모이 레이아웃을 가정합니다. 추가 형식이나 대안 메모이 레이아웃을 지원하려면 커널 랜처를 수정해야 합니다.

이제 우리는 고성능 ROCm 커널을 확보했으므로, 이제 자연스러운 질문은: 이를 실제 PyTorch 워크플로우에 어떻게 통합하고 다른 사람들과 공유할 것인가? 바로 다음 장에서 kernels 를 사용하여 이를 다루겠습니다.

이 가이드는 상당히 기술적이지만, 모든 세부 사항을 이해하지 않고도 단계별로 따라갈 수 있으며 정상적으로 작동합니다. 호기심이 있으시다면 나중에 개념을 더 깊이 파고들 수 있습니다.

Hugging Face Kernel Builder 는 파일을 다음과 같이 조직화되도록 기대합니다:

gemm/
├── build.toml
├── gemm
...

build.toml: 프로젝트 매니페스트; 빌드 프로세스의 핵심입니다.gemm/: GPU 마법이 일어나는 원본 CUDA 소스 코드입니다.flake.nix: 완벽하게 재현 가능한 빌드 환경을 위한 열쇠입니다.torch-ext/gemm/: 원본 PyTorch 연산자를 위한 Python 래퍼입니다.

때로는 프로젝트가 다른 파일에 의존할 수 있으며, 테스트나 헬퍼 스크립트와 같은 경우에도 문제가 없이 추가할 수 있습니다. 우리 프로젝트의 경우 구조는 다음과 같습니다:

gemm/
├── build.toml
├── gemm
...

RadeonFlow Kernels 의 gemm 커널 원본 파일을 살펴보면, HIP 소스 파일로 .cpp 확장자를 사용합니다. 첫 단계로, 이 확장자를 내용과 사용법에 따라 .h 또는 .hip 로 변경해야 합니다:

• 헤더 파일에 커널 선언, 인라인 함수, 다른 파일에 포함될 템플릿 코드가 포함되어 있다면 .h 를 사용
• HIP/GPU 코드를 별도로 컴파일해야 하는 구현 파일에 포함되어 있다면 .hip 을 사용 (예: 커널 라운чер, 복잡한 구현을 가진 장치 함수)

우리의 예에서 gemm_kernel.h, gemm_kernel_legacy.h, 및 transpose_kernel.h 는 헤더 파일이고, gemm_launcher.hip 은 HIP 구현 파일입니다. 이 네이밍 규칙은 커널 빌더 (kernels/builder) 가 각 파일 유형을 올바르게 식별하고 컴파일하도록 돕습니다.

이 파일은 전체 빌드를 조정합니다. 커널 빌더가 무엇을 컴파일해야 하고 어떻게 모든 것이 연결되는지 알려줍니다.

[general]
name = "gemm"
universal = false
...

general
이 섹션에는 일반적인 프로젝트 설정이 포함되어 있습니다.

name(필수): 프로젝트의 이름입니다. 이는 커널 이름과 일치해야 하며 Python 패키지에 사용됩니다.universal(선택적): true 로 설정되면 커널은 universal 커널이 됩니다. Universal 커널은 순수 Python 패키지 (컴파일된 파일 없음) 입니다. Universal 커널은 아래에 설명된 다른 섹션을 사용하지 않습니다. 좋은 예는 Triton 커널입니다. 기본값: false

torch
이 섹션은 Torch 확장자 구성을 설명합니다. PyTorch 에 커널을 노출할 Python 바인딩을 정의합니다.

src(필수): PyTorch 확장자를 위한 소스 파일 및 헤더 목록입니다. 우리 경우에는 C++ 바인딩 파일을 포함하며, 이는 Python 인터페이스를 만듭니다.

kernel.gemm
이름 "gemm" 을 가진 커널을 지정합니다. 여러 커널 섹션을 정의할 수 있으므로 같은 build.toml 파일에 여러 커널 섹션을 정의할 수 있습니다.

backend(필수): 커널을 실행하는 컴퓨팅 백엔드입니다. AMD GPU 지원을 위해

Hugging Face Hub 에서 ROCm 커널을 빌드하고 공유하는 과정은 이제 Nix 를 기반으로 한 깔끔하고 재현 가능한 워크플로우와 PyTorch 와의 원활한 통합 덕분에 개발자가 설정에 집중하기보다 성능 최적화에 집중할 수 있게 되었습니다. 빌드가 완료되면 사용자 정의 커널은 Hugging Face Hub 에서 공유될 수 있으며, 몇 줄의 코드로 커뮤니티에 즉시 액세스 가능하고 프로젝트 전반에 걸쳐 사용 가능합니다. 🚀

• kernels – Hub 에서 커널을 빌드하고 관리하며 로드하는 라이브러리입니다. kernel-builder 도구를 포함합니다.
• Kernels Community Hub – 커뮤니티에서 커널을 공유하고 발견합니다.