PyTorch Helion, TPU 백엔드로 이기종 커널 작성 문턱 낮춘다

들어가며

PyTorch 생태계가 고성능 커스텀 커널 개발을 더 높은 추상화 수준으로 끌어올리고 있다. PyTorch Blog가 공개한 Helion의 TPU 백엔드는 개발자가 PyTorch에 가까운 방식으로 커널을 작성하면, 컴파일러가 이를 TPU용 Pallas 코드로 변환하는 흐름을 제시한다.

이는 단순한 문법 편의 기능이 아니다. TPU는 GPU와 달리 명시적인 메모리 공간, HBM과 온칩 VMEM 사이의 데이터 이동, 계산과 전송의 겹치기 같은 요소가 성능에 큰 영향을 준다. 기존에는 Pallas와 TPU 메모리 계층에 대한 깊은 이해가 필요했지만, Helion은 이 부담의 일부를 컴파일러와 자동 튜너로 옮기려 한다.

핵심 내용

• Helion은 성능 이식성을 지향한다: 하나의 고수준 커널 표현에서 GPU와 TPU 같은 서로 다른 하드웨어에 맞는 코드를 생성하는 것이 목표다.
• TPU 백엔드는 Pallas를 대상으로 한다: 개발자는 Helion으로 작성하고, 백엔드는 TPU 프로그래밍을 위한 저수준 DSL인 Pallas 코드를 생성한다.
• 파이프라이닝이 최적화의 중심이다: Helion은 HBM에서 VMEM으로 데이터를 가져오는 과정과 행렬·벡터 연산을 최대한 겹치도록 코드를 구성한다.
• Flash Attention은 더 복잡한 사례다: Q 타일을 처리하는 외부 루프와 K/V 타일을 순회하는 내부 루프가 있으며, Helion은 emit_pipeline 사용, 사전 적재, 루프 전개 등 여러 전략을 자동 튜닝한다.
• 제시된 성능 수치가 눈에 띈다: 원문에 따르면 Flash Attention 워크로드에서 Helion 생성 커널은 TPU v7에서 838 TFLOPs를 기록했고, 이는 단일 tensor core 기준 약 79% MFU다.

의미와 영향

이번 발표는 PyTorch가 TPU를 고성능 ML 워크로드의 주요 타깃으로 본격적으로 다루고 있음을 보여준다. 최신 TPU와 고성능 GPU는 BF16 연산 성능과 HBM 대역폭 같은 핵심 지표에서 비교 가능한 수준에 있지만, 개발자가 성능을 끌어내는 방식은 상당히 다르다.

Helion이 제공하려는 가치는 세 가지로 요약할 수 있다. 첫째, Pallas 전문가가 아니어도 성능이 중요한 커널 최적화에 접근할 수 있다. 둘째, TPU와 GPU를 모두 쓰는 팀이 동일한 고수준 커널 코드를 유지할 가능성이 커진다. 셋째, 입력 shape에 따라 달라지는 파이프라인과 버퍼 선택을 자동 튜닝으로 탐색할 수 있다.

다만 원문에 제시된 결과는 특정 워크로드와 설정에 대한 것이므로, 모든 TPU 커널에서 같은 수준의 효율을 보장한다고 볼 수는 없다. 그럼에도 방향성은 분명하다. AI 인프라 경쟁은 더 빠른 칩뿐 아니라, 그 칩의 성능을 개발자가 얼마나 쉽게 활용할 수 있게 만드는 도구 경쟁으로 확장되고 있다.

출처: PyTorch Blog