Flash-MSA: Sparse Attention Kernels을 통한 Million-Token Training 가속화

DeepSeek-V4, GLM-5.2 등 최신 프론티어 모델들이 추론 속도 향상을 위해 Sparse Attention을 사용하고 있지만, 이를 효율적으로 학습할 수 있는 코드는 부족한 실정입니다. Flash-MSA는 이를 해결하기 위해 Hopper 및 Blackwell GPU를 지원하는 고성능 오픈소스 학습 커널을 제공합니다.

MSA(MiniMax Sparse Attention)의 주요 특징:

• Blockwise Sparsity: 개별 KV 토큰 대신 128 단위의 블록을 선택하여 캐싱 효율성을 높였습니다.
• GQA(Grouped-Query Attention) 적용: 기존 MLA 방식 대신 GQA를 사용하여 다양한 모델 구조에서도 접근 가능하도록 설계했으며, 어텐션 헤드의 표현력을 높였습니다.
• Group-wise Specialization: 각 프록시 헤드가 서로 다른 KV 서브셋을 선택할 수 있도록 하여 어텐션의 다양성을 확보했습니다.

커널 설계 핵심:

• Forward Pass: 프록시 어텐션 수행 후 Sparse Main Attention을 실행하며, 메모리 오버헤드를 줄이기 위해 프록시 어텐션 시 top-k 인덱스만을 캐싱합니다.
• Backward Pass: 메인 어텐션과 프록시 어텐션의 그래디언트를 계산하기 위해 Double-attention combined pass를 수행하여 효율성을 극대화했습니다.
• 메모리 최적화: 블록 단위 희소성(Block-sparsity) 덕분에 전체 학습 단계에서 프록시 포워드 결과인 블록 인덱스만 저장하면 되므로 메모리 사용량을 크게 절감했습니다.