Helion, 컴파일 타임 휴리스틱으로 GPU 커널 튜닝 비용 낮춘다

들어가며

GPU 커널의 성능은 연산량만으로 결정되지 않는다. 타일의 블록 크기, 워프 수, 파이프라인 단계, 메모리 접근 방식과 같은 설정이 실행 시간에 큰 영향을 준다. 가능한 설정을 차례로 컴파일하고 벤치마크하는 자동 튜닝은 좋은 구성을 찾을 수 있지만, 커널과 입력 형태마다 검색을 반복해야 하므로 대기 시간이 빠르게 늘어난다. PyTorch의 머신러닝 커널용 DSL인 Helion은 그동안 성능을 자동 튜닝에 크게 의존해 왔다. 새 컴파일 타임 휴리스틱은 프로그램 구조를 먼저 분석해, 긴 검색 없이도 합리적인 설정을 생성하는 선택지를 제공한다.

핵심 작동 방식

• 연산 유형별로 전략을 선택한다. Helion은 디바이스 IR을 훑어 행렬 곱이 포함된 커널을 matmul 휴리스틱으로, 행렬 곱은 없지만 리덕션이 있는 커널을 reduction 휴리스틱으로 보낸다. 나머지는 pointwise 전략을 사용한다. matmul은 일반적인 GEMM에 한정되지 않고 hl.dot이나 torch.baddbmm 같은 연산도 포함한다.
• IR과 데이터 생존 기간을 분석한다. 어떤 축이 타일링되는지, 어떤 블록 크기를 조정할 수 있는지, 각 타일이 커널의 어느 시점까지 살아 있는지를 기록한다. 값이 마지막으로 사용되는 위치를 추적하고, 중간 결과가 레지스터나 공유 메모리에 놓일 가능성을 근사한다.
• 하드웨어 자원 사용량을 모델링한다. 후보 블록 크기마다 하나의 CTA가 사용할 레지스터, 공유 메모리, 그리고 Blackwell에서는 텐서 메모리의 양을 추정한다. 이를 바탕으로 SM당 동시 CTA 수, 레지스터 스필 위험, 해당 구성이 실제로 실행 가능한지를 판단한다.
• 상충하는 설정을 조정한다. 파이프라인 단계를 늘리면 메모리 지연을 숨길 수 있지만 공유 메모리를 더 사용한다. 워프 수를 늘리면 스레드별 레지스터 압력은 줄어들 수 있으나, 연산량이 작은 경우에는 작은 행렬 곱을 더 많은 CTA로 실행하는 편이 유리할 수 있다. 그리드가 GPU의 SM 수보다 적은 CTA를 만들면 타일 크기를 줄여 병렬성을 높이는 방식도 사용한다.

결과와 의미

보고된 자료에 따르면 휴리스틱 분석을 켰을 때의 추가 컴파일 비용은 수십 밀리초 수준으로, 일반적으로 Triton 컴파일 시간에 묻힌다. H100에서 9개 linear attention 변형과 96개를 넘는 커널·형상 조합을 평가한 결과, Flash Linear Attention의 손작성 Triton 구현 대비 forward는 1.18배, backward를 포함하면 1.12배 빨랐다. Helion의 기존 자동 튜닝 비활성화 기본값과 비교하면 2.13배였다.

vLLM의 운영용 리덕션 커널 6종과 153개를 넘는 커널·형상 조합에 대한 평가에서는 휴리스틱 구성이 손작성 CUDA/C++ 연산자보다 1.27배, 기존 기본값보다 2.2배 빨랐다. 이 수치가 컴파일 타임 규칙이 항상 철저한 자동 튜닝을 이긴다는 뜻은 아니다. 자동 튜닝은 훨씬 넓은 후보 공간을 실제로 시험할 수 있다. 핵심은 성능과 대기 시간 사이에 실용적인 절충점을 제공한다는 데 있다.

자동 튜닝을 끄면 주 휴리스틱이 기본 구성을 직접 만들고, 켜면 보조 휴리스틱이 약 5~20개의 관련 후보를 검색의 출발점으로 제공한다. 하나의 커널에 여러 행렬 곱이 들어갈 때는 루프를 넘어 유지되는 누산기, 동적 연산량, 출력 영역 등을 기준으로 우선순위를 매긴 뒤 공유 블록 크기와 워프 설정의 충돌을 해결한다. Helion의 변화는 GPU 최적화를 무작정 후보를 시험하는 작업에서, 프로그램 구조와 하드웨어 자원 모델을 활용하는 컴파일러 단계의 판단으로 옮기려는 흐름을 보여준다.

출처: PyTorch Blog