HipKittens로 본 AMD GPU 커널 최적화, 왜 어셈블리까지 내려가야 했나

핵심 메시지

• AMD의 MI355는 사양표상 엔비디아 B200에 뒤지지 않지만, 성숙한 커널 소프트웨어가 없어 그 성능이 AI 워크로드에서 잠겨 있다는 것이 강연의 문제의식이다.
• HipKittens는 스탠퍼드의 ThunderKittens가 쓰던 타일 중심 추상화가 AMD 하드웨어에도 통하는지 확인하려는 시도이며, 인터페이스는 유지하되 내부는 처음부터 다시 구현했다.
• 최적화는 메모리 이동, 프로세서 내부 점유율, 캐시 재사용이라는 세 축으로 정리된다.
• 엔비디아에서 표준이 된 생산자-소비자 워프 특화가 AMD에서는 그대로 통하지 않는다. 레지스터 재할당이 없다는 하드웨어 차이 때문이며, 대신 두 웨이브가 번갈아 일하는 핑퐁 스케줄을 찾아냈다.
• 칩렛 구조 탓에 L2 캐시가 분산돼 있어, 스레드 블록을 어떤 순서로 배치하느냐가 곧 성능이 된다.

쉽게 이해하기

GPU MODE 강연에 나온 HipKittens의 주 개발자는 문제를 돈으로 환산하며 시작한다. 분기 GPU 매출이 500억 달러를 넘고 이 하드웨어 위에서 학습과 추론이 몇 달씩 돌아가는데, 하드웨어에 맞춰 잘 짜인 알고리즘도 다음 세대 하드웨어에서는 성능이 무너진다는 것이다. 널리 쓰이는 플래시 어텐션 커널조차 A100에서 H100으로 옮겨 가며 성능이 크게 떨어졌고, 제대로 된 커널이 나오기까지 H100 출시 후 2년이 걸렸다고 그는 말한다. 결국 나쁜 소프트웨어가 수십억 달러어치 연산을 낭비하는 셈이다.

AMD 쪽 사정은 더 나쁘다. 강연자는 AMD 소프트웨어 생태계에 존재하는 선택지를 하나씩 검토한다. 파이토치는 쓰기 쉽지만 새 아키텍처를 즉시 지원하지 못하고 AMD 백엔드가 더 느리다. 트라이톤·모조·타일랭 같은 컴파일러는 사용성을 얻는 대신 성능을 내주는데, 특히 AMD에서는 레지스터 수명 추적과 메모리 접근 하강이 잘 되지 않아 레지스터 스필이 생기고 기본적인 행렬 곱에서도 성능이 나오지 않았다고 한다. 남은 선택지는 AMD가 내부적으로 쓰는 방식, 즉 순수 어셈블리다. 컴파일러가 방해하지 않아 빠르지만 읽기도 고치기도 어렵고 전문가가 있어야 하며, 무엇보다 모든 워크로드로 확장되지 않는다.

하드웨어 차이도 구체적이다. AMD에서는 워프 대신 웨이브라 부르고 한 웨이브가 32개가 아닌 64개 스레드로 이뤄지며, 레지스터 파일은 두 배 크지만 공유 메모리는 상대적으로 작다. 다이 하나가 여덟 개의 작은 칩렛으로 나뉘고 각 칩렛이 자기 L2 캐시를 가지며 마지막 레벨 캐시를 공유한다. 엔비디아에 있는 텐서 메모리 가속기나 세밀한 비동기 제어, 레지스터 재할당 같은 장치는 없거나 다르게 동작한다. 강연자는 숫자가 크다고 더 좋은 하드웨어라는 뜻은 아니며, AMD가 소프트웨어를 못 만든다기보다 하드웨어 자체가 프로그래밍하기 더 까다롭다는 쪽에 가깝다고 정리한다.

메모리 이동에서 가장 성가셨던 것은 레지스터 배치다. 엔비디아에는 16×16 코어 행렬이라는 규칙적인 단위가 있어 큰 타일을 그 단위의 반복으로 다룰 수 있지만, AMD의 행렬 코어 명령에는 그런 단위가 없다. 그래서 HipKittens에서는 레지스터 타일을 만들 때 어떤 기본 모양으로 쌓을지 프로그래머가 지정해야 한다. 여기에 컴파일러 문제까지 겹친다. 512개 레지스터가 누산기용과 벡터용으로 갈리는데 컴파일러가 누산기 레지스터를 행렬 코어 입력으로 코드 생성하지 못해 스필이 발생했고, 개발자가 레지스터 번호 범위를 직접 지정하는 '명시적 레지스터 고정'을 도입해 우회했다.

공유 메모리에서는 문서화되지 않은 동작을 직접 역공학해야 했다. 같은 공유 메모리 읽기라도 스레드당 128비트를 읽는 명령은 뱅크가 64개처럼 동작하고, 96비트나 64비트 명령에서는 32개처럼 동작해 이유 없는 뱅크 충돌이 생겼다. AMD 엔지니어에게 물어도 답이 오지 않아, 스레드를 하나씩 마스킹해 어떤 스레드들이 같은 단계에 참여하는지 찾고 뱅크 충돌이 나는 지점을 세어 뱅크 수를 알아내는 방식으로 밝혀냈다. 이 분석 스크립트는 저장소에 함께 공개돼 있다.

주요 인사이트

• 점유율 확보 방식이 플랫폼마다 다르다는 점이 이 강연의 핵심 발견이다. 엔비디아에서는 데이터를 나르는 생산자 워프와 계산하는 소비자 워프를 나누는 방식이 표준이지만, AMD에서는 레지스터를 재할당할 수 없어 생산자가 필요 이상으로 레지스터를 붙들고 소비자가 굶는 비효율이 생긴다.
• 대안으로 제시된 것은 SIMD당 두 웨이브를 두고 한쪽이 대량 메모리 연산을 하는 동안 다른 쪽이 대량 계산을 하다가 서로 역할을 바꾸는 8웨이브 핑퐁 스케줄이다. 반복문 진입 전에 한 웨이브를 조건부 배리어로 붙잡아 두어 계속 어긋난 상태를 유지하는, 강연자 표현으로 다소 편법적인 방식이다.
• 칩렛 구조에서는 작업 배치 순서 자체가 최적화 대상이다. 스레드 블록이 행 우선 순서로 배정되면 같은 L2를 공유하는 블록들이 서로 겹치지 않는 타일을 읽어 재사용이 사라진다. HipKittens는 하나의 칩렛에 연속으로 배정할 블록 수와 창 크기를 조절하는 방식으로 이 순서를 조정한다.
• 성능 결과로는 500줄 남짓한 코드로 쓴 어텐션 순방향 커널이 좋은 성능을 냈고, 순수 어셈블리 기준선과 비교한 BF16 행렬 곱도 경쟁력이 있었으며, 어셈블리로는 확장이 어려운 어텐션 역방향에서는 오히려 더 나은 경우가 있었다고 소개된다.
• 라이브러리를 별도로 만든 이유에 대해 강연자는 추상화 수준이 어셈블리 바로 위라 하드웨어 중립적일 수 없기 때문이라고 답한다. 다만 레지스터 타일·공유 타일 같은 인터페이스는 ThunderKittens와 거의 같게 유지해, 두 코드를 서로 알아볼 수 있게 했다.

자주 묻는 질문

관련 AI 소식