Gluon 커널 DSL과 선형 레이아웃, Triton으로 부족했던 Blackwell GPU 최적화의 해법

핵심 메시지

• Triton은 최소한의 노력으로 Blackwell에서 이론 최대 성능의 약 80%를 내주지만, 하드웨어가 워프 특화를 요구하면서 컴파일러가 떠안아야 할 선택지가 감당하기 어려울 만큼 늘어났다.
• Gluon은 Triton의 텐서 추상화는 그대로 두면서 레이아웃, 워프 특화, 공유 메모리 관리를 프로그래머의 손에 되돌려준다.
• 레이아웃은 0과 1로 이루어진 체 F2 위의 선형 사상으로 정식화되며, 덧셈이 XOR로 바뀌는 순간 스위즐링 공식이 기저에서 자연스럽게 유도된다.
• 2 CTA 모드, TMA 멀티캐스트, 클러스터 런치 컨트롤을 겹겹이 쌓으면 철저한 오토튜닝 없이도 cuBLAS와 대등한 행렬곱이 나온다.
• 대가는 이식성이다. 하드웨어 세대에 직접 대고 프로그래밍하는 방식이라 다른 세대로 커널을 그대로 옮길 수 없다.

쉽게 이해하기

GPU MODE 104회 강의는 Triton을 오늘날의 표준 커널 언어로 키운 개발자들이 직접 나와 새 언어 Gluon을 소개하는 자리였다. 발표는 피터 벨, 케런 저우, 마리오 레스카노 세 사람이 나눠 맡았다. Triton은 2023년 무렵 사실상 유일한 커널 DSL이라 불릴 만큼 자리를 굳혔고, 그 뒤로 비슷한 목표를 가진 언어들이 쏟아져 나왔다.

Triton의 매력은 단순함과 성능의 균형이었다. tl.dot 하나가 텐서 코어로 들어가는 관문이었고, num_stages를 4로 두면 로드가 알아서 파이프라이닝됐다. 슬라이드 한 장에 들어가는 행렬곱 커널이 Blackwell에서 이론 최대 성능의 약 80%를 냈고, 조금 더 공을 들이면 95% 근처까지 올라갔다.

문제는 하드웨어가 변했다는 점이다. Ampere 시절에는 네 개의 워프가 나란히 데이터를 나눠 싣고 동기화한 뒤 함께 행렬 연산을 수행하는 식이라 Triton의 프로그래밍 모델과 잘 맞았지만, Blackwell에서는 워프 하나가 TMA 명령 한 번으로 거대한 데이터 덩어리를 끌어와 대역폭을 채우고 행렬 연산도 워프 하나가 텐서 코어를 포화시킨다. 첫 네 워프는 에필로그만 처리하고 나머지가 로드와 연산을 전담하는 구조로 바뀌면서 컴파일러가 내려야 할 판단이 함께 폭증했고, 발표팀은 플래시 어텐션처럼 복잡한 커널까지 모두 최적으로 다루는 것은 현실적으로 불가능하다고 봤다. 다른 진영의 입장을 묻자 돌아온 답은 경험담이었는데, 벤치마크에서 좋던 도구에 실제 프로덕션 커널을 던졌더니 컴파일되는 것이 하나도 없었고 레지스터를 적게 잡으면 스필링으로 성능이 무너지고 넉넉히 잡으면 PTXAS가 오류를 내며 포기했다고 한다.

그래서 Gluon은 모든 텐서 연산에 레이아웃 인자를 붙여, 추상적인 텐서가 하드웨어의 어느 자리에 놓일지를 사람이 지정하게 한다. 규칙 하나가 인상적인데, 컴파일러는 결코 새 레이아웃을 지어내지 않고 코드에 적힌 것에서만 가져온다. 그 위에 fork-join 형태의 워프 특화, 공유 메모리 직접 할당, 배리어와 비동기 로드 같은 하드웨어 명령이 그대로 노출된다. 그렇다고 컴파일러가 놀지는 않아서, 포인터 정렬을 분석해 로드를 벡터화하고 임의의 레이아웃 사이 변환을 뱅크 충돌이 적게 처리하며 CUDA로는 50줄이 될 합산을 한 줄로 만든다.

마리오가 맡은 선형 레이아웃 설명은 이 언어의 수학적 뼈대다. 선형 레이아웃은 0과 1만 있는 체 F2 위의 선형 사상, 즉 0과 1로 채워진 행렬이고 벡터에 적용할 때 곱셈은 AND, 덧셈은 XOR로 바뀌는데, XOR이 등장하는 순간 그동안 왜 맞는지 모른 채 외우던 스위즐링 공식이 정준 기저를 읽어 유도되는 대상이 된다. 케런의 사례 연구는 이것이 실제 성능으로 이어지는지를 보여주는데, 비동기 명령인 tcgen05 행렬 연산에서 출발해 B 타일을 두 CTA가 나눠 갖는 2 CTA 모드, CTA 하나만 로드 명령을 내보내 L2 트래픽을 줄이는 TMA 멀티캐스트, 작업량이 치우쳤을 때 다음 타일을 밀어 넣는 클러스터 런치 컨트롤을 차례로 얹는다. 최종 형태는 스케줄링·에필로그·행렬 연산·로드가 각각 워크 그룹을 맡아 완전히 겹쳐 도는 파이프라인이고, 블록 크기와 스테이지 수를 전부 훑는 오토튜닝 없이도 cuBLAS와 대등한 성능이 나왔다. 여기에 구간을 골라 재는 Proton 프로파일러와 레이아웃을 눈으로 확인하는 웹 시각화 도구가 함께 제공된다.

주요 인사이트

• '컴파일러를 더 똑똑하게 만들면 된다'는 입장과 '하드웨어가 원하는 것을 그대로 노출하라'는 입장이 맞서는 자리에서, 발표팀은 벤치마크 점수가 아니라 프로덕션 커널이 컴파일조차 되지 않았다는 경험을 근거로 후자를 택했다.
• 레이아웃을 선형대수로 정식화한 것이 이 강의의 진짜 성과다. 마법 공식처럼 전해지던 스위즐링이 기저를 읽어 유도할 수 있는 대상이 되면서, 최적의 스위즐을 찾는 일도 사람의 감이 아니라 계산의 문제가 됐다.
• 저수준으로 내려간다고 컴파일러의 몫이 사라지지는 않는다. 동기화와 펜스 삽입, 뱅크 충돌 최소화, 공유 메모리 정적 패킹처럼 사람이 놓치기 쉬운 일은 여전히 자동으로 처리된다.
• CUDA에서 특정 명령을 스레드 하나로 내보내기 위해 쓰는 관례는 하드웨어의 본질이 아니라, 스레드 단위 코드를 SIMD로 바꾸는 과정에서 새어 나온 추상화의 틈이라는 지적이 나왔다. Gluon은 그 판단을 언어 내부로 감춘다.
• 성능과 이식성을 맞바꾸는 선택이므로 모든 커널을 이 언어로 쓸 이유는 없다. 개선점은 다시 Triton으로 되먹임된다는 점도 함께 언급됐다.

자주 묻는 질문

관련 AI 소식