HazyResearch/HipKittens

HipKittens란 무엇인가?

HipKittens는 AMD GPU에서 고성능 AI 커널(CDNA3, CDNA4, CDNA5)을 작성하기 위한 저수준 빌딩 블록을 제공하는 오픈 소스 C++ 라이브러리입니다. ThunderKittens 생태계에 속하며 Triton이나 Torch‑Inductor와 같은 프로젝트가 사용하는 것과 동일한 종류의 텐서 코어 스타일 프로그래밍을 대상으로 하지만 AMD의 ROCm 스택(HIP, hipcc)용으로 작성되었습니다.

핵심 아이디어

아이디어	의미
타일 프리미티브	AMD 텐서 코어 유닛에 직접 매핑되는 작고 고정된 크기의 블록. 라이브러리는 병합되고, 뱅크 충돌을 피하며, 네이티브 텐서 코어 레이아웃을 사용하는 메모리 액세스 헬퍼를 제공합니다.
Python 스타일의 대량 계산	Python 함수(예: `gemm`, `attention`)처럼 보이는 얇은 C++ 래퍼이지만 인라인 어셈블리와 HIP 호출로 컴파일됩니다.
비동기 로드/스토어	공유 메모리로 비차단 로드를 발행하여 하드웨어가 대기 시간을 숨길 수 있게 하는 헬퍼.
스케줄링 패턴	계산과 메모리를 겹쳐 GPU 파이프라인을 바쁘게 유지하는 8‑wave ping‑pong 및 4‑wave interleave의 두 가지 재사용 가능한 패턴.

무엇을 할 수 있나요?

BF16 GEMM 커널(예: 8192×8192×8192 행렬 곱셈)로 AITER 및 PyTorch와 벤치마크할 수 있습니다.
Attention 커널(MHA, GQA, 인과/비인과)로 순방향 및 역방향 패스를 모두 지원하며 배치 크기, 시퀀스 길이, 헤드 수 및 헤드 차원을 구성할 수 있습니다.
메모리 바운드 커널(회전 위치 임베딩 및 융합 레이어 정규화 등).
BERT 또는 LLaMA용 엔드투엔드 훈련 스크립트로 HipKittens attention 커널을 PyTorch 훈련 루프에 연결할 수 있습니다.

이 모든 것은 ROCm의 hipcc로 컴파일되며 제공된 AMD Docker 이미지 내에서 실행할 수 있습니다.

누가 사용할까요?

연구자: AMD 전용 또는 멀티 실리콘 AI 성능을 탐구하고 가볍고 "메탈에 가까운" 커널 라이브러리가 필요한 사람.
엔지니어: AMD GPU용 맞춤형 트랜스포머 백엔드를 구축하는 사람, 특히 AITER, Triton 또는 HipBLASLT와 같은 기존 백엔드와 비교하려는 경우.
학생: AMD 하드웨어에서 텐서 코어 프로그래밍을 배우는 사람. 라이브러리의 단위 테스트와 예제 make 파일은 구체적인 학습 자료로 사용할 수 있습니다.

시작하는 방법

저장소를 클론하고 일치하는 ROCm Docker를 가져옵니다(README에 MI35x 및 MI30x GPU용 이미지 두 개가 나열되어 있습니다).
컨테이너에 들어가 env.src를 소스로 지정하여 환경 변수를 설정하고 비교 벤치마크를 실행하려면 AITER 기준선을 설치합니다.
커널을 빌드합니다(예: cd kernels/cdna4/gemm/bf16fp32 && make && python bench.py). bench.py 스크립트는 HipKittens 커널, AITER 참조 및 PyTorch 구현을 실행하고 타이밍 숫자를 출력합니다.
단위 테스트를 실행하여 GPU 세대를 테스트합니다(cd HipKittens/tests/unit/cdna4 && make && ./unit_tests).
분석 스크립트를 HipKittens/analysis에서 탐색하여 논문의 성능 플롯을 재현하거나 training/을 참조하여 빠른 BERT/LLaMA 훈련 실행을 수행합니다.

언제 유용한가요?

AMD GPU에서 훈련 또는 추론(MI300X, MI350X, MI455 등)을 수행하고 일반 ROCm 라이브러리보다 빠른 커널이 필요한 경우.
새로운 스케줄링 또는 메모리 액세스 패턴을 실험하고 싶지만 직접 원시 HIP 어셈블리를 작성하고 싶지 않은 경우.
멀티 실리콘 AI 스택을 구축하고 있으며 다른 DSL(예: Triton)이 사용하는 것과 동일한 타일 기반 추상화를 따르지만 AMD 하드웨어에서 작동하는 참조 구현이 필요한 경우.

연락처 및 커뮤니티: William Hu(willhu@stanford.edu)와 Simran Arora(simran@cs.stanford.edu)가 HipKittens 사용자를 위한 Discord 채널을 운영합니다. 저장소에는 더 깊은 배경을 위한 arXiv 논문, 블로그 게시물 및 YouTube 강의 링크도 있습니다.