GPU 고속 동작 돕는 AI 도메인 특화 언어 개발했다

스탠퍼드 대학 연구팀이 GPU를 최대한 활용해 단위 시간당 연산량을 최대화하기 위한 도메인 특화 언어(DSL)인 썬더키튼(ThunderKittens)을 출시했다.

연구팀은 엔비디아 H100을 사용해 GPU 활용률 최대화에 주력했다. H100은 텐서코어를 사용하는 반정밀도 행렬곱 연산 성능이 989TFLOPS로 다른 모든 연산 능력 합계인 60TFLOPS를 크게 웃돈다. 그러니까 H100 GPU 활용률은 대부분 텐서코어 사용률에 의존한다. 연구팀은 WGMMA 명령, 주소 생성, 공유 메모리, 점유율 등 4가지 부분을 집중적으로 개선해 모든 GPU 사이클에서 텐서코어가 작업할 수 있도록 했다.

먼저 WGMMA 명령. H100에는 WGMMA(warp group matrix multiply accumulate)라는 새로운 명령어 세트가 추가됐다. WGMMA 명령어를 사용하면 스트리밍 멀티프로세서(SM) 128개 스레드가 협력적으로 동기화해 공유 메모리에서 직접 행렬 연산을 수행한다. 연구팀이 진행한 마이크로벤치마크에 따르면 WGMMA 명령어를 사용하지 않으면 GPU 활용률이 63%에서 머무르게 된다. 하지만 WGMMA 명령어를 사용할 때 공유 메모리에 데이터를 어떻게 배치해야 하는지가 복잡한 문제였고 엔비디아 문서에 오류가 있어 연구팀이 적절한 데이터 배치를 위해 많은 노력을 기울였다고 한다. 그렇지만 WGMMA 명령어를 사용하지 않으면 GPU 활용률 37%를 잃게 되어 피할 수 없는 문제였다.

다음은 주소 생성. H100은 텐서코어와 메모리 모두 빠르게 작동하므로 데이터를 가져오기 위한 메모리 주소를 생성하는 것만으로도 칩 리소스를 상당히 소모한다. 엔비디아가 제공하는 TMA(Tensor Memory Accelerator) 명령어를 사용하면 글로벌 메모리나 공유 메모리에서 다차원 텐서 레이아웃을 지정하고 텐서 일부를 비동기적으로 가져올 수 있다. TMA를 사용하면 주소 생성 비용을 크게 절감할 수 있다.

이어 공유 메모리. 공유 메모리의 단일 액세스 레이턴시는 30 사이클로 비교적 작지만 과거에는 다른 부분이 병목 지점이어서 무시됐다. 하지만 이번 같은 최대 최적화 작업에선 이런 작은 레이턴시도 고려해야 했다. 연구팀은 레지스터와 공유 메모리간 데이터 이동을 최소화하고 데이터를 이동해야 할 때는 WGMMA나 TMA 명령어를 사용해 비동기적으로 공유 메모리와 레지스터간 데이터를 이동시켰다.

다음은 점유율. 점유율은 GPU가 실행할 수 있는 최대 Warp 수에 비해 실제로 실행한 Warp 수 비율이다. H100은 칩 비동기 기능이 강화되어 메모리 페치, 행렬곱 실행, 공유 메모리 제거, 레지스터 연산 등을 동시에 실행할 수 있으므로 이전 세대 하드웨어보다 낮은 점유율에서도 성능을 높일 수 있다. 그래도 점유율이 높을수록 하드웨어 실제 성능이 향상되기 쉽다. A100이나 RTX 4090 등 다른 하드웨어는 H100보다 동기 명령 디스패치에 더 의존하므로 점유율 향상이 중요하다.

이런 요소를 개선하기 위해 연구팀은 CUDA에 내장할 도메인 특화 언어(DSL)로 썬더키튼을 설계하고 출시했다. 실제로 파이토치(PyTorch) FA2와 썬터키튼(TK)를 사용했을 때 플래시 어텐션(Flash Attention) 계산 능력을 측정한 결과 평균적으로 썬더키튼이 30% 성능 향상을 이뤘다. 또 라이너 어텐션(Linear Attention) 계산에서 썬더키튼은 215TFLOPS 계산 성능을 발휘해 기존 방식보다 큰 폭으로 고속화됐다고 한다. 썬더키튼 코드는 깃허브에서 오픈소스 라이선스로 공개되어 있다. 관련 내용은 이곳에서 확인할 수 있다.