FLAMIA, Blackwell GPU에서 어텐션 연산 1.1~1.4배 가속
FLAMIA: FLAsh MIcroscaling Attention for Fast and Causally Safe Training on Blackwell
- 매체
- Hugging Face
- 게시
- 중요도
- 3/5
- 화제 지수
- 18
한국어 요약
NVIDIA Blackwell GPU의 저정밀 행렬 곱셈을 활용한 어텐션 커널 FLAMIA를 제시했다. MXFP8과 혼합 NVFP4/FP8 학습 경로로 양자화 오버헤드를 줄이고, 인과 정보 누수를 방지하며, 정방향과 역방향 계산의 수치 일관성을 확보한다.
원문 발췌 · English
Attention remains a substantial computational cost in Transformer training, particularly at long context lengths. NVIDIA Blackwell GPUs provide high-throughput low-precision matrix multiplication, but exploiting it in attention requires addressing quantization overhead, causal information leakage through shared scales, and numerical inconsistencies between the forward and backward passes. We present FLAMIA, a family of attention kernels for Blackwell with MXFP8 and mixed NVFP4/FP8 training paths. Our design combines token-block scaling that reuses quantized operands across matrix orientations, targeted safeguards against scale-induced leakage, and consistent forward and backward computations. Fused quantization and architecture-specific kernel optimizations reduce overhead, while document-local scaling and workload-aware scheduling support variable-length documents directly in packed layouts without repacking. Across various causal packed-document workloads, the MXFP8 path accelerates the combined forward and backward attention call, including quantization, by factors of 1.10--1.40times on B200 and 1.06--1.42times on B300 relative to the fastest evaluated baseline for each workload. We validate training stability and quality preservation on pretraining Llama-style models. These results show how jointly designing quantization, training numerics, and document scheduling translates low-bit hardware throughput into practical attention acceleration.