PyTorch Helion: ML 기반 LFBO 알고리즘으로 커널 오토튜닝 속도와 성능 동시 개선

섹션별 상세

01

Helion DSL의 오토튜닝 엔진은 블록 크기, 루프 순서 등 방대한 하이퍼파라미터 공간을 탐색하여 최적의 커널 구성을 찾지만, 이 과정에서 발생하는 긴 컴파일 및 벤치마킹 시간이 개발자의 주요 병목 현상으로 작용한다.

02

새로운 LFBO Pattern Search 알고리즘은 베이지안 최적화 기법을 변형하여, Random Forest 분류기를 통해 성능이 우수할 것으로 예측되는 후보군만 지능적으로 필터링함으로써 불필요한 컴파일 횟수를 획기적으로 줄인다.

03

회귀(Regression) 대신 분류(Classification) 모델을 선택함으로써, 모든 구성의 지연 시간을 예측하는 대신 상위 10%의 고성능 구성을 식별하는 데 모델 용량을 집중하고 컴파일 오류나 타임아웃이 발생하는 부적합한 구성을 효과적으로 배제한다.

04

탐색의 다양성을 확보하기 위해 Random Forest의 리프 노드 동시 발생 빈도를 기반으로 유사도 점수를 계산하고, 기존에 탐색된 구성과 너무 유사한 후보에는 페널티를 부여하여 더 넓은 파라미터 공간을 탐색하도록 설계되었다.

05

NVIDIA B200 및 AMD MI350 하드웨어 테스트 결과, LayerNorm 및 FlashAttention 등 주요 커널에서 튜닝 시간은 최대 50% 단축되었으며, 커널 지연 시간은 오히려 개선되는 등 효율성과 성능의 트레이드오프 문제를 해결했다.

도메인 특화 언어(DSL): — 특정 도메인의 문제를 해결하기 위해 설계된 프로그래밍 언어이다. Helion은 고성능 ML 커널 작성을 위해 PyTorch와 유사한 구문을 제공하며 최적화 작업을 자동화 엔진에 위임한다.
오토튜닝(Autotuning): — 하드웨어 특성에 맞춰 커널의 하이퍼파라미터인 블록 크기, 루프 순서, 메모리 패턴 등을 자동으로 탐색하여 최적의 성능을 찾는 과정이다. 수만 개의 조합 중 최적점을 찾는 것이 핵심이다.
가능도 미사용 베이지안 최적화(LFBO): — 복잡한 확률 밀도 함수를 직접 계산하는 대신 분류 모델을 사용하여 유망한 영역을 식별하는 경량화된 베이지안 최적화 방식이다. 계산 비용이 낮아 실시간 탐색 엔진에 적합하다.
주성분 분석(PCA): — 고차원의 데이터를 정보 손실을 최소화하면서 저차원으로 투영하는 차원 축소 기법이다. 이 아티클에서는 수만 개의 커널 구성 분포를 2차원 평면에 시각화하여 탐색 범위를 비교하는 데 사용되었다.