Blackwell GPU용 저정밀도 Attention 커널 FLAMIA 공개
원제 FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell
추천 13댓글 1
Key Point
Transformer 학습의 가장 비용이 높은 부분인 Attention 연산을 최신 Blackwell GPU의 저정밀도 연산으로 1.1~1.4배 빠르게 만드는 기법이 개발되어, 대규모 모델 학습 효율에 직접적인 영향을 미칠 수 있다.
핵심 요약
- NVIDIA Blackwell GPU의 고처리량 저정밀도 행렬 곱셈을 Attention 연산에 활용하기 위해 FLAMIA라는 attention 커널 패밀리를 제시했다.
- 양자화 오버헤드, 공유 스케일을 통한 인과 정보 누출, 정방향과 역방향 계산 간 수치 불일치 등 세 가지 주요 문제를 해결했다.
- MXFP8과 혼합 NVFP4/FP8 학습 경로를 지원하며, 토큰-블록 스케일링으로 양자화된 피연산자를 행렬 방향 간에 재사용한다.
- 스케일 유도 누출 방지를 위한 타겟 세이프가드와 정방향·역방향 계산의 일관성 보장 메커니즘을 포함했다.
- 융합 양자화와 아키텍처별 커널 최적화로 오버헤드를 줄였고, 문서 로컬 스케일링과 워크로드 인식 스케줄링으로 가변 길이 문서를 재패킹 없이 압축 레이아웃에서 직접 처리한다.
- B200에서 정방향·역방향 attention 호출(양자화 포함)을 1.10~1.40배 가속화하고, B300에서는 1.06~1.42배 가속화했다.
- Llama 스타일 모델 사전학습에서 학습 안정성과 모델 품질이 보존됨을 검증했다.
- 양자화, 학습 수치, 문서 스케줄링을 공동으로 설계하여 저비트 하드웨어 처리량을 실제 attention 가속화로 전환할 수 있음을 보였다.