▲

GN⁺: GPU의 폭발적인 성능 향상

 (hazyresearch.stanford.edu)
1P by neo 6달전 | favorite | 댓글 1개

H100 GPU의 특징

  • 80GB HBM3 메모리와 3TB/s 대역폭 제공 (실제로는 조금 더 낮음)
  • 50MB L2 캐시와 12TB/s 대역폭 제공. GPU에 두 개의 25MB 섹션으로 분할되어 크로스바로 연결됨 (크로스바는 성능 저하 요인)
  • 132개의 Streaming Multiprocessor(SM)로 구성되며, 각 SM은 아래 구성을 가짐:
    • 256KB L1 캐시 내에 최대 227KB의 공유 메모리 제공 (합쳐서 약 33TB/s 대역폭)
    • 비동기 주소 생성 및 메모리 fetching이 가능한 Tensor Memory Accelerator(TMA) 제공. 온칩 메모리 네트워크 지원 등의 기능도 제공하지만 이번 포스트에서는 다루지 않음
    • 4개의 쿼드런트로 나뉘며, 각 쿼드런트는 warp scheduler, 512개 vector register(각각 32개의 4-byte 워드 포함), matrix multiply용 tensor core, sum/multiply 등의 병렬 연산을 지원하는 내장 instruction으로 구성됨

H100 GPU 성능 최적화 팁

  • Tensor Core를 최대한 활용하는 것이 중요. H100은 989 TFLOPs의 FP16 행렬곱 연산 성능을 가지므로, Tensor Core 활용도에 따라 전체 GPU 활용률이 크게 좌우됨
  • 단, Tensor Core를 최대한 활용하기 위해서는 아래 사항들을 고려해야 함:
    • Warp Group Matrix Multiply Accumulate(WGMMA) 명령어 활용이 필수적이지만 사용이 까다로움
    • Shared Memory가 생각보다 빠르지 않고, 사용시 세심한 주의가 필요함
    • 주소 생성이 비용이 크므로 Tensor Memory Accelerator(TMA) 등을 활용해 최적화해야 함
    • Occupancy를 높이는 것이 여전히 도움이 되며, register가 주요 리소스임
  • 이런 특징들은 H100 뿐만 아니라 다른 GPU에도 어느정도 적용되지만, 특히 H100에서 Tensor Core 활용이 중요하고 까다로운 편

ThunderKittens : H100에 최적화된 CUDA 임베디드 DSL

  • NVIDIA H100 등 최신 GPU의 성능을 최대한 끌어내기 위해 개발된 CUDA 기반 임베디드 DSL
  • 아래 4개의 타일 기반 템플릿 타입을 제공:
    • Register Tile (레지스터에 저장된 2D 텐서)
    • Register Vector (레지스터에 저장된 1D 텐서)
    • Shared Tile (Shared Memory에 저장된 2D 텐서)
    • Shared Vector (Shared Memory에 저장된 1D 텐서)
  • 또한 타일 조작을 위한 다양한 연산자 (exp, mul, sum 등)를 warp 또는 warp 그룹 레벨에서 제공
  • 기존 Flash Attention, Flash Attention 2 커널을 ThunderKittens으로 구현하니 코드가 크게 간결해지고, H100에서 최대 30% 성능 향상
  • Based Linear Attention 커널도 ThunderKittens으로 구현하여 215 TFLOPs 성능 달성 (알고리즘 특성 상 recompute 포함하면 300 TFLOPs 이상)

철학적 관점에서의 고찰

  • ThunderKittens이 잘 작동하는 이유는 모든 것을 다 지원하려 하지 않고, 단순하면서도 GPU 아키텍처에 잘 들어맞는 타일 기반 추상화를 제공하기 때문
  • 전통적인 32-bit 워드 대신 1024-bit 벡터 레지스터를 사용하는 것도 발전이지만, 16x16 타일을 레지스터의 단위로 보는 패러다임의 전환이 필요
  • AI 워크로드가 결국 행렬곱, reduction, reshape 위주라는 점에서 타일 기반 접근이 타당하며, 하드웨어 관점에서도 systolic array 외에 작은 행렬곱을 지원하는 방향으로 진화할 것
  • 더 나아가 하드웨어에 최적화된 형태로 AI 알고리즘을 설계하는 방식으로 사고를 전환할 필요가 있음. 예를 들어 RNN의 상태 크기를 SM에 들어갈 만큼으로 제한하고, 연산 밀도도 하드웨어가 요구하는 수준으로 맞추는 식의 접근이 필요

GN⁺의 의견

  • ThunderKittens은 CUDA에 익숙한 개발자들에게는 매력적인 선택지가 될 수 있음. 기존 커널 코드를 크게 건드리지 않고도 쉽게 성능 향상을 누릴 수 있기 때문
  • 다만 초보자에게는 여전히 진입장벽이 높을 수 있음. 앞으로 더 다양한 예제 코드와 학습 자료가 뒷받침되어야 할 듯
  • H100 뿐만 아니라 AMD 등 다른 GPU에도 ThunderKittens 지원을 확대한다는 계획이 흥미로움. 벤더 종속성을 낮추는데 기여할 수 있을 것으로 보임
  • 궁극적으로는 AI 모델/알고리즘 자체를 하드웨어에 최적화된 형태로 설계하는 것이 매우 중요한 포인트임. 이를 위해서는 하드웨어 특성에 대한 깊은 이해가 선행되어야 하는데, ThunderKittens은 개발자들이 그런 인사이트를 얻는데 도움이 될 수 있을 것
  • Hazy Research의 지속적인 연구개발과 오픈소스 기여가 CUDA 생태계 활성화에 큰 도움이 될 것으로 기대됨. 다만 장기적으로는 보다 추상화 레벨이 높은 프레임워크의 등장도 필요해 보임
▲
neo 6달전  [-]
Hacker News 의견

  • AI 하드웨어의 요구 사항이 점점 명확해지고 있음. GPU는 원래 다른 목적으로 설계되었지만, 좋은 행렬 곱셈 하드웨어를 가지고 있어 AI에 사용되고 있음. "AI GPU"는 실제 GPU의 일부 기능을 제외할 수 있으며, 더 짧은 숫자(16비트, 8비트, 2비트, 1비트 부동 소수점)로 향하는 추세가 있음. 이 논문은 16x16 타일을 좋아하는 하드웨어가 많은 이점이 있음을 시사함.

  • AI 전용 보조 프로세서(NPU)가 필요함. 특히 개발자, 전문가, 게이머 등을 위한 프로슈머급 데스크톱 시스템에서 필요함. GPU는 기업에서는 작동하지만 개인 컴퓨팅 측면에서 AI에 사용하기에는 불편함. 특히 VRAM 제한과 Vulkan 이외의 표준 개방형 API 부재가 문제임.

  • AI 연구를 발전시키려면 신경과학, 심리학을 더 잘 연구해야 함. 또한 신경망의 그래프 토폴로지와 관련된 것들도 관련이 있을 수 있음.

  • CUTLASS를 사용자 친화적으로 만든 것인가?

  • ThunderKittens 마스코트는 고양이/소니 아이보 분위기가 있음. AI로 잘 생성된 것으로 보임.

  • 유니버설 베이직 컴퓨트(UBC)가 유니버설 베이직 인컴의 대체재로 고려된다면, 그것은 매우 디스토피아적인 미래가 될 것임. 엔비디아 같은 한 회사가 모든 컴퓨팅을 만든다고 상상해보라.

답변달기