PyTorch용 고성능 TPU 커널 DSL, Helion 공개

Google과 협력하여 구축된 Helion은 성능 이식성을 갖춘 ML 커널 작성을 위한 PyTorch의 고수준 DSL입니다. 기존의 저수준 DSL인 Pallas의 높은 진입 장벽을 해결하여, 사용자가 익숙한 PyTorch 스타일의 코드를 작성하면 최적화된 TPU 코드로 컴파일할 수 있도록 지원합니다.

주요 특징 및 성과는 다음과 같습니다:

• 고성능 달성: Flash Attention 워크로드 기준, TPU v7에서 838 TFLOPs(텐서 코어 하나 기준 약 79% MFU)를 달성했습니다.
• 오토튜닝(Autotuning): 입력 형태(shape)에 따라 다양한 코드 생성 전략을 탐색하여 최적의 파이프라이닝 방식을 선택하며, TPU의 VMEM과 연산 자원을 극대화합니다.
• 하드웨어 이식성: TPU와 GPU 간의 커널 집합을 동일하게 유지할 수 있어 크로스 하드웨어 사용자의 워크플로를 단순화합니다.
• 사용 사례: 성능 최적화가 필요한 경우, Pallas 비전문가, 그리고 여러 하드웨어 간 동일한 커널을 유지하려는 사용자에게 적합합니다.