pytorch/helion

Helion – Python에 내장된 DSL로 GPU 커널 작성하기

개념 – Helion을 사용하면 GPU에서 실행되는 작은 프로그램인 커널 코드를 일반 PyTorch와 같은 구문으로 Python에서 직접 작성할 수 있습니다. 내부적으로 이 코드는 Triton 커널(빠른 GPU 기계어로 컴파일되는 저수준 언어)로 변환됩니다.

중요성 – 효율적인 GPU 커널을 작성하는 것은 어렵습니다. 타일 크기, 그리드 차원, 메모리 접근 패턴 등 많은 저수준 설정을 결정해야 합니다. Helion은 이러한 작업의 대부분을 자동화합니다. 단일 Helion 함수는 하나의 GPU 커널로 컴파일되지만, 첫 실행 전에 수백 가지 Triton 변형을 시도하고 가장 빠른 것을 선택하는 오토튜너를 실행합니다. 그 결과 작성하기 쉽고 수동 튜닝에 가까운 성능을 제공하는 커널이 탄생합니다.

주요 기능

• Python 우선 API – @helion.kernel()로 커널을 정의하고 내부에서 일반 PyTorch 텐서를 사용합니다. 외부 Python 코드는 CPU에서 실행되고(할당, shape 검사), 내부 루프는 GPU 코드가 됩니다.
• 자동 타일링 및 그리드 선택 – hl.tile은 반복 공간을 분할하며, Helion이 타일 차원, 루프 순서, 그리드 크기를 결정합니다.
• 풍부한 오토튜닝 – 다양한 구성 공간(블록 크기, 루프 순서, 인덱싱 전략, 워프 특수화 등)을 검색하고 최적 구성을 캐시하여 이후 실행에 사용할 수 있습니다. 또한 자체 구성을 제공하여 비용이 많이 드는 검색을 건너뛸 수도 있습니다.
• 메모리 접근 전략 – 포인터 기반, 블록 포인터 또는 TensorDescriptor(TMA) 인덱싱을 자동으로 선택하며, 재정의할 수 있습니다.
• 리덕션, 마스킹, 영구 커널 지원 – 추가 코드 없이 softmax, sum 또는 대규모 리덕션과 같은 일반적인 패턴을 처리합니다.
• 백엔드 – 주로 Triton이지만, 실험적인 TileIR(NVIDIA Blackwell GPU용)과 CuTe 백엔드(CUDA 13 이상)도 있습니다.
• 디버그 및 개발 도구 – 생성된 Triton 코드 출력, 재현 가능한 스크립트 생성, 단계별 디버깅을 위한 인터프리터 모드에서 커널 실행.

일반적인 작업 흐름

1. @helion.kernel()과 hl.tile을 사용하여 커널 함수를 작성합니다.
2. CUDA 텐서로 함수를 호출합니다. 첫 호출은 오토튜닝을 트리거합니다(행렬 곱 예제의 경우 약 10분).
3. 선택적으로 최적의 helion.Config를 추출하고 하드코딩하여 향후 튜닝을 방지합니다.
4. 저장된 구성으로 커널을 배포하여 예측 가능한 시작 시간을 보장합니다.

사용 대상 – 사용자 지정 GPU 커널이 필요하지만 Triton이나 CUDA를 직접 작성하고 싶지 않은 연구자와 엔지니어. PyTorch 코드베이스에 자연스럽게 통합되며, 기존 라이브러리에 특수 구현이 없는 워크로드를 가속화할 수 있습니다.

설치 – pip install helion (Python 3.10-3.14, PyTorch ≥ 2.9, Triton ≥ 3.5 필요). 개발 설정은 uv 도구와 CuTe 백엔드용 선택적 확장을 사용합니다.

커뮤니티 – GPU MODE Discord(#helion)에서 논의가 진행됩니다. 이 프로젝트는 거버넌스 모델을 따르며 BSD 라이선스로 배포됩니다.