GPU Programming Fundamentals + ThunderKittens

2024 년 시점에서 GPU 프로그래밍의 풍경은 이미 복잡하다 — CUDA, Triton, CUTLASS, cuDNN, torch.compile 의 Inductor. 그 위에 또 하나의 라이브러리(ThunderKittens) 가 왜 필요한지부터 답하지 않으면 강의는 시작될 수 없다.

강의는 두 명의 화자가 두 단계로 답을 나눈다.

1. William Brandon — fundamentals. H100 의 메모리 계층(HBM ↔ L2 ↔ SMEM ↔ register), warp 와 warpgroup, TMA(Tensor Memory Accelerator), WGMMA(Warpgroup Matrix Multiply-Accumulate) 같은 하드웨어 사실들을 정리한다. 이 사실들이 곧 DSL 이 어떤 형태가 되어야 하는지를 결정한다.
2. Simran Arora — ThunderKittens. 위 fundamentals 위에 얹는 mini-DSL 의 디자인 — tile 을 1급 시민으로, producer/consumer 를 명시적 abstraction 으로, 그러면서도 CUDA 처럼 직접 PTX 까지 내려갈 수 있게.

A100(Ampere) 에서 H100(Hopper) 로 넘어가면서 GPU 프로그래밍 모델 자체가 변했다. async copy 가 first-class 가 되고, warpgroup (4 warp = 128 thread) 단위의 행렬곱 명령(WGMMA)이 등장했다. 기존 CUDA / Triton 의 추상은 이 변화에 자연스럽게 맞지 않는다.

왜 기존 도구로는 부족한가 — 세 가지 구체적 부담.

• WGMMA 의 메모리 layout 요구 — 결과가 register tile 안에 특정 swizzling 으로 떨어진다. 이 layout 을 사람이 직접 들고 다니면 코드가 PTX 같아지고, 컴파일러에 맡기면 layout 변환 비용이 새는 자리에서 발생.
• TMA 의 box copy + barrier 모델 — async memcpy 가 mbarrier 의 arrival 로 동기화된다. CUDA stream 모델보다 fine-grained, Triton 의 자동화 모델보다 explicit.
• warpgroup specialization — 4 warp 묶음이 다른 역할을 동시에 한다 (예: producer warp 와 consumer warp). 이게 FA3 의 ping-pong, FA4 의 5-way specialization 의 핵심.

ThunderKittens 의 모든 코드는 tile 위에서 돈다. 최소 단위 16×16, 자료형은 register 에 있는 rt 와 shared memory 에 있는 st. 그리고 row/column vector 는 tile 의 reduction 결과로 자연스럽게 따라온다.

// ThunderKittens 의 GEMM inner loop (개념)
using a_t = st<bf16,128,64>;
using b_t = st<bf16,64,128>;
using c_t = rt<float,128,128>;

c_t acc; kittens::zero(acc);

for (int k = 0; k < K; k += 64) {
  // async TMA copy
  kittens::load_async(a_smem, A.tile_at(m,k));
  kittens::load_async(b_smem, B.tile_at(k,n));
  kittens::wait(barrier, phase);

  // WGMMA on warpgroup of 4 warps
  warpgroup::mma_AB(acc, a_smem, b_smem);
  warpgroup::mma_commit_group();
}

warpgroup::mma_async_wait();
kittens::store(C.tile_at(m,n), acc);

CUDA 의 전통은 모든 thread 가 동일한 일을 한다 (SIMT). H100 시대에는 그 모델이 깨진다 — 일부 warp 는 TMA 만 발사하고, 다른 warp 는 MMA 만 돌린다. 이게 warp specialization, ThunderKittens 가 가장 잘 잡아주는 패턴.

강의에서 가장 진하게 다뤄진 예제. FlashAttention-3 의 핵심 트릭(warpgroup 두 그룹이 ping-pong 으로 softmax 와 MMA 를 번갈아 돌린다)을 ThunderKittens 의 producer/consumer 추상으로 풀어 본다.

// ThunderKittens FA3-style attention (개념)
template<int NUM_WORKERS>
__global__ void attn_kernel(...) {
  if (warpid < 2) {
    // PRODUCER warps
    for (int k=0; k<Tk; ++k) {
      tma::load_async(K_smem[stage], K_gmem[k]);
      tma::load_async(V_smem[stage], V_gmem[k]);
      tma::arrive(barrier[stage]);
      stage = (stage+1) % STAGES;
    }
  } else {
    // CONSUMER warpgroups (G0 / G1 ping-pong)
    for (int k=0; k<Tk; ++k) {
      tma::wait(barrier[stage], phase);
      warpgroup::mma_AB(s_acc, q_reg, K_smem[stage]);
      warpgroup::mma_commit_group();
      // softmax on the OTHER group while this one mma's
      if (group_id == 0) other.softmax_async();
      warpgroup::mma_async_wait();
      kittens::row_max(m_new, s_acc);
      kittens::exp(s_acc);
      // online softmax accumulator update ...
    }
  }
}

실제 ThunderKittens repo 의 examples/attn/h100/ 안에 FA3 와 호환되는 reference 구현이 있고 — 강의에서는 그 코드의 핵심 30 여 줄이 어떻게 FA3 의 trick 들을 그대로 노출하는지 line-by-line 로 짚는다 (확인 필요).

ThunderKittens 의 가장 자주 인용되는 숫자 — H100 위 BF16 GEMM 에서 855 TFLOPs (peak 989 의 86%) 를 100 줄 미만 코드로. 비교 기준을 어떻게 잡느냐에 따라 의미가 달라진다.

왜 ThunderKittens 가 86% 를 받는지 두 가지 핵심.

• WGMMA 를 직접 부른다 — 컴파일러가 dot product 를 보고 lowering 하는 게 아니라 사람이 warpgroup::mma_AB 를 명시적으로 적는다. layout 변환 비용 0.
• TMA 의 box copy 추상 — global → SMEM 이 single instruction 으로 발사되고 mbarrier 로 완료. async pipeline 의 fall-through case 가 없다.

ThunderKittens 는 Triton 의 대체재가 아니라 다른 추상 레벨의 도구다. 강의에서 두 저자가 명시적으로 짚은 비교.

• Together AI — inference 백엔드의 H100 attention/GEMM 핫패스. Together 의 enginereing 블로그에서 ThunderKittens 기반 커널 사용 사실을 명시.
• Jump Trading — quant trading 인프라. 대용량 행렬 연산의 latency-critical 자리.
• Cursor — code completion 모델 inference. (확인 필요 — 어떤 kernel 인지 정확히는 미상)
• HazyResearch internal — Hyena, BASED, LoLCATs 같은 efficient attention 변형들의 reference 구현.

채택의 공통 패턴 — 모두 "기존 cuBLAS / FlashAttention 으로는 풀기 어려운 한 자리" 가 있는 곳. 새 attention variant, 비표준 GEMM shape, 또는 fused custom operator.

• HipKittens — AMD MI300 위로의 포팅. matrix core 와 LDS 의 양상이 다르므로 1:1 mapping 은 아니지만 tile-first 추상은 보존.
• ThunderMittens — Apple Silicon (M-series) 의 GPU. Metal performance shader 위에 비슷한 abstraction 시도.
• Blackwell (B200) 대응 — TMA 의 5th-gen tensor core 인 tcgen05.mma, async warp specialization 의 5-way 분리. ThunderKittens 가 FA4 같은 커널의 reference 구현 자리로 다시 호출됨 (L080 참조).
• compiler 레이어 추가 — TK 코드를 기반으로 하는 더 높은 level 의 schedule autotuner. 강의에서 언급되었는지 확인 필요.

큰 그림 — tile-first DSL 이 단일 회사의 single-arch 도구가 아니라 크로스 아키텍처 추상 으로 자리잡으려는 시도. 성공 여부는 AMD/Apple 위 성능과 채택 곡선이 결정.

손에 새기기 — 실습 시퀀스

1. TK 의 hello-world GEMM — 256³ BF16 matmul 을 TK 로 짠다. warpgroup::mma_AB 한 줄이 어떤 PTX 로 lowering 되는지 ncu 의 SASS view 로 확인.
2. 2-stage producer/consumer — TK 의 stage buffer 갯수를 1, 2, 3, 4 로 바꿔가며 GEMM 의 throughput 을 측정. stage=2 부근에서 plateau 가 보이는지 확인.
3. WGMMA layout 학습 — accumulator register tile 의 element 가 어떤 thread/warp 에 분산되는지 직접 print 해본다. swizzling 의 그림이 손에 잡힐 때까지.
4. TMA 의 box copy 직접 발사 — descriptor 만 만들어서 cp.async.bulk.tensor 한 번 발사. mbarrier::wait 로 완료 확인. async copy 의 model 을 손에 익힌다.
5. FA3 의 ping-pong 패턴 직접 짜기 — TK 의 examples/attn 을 copy 해서 한 줄씩 분해. 두 warpgroup 이 어떻게 번갈아 도는지 print debug 로 직접 본다.
6. Triton 과의 동일 op 비교 — 같은 GEMM shape 을 Triton 과 TK 로 짜고 throughput / SASS / register usage 를 직접 diff.

• FA4 reference 가 ThunderKittens 위에 올라왔는가 — Modal 블로그가 reverse-engineer 한 FA4 가 TK 의 어느 commit 부근에서 reference 가 됐는지. 강의가 영상 시점 (2025) 에 FA4 를 다뤘는지 확인.
• autotuner 의 부재 — TK 는 template parameter 를 사람이 고른다. autotuner 가 없는 게 의도인가, 추후 추가될 것인가.
• HipKittens 의 성능 격차 — AMD MI300 위 GEMM 의 peak% 가 H100 의 86% 만큼 잘 나오는지. (확인 필요)
• compile time — C++ template 기반이라 nvcc 컴파일 시간이 무거울 수 있다. flash-attention build 시간만큼 부담스러운지.
• backward pass 의 자리 — forward pass 의 ping-pong 이 fancy 한데 backward 는 같은 추상에 자연스럽게 들어가는지. 강의가 다뤘는지 확인.
• industry adoption 의 진짜 깊이 — Together / Jump / Cursor 가 어떤 커널을 정확히 어디에 쓰는지의 구체.