HazyResearch/ThunderKittens

ThunderKittens – 고성능 AI 커널을 위한 CUDA-DSL

무엇인가요 – ThunderKittens는 헤더 온리 C++/CUDA 라이브러리로, 깊이 있는 학습 커널(예: GEMM, FlashAttention)을 컴팩트하고 타일 기반의 코드로 작성할 수 있게 해줍니다. 낮은 수준의 GPU 세부 사항(텐서 코어 호출, 비동기 복사, 공유 메모리 벤딩)을 추상화하면서도 하드웨어에 가까운 상태를 유지하여, 최신 NVIDIA GPU(H100, Blackwell, Vera Rubin)에서 이론적 최고 성능에 근접한 실행 속도를 제공합니다.

왜 중요한가요 – 대규모 언어 모델의 학습과 추론은 행렬 곱셈 및 어텐션 커널의 실행 속도에 제한됩니다. 이러한 커널을 수작업으로 작성하는 것은 실수를 유발하며 NVIDIA 프로그래밍 모델에 대한 깊은 지식이 필요합니다. ThunderKittens는 다음 특징을 가진 작고 확장 가능한 DSL을 제공합니다:

• 많은 커널에서 소스 코드가 100줄 미만입니다.
• 컴파일 시점에 정확한 레이아웃 처리를 보장합니다.
• 어셈블리 코드를 작성하지 않고도 최신 텐서 코어 명령어(WGMMA, TCGEN05, MXFP8, NVFP4)를 사용할 수 있습니다.
• C++에서 직접 사용하거나, PyBind11를 통해 PyTorch용으로 래핑할 수 있습니다.

핵심 개념

개념	기능
타일 프리미티브	16×16 이상의 블록(레지스터 타일, 공유 타일, 벡터)을 조작하며, 텐서 코어 레인에 자연스럽게 매핑됩니다.
워프 / 워프 그룹	기본적으로 단일 워프(32스레드)용으로 함수를 작성합니다. 4워프의 협업 그룹(워프 그룹)은 비동기 행렬 곱셈-누산 명령어를 노출합니다.
TMA / 비동기 복사	NVIDIA의 TMA(Tensor Memory Access) 메커니즘을 통해 로드/스토어를 숨기며 지연을 감추는 내장 헬퍼가 있습니다.
정적 레이아웃 체크	템플릿이 데이터 타입, 형태, 메모리 레이아웃을 인코딩하여, 불일치한 연산은 컴파일 시점에 감지됩니다.
Load-Store-Compute-Finish 템플릿	메모리 이동과 계산을 겹치는 추천 패턴으로, 오쿠파시를 최대화합니다.

일반적인 워크플로우

1. 리포지토리 복제 후 CUDA 소스에 kittens.cuh 포함.
2. 제공된 matmul_layout / matmul_template (또는 기타 프리미티브)를 사용해 커널 정의 – producer, consumer, 선택적 common_setup 콜백만 채우면 됩니다.
3. 제공된 Makefile로 커널 컴파일 (CUDA 12.8+, C++20). 각 커널은 kernels/ 아래의 별도 폴더에 있으며, 독립적으로 빌드 가능.
4. Python에서 호출 (옵션) – make 후, 작은 PyBind11 래퍼를 통해 PyTorch 2.8+에서 커널을 호출할 수 있습니다.
5. 벤치마크 / 테스트 – 정확성 테스트 및 성능 스크립트는 각 커널과 동일한 폴더에 위치해 있습니다.

예시: H100에서 855 TFLOPs 행렬 곱셈

#include "kittens.cuh"
#include "prototype.cuh"
using namespace kittens;
using namespace kittens::prototype;

// 레이아웃 정의 (타일, 글로벌 포인터 등)
template<int M_BLOCK, int N_BLOCK>
struct matmul_layout { … };

// producer, consumer, common_setup를 연결하는 커널 템플릿
template<int _M_BLOCK=2, int _N_BLOCK=4, int _SUPER_M=12>
struct matmul_template { … };

전체 소스(약 100줄)는 README에 표시되어 있으며, H100의 이론적 최고 성능의 약 86%에 도달하는 커널로 컴파일됩니다.

지원되는 하드웨어 – 주로 NVIDIA Hopper(H100) 및 Blackwell(B200) GPU. 2026년 9월 이후 새롭게 출시된 Vera Rubin GPU도 지원됩니다. Ampere는 작동하지만, 향후 업데이트는 제공되지 않습니다. AMD 사용자는 자매 프로젝트 HipKittens 를 참고하세요.

설치 체크리스트

• CUDA 12.8+ (CUDA_HOME 설정, PATH/LD_LIBRARY_PATH 업데이트).
• C++20 컴파일러 (gcc-11 또는 clang-11).
• (옵션) Python 바인딩이 필요하면 PyTorch 2.8+ 및 PyBind11.
• 리포지토리 복제, 헤더 포함, 제공된 Makefile로 커널 컴파일.

사용 사례 – Together AI, Jump Trading, Cursor 등에서 프로덕션 규모의 학습 및 추론에 사용되고 있습니다. 스탠포드 대학 Hazy Research Lab에서도 내부적으로 사용 중입니다.

학습 리소스

• ThunderKittens 매뉴얼 – 타일, 스코프, API 규약을 다루는 짧은 가이드.
• 교육용 커널 시리즈 – kernels/gemm/educational_h100은 단계별 GEMM 구현을 설명합니다.
• 딥다이브 블로그 – Hamza Elshafie의 분석(2026년 5월)은 DSL 내부 구조를 설명합니다.
• 온보딩 문서 – 새 사용자를 위한 Google 문서(README 내 링크).

데모 – demos/ 폴더에는 LLM 학습 및 추론용 실행 가능한 예제(예: Qwen, Llama, LoLCATS)가 포함되어 있습니다. ThunderKittens 커널을 PyTorch 워크플로우에 통합하고 Hugging Face의 8B 모델을 실행하는 방법을 보여줍니다.

---

결론 – NVIDIA 텐서 코어의 성능을 최대한 끌어내야 하는 커스텀 어텐션, GEMM, 또는 기타 행렬 중심 연산을 위한 커널을 작성해야 한다면, ThunderKittens는 CUDA 위에 직접 위치한 작고 타입 안전한 DSL을 제공하여, 일반적인 어셈블리 수준의 번거로움 없이 프로덕션 수준의 커널을 작성할 수 있게 해줍니다.