제로부터 GPU: 프로덕션용 CUDA 커널 구축 및 확장 가이드

우리는 이러한 목적을 위해 kernel-builder 라이브러리를 만들었습니다. 커스텀 커널을 로컬에서 개발한 후, 여러 아키텍처를 위한 빌드를 생성하고 전 세계에 사용 가능하게 할 수 있습니다.

이 가이드에서는 지점에서부터 완전하고 현대적인 CUDA 커널을 구축하는 방법을 보여드리겠습니다. 그리고 실제 엔지니어링 전략을 바탕으로 프로덕션 및 배포의 어려운 과제를 해결하여, 단순히 빠르지만 효율적이고 유지보수 가능한 시스템을 구축하는 방법을 보여드리겠습니다.

작업이 끝나면 다른 개발자들이 다음과 같이 허브에서 직접 커널을 사용할 수 있게 될 것입니다:

import torch
from kernels import get_kernel
# Hugging Face Hub 에서 커스텀 커널 다운로드
...

영상으로 보는 것이 더 좋을까요? 이 가이드를 따르는 유튜브 영상을 확인해 보세요.

시작해 보겠습니다! 🚀

RGB 이미지를 회색조로 변환하는 실용적인 커널을 구축하겠습니다. 이 예제는 PyTorch 의 현대적 C++ API 를 사용하여 함수를 일급 시민인 네이티브 연산자로 등록합니다.

좋은 프로젝트의 기초는 깔끔하고 예측 가능한 구조입니다. Hugging Face Kernel Builder 는 파일이 다음과 같이 조직되어야 한다고 기대합니다:

img2gray/
├── build.toml
├── csrc
...

: 빌드 프로세스의 뇌인 프로젝트 매니페스트.build.toml

: GPU 마법이 일어나는 원시 CUDA 소스 코드.csrc/

: 완벽하게 재현 가능한* 빌드 환경의 열쇠.flake.nix

: 원시 PyTorch 연산자를 위한 Python 래퍼.torch-ext/img2gray/

이 파일은 전체 빌드를 조정합니다. kernel-builder 에 무엇을 컴파일하고 모든 것이 어떻게 연결되는지 알려줍니다.

# build.toml
[general]
name = "img2gray"
...

누구나 어떤 머신에서든 커널을 빌드할 수 있도록 보장하기 위해, 우리는 flake.nix 파일을 사용합니다. 이는 kernel-builder 와 그 의존체의 정확한 버전을 잠금하여 "내 머신에서는 작동하지만" 문제를 제거합니다.

# flake.nix
{
description = "img2gray kernel 의 Flake";
...

이제 GPU 코드입니다. ** csrc/img2gray.cu** 내부에서, 2D 그리드 스레드를 사용하는 커널을 정의할 것입니다—이미지를 처리하는 데 자연스럽고 효율적인 구성입니다.

// csrc/img2gray.cu
#include <cstdint>
#include <torch/torch.h>
...

이것이 가장 중요한 단계입니다. 우리는 단순히 Python 에만 바인딩하지 않습니다. 우리는 함수를 PyTorch 의 네이티브 연산자로 등록합니다. 이는 PyTorch 생태계에서 일급 시민이 되며, torch.ops 네임스페이스 아래에서 표시됩니다.

파일 ** torch-ext/torch_binding.cpp** 는 이 등록을 처리합니다.

// torch-ext/torch_binding.cpp
#include <torch/library.h>
#include "registration.h" // 빌드에서 포함됨
...

간단히 말하면, TORCH_LIBRARY_EXPAND 는 우리가 연산자를 정의할 수 있게 해줍니다. 이는 향후 쉽게 확장하거나 수정할 수 있습니다.

이 접근 방식은 두 가지 주요 이유로 중요합니다:

Compatiblity with: 이 방식으로 커널을 등록함으로써, torch.compile 는 이를 "볼" 수 있습니다. 이는 PyTorch 가 커스텀 연산자를 더 큰 계산 그래프로 융합하여 오버헤드를 최소화하고 성능을 극대화할 수 있게 합니다. 이는 커스텀 코드가 PyTorch 의 광범위한 성능 생태계와 원활하게 작동하는 열쇠입니다.Hardware-Specific Implementations: 이 시스템은 동일한 연산자를 위한 다른 백엔드를 제공하도록 허용합니다. 다른 TORCH_LIBRARY_IMPL(img2gray, CPU, ...) 를 추가할 수 있습니다.

C++ CPU 함수를 지시하는 block 을 나타냅니다. PyTorch 의 dispatcher 는 입력 텐서 (tensor) 의 device 에 따라 자동으로 올바른 구현체 (CUDA 또는 CPU) 를 호출하여 코드를 강력하고 포터블하게 만듭니다.**__init__.py 설정

wrapper

torch-ext/img2gray/ 디렉터리에서는 사용자 친화적인 방식으로 커스텀 연산자를 노출하기 위해 Python 패키지로 이 디렉터리를 만들기 위한 __init__.py 파일을 필요로 합니다.

_ops 모듈은 템플릿에서 kernel-builder 가 자동 생성하여 등록한 C++ 함수를 위한 표준 네임스페이스를 제공합니다.

# torch-ext/img2gray/__init__.py
import torch
from ._ops import ops
...

이제 우리의 커널과 바인딩이 준비되었으므로 이제 빌드해야 합니다. kernel-builder 도구는 이 과정을 간소화합니다. 단일 명령어, nix build . -L 으로 커널을 빌드할 수 있지만, 개발자로서 우리는 더 빠르고 반복적인 워크플로우를 원합니다. 이를 위해 모든 필요한 의존성이 사전 설치된 개발 쉘 (shell) 에 진입하기 위한 nix develop 명령어를 사용합니다.

더 구체적으로 말하면, 우리가 사용하고 싶은 정확한 CUDA 와 PyTorch 버전을 선택할 수 있습니다. 예를 들어, PyTorch 2.7 과 CUDA 12.6 를 사용하여 커널을 빌드하려면 다음 명령어를 사용할 수 있습니다:

# Drop into a Nix shell (all dependencies isolated sandbox)
nix develop .#devShells.torch28-cxx11-cu128-x86_64-linux

위 devShell 이름은 다음과 같이 해석할 수 있습니다:

nix develop .#devShells.torch27-cxx11-cu126-x86_64-linux
│ │ │ │
│ │ │ └─── Architecture: x86_64 (Linux)
...

이 시점에서 우리는 모든 의존성이 설치된 Nix 쉘 안에 있습니다. 이제 이 특정 아키텍처를 위한 커널을 빌드하고 테스트할 수 있습니다. 나중에 여러 아키텍처를 다루고 커널의 최종 버전을 배포할 것입니다.

build2cmake generate-torch build.toml
# this creates: torch-ext/img2gray/_ops.py, pyproject.toml, torch-ext/registration.h, setup.py, cmake/hipify.py, cmake/utils.cmake, CMakeLists.txt

이 명령어는 커널을 빌드하는 데 사용되는 몇 가지 파일을 생성합니다: CMakeLists.txt, pyproject.toml, setup.py 및 cmake 디렉터리입니다. CMakeLists.txt 파일은 커널을 빌드하기 위한 CMake 의 주요 진입점입니다.

python -m venv .venv
source .venv/bin/activate

이제 수정 모드 (editable mode) 에서 커널을 설치할 수 있습니다.

pip install --no-build-isolation -e .

🙌 놀랍습니다! 이제 우리는 PyTorch 바인딩의 모범 사례를 따르는 완전히 재현 가능한 빌드 프로세스를 가진 커스텀 빌드 커널을 가지고 있습니다.

모든 것이 제대로 작동하는지 확인하기 위해 커널이 등록되었는지 그리고 예상대로 작동하는지 확인하기 위해 간단한 테스트를 실행할 수 있습니다. 만약 그렇지 않다면, nix 환경을 재사용하여 소스 파일을 편집하고 빌드를 반복하여 수정할 수 있습니다.

# scripts/sanity.py
import torch
import img2gray
...

이제 우리는 작동하는 커널을 가지고 있으므로 다른 개발자와 세계와 공유할 차례입니다!

공유하기 전에 우리가 하고 싶은 작은 일은 빌드 프로세스 동안 생성된 모든 개발 아티팩트를 정리하여 불필요한 파일을 업로드하지 않도록 하는 것입니다.

build2cmake clean build.toml

이전에 우리는 PyTorch 와 CUDA 의 특정 버전으로 커널을 빌드했습니다. 그러나 더 넓은 대중에게 제공하기 위해 모든 지원되는 버전으로 빌드해야 합니다. kernel-builder 도구가 이를 도와줄 수 있습니다.

이는 또한 compliant kernel 개념의 위치입니다.

이제 적용됩니다. 호환성 있는 커널은 지원되는 PyTorch 및 CUDA 버전 모두에서 빌드되고 실행될 수 있는 것입니다. 일반적으로 이는 사용자 정의 구성을 필요로 하지만, 저희의 경우 kernel-builder 도구가 프로세스를 자동화합니다.

# dev shell 외부에서 다음 명령어를 실행하세요
# 샌드박스 내부에 있다면 `exit` 로 나설 수 있습니다
nix build . -L

이 프로세스는 지원되는 PyTorch 및 CUDA 버전 모두를 위해 커널을 빌드하기 때문에 시간이 걸릴 수 있습니다. 출력은 result 디렉토리에 표시됩니다.

kernel-builder 팀은 지원되는 빌드 변형을 능동적으로 유지 관리하며, 최신 PyTorch 및 CUDA 릴리스와 동기화되고 더 넓은 호환성을 위해 트레일링 버전을 지원합니다.

마지막 단계는 결과를 예상되는 build 디렉토리로 이동하는 것입니다 (이것이 kernels 라이브러리가 이를 찾기 위한 곳입니다).

mkdir -p build
rsync -av --delete --chmod=Du+w,Fu+w result/ build/

빌드 아티팩트를 Hub 로 푸시하면 이전 포스트에서 보았듯이 다른 개발자가 커널을 사용하도록 간단해집니다. 우리는 kernels upload 명령어를 사용할 수 있습니다:

kernels upload <path_to_kernel> --repo_id hub-username/img2gray

표준 git 기반 프로세스를 사용하여 업로드도 가능합니다.

먼저 새 레포지토리를 생성합니다:

hf repo create img2gray

Hugging Face Hub 에 huggingface-cli login 으로 로그인했는지 확인하세요.
이제 프로젝트 디렉토리에서 새 레포지토리에 프로젝트를 연결하고 코드를 푸시합니다:

# git 초기화 및 Hugging Face Hub 연결
git init
git remote add origin https://huggingface.co/<your-username>/img2gray
...

완벽합니다! 이제 커널은 Hugging Face Hub 에 있으며, 다른 사람들이 사용하도록 준비되어 있고 kernels 라이브러리와 완전히 호환됩니다. 저희의 커널과 모든 빌드 변형은 drbh/img2gray 에서 이용 가능합니다.
kernels 라이브러리를 사용하면 커널을 전통적인 의미로 "설치"하지 않습니다. Hub 레포지토리로 직접 로드하여 새로운 연산자를 자동으로 등록합니다.

# /// script
# requires-python = ">=3.10"
# dependencies = [
...

사용 가능한 커널이 준비되면, 커널을 배포하는 것을 더 쉽게 만들 수 있는 몇 가지 작업을 수행할 수 있습니다. 버전 관리를 사용하여 커널의 하류 사용이 깨지지 않으면서 API 변경을 하는 도구로 사용하는 방법을 논의하겠습니다. 그 후, Python wheels 를 만들어야 하는 방법을 보여드리겠습니다.

일단 시간이 지나면 커널을 업데이트하고 싶을 수 있습니다. 새로운 성능 개선 방법을 찾았거나 커널의 기능을 확장하고 싶을 수도 있습니다. 일부 변경 사항은 커널의 API 를 변경해야 합니다. 예를 들어, 더 높은 버전은 하나의 공중 함수에 새로운 필수 인자를 추가할 수 있습니다. 이는 하류 사용자에게 불편할 수 있으며, 코드가 이 새로운 인자를 추가하기까지 깨질 수 있기 때문입니다.

커널의 하류 사용자는 이를 피하기 위해 사용하는 커널을 특정 리비전에 고정 (pinning) 할 수 있습니다. 예를 들어, 각 Hub 레포지토리는 Git 레포지토리이기도하므로, 그들은 git commit shorthash 를 사용하여 커널을 리비전에 고정할 수 있습니다:

from kernels import get_kernel
img2gray_lib = get_kernel("drbh/img2gray", revision="4148918")

Git shorthash 를 사용하면 분해의 가능성이 줄어듭니다; 그러나 해석하기 어렵고 버전 범위 내에서 부드러운 업그레이드를 허용하지 않습니다. 따라서 Hub 커널을 위해 친숙한 세마틱 버전을 사용하는 것을 권장합니다. 커널에 버전을 추가하는 것은 쉽습니다: 단순히 vx.y.z 형식의 Git 태그를 추가하면 됩니다

여기서 x.y.z 는 버전입니다. 예를 들어, 현재 커널의 버전이 1.1.2 인 경우, 이를 v1.1.2 로 태그할 수 있습니다

그런 다음 get_kernel 을 사용하여 해당 버전을 가져올 수 있습니다:

from kernels import get_kernel
img2gray_lib = get_kernel("drbh/img2gray", revision="v1.1.2")

버전 범위를 사용하면 버전 관리가 더욱 강력해집니다. 세마틱 버전에서 버전 1.y.z 는 다음 각 succeeding x 와 y 에 대해 공개 API 에서 후방 호환성 변경이 없어야 합니다. 따라서, 커널의 버전이 작성 시점에 1.1.2 인 경우, 버전을 적어도 1.1.2 이고 2.0.0 보다 낮게 요청할 수 있습니다:

from kernels import get_kernel
img2gray_lib = get_kernel("drbh/img2gray", version=">=1.1.2,<2")

이는 코드가 항상 1.y.z 시리즈에서 최신 커널을 가져올 것을 보장합니다. 버전 범위는 Python 스타일의 버전 지정자일 수 있습니다.

버전을 huggingface-cli 로 태그할 수 있습니다:

$ huggingface-cli tag drbh/img2gray v1.1.2

대규모 프로젝트에서는 각 get_kernel 호출이 아닌 전역으로 커널 버전을 조정하고 싶을 수 있습니다. 또한 모든 사용자가 동일한 커널 버전을 갖도록 커널을 잠금하는 것이 종종 유용하며, 이는 버그 보고서 처리에 도움이 됩니다.

kernels 라이브러리는 프로젝트 수준에서 커널을 관리하는 훌륭한 방법을 제공합니다. 이를 위해 pyproject.toml 파일의 프로젝트 빌드 시스템 요구 사항에 kernels 패키지를 추가합니다. 이렇게 하면 tools.kernels 섹션에서 프로젝트의 커널 요구 사항을 지정할 수 있습니다:

[build-system]
requires = ["kernels", "setuptools"]
build-backend = "setuptools.build_meta"
...

버전은 Python 의존성과 동일한 유형의 버전 지정자로 지정할 수 있습니다. 이는 버전 태그 (va.b.c) 를 사용할 수 있는 또 다른 곳입니다 -- kernels 는 프로젝트의 버전 태그를 사용하여 어떤 버전을 사용할 수 있는지 조회합니다. pyproject.toml 에서 커널을 지정하고 나면, kernels 명령줄 유틸리티를 사용하여 특정 버전으로 잠금할 수 있습니다. 이 유틸리티는 kernels Python 패키지의 일부입니다:

$ kernels lock .

이는 pyproject.toml 에서 지정된 범위와 호환되는 최신 커널 버전을 가진 kernels.lock 파일을 생성합니다. kernels.lock 은 프로젝트의 Git 저장소에 제출해야 하므로, 프로젝트의 모든 사용자가 잠금된 커널 버전을 얻을 수 있습니다. 새로운 커널 버전이 출시되면 kernels lock 을 다시 실행하여 잠금 파일을 업데이트할 수 있습니다.

프로젝트에서 잠긴 커널을 완전히 구현하려면 마지막 부분이 필요합니다. get_locked_kernel 은 get_kernel 의 대조품이며 잠긴 커널을 사용합니다. 따라서 잠긴 커널을 사용하려면, 모든 get_kernel 발생을 get_locked_kernel 으로 교체합니다:

from kernels import get_kernel
img2gray_lib = get_locked_kernel("drbh/img2gray")

그것이 전부입니다! 프로젝트의 get_locked_kernel("drbh/img2gray") 호출은 이제 kernels.lock 에서 지정된 버전을 사용합니다.

get_locked_kernel