- AI 연산 비용이 커지는 상황에서 Hazy Research는 NVIDIA H100의 tensor core를 쉬지 않게 유지하는 것이 GPU 성능 최적화의 핵심이라고 정리함
- H100은 half-precision 행렬곱에서 989 TFLOPs를 내지만 일반 연산은 약 60 TFLOPs에 그쳐, tensor core가 멈추는 순간 활용률이 크게 떨어짐
- 최대 성능에 가까워지려면 WGMMA, shared memory 배치, 주소 생성, occupancy를 함께 다뤄야 하며
wgmma.mma_async없이는 마이크로벤치마크에서 피크의 약 63%에 머무름 - 공개된 CUDA 내장 DSL ThunderKittens는 tile·vector 추상화로 swizzling과 register layout 같은 복잡성을 감싸 FlashAttention 계열 커널 작성을 단순화함
- H100용 FlashAttention-2 forward 커널은 약 100줄로 작성되며 FlashAttention-2보다 약 30% 빠르고, Based linear attention 커널은 215 TFLOPs로 동작함
H100 성능을 좌우하는 조건
- AI는 많은 컴퓨트를 사용하며, Hazy Research는 최근 몇 년간 AI가 더 적은 컴퓨트를 쓰거나 주어진 컴퓨트에서 더 효율적으로 돌도록 하는 작업을 해옴
- 컴퓨트 절감 예시로 Based, Monarch Mixer, H3, Hyena, S4가 있음
- 효율적 실행 예시로 FlashAttention, FlashAttention-2, FlashFFTConv가 있음
- 실용적 목표는 GPU를 빠르게 만드는 데서 배운 내용을 정리하고, 빠른 커널 작성을 돕는 CUDA 내장 DSL ThunderKittens를 공개하는 것임
- 더 넓게는 하드웨어 이해가 AI 컴퓨트를 바라보는 방식을 어떻게 바꿨는지 다룸
NVIDIA H100의 구조와 병목
- H100 SXM GPU는 다음 구성을 기준으로 논의됨
- 80GB HBM3, 대역폭 3TB/s
- 50MB L2 캐시, 대역폭 12TB/s, GPU 전반에 25MB 섹션 2개로 나뉘고 crossbar로 연결됨
-
132개 SM
- 각 SM은 최대 227KB shared memory를 포함한 256KB L1 캐시를 가지며, 함께 약 33TB/s 대역폭을 가짐
- Hopper의 새 하드웨어인 Tensor Memory Accelerator(TMA) 가 비동기 주소 생성과 메모리 fetch를 담당함
- 각 SM은 4개 quadrant로 구성되고, 각 quadrant에는 warp scheduler, 512개 vector register, 행렬곱용 tensor core, 병렬 내장 명령들이 있음
- 모든 컴퓨트는 SM에서 일어나며, 대부분은 register에서 처리됨
- H100에서 성능을 내는 핵심은 tensor core를 계속 fed 상태로 유지하는 것임
- H100은 half-precision 행렬곱에서 989 TFLOPs, “그 외” 연산에서 약 60 TFLOPs를 제공함
- tensor core가 사용되는 사이클에는 최소 94%의 하드웨어 활용률에 도달함
- tensor core가 사용되지 않는 사이클에는 최대 6% 활용률에 머묾
WGMMA: 필요하지만 까다로운 명령
- H100에는 warp group matrix multiply accumulate 명령인
wgmma.mma_async가 있음- PTX에서는
wgmma.mma_async - SASS에서는
HGMMA/IGMMA/QGMMA/BGMMA
- PTX에서는
- 이전 GPU의
wmma.mma.sync,mma.sync는 32개 thread로 구성된 warp 하나가 tensor core에 데이터를 넣고 결과를 기다리는 동기 방식이었음 wgmma.mma_async는 연속된 128개 thread가 SM의 모든 quadrant에 걸쳐 협력적으로 동기화하고, shared memory에서 직접 비동기 행렬곱을 시작함- warp들은 행렬곱이 진행되는 동안 register로 다른 작업을 할 수 있음
- 결과는 원하는 시점에 기다릴 수 있음
- 마이크로벤치마크에서 H100의 전체 compute를 끌어내려면 이 명령들이 필요했음
- 사용하지 않으면 GPU가 피크 활용률의 약 63% 에서 머무르는 것으로 관찰됨
- tensor core가 로컬 리소스에서도 깊은 하드웨어 파이프라인을 요구하기 때문일 수 있음
- 가장 큰 난점은 memory layout의 복잡성임
- unswizzled shared memory layout은 coalescing이 매우 나빠 L2 대역폭을 많이 요구함
- swizzled layout은 문서가 잘못되어 있어 파악하는 데 시간이 걸렸음
- swizzled layout은 특정 행렬 shape에서만 동작하는 것으로 보이고,
wgmma.mma_async의 다른 기능과 잘 맞지 않음 - 하드웨어는 tensor core로 가는 중 sub-matrix transpose를 할 수 있지만, layout이 swizzled가 아닐 때만 가능함
- FlashAttention 같은 커널에서는 TMA와 L2 캐시가 충분히 빨라 이 문제를 어느 정도 숨길 수 있음
- 하드웨어를 완전히 쓰려면 memory request를 coalescing하고 bank conflict를 피해야 하므로 layout 제어가 중요함
Shared memory와 bank conflict
- Shared memory의 single-access latency는 약 30 cycles로 보이며, 이 시간 동안 SM의 tensor core는 거의 두 번의 32x32 정방 행렬곱을 수행할 수 있음
- FlashAttention 같은 이전 작업에서는 주로 HBM-SRAM 병목에 집중했고, 과거에는 이 병목이 실제로 중요했음
- HBM이 빨라지고 tensor core가 칩의 다른 부분보다 더 빠르게 커지면서, shared memory의 작은 latency도 제거하거나 숨겨야 하는 대상이 됨
- Shared memory는 32개 bank로 나뉘어 있어 조심하지 않으면 bank conflict가 발생함
- 같은 memory bank에 여러 다른 memory 조각을 동시에 요청하면 요청이 직렬화됨
- 경험상 커널이 불균형하게 느려질 수 있음
- WGMMA와 MMA 명령이 요구하는 register layout은 단순하게 쓰면 bank conflict를 겪을 수 있음
- 해결책은 여러 swizzling 패턴으로 shared memory를 재배치해 conflict를 피하는 것임
- 가능하면 register와 shared memory 사이 이동을 피하고, 필요할 때는 WGMMA와 TMA 같은 내장 하드웨어로 비동기 데이터 이동을 수행하는 편이 유리함
- 실제 warp를 사용한 동기 이동은 가장 일반적이지만 최악의 fallback에 가까움
주소 생성과 TMA
- H100은 tensor core와 memory가 모두 빨라서, fetch할 memory address를 생성하는 작업 자체가 칩 리소스의 상당 부분을 차지함
- 복잡한 interleaved pattern이나 swizzling pattern이 추가되면 더 두드러짐
- NVIDIA의 Tensor Memory Accelerator(TMA) 는 global/shared memory의 다차원 tensor layout을 지정하고, 해당 tensor의 subtile을 비동기로 fetch한 뒤 완료 시 barrier를 트리거할 수 있게 함
- TMA는 주소 생성 비용을 줄이고 pipeline 구성도 쉽게 만듦
- TMA는
wgmma.mma_async처럼 H100의 잠재력을 끌어내는 데 필수로 평가됨- 경험상 WGMMA보다 더 중요할 수도 있음
- register 리소스와 instruction dispatch를 절약함
- global memory에 비동기 reduction을 수행하는 기능도 있어 복잡한 backward kernel에서 유용함
- TMA 역시 swizzling mode를 이해하려면 일부 reverse engineering이 필요했지만, WGMMA보다는 덜 고통스러웠음
Occupancy가 숨겨주는 비용
- CUDA에서 occupancy는 같은 실행 하드웨어에 co-scheduled된 thread 수를 뜻함
- SM quadrant의 warp scheduler는 매 cycle마다 명령을 받을 준비가 된 warp에 instruction을 issue하려고 함
- H100은 이전 세대보다 occupancy에 덜 의존하는 면이 있음
- 비동기 기능 덕분에 단일 instruction stream도 memory fetch, matrix multiply, shared memory reduction, register math를 동시에 바쁘게 만들 수 있음
- 하지만 occupancy는 실수와 동기화 비용을 숨기는 데 매우 유용함
- 완벽히 설계된 pipeline은 추가 occupancy 없이도 빠를 수 있음
- 실제 관찰에서는 NVIDIA GPU가 occupancy를 염두에 두고 설계된 것으로 보였음
- synchronization과 실수 가능성이 많기 때문에 occupancy를 높이면 실현 하드웨어 활용률이 좋아지는 경우가 많았음
- H100에서는 occupancy가 유용한 수준이지만, A100과 RTX 4090에서는 각각 더 중요해졌다고 봄
- H100 대비 동기 instruction dispatch에 더 의존하기 때문일 가능성을 언급함
ThunderKittens: CUDA 안의 작은 DSL
- ThunderKittens는 H100에서 빠른 커널을 쉽게 작성하기 위해 만든 CUDA 내장 DSL임
- 초기에는 연구실 내부 사용을 위해 만들었고, 이후 공개됨
- 이름은 kittens가 귀엽고 코드에서
kittens::를 입력하게 하는 것이 재미있다고 생각해 붙임 - ThunderKittens는 단순함을 목표로 하며 네 가지 templated type을 제공함
- Register tiles: register file 위의 2D tensor
- Register vectors: register file 위의 1D tensor
- Shared tiles: shared memory 안의 2D tensor
- Shared vectors: shared memory 안의 1D tensor
- Tile은 height, width, layout으로 parameterized됨
- Register vector는 length와 layout으로 parameterized되고, shared vector는 length만 사용함
- shared vector는 일반적으로 bank conflict를 겪지 않음
- 제공 연산은 warp level 또는 협력 warp group level에서 tile·vector를 조작함
- initializer: shared vector를 zero로 만드는 작업 등
- unary op:
exp등 - binary op:
mul등 - row/column op:
row_sum등
- ThunderKittens는 CUDA 안에 내장되어 있어 Triton 같은 라이브러리와 달리 추상화가 “gracefully” 실패한다고 설명함
- 빠진 기능이 있으면 원하는 방식으로 확장할 수 있음
FlashAttention 예시와 성능
- ThunderKittens 예시로 RTX 4090용 간단한 forward FlashAttention 커널이 제시됨
- headdim=64만 다룸
n은 256의 배수여야 함- 약 60줄의 CUDA 코드로 작성됨
- 하드웨어 활용률은 75%
- 복잡성 대부분은 swizzling pattern이나 register layout이 아니라 알고리듬 자체에 있음
- H100용 FlashAttention-2 forward pass도 ThunderKittens로 작성됨
- TMA, WGMMA, swizzling mode, descriptor의 복잡성을 ThunderKittens가 감쌈
- 커널은 약 100줄
- H100에서 FlashAttention-2보다 약 30% 빠름
- ThunderKittens는 GPU에서 쓸 수 있는 “mini-pytorch”처럼 layout과 instruction을 감싸고 primitive를 제공함
- Based linear attention과 앞으로 공개될 다른 architecture용 kernel도 함께 공개됨
- Based linear attention kernel은 215 TFLOPs로 동작함
- 알고리듬 자체의 recompute를 고려하면 300 TFLOPs를 넘음
- Linear attention은 이론적으로 더 효율적이지만, 실제 하드웨어에서는 역사적으로 효율이 크게 낮았음
- 이 결과가 높은 throughput application 범위를 넓힐 수 있다고 봄
Tile 중심 사고
- ThunderKittens가 잘 작동한 이유는 모든 것을 하려 하지 않기 때문이라고 봄
- CUDA는 ThunderKittens보다 훨씬 표현력이 높음
- ThunderKittens는 작고 단순한 DSL임
- 핵심 추상화는 small tile이며, 이는 AI와 하드웨어가 향하는 방향과 맞는다고 봄
- ThunderKittens는 16보다 작은 차원을 지원하지 않음
- 하드웨어도 그런 작은 차원을 특별히 원하지 않는다고 봄
- “matrix multiply가 16x16보다 작다면 그것이 AI인지 확신할 수 있느냐”는 식으로 문제를 제기함
- CPU 시대의 32-bit word를 register로 보는 관점은 AI 하드웨어에는 맞지 않는다고 봄
- CUDA의 1024-bit vector register는 올바른 방향의 한 단계로 봄
- 여기서 register는 16x16 tile의 데이터임
- AI는 여전히 matrix multiply, reduction, reshape가 중심이므로 tile 추상화가 AI와 하드웨어 모두에 맞는다고 봄
- 앞으로는 AI 아이디어를 하드웨어에 잘 매핑되는 방식으로 재정렬해야 함
- recurrent state 크기는 SM에 들어갈 수 있을 만큼 커야 함
- compute density는 하드웨어가 요구하는 수준보다 낮아서는 안 됨
- 하드웨어에서 배운 내용을 AI 설계에 맞추는 것이 앞으로의 중요한 방향임
AMD 지원 계획
- ThunderKittens의 AMD hardware 지원이 곧 나올 예정임