Flash-MSA: 세계 최초 오픈소스 스파스 어텐션 학습 커널, 백만 토큰 훈련 가속화

Flash-MSA: Accelerating Million-Token Training with Sparse Attention Kernels

프론티어 모델들이 추론 속도를 높이기 위해 스파스 어텐션을 사용하지만, 이를 효율적으로 학습시킬 수 있는 코드는 공개된 적이 없었습니다. 이 글은 MiniMax의 스파스 어텐션(MSA)을 위한 세계 최초의 고성능 오픈소스 학습 커널을 소개합니다. CuTeDSL로 작성된 이 커널은 Hopper 및 Blackwell GPU에서 동작하며, 블록 단위 희소성, GQA 기반 메인 어텐션, 그룹별 프록시 헤드 특성화를 지원합니다. 또한 KL 발산 손실을 메모리 효율적으로 계산하는 기법과 정방향/역방향 커널의 정확성을 검증한 결과를 제시합니다.