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

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

ML 커널 최적화를 위한 DSL인 Helion은 고성능 커널을 생성하지만, 수천 개의 구성을 탐색하는 오토튜닝 과정에서 막대한 시간이 소요되는 한계가 있었다. 이를 해결하기 위해 연구팀은 기존의 패턴 검색 방식에 ML 분류 모델을 결합한 'LFBO Pattern Search' 알고리즘을 개발했다. 이 알고리즘은 수집된 데이터를 바탕으로 상위 10%의 유망한 후보만 선별하여 벤치마킹함으로써 탐색 효율을 극대화한다. 결과적으로 NVIDIA B200 하드웨어에서 튜닝 시간을 36.5% 줄이면서도 커널 성능을 2.6% 향상시키는 성과를 거두었으며, 현재 Helion의 기본 검색 알고리즘으로 채택되었다.

배경

PyTorch 및 커널 최적화 기본 지식, Bayesian Optimization 개념, GPU 아키텍처(NVIDIA B200, AMD MI300 시리즈)에 대한 이해

대상 독자

고성능 ML 커널을 개발하거나 GPU 최적화 엔진을 연구하는 엔지니어

의미 / 영향

이 기술은 ML 모델의 추론 및 학습 속도를 결정짓는 커널 최적화 주기를 획기적으로 단축합니다. 특히 하드웨어 가속기마다 다른 최적 구성을 찾는 비용을 낮춰 새로운 AI 하드웨어 도입 시 소프트웨어 스택 최적화 속도를 높이는 데 기여할 것입니다.

섹션별 상세

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

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

LFBO Pattern Search 알고리즘의 작동 워크플로우 다이어그램이다. — Diagram이웃 생성, ML 모델을 통한 필터링, 벤치마킹, 모델 재학습, 검색 복사본 업데이트로 이어지는 순환 구조를 시각화하여 알고리즘의 핵심 메커니즘을 설명한다.

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

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

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

NVIDIA B200 하드웨어에서 LFBO와 Pattern Search의 성능 및 시간 비교 차트이다. — ChartLFBO가 기존 방식 대비 여러 커널에서 성능 가속(Speedup)을 달성함과 동시에 실제 실행 시간(Wallclock Time)을 유의미하게 단축했음을 보여준다. 특히 FlashAttention과 LayerNorm에서 시간 절감 효과가 두드러진다.

AMD MI350 하드웨어에서의 LFBO 성능 및 시간 벤치마크 결과이다. — ChartNVIDIA뿐만 아니라 AMD 하드웨어에서도 LFBO 알고리즘이 일관되게 튜닝 시간을 줄이고 커널 지연 시간을 개선할 수 있음을 입증한다. 하드웨어 범용성을 확인할 수 있는 지표이다.

시간 경과에 따른 최소 커널 지연 시간의 변화를 나타낸 그래프이다. — ChartLFBO가 Pattern Search보다 훨씬 이른 시점에 더 낮은 지연 시간을 가진 구성을 찾아내며, 성능 향상 폭(Jump) 또한 더 크다는 것을 시각적으로 보여준다.

PCA 공간에서 시각화된 LFBO와 Pattern Search의 탐색 분포 비교도이다. — ChartLFBO가 기존 방식보다 절반 이하의 평가 횟수만으로도 훨씬 더 넓고 다양한 파라미터 영역을 탐색했음을 증명한다. 이는 ML 가이드에 의한 효율적인 탐색 능력을 뒷받침한다.

실무 Takeaway

커널 오토튜닝 시 모든 후보를 전수 조사하는 대신 ML 분류기를 활용해 상위 10% 유망 후보만 선별하면 튜닝 시간을 30% 이상 단축할 수 있다.
시스템 비용 모델링 시 단순 회귀보다 분류 모델을 사용하면 데이터가 없는 오류 케이스를 더 잘 처리하고 고성능 영역에 최적화 역량을 집중할 수 있다.
탐색 알고리즘에 다양성 페널티를 도입하여 PCA 공간상에서 더 넓은 범위를 탐색하게 함으로써 국소 최적점(Local Optima) 탈출 가능성을 높일 수 있다.

언급된 리소스

문서Helion DSL Introduction Blog Post