FLAMIA, Blackwell에서 인과적으로 안전한 저정밀 Attention 가속

Transformer 학습에서 Attention은 특히 컨텍스트 길이가 길어질수록 큰 계산 비용을 차지합니다. NVIDIA Blackwell은 FP8과 FP4 같은 저정밀 행렬 연산을 높은 처리량으로 실행하지만, 이를 Attention에 적용하려면 단순히 데이터 형식을 낮추는 것만으로는 부족합니다. 양자화 자체가 추가 비용을 만들고, 공유 스케일이 인과성을 훼손할 수 있으며, 순전파와 역전파가 서로 다른 수치 경로를 사용하면 학습 안정성에도 영향을 줄 수 있습니다.

Yandex Research가 제안한 FLAMIA(FLAsh MIcroscaling Attention)는 FlashAttention을 Blackwell의 저정밀 하드웨어에 맞게 설계한 커널 제품군입니다. 주요 경로는 MXFP8과 혼합 NVFP4/FP8 방식입니다. 혼합 방식은 Query-Key 점수 계산과 역전파 재계산에 NVFP4를 사용하고, 나머지 곱셈에는 FP8을 사용합니다. 모든 연산을 FP4로 처리하면 양자화 오버헤드가 커지고 손실이 소폭 악화될 수 있기 때문에, FLAMIA는 단계별로 정밀도를 선택합니다.

핵심 설계

• 양자화 표현 재사용. 32개 토큰 블록 단위로 2의 거듭제곱 스케일을 공유하고, 점수 그래디언트에는 32×32 블록을 사용합니다. 같은 저정밀 피연산자를 서로 다른 행렬 방향에서 재사용해 중복 양자화와 메모리 이동을 줄이는 방식입니다.
• 인과성을 고려한 스케일링. 인과 Attention에서는 현재 위치가 미래 토큰을 볼 수 없어야 합니다. 그러나 미래 위치가 포함된 블록에서 공유 스케일을 계산하면, 이후 마스크를 적용하더라도 스케일 자체를 통해 정보가 간접적으로 전달될 가능성이 있습니다. FLAMIA는 문서 단위 스케일, 포화 방지, key별 스케일 접기 등을 사용해 위험을 줄이고, 지원되는 배치 호출에서는 더 엄격한 인과 옵션도 제공합니다.
• 순전파와 역전파의 수치 일치. 역전파에서 확률을 다시 계산할 때 순전파와 호환되는 양자화, 타일링, 스케일링, 지수 계산 규칙을 적용합니다. 한 번의 그래디언트가 비슷하다는 사실만으로 장시간 학습의 안정성이 보장되는 것은 아니기 때문입니다.
• 하드웨어 인식형 실행. 양자화를 커널 내부에 융합하고 softmax와 행렬 곱을 겹쳐 실행합니다. 저정밀 역전파에는 협력 스레드 배열을 사용해 GPU 활용도를 높입니다. 길이가 서로 다른 문서를 packed layout으로 처리할 때도 재패킹 없이 문서별 양자화와 작업량 기반 스케줄링을 적용합니다.

결과와 의미

여러 인과 packed-document 워크로드에서 양자화 시간을 포함한 순전파·역전파 Attention 전체를 측정한 결과, MXFP8 경로는 각 워크로드에서 가장 빠른 평가 기준선보다 B200에서 약 1.101.40배, B300에서 약 1.061.42배 빠른 결과를 보였습니다. Llama 계열 모델의 사전학습 실험에서는 학습 안정성과 모델 품질 유지도 확인했습니다.

FLAMIA의 핵심 시사점은 저정밀 가속이 단순히 더 작은 데이터 형식을 선택하는 문제가 아니라는 것입니다. 양자화 단위, 인과 의미, 순전파와 역전파의 수치 동작, GPU 스케줄링을 함께 설계해야 이론적인 하드웨어 처리량을 실제 학습 속도로 전환할 수 있습니다. 다만 보고된 수치는 평가된 워크로드에 대한 결과이므로 모든 모델과 시퀀스 길이 분포에 동일한 가속률이 적용된다고 볼 수는 없습니다.

출처: Hugging Face Daily Papers