TIRx 커널 언어 공개, CUDA 수준 스레드 제어와 트라이톤식 타일 문법을 얕은 컴파일러로 묶다

핵심 메시지

• 홉퍼와 블랙웰이 워프 특수화, 비동기 연산, 새 메모리 공간을 들고 오면서 고수준 컴파일러를 확장하는 일이 차라리 CUDA로 직접 짜는 것보다 어려워졌다는 문제의식에서 출발한다.
• TIRx는 기본 작성 수준을 스레드에 두되 자주 쓰는 연산에는 타일 인터페이스를 제공하고, 중간 표현과 백엔드 사이 단계를 최소로 유지하는 '얕고 열린' 컴파일러를 지향한다.
• 레이아웃은 파이토치의 모양·보폭 개념에서 출발해, 텐서 메모리의 레인과 칼럼 같은 하드웨어 축과 어느 스레드가 값을 갖는지, 어디에 복제되는지까지 하나의 매핑으로 표현한다.
• 어떤 타일 연산을 어떤 하드웨어 명령으로 바꿀지 정하는 규칙이 파이썬 코드로 등록돼 있어, 개발자가 중단점을 걸고 들여다보며 고칠 수 있다는 점이 확장의 진입장벽을 낮춘다.
• 발표 말미의 문답에서는 에이전트가 커널을 대량으로 쓰는 시대에 컴파일러의 역할이 코드 생성기가 아니라 동기화 오류를 잡아내는 검증 도구로 옮겨갈 것이라는 전망이 나왔다.

쉽게 이해하기

GPU MODE 109회 강연에 선 발표자는 카네기멜런대학교 박사과정 보한 허우(Bohan Hou)로, TVM 계열 연구로 알려진 톈치 첸 교수 연구실 소속이다. 그가 소개한 TIRx는 고성능 머신러닝 커널을 위한 프로그래밍 언어이자 컴파일러 기반 구조다. 목표는 CUDA 수준의 직접적인 제어와 트라이톤 같은 타일 문법의 간결함을 한자리에 두되, 그 사이에 깊고 경직된 컴파일러 층을 쌓지 않는 것이다.

문제의식은 하드웨어 세대교체에서 나온다. 고수준 DSL은 자기가 전제한 하드웨어 모델에 맞을 때는 생산성이 높지만, 새 세대가 다른 실행 모델을 들고 오면 이야기가 달라진다. 워프 특수화와 비동기 연산, 새로운 메모리 공간이 추가되면 손볼 곳이 장치 명령어에 그치지 않고 언어와 스케줄링 모델, 분석과 하강 파이프라인까지 번진다. 그래서 최근 저수준 DSL이 여럿 등장했다는 것이 발표자의 진단이다.

TIRx가 세운 네 가지 원칙은 명확하다. 백엔드마다 자기 방식의 하강 경로를 쓰게 하고, 기본 작성 수준을 스레드에 두며, 중간 표현과 네이티브 백엔드 사이 단계를 몇 개로 줄이고, 같은 중간 표현과 컴파일러 패스를 파이썬·C++·러스트 등 여러 언어에서 그대로 다루게 하는 것이다. 실제로 발표에서는 256개 실수를 두 배로 만드는 짧은 예제부터 블랙웰 텐서 메모리를 쓰는 행렬곱 커널까지, 모든 하드웨어 동작이 코드 위에 그대로 보이는 방식으로 전개된다.

발표의 절반 가까이는 레이아웃 표현에 쓰인다. 출발점은 익숙한 모양과 보폭이지만, TIRx는 텐서의 모양과 레이아웃의 모양을 분리해 사용자에게는 의미 있는 형태를 남기면서 내부적으로는 타일 분할 순서를 드러낸다. 여기에 각 보폭이 어떤 물리적 축에 기여하는지 이름을 붙이면, 블랙웰 텐서 메모리의 레인과 칼럼 같은 구조적 좌표는 물론 어느 스레드가 값을 소유하는지, 하나의 값이 몇 군데로 복제되는지까지 같은 식 하나로 적을 수 있다.

성능 면에서는 블랙웰 세대 GPU에서 인과·비인과 플래시 어텐션, 희소 MoE 디코드, 텐서 병렬 환경의 행렬곱-통신 융합 커널 결과를 제시하며 고도로 최적화된 기준선과 대등하다고 밝혔다. 함께 만들고 있는 오픈 커널 라이브러리에는 밀집·블록 스케일 행렬곱, 그룹 행렬곱, 플래시 어텐션, 정규화, 융합 MoE 커널 등이 들어 있다. 개발 도구로는 커널 내부의 시간 구간을 표시해 워프의 생명주기와 파이프라인 사이 빈틈을 눈으로 확인하는 추적 기능이 소개됐다.

주요 인사이트

• 이 강연의 핵심 주장은 '추상화를 높이는 대신 컴파일러를 얇게 만들자'는 것이다. 새 하드웨어 기능이 나왔을 때 컴파일러 여러 층을 뜯어고치는 대신, 커널 코드에서 곧바로 그 명령을 부를 수 있게 열어두는 쪽을 택했다.
• 레이아웃을 '논리 원소에서 물리 위치로 가는 하나의 매핑'으로 통일한 점이 설계의 중심이다. 스레드 소유권, 메모리 배치, 장치 간 분산이 서로 다른 개념처럼 보이지만 같은 표기로 적히면 커널 엔지니어가 이미 하던 판단과 코드가 가까워진다.
• 확장 지점을 파이썬 패스로 둔 선택은 성능이 아니라 사람에 대한 배려다. 거대한 C++ 컴파일러를 고치는 대신 중단점을 걸고 디스패치 규칙을 들여다볼 수 있다면, 기여자의 수가 달라진다.
• 빠진 명령어 하나가 개발을 막지 않도록 CUDA 함수를 소스 문자열로 넣어 부르는 탈출구를 남겨둔 점도 실용적이다. 라이브러리가 완성되기를 기다리지 않아도 되는 구조다.
• 문답에서 나온 전망이 인상적이다. 에이전트가 커널을 아무리 많이 써내도 동기화 버그는 수만 번에 한 번, 수십 분 실행 뒤에야 드러나기 때문에, 결국 사람이 검토해야 하는 병목이 남는다는 것이다.

자주 묻는 질문

관련 AI 소식