ThunderKittens 해부: 고성능 AI 커널을 위한 컴팩트한 DSL의 구조

서론 (Introduction)

현대의 머신러닝 (ML) 워크로드는 커스텀 GPU 커널 (GPU kernels)에 크게 의존합니다. 모델이 깔끔한 텐서 연산 (tensor operations)으로 표현되더라도, 성능은 거의 항상 그 밑단에 있는 특화된 구현체로부터 나옵니다. 이에 대한 좋은 예로는 다양한 어텐션 메커니즘 (attention mechanisms), 다양한 정밀도 (precisions)에 걸친 GEMM, 그리고 오늘날 최첨단 모델에서 상당히 흔한 아키텍처적 선택이 된 MoE 스타일의 그룹화된 GEMM (grouped GEMMs) 등이 있습니다.

스케일링 법칙 (scaling laws)의 관점에서 본다면 이는 매우 중요합니다. 더 나은 모델은 일반적으로 더 나은 알고리즘, 더 많은 데이터, 그리고 더 많은 컴퓨팅 자원의 조합을 통해 만들어졌습니다. 우리가 이를 계속해서 밀어붙이고자 한다면, 알고리즘의 품질뿐만 아니라 해당 알고리즘이 하드웨어에서 실제로 얼마나 효율적으로 실행되는지에도 관심을 가져야 합니다. Tri Dao가 언급했듯이, 이를 명확하게 표현하는 한 가지 방법은 다음과 같습니다:

우리는 두 항을 모두 개선하고 싶습니다. 알고리즘 측면에서 연구자들은 새로운 아키텍처와 새로운 학습 및 추론 레시피 (training and inference recipes)를 빠르게 반복해야 합니다. 하지만 이 모든 것이 대규모로 의미를 갖기 위해서는, 실제 하드웨어에서 실제로 빠르게 실행되는 코드로 변환되어야 합니다. 여기에는 지속적인 긴장 관계가 존재합니다. 우리는 연구에 충분히 생산적인 프로그래밍 환경을 원하면서도, 동시에 본격적인 성능을 얻고 잘 확장될 수 있도록 하드웨어에 충분히 가까운 (close to the metal) 환경을 원합니다.

이것이 바로 GPU 프로그래밍 DSL (DSLs)이 차지하는 영역이며, 이들은 상당히 넓은 범위를 아우릅니다.

최상위 계층에서는 PyTorch와 같은 프레임워크를 통해 연구자들이 GPU를 전혀 신경 쓰지 않고도 텐서 표현식 (tensor expressions)을 작성할 수 있습니다. 프레임워크가 커널 디스패치 (kernel dispatch)를 처리하며, PyTorch 2에서는 TorchDynamo + TorchInductor가 종종 Triton을 생성함으로써 경쟁력 있는 GPU 코드를 만들어낼 수 있습니다. 한 단계 낮은 수준에서는 Triton이 대부분의 CUDA 복잡성을 숨기면서도 타일링 (tiling), 메모리 액세스 패턴 (memory access patterns), 그리고 프로그램 구조에 대해 더 명시적인 제어를 제공합니다. 더 낮은 수준인 CUDA C++, CUTLASS/CuTe, 또는 PTX로 내려가면 직접적인 하드웨어 제어가 가능해지지만, 이제는 메모리 레이아웃 (memory layouts), 워프 동기화 (warp synchronisation), 텐서 코어 스케줄링 (tensor core scheduling), 그리고 수많은 상용구 코드 (boilerplate)를 직접 관리해야 합니다.

더 깊이 내려갈수록 GPU 계층 구조 (GPU hierarchy)에 대해 더 많이 추론할 수 있지만, 무언가를 수행하기 위해 더 많은 전문 지식과 코드가 필요하게 됩니다.

Stanford의 Hazy Research Lab에서 개발한 CUDA 내부의 임베디드 DSL (embedded DSL)인 ThunderKittens는 이 스펙트럼 상에서 진정으로 흥미로운 지점에 위치합니다. 이 기술의 이면에 있는 연구 질문은 명확합니다: 광범위한 AI 워크로드 전반에 걸쳐 빠른 커널을 지원하면서도, 프로그래밍 추상화 (programming abstraction)를 얼마나 작게 유지할 수 있는가?

하드웨어를 완전히 숨기거나 모든 것을 노출하는 대신, TK는 중간 지점을 찾았습니다. 즉, 타일 레이아웃 (tile layouts), 공유 메모리 할당 (shared memory allocation), 레지스터 프래그먼트 (register fragments), TMA 텐서 맵 (TMA tensor maps), 텐서 코어 디스크립터 (tensor core descriptors)와 같은 반복적인 배관 작업 (plumbing)을 추상화하는 동시에, 어떤 데이터가 어디로 이동하는지, 파이프라인 단계 (pipeline stages)가 어떻게 구성되는지, 그리고 작업이 어떻게 스케줄링되는지를 신중하게 추론할 수 있을 만큼 하드웨어에 가깝게 유지합니다. 또한 CUDA에 임베디드되어 있기 때문에, 라이브러리가 노출하지 않는 기능이 필요할 때는 언제든지 로우 CUDA (raw CUDA)나 인라인 PTX (inline PTX)로 내려갈 수 있습니다.

이것이 제가 이 포스트에서 사용하고자 하는 프레임워크입니다. 저는 TK가 어떤 추상화를 노출하는지, 왜 그런 형태를 띠고 있는지, 그리고 이것이 Hopper 및 Blackwell GPU 하드웨어에 어떻게 매핑되는지를 이해하고자 합니다.

우리는 핵심 프로그래밍 모델인 글로벌 레이아웃 (global layouts), 공유 타일 (shared tiles), 레지스터 타일 (register tiles), 벡터 (vectors), 연산 래퍼 (compute wrappers), 그리고 메모리 이동 (memory movement)부터 시작하겠습니다. 그런 다음 최신 Blackwell 전용 추가 기능인 tcgen05를 살펴보겠습니다.

, 2xSM MMA, 텐서 메모리 (tensor memory), 그리고 클러스터 런칭 컨트롤 (Cluster Launch Control)에 대해 알아보겠습니다. 마지막으로, TK의 lcf 파이프라인 템플릿을 사용하여 어텐션 프리필 (attention prefill) 커널을 구축하고, 이를 FlashAttention-2 및 3와 벤치마킹함으로써 모든 내용을 구체화하겠습니다.

ThunderKittens 프로그래밍 모델 및 핵심 추상화 (Core Abstractions)

더 깊이 들어가기 전에, 먼저 한 걸음 물러나 ThunderKittens가 실제로 무엇을 하려고 하는지 질문해 보는 것이 도움이 됩니다.

높은 수준에서 보면, TK는 AI 커널에 잘 매핑되는 주관적인(opinionated) 추상화 세트를 제공합니다. 모든 것을 로우 CUDA (raw CUDA)로 직접 작성하는 대신, TK를 사용하면 타일 (tiles) 단위로 작업하고 해당 타일들에 대한 상위 수준의 연산을 수행할 수 있습니다. 이러한 연산 중 다수는 PyTorch 프리미티브 (primitives)와 다소 유사하게 느껴지며, 이는 머신러닝 (ML) 배경을 가진 사람들에게 커널 개발에 익숙한 느낌을 줍니다.

이것이 바로 프레임워크의 핵심 아이디어입니다. 즉, 현대적인 GPU를 효율적으로 활용하는 데 필요한 제어력을 포기하지 않으면서도, 고성능 커널을 작성하는 복잡성을 줄이는 것입니다.

논문에서는 다음과 같이 설명합니다:

“이러한 모든 하드웨어 기능을 활용하기 위해 수많은 기술이 분명히 필요함에도 불구하고, 우리의 핵심적인 기술적 발견은 실제로 많은 AI 커널에 있어 고성능 커널 작성을 단순화할 수 있는 소수의 핵심 추상화가 존재한다는 것입니다.”

이 주장을 진정으로 이해하기 위해서는 GPU 하드웨어 자체에 대한 견고한 멘탈 모델 (mental model)을 갖는 것이 중요합니다. 워프 (warps), 스레드 블록 (thread blocks), 공유 메모리 (shared memory), 텐서 코어 (tensor cores), 그리고 점유율 (occupancy)과 같은 개념이 아직 직관적으로 느껴지지 않는다면, 제 H100 GEMM 최적화 포스트의 도입부를 먼저 읽어보실 것을 강력히 권장합니다. 최신 Blackwell 구성 요소에 대해서는 이 블로그에서 다룰 예정이니, 그 부분은 저를 믿고 따라와 주시기 바랍니다. 이러한 기초가 갖춰지면, TK 프로그래밍 모델을 이해하기가 훨씬 쉬워집니다.

그 지점부터 우리는 ThunderKittens가 구축하는 핵심 추상화(abstractions)가 무엇인지, 그리고 왜 이것들이 현대 GPU 하드웨어에 매우 자연스럽게 매핑되는지를 살펴보기 시작할 수 있습니다. 간단한 복습을 위해, 여기 GPU 메모리 계층 구조(memory hierarchy)와 그에 대응하는 CUDA 프로그래밍 모델을 TK의 타일 추상화 (tile abstractions) (제 어설픈 삼각형 그림을 양해해 주세요)와 함께 정리했습니다:

[IMG:1]

이제 TK의 타일 추상화가 이러한 메모리 계층 구조의 관점에 어떻게 부합하는지 알 수 있습니다. 타일 추상화는 TK 프로그래밍 모델의 근본적인 빌딩 블록(building blocks)이며, 다른 모든 구성 요소는 그 위에 계층적으로 쌓여 있습니다. 높은 수준에서 보면, 주요 ThunderKittens 추상화는 다음과 같습니다:

타일 추상화 (Tile Abstractions)

TK는 모든 것이 GPU 계층 구조에 깔끔하게 매핑되는 타일(tiles) 단위로 표현되어야 한다는 아이디어를 중심으로 구축되었습니다. 가장 기본적인 수준에서 TK는 **높이가 16으로 고정된 베이스 타일 (base tile)**을 사용하며, 너비는 데이터 타입 (datatype)에 따라 달라집니다. fp16, bf16 및 기타 16개 컬럼 케이스의 경우, 해당 베이스 타일은 16×16입니다. fp8과 같은 1바이트 타입의 경우, 16×32가 됩니다. 이는 임의적인 선택이 아니었습니다. 이는 텐서 코어(tensor core) 명령어가 프래그먼트(fragments)를 노출하는 방식과, 해당 프래그먼트들이 공유 메모리(shared memory)를 통해 스테이징(staged)되어야 하는 방식에서 직접적으로 기인합니다.

우리는 이미 이전의 텐서 코어 명령어에서 이를 확인할 수 있습니다. SASS에서 HMMA 명령어는 종종 16×8×16 또는 16×8×8과 같은 형상(shapes)으로 동작하므로, 하드웨어 명령어 자체가 문자 그대로 정사각형인 16×16 곱셈이 아닐지라도, 프래그먼트 레이아웃(fragment layout)에서 16이라는 범위는 여전히 매우 실질적입니다. Hopper 아키텍처에서는 워프 그룹(warp group) 버전이 HGMMA로 나타나며, 64×256×16과 같이 훨씬 더 큰 형상을 가집니다. 따라서 핵심은 모든 텐서 코어 명령어가 정확히 16×16이라는 것이 아닙니다. 핵심은 하드웨어가 매우 강력한 16 기반 구조를 노출하며, 더 큰 프래그먼트들은 해당 구조의 반복되는 조각들을 조합함으로써 구축된다는 점입니다.

데이터 타입 의존성은 Hopper의 WGMMA 입력 프래그먼트를 살펴보면 특히 명확해집니다. 16비트 입력의 경우, m64nNk16 제품군은 자연스럽게 64×16을 노출합니다.

TK가 반복되는 16×16 조각으로 취급할 수 있는 input slab입니다. fp8의 경우, 이에 대응하는 제품군은 m64nNk32이며, 이제 동일한 입력 측의 너비가 64×32로 확장됩니다. 이것이 바로 TK가 행의 입도(granularity)를 16으로 고정하면서도, 1바이트 타입에 대해서는 열의 입도를 확장할 수 있도록 설계한 정확한 이유입니다. 추상화(abstraction)는 여전히 동일하지만, 이를 기반으로 구축된 하드웨어 파편(fragment)이 더 넓어진 것입니다.

따라서 TK가 더 큰 타일(tile)을 빌드할 때, 이를 가장 쉽게 이해하는 방법은 전체 타일 형상이 채워질 때까지 이러한 기본 파편들을 쌓아 올린다고 생각하는 것입니다. st_bf<64,64>와 같은 더 큰 16비트 타일은 16×16 파편들의 4×4 배열로 볼 수 있습니다. 반면 st_fp8e4m3<64,64>와 같은 fp8 타일은 16×32 파편들의 4×2 배열이 됩니다. 공유된 타일 자체는 여전히 하나의 공유 메모리 객체로 저장되지만, 이는 TK가 다루는 형상을 추론하는 데 유용한 방식입니다.

다음으로 넘어가기 전 한 가지 작은 주의사항이 있습니다. TK는 Hopper나 Blackwell이 하드웨어에서 지원하는 모든 수치 포맷(numerical format)을 노출하려고 시도하지 않습니다. 제가 살펴보고 있는 현재 리포지토리 상태(commit #01cb68c)에서, 주요 타일 타입은 TK가 실제로 기반을 두고 빌드하는 포맷들을 다룹니다:

bf16, half (FP16), float (FP32), 일반적인 FP8 포맷인 fp8e4m3 및 fp8e5m2, Blackwell의 fp8e8m0, 그리고 fp4e2m1_2와 같은 패킹된 (packed) FP4 저장 방식입니다.

FP4는 우리의 사고 모델에서 약간의 추가적인 주의가 필요한 부분입니다. TK는 스칼라(scalar) fp4e2m1 타입을 정의하지만, 타일 타입은 패킹된 저장 방식을 사용합니다. 단일 FP4 값은 단 4비트뿐이기 때문입니다. 두 개의 FP4 값이 자연스럽게 1바이트에 들어가므로, 타일 대상 타입은 fp4e2m1_2와 같은 형태가 됩니다. 이는 패킹된 FP4가 TK의 주소 지정 가능한 타일 단위에서는 여전히 16×32 기본 타일처럼 보이지만, 실제 개별 FP4 값을 세어보면 동일한 타일이 16×64 스칼라를 나타낸다는 것을 의미합니다. 따라서 일반적인 규칙은 여전히 유효하지만, 데이터 타입 자체가 패킹되었을 때

이것이 핵심 아이디어입니다. TK의 타일 형태는 텐서 코어 (Tensor Core) 스타일의 행렬 조각 (matrix fragments)과 자연스럽게 부합합니다. 이는 모든 명령어가 문자 그대로 하나의 정사각형 타일이기 때문이 아니라, 하드웨어가 실제로 처리하는 더 큰 조각들을 구성하기 위한 깔끔한 소프트웨어 단위 (software unit)를 제공하기 때문입니다. 이는 이미 오래된 HMMA 스타일의 명령어에서도 확인할 수 있는데, 여기서 16 단위의 세분성 (granularity)은 명확하며, 전체 워프 그룹 (warp group) 연산이 훨씬 더 큼에도 불구하고 Hopper HGMMA 명령어에서도 동일하게 적용됩니다. 따라서 TK가 더 큰 무언가를 구축할 때, GPU에 어색한 추상화 (abstraction)를 강요하는 것이 아닙니다. 단지 하드웨어가 작동하고자 하는 방식과 이미 일치하는 더 작은 조각들로부터 더 큰 타일을 구성하는 것뿐입니다.

실제로 TK는 세 가지 핵심 타일 수준 추상화 (tile-level abstractions)를 중심으로 프로그래밍 모델을 구축합니다: 글로벌 레이아웃 기술자 (global layout descriptors, gl), 공유 타일 (shared tiles, st), 그리고 **레지스터 타일 (register tiles, rt)**입니다. 이와 더불어, LayerNorm이나 RMSNorm처럼 벡터 형태로 표현하는 것이 더 자연스러운 커널을 위해 공유 및 레지스터 **벡터 추상화 (vector abstractions, sv, rv)**도 제공합니다. 단, TK 2.0은 현재 주로 Hopper 및 Blackwell GPU를 위해 구축되고 테스트됩니다. 이 프로젝트는 더 이상 Ampere를 적극적으로 지원하지 않는다고 명시하고 있습니다.

개별 추상화를 더 자세히 살펴보기 전에, TK 코드베이스 전반에 걸쳐 나타나는 몇 가지 일반적인 상수 (constants)와 스레드 인덱싱 도우미 (thread indexing helpers)를 빠르게 언급해 두는 것이 유용합니다. 이것들은 보통 kittens::xxx로 직접 접근합니다.

상수 / 헬퍼 (Constant / Helper)	값 (Value)	의도 (Intent)
BASE_TILE_DIM	16	근본적인 16 기반의 입도 (granularity).
TILE_COL_DIM<T>	16, 또는 FP8과 같은 1바이트 타입의 경우 32	타입 T에 대한 기본 타일 너비 (width).
TILE_ROW_DIM<T>	16	타입 T에 대한 기본 타일 높이 (height).
TILE_ELEMENTS<T>	TILE_COL_DIM<T> * TILE_ROW_DIM<T>	하나의 기본 타일에 포함된 요소 (elements) 수.
WARP_THREADS	32	하나의 워프 (warp) 내 스레드 수.
WARPGROUP_THREADS	128	하나의 워프 그룹 (warpgroup) 내 스레드 수.
WARPGROUP_WARPS	4	하나의 워프 그룹 내 워프 수.
warpid()	threadIdx.x >> 5	블록 내 워프 인덱스 (warp index).
warpgroupid()	threadIdx.x >> 7	블록 내 워프 그룹 인덱스 (warpgroup index).
laneid()	threadIdx.x & 0x1f	워프 내 레인 인덱스 (lane index).

이것은 섹션의 나머지 부분을 위한 훨씬 더 깔끔한 토대가 됩니다. 왜냐하면 이제 이후의 추상화(abstractions)들을 동일한 기저 아이디어 관점에서 읽을 수 있기 때문입니다. 즉, TK는 행(row)의 입도를 16으로 고정하고, 너비(width)는 데이터 타입(datatype)을 따르도록 하며, 하드웨어의 프래그먼트(fragment) 구조와 이미 일치하는 형상(shapes)들을 조합함으로써 더 큰 공유 메모리(shared) 및 레지스터(register) 수준의 객체들을 구축합니다.

레지스터 타일 (Register tiles): 레지스터 타일은 TK가 연산 중에 레지스터에 상주하는 값들을 위해 사용하는 주요 추상화 단위입니다. GEMM 스타일의 커널에서, 이것들은 누산기 프래그먼트(accumulator fragments)를 보유하는 타일인 경우가 매우 많으며, 이것이 레지스터 타일이 텐서 코어(tensor core) 명령 주변에서 매우 두드러지게 나타나는 이유입니다. 소스 코드에서 일반적인 형태는 rt<T, rows, cols, layout>이며, 따라서 타입은 데이터 타입(datatype), 형상(shape), 그리고 레이아웃(layout)에 의해 매개변수화됩니다. 하지만 실제로는 FP32 레지스터 타일을 위한 rt_fl<M, N> 또는 BF16 레지스터 타일을 위한 rt_bf<M, N>와 같은 더 짧은 별칭(aliases)을 주로 보게 될 것입니다.

내부적으로 레지스터 타일은 위에서 소개한 것과 동일한 "빌딩 블록 (building blocks)" 이야기를 따릅니다. 행의 입도는 16으로 고정된 상태를 유지하며, 너비는 데이터 타입에 따라 달라집니다. rt_base.cuh에서 기본 레지스터 프래그먼트는 TILE_ROW_DIM<T>와 TILE_COL_DIM<T>를 직접 사용하며, 그 후 rt.cuh에서