트라이톤(Triton) 콘퍼런스 정리: CUDA 대신 파이썬으로 GPU 커널 최적화하기

핵심 메시지

• 트라이톤은 파이썬에 가까운 문법으로 GPU 커널을 작성하게 해주는 DSL로, CUDA를 배우는 부담을 덜어주는 대신 약간의 성능을 양보한다.
• CUDA가 개별 스레드 단위로 동작하는 반면 트라이톤은 여러 스레드를 묶은 '블록'을 기본 단위로 삼아 스레드 간 상호작용 관리를 감춰준다.
• 파이토치부터 CUDA까지 이어지는 여러 계층의 도구들은 개발 생산성과 성능 사이의 파레토 프론티어 위에 서로 다른 위치를 차지한다.
• torch.compile이 자동으로 만들어내지 못하는 최적화가 있어, 링크드인에서 시작된 라이거 커널 같은 '기성품 커널' 프로젝트가 의미를 가진다.
• AMD와 엔비디아 모두 자사 하드웨어에서 트라이톤이 잘 돌도록 작업하고 있어, 한 번 작성한 커널 코드가 여러 하드웨어에서 동작할 여지가 커진다.

쉽게 이해하기

영상은 마이크로소프트 실리콘밸리 캠퍼스에서 열린 트라이톤 콘퍼런스를 하루 동안 참관한 머신러닝 엔지니어의 기록이다. 트라이톤은 파이썬 위에 얹힌 DSL로, 사실상 CUDA를 대체하는 것을 목표로 한다. GPU에서 매우 빠르게 도는 커스텀 커널을 쓰고 싶지만 워프와 스케줄링, 스레드 블록의 상호작용까지 익혀야 하는 CUDA의 학습 부담이 크다는 문제의식에서 출발했다.

발표자가 꼽은 가장 큰 단순화는 연산의 기본 단위다. CUDA는 개별 스레드 위에서 동작하지만 트라이톤은 스레드를 추상화하고 여러 스레드를 한꺼번에 묶은 블록을 다룬다. 그 덕분에 CUDA에서라면 직접 관리해야 할 스레드 간 상호작용이 상당 부분 감춰지고, 커스텀 마스크가 들어간 플래시 어텐션이나 더 최적화된 행렬 곱 커널을 훨씬 수월하게 작성할 수 있다.

콘퍼런스 전반부의 주제는 생태계였다. GPU 모드(이전 이름은 CUDA 모드) 커뮤니티의 발표자는 C++ 빌드 과정의 번거로움을 짚은 뒤, 파이썬다운 문법을 가진 트라이톤이 왜 그렇게 빠르게 인기를 얻었는지 설명했다. 이어진 발표에서는 파이토치보다 아래, 트라이톤보다는 위에 놓이는 또 다른 DSL 헬리온이 소개됐다. 헬리온은 오토튜너가 수천 개의 트라이톤 커널을 생성한 뒤 입력 형태와 지정한 파라미터에 맞춰 가장 좋은 것을 자동으로 고른다.

발표자가 하루 종일 관찰한 공통된 흐름은 개발 생산성과 GPU 활용률 사이의 파레토 프론티어였다. 파이토치, 헬리온, 트라이톤, TLX, 글루온, CUDA는 각각 이 곡선 위의 다른 지점에 서 있다. 헬리온이 트라이톤보다 한 층 위에서 더 단순하게 쓰이는 도구라면, 글루온은 트라이톤 컴파일러의 일부를 노출해 더 세밀한 제어권을 주는 아래 계층이다. TLX는 그 사이에서 엔비디아·AMD 카드에 특화된 API를 추가로 열어준다. 바이트댄스는 통신 프리미티브를 다루는 '트라이톤 디스트리뷰티드'를 발표했다.

마지막 발표는 링크드인에서 시작된 라이거 커널이었다. 발표자는 'torch.compile이 알아서 트라이톤 코드를 만들어주는데 왜 라이거가 필요한가'라는 질문에 대한 답이 인상적이었다고 전한다. 컴파일러는 똑똑하지만, 피크 메모리 사용량을 줄이는 아주 구체적이고 단계적인 메모리 조작까지 찾아내지는 못한다는 것이다. 소프트맥스나 어텐션 연산에 torch.compile을 돌린다고 플래시 어텐션이 나오지는 않는다는 설명이 그 근거로 제시됐다.

주요 인사이트

• 도구가 여러 계층으로 늘어난 것은 혼란이 아니라 선택지의 확대다. 발표자는 모든 계층을 다 쓰는 것이 목표가 아니라, 감당할 수 있는 비용과 보유한 기술, 투자할 시간에 따라 어느 계층에서 일할지 고를 수 있게 되는 것이 핵심이라고 정리한다.
• 성능이 정말 중요한 극히 일부 구간만 아래 계층으로 내려가는 전략이 현실적이다. 대부분의 커스텀 커널은 트라이톤이나 헬리온으로 쓰고, 임무상 결정적인 부분만 TLX나 글루온으로, 그보다 더 극단적인 경우에만 CUDA나 PTX로 내려가는 식이다.
• AMD와 엔비디아 엔지니어가 각각 자사 하드웨어에서 트라이톤을 최적화하는 작업을 소개했다는 점은 중요하다. 발표자는 다른 카드 공급자들이 트라이톤 백엔드를 지원하게 되면 컴퓨팅 공급이 늘어 연산 비용이 내려갈 수 있다고 내다봤다.
• 영상 말미의 대화는 생성 모델의 학습 데이터 편향을 드러낸다. 로보틱스 종사자는 이미지·영상 생성 모델이 창고처럼 어수선한 장면에서 잘 작동하지 않는다고 지적했고, 발표자는 이들 모델의 목표가 '미적으로 보기 좋은 결과'에 맞춰져 있어 로봇 학습에 필요한 분포와 직교하기 때문이라고 해석했다.

자주 묻는 질문

관련 AI 소식