🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

허깅페이스의 kernels와 kernel-builder를 이용해 AMD MI300X용 ROCm FP8 GEMM 커널을 재현 가능한 형태로 구성하고, 빌드 설정부터 PyTorch 네이티브 연산자 연결까지 구현하는 방법을 설명한 기술 안내서다.

📌 핵심 요약

• 사용자 정의 GPU 커널은 딥러닝 연산을 특정 작업에 맞게 최적화할 수 있지만, 아키텍처별 컴파일 설정과 빌드 플래그, PyTorch 확장 연동 과정에서 복잡한 오류가 발생하기 쉽다.
• 허깅페이스 kernels 라이브러리와 kernel-builder는 CUDA, ROCm, Metal, XPU 등 여러 백엔드의 커널을 일관된 구조로 빌드하고 공유할 수 있게 하며, 이 글은 그중 ROCm과 AMD GPU에 집중한다.
• 예제로 사용된 RadeonFlow GEMM은 AMD Instinct MI300X에 최적화된 FP8 블록 단위 행렬 곱셈 커널로, e4m3fnuz 입력과 블록별 스케일링을 사용해 처리량과 수치 정확성의 균형을 맞춘다.
• 프로젝트는 build.toml, flake.nix, HIP 소스와 헤더, PyTorch 바인딩 및 Python 래퍼로 구성되며, 파일 역할에 따라 .h와 .hip 확장자를 구분해야 kernel-builder가 올바르게 처리할 수 있다.
• build.toml에서 ROCm 백엔드와 gfx942 아키텍처, Torch 의존성, 소스 및 include 경로를 선언하고, C++ 바인딩에서 텐서의 장치·연속성·차원을 검증한 뒤 HIP 실행 함수에 연결해 torch.ops에서 사용할 수 있는 네이티브 PyTorch 연산자로 만든다.

🧩 주요 포인트

1. 사용자 정의 GPU 커널은 딥러닝 연산을 특정 작업에 맞게 최적화할 수 있지만, 아키텍처별 컴파일 설정과 빌드 플래그, PyTorch 확장 연동 과정에서 복잡한 오류가 발생하기 쉽다.
2. 허깅페이스 kernels 라이브러리와 kernel-builder는 CUDA, ROCm, Metal, XPU 등 여러 백엔드의 커널을 일관된 구조로 빌드하고 공유할 수 있게 하며, 이 글은 그중 ROCm과 AMD GPU에 집중한다.
3. 예제로 사용된 RadeonFlow GEMM은 AMD Instinct MI300X에 최적화된 FP8 블록 단위 행렬 곱셈 커널로, e4m3fnuz 입력과 블록별 스케일링을 사용해 처리량과 수치 정확성의 균형을 맞춘다.
4. 프로젝트는 build.toml, flake.nix, HIP 소스와 헤더, PyTorch 바인딩 및 Python 래퍼로 구성되며, 파일 역할에 따라 .h와 .hip 확장자를 구분해야 kernel-builder가 올바르게 처리할 수 있다.
5. build.toml에서 ROCm 백엔드와 gfx942 아키텍처, Torch 의존성, 소스 및 include 경로를 선언하고, C++ 바인딩에서 텐서의 장치·연속성·차원을 검증한 뒤 HIP 실행 함수에 연결해 torch.ops에서 사용할 수 있는 네이티브 PyTorch 연산자로 만든다.

🧠 상세 정리

1. 사용자 정의 커널 빌드가 어려운 이유

사용자 정의 커널은 이미지 처리, 텐서 변환, 행렬 연산처럼 계산량이 큰 작업을 실제 워크로드에 맞게 최적화하는 핵심 수단이다. 그러나 커널을 대상 GPU 아키텍처에 맞춰 컴파일하고, 필요한 빌드 플래그를 연결하며, PyTorch 확장 모듈에 안정적으로 통합하는 과정은 빠르게 복잡해질 수 있다. 이 과정에서는 CMake나 Nix 구성뿐 아니라 컴파일러 오류와 ABI 호환성 문제까지 함께 다뤄야 한다. 허깅페이스의 kernels 라이브러리는 kernel-builder를 이용해 이러한 빌드 절차를 구조화하고, 완성된 커널을 kernels-community와 공유할 수 있도록 설계됐다. CUDA, ROCm, Metal, XPU 등 여러 GPU 및 가속기 백엔드를 지원하므로 커널의 성능뿐 아니라 이식성과 PyTorch 통합 편의성도 함께 확보하는 것이 목표다.

2. ROCm 안내서의 범위와 예제

이 글은 kernels가 지원하는 여러 백엔드 가운데 ROCm 호환 커널만을 대상으로 빌드, 테스트, 공유 과정을 설명한다. 독자는 AMD GPU에서 효율적으로 실행되는 커널을 구성하는 방법과 함께 재현성, 패키징, 배포에 필요한 기본 원칙을 따라가게 된다. 전체 내용은 기존 kernel-builder 안내서를 ROCm 중심으로 간소화한 형태이며, CUDA 전반을 다루는 별도의 안내서와 범위가 구분된다. 실습 예제로는 RadeonFlow_Kernels 프로젝트의 GEMM 커널을 사용한다. 세부 개념을 모두 이해하지 않아도 단계별 절차를 따라갈 수 있도록 구성됐지만, 실제 코드는 HIP 실행 방식과 PyTorch의 C++ 확장 구조를 함께 다루는 기술적인 내용이다.

3. RadeonFlow FP8 GEMM 커널의 특징

RadeonFlow GEMM은 AMD Instinct MI300X GPU에 최적화된 고성능 FP8 블록 단위 행렬 곱셈 구현이다. GEMM은 두 행렬 A와 B를 곱해 C를 만드는 연산으로, 대부분의 딥러닝 워크로드에서 핵심 구성 요소로 사용된다. 이 커널은 더 높은 처리량과 더 낮은 메모리 대역폭 사용을 위해 FP8 저정밀 형식을 채택하며, AMD Developer Challenge 2025를 위해 개발되어 2025년 6월 대상을 받았다. 입력에는 지수 4비트와 가수 3비트로 구성된 e4m3fnuz 형식을 사용한다. FP8은 FP16이나 FP32보다 표현 범위가 좁기 때문에 블록별 스케일 계수인 a_scale과 b_scale을 적용해 값이 계산에 적합한 범위에 놓이도록 조정하고, 저정밀 연산에서 정확성 손실을 줄인다.

4. 입출력 형식과 지원 범위

커널의 실행 인자는 a, b, a_scale, b_scale, c이며, a와 b는 양자화된 입력 행렬이고 c는 결과 행렬이다. 원문에 제시된 커널 형식에서 a는 K×M, b는 K×N 크기의 e4m3fnuz 데이터이며, a_scale은 (K // 128)×M, b_scale은 (K // 128)×(N // 128) 크기의 FP32 스케일 행렬이다. 출력 c는 M×N 크기의 BF16 행렬로 생성된다. 이 구현은 대회에서 요구한 전치 메모리 배치를 전제로 하며, 특정 행렬 크기에 맞춰 미리 컴파일되어 있다. 따라서 새로운 행렬 크기나 다른 메모리 배치를 지원하려면 단순히 입력만 바꾸는 것이 아니라 커널 실행기를 수정해야 한다. 이러한 제약은 커널이 범용 GEMM 구현이라기보다 정해진 조건에서 높은 성능을 내도록 특화됐음을 보여준다.

5. 프로젝트 구조와 HIP 파일 구분

kernel-builder가 기대하는 기본 프로젝트에는 빌드 명세인 build.toml, 커널 소스 디렉터리, 재현 가능한 환경을 정의하는 flake.nix, C++ 바인딩과 Python 래퍼를 담은 torch-ext 디렉터리가 포함된다. RadeonFlow 예제는 여기에 공통 GPU 헤더, 타이머와 산술 유틸리티, 성능 측정 및 검사 관련 파일을 추가한 확장 구조를 사용한다. 원본 RadeonFlow 파일 중 일부 HIP 소스는 .cpp 확장자를 사용하므로, 실제 내용과 컴파일 방식을 기준으로 .h 또는 .hip으로 변경해야 한다. 다른 파일에 포함되는 선언, 인라인 함수, 템플릿 코드는 .h로 두고, 별도로 컴파일해야 하는 HIP 또는 GPU 구현은 .hip으로 둔다. 예제에서는 gemm_kernel.h, gemm_kernel_legacy.h, transpose_kernel.h가 헤더이고 gemm_launcher.hip가 구현 파일이며, 이 구분을 통해 kernel-builder가 각 파일을 올바른 방식으로 식별하고 컴파일한다.

6. build.toml과 flake.nix 설정

build.toml은 무엇을 어떤 백엔드용으로 컴파일하고 각 구성 요소를 어떻게 연결할지 정의하는 빌드의 중심 명세다. general 섹션에서는 프로젝트와 Python 패키지에 사용될 이름을 gemm으로 지정하고, 컴파일 파일이 없는 순수 Python 패키지가 아니므로 universal을 false로 설정한다. torch 섹션에는 Python 인터페이스를 구성하는 C++ 바인딩 소스와 헤더를 등록한다. kernel.gemm 섹션에서는 backend를 rocm으로, rocm-archs를 MI300 계열 대상인 gfx942로 지정하고, PyTorch 텐서 연산을 위해 torch 의존성을 선언한다. 커널과 유틸리티 및 검사 관련 소스 목록과 include 디렉터리도 이곳에 명시한다. flake.nix는 허깅페이스 kernels 저장소의 kernel-builder와 의존성 버전을 고정하고 genFlakeOutputs를 호출함으로써 다른 환경에서도 동일한 빌드 구성을 재현할 수 있게 한다.

7. HIP 실행기와 형상별 디스패치

실제 GPU 실행 로직은 gemm/gemm_launcher.hip에 정의되며, 설정에 따라 새롭게 최적화된 GEMM 구현을 호출하거나 기존 구현으로 대체할 수 있다. 외부에서 호출되는 run 함수는 입력과 스케일 및 출력 버퍼를 void 포인터로 받은 뒤, 각각 FP8 입력 포인터, FP32 스케일 포인터, BF16 출력 포인터로 변환한다. KernelTimerScoped는 2×M×N×K에 해당하는 연산량을 기준으로 실행 시간과 GFLOPS 같은 성능 지표를 기록할 수 있게 한다. 이후 pack_shape로 M, N, K 조합을 묶고 switch 문에서 미리 준비된 DISPATCH_GEMM 구성을 선택한다. 원문 예제에는 1024·1536·7168 조합과 6144·7168·2304 조합에 대응하는 디스패치가 제시된다. 등록되지 않은 형상이 들어오면 지원하지 않는 M, K, N 값을 출력한 뒤 실행을 중단하므로, 지원 범위는 컴파일된 디스패치 항목으로 명확히 제한된다.

8. PyTorch 네이티브 연산자 연결과 검증

torch-ext/torch_binding.cpp는 HIP 실행 함수를 단순한 Python 함수가 아니라 torch.ops를 통해 접근할 수 있는 네이티브 PyTorch 연산자로 연결하는 역할을 한다. 바인딩 코드는 먼저 a, b, 두 스케일 텐서와 출력 텐서가 모두 GPU 장치에 있는지 TORCH_CHECK로 검사하고, 각 텐서가 연속 메모리인지도 확인한다. 이어 a를 [M, K], b를 [K, N], 출력을 [M, N]으로 해석해 차원을 구하고, b의 첫 번째 차원이 K와 일치하는지와 출력의 두 차원이 M, N에 맞는지를 검증한다. 검증을 통과하면 텐서의 data_ptr을 이용해 원시 데이터 주소를 얻고 HIP 실행 함수 run에 전달한다. 제시된 코드는 기본 HIP 스트림인 스트림 0을 사용하므로, 실행기까지 전달되는 스트림도 해당 값이다. 이 계층은 PyTorch 텐서 수준의 안전성 검사와 저수준 ROCm 커널 실행을 분리하면서, 고성능 커널을 일반적인 PyTorch 작업 흐름 안에서 호출할 수 있도록 만든다.

🧾 핵심 주장 / 시사점

• kernel-builder의 핵심 가치는 커널 코드 자체를 새로 만드는 데 있기보다, 소스 분류와 아키텍처 지정, PyTorch 바인딩, 재현 가능한 빌드 환경을 하나의 규약으로 묶어 통합 과정의 복잡성을 줄이는 데 있다.
• RadeonFlow GEMM의 성능은 FP8 사용만으로 결정되지 않으며, 블록별 스케일링과 특정 행렬 형상용 디스패치, MI300X의 gfx942 아키텍처를 겨냥한 사전 컴파일이 함께 작동한 결과다.
• 실제 배포 가능한 커널 패키지를 만들려면 GPU 코드의 속도뿐 아니라 지원 형상과 메모리 배치의 제약을 명시하고, PyTorch 바인딩 계층에서 장치·연속성·차원을 엄격히 검증하는 과정이 필요하다.

✅ 액션 아이템

• 허깅페이스 kernels와 kernel-builder로 MI300X용 ROCm FP8 커널을 구성할 때 build.toml의 ROCm 백엔드, gfx942, Torch 의존성, 경로 선언을 우선 점검한다.
• RadeonFlow 방식의 e4m3fnuz 입력 블록별 스케일링 구현을 반영해 처리량과 수치 정확도 균형 목표를 수치 기준으로 정의한다.
• 프로젝트 구성요소를 build.toml·flake.nix·HIP 소스·헤더·C++ 바인딩·Python 래퍼로 분리해 .h/.hip 구분 규칙을 고정하고 kernel-builder 처리 재현성을 유지한다.

❓ 열린 질문

• e4m3fnuz 입력과 블록별 스케일링은 MI300X FP8 GEMM에서 처리량 대비 정확도 개선 효과를 어떤 지표로 판단할 것인가?
• ROCm 아키텍처별 플래그 설정에서 gfx942 이외 대상 추가 시 어떤 오류 패턴이 먼저 드러나며 범위별 원인 구분은 어떻게 할 것인가?
• torch.ops 노출 전 C++ 바인딩에서 텐서 장치·연속성·차원 검증 실패를 어디서 분기해야 호출 안정성을 보장할 수 있는가?