B200 어텐션 커널을 60개 다이어그램으로 처음부터 SOTA급 성능까지
B200 Attention Kernel from Scratch to Near-SOTA in 60 Diagrams
CUDA와 약간의 PTX로 작성된 B200 어텐션 커널을 처음부터 구현하여 FlashAttention-4 성능의 94.4%에 도달하는 과정을 60개의 다이어그램으로 설명합니다. 초보자도 이해할 수 있도록 순진한 커널에서 시작해 한 번에 하나씩 최적화를 추가하며, 각 단계의 코드와 다이어그램을 제공합니다. Blackwell의 비대칭 하드웨어 확장으로 인해 소프트맥스가 주요 병목이 되는 문제를 다루고, 이를 해결하기 위한 최적화 기법을 소개합니다. 마지막으로 최종 커널을 비디오 생성 모델에 적용하는 캡스톤 프로젝트를 포함합니다.
소프트맥스는 tensor core가 아닌 ALU와 MUFU에서 실행되는데, Blackwell B200에서 tensor core 처리량은 약 두 배로 증가한 반면 exp 유닛은 거의 그대로여서, 우리 타일 크기에서는 소프트맥스가 매우 강력한 MMA와 거의 비슷한 사이클을 소비하며 이 커널의 주요 병목이 됩니다.