Blackwell용 FlashAttention-4, 엔드투엔드 MXFP8 지원

들어가며

FP8을 어텐션에 적용하는 일은 데이터 타입을 단순히 바꾸는 것보다 어렵습니다. 블록 단위 스케일링을 사용하는 행렬곱에서는 스케일 팩터를 계산하고, 적절한 축에 맞춰 배치하며, 메모리와 동기화 제약을 지키면서 텐서 코어에 공급해야 합니다. PyTorch 팀은 NVIDIA Blackwell GPU용 FlashAttention-4를 확장해 MXFP8 기반 순전파와 역전파를 모두 지원한다고 설명했습니다.

핵심 변화

• 순전파와 역전파를 모두 MXFP8로 처리합니다. LLM 형태의 입력에서 순전파는 최대 2.85 PFLOPS, 역전파는 2 PFLOPS를 기록했습니다. 내부 테스트 형상에서는 각각 2.54 PFLOPS와 1.58 PFLOPS를 달성했으며, BF16 대비 최대 향상 폭은 순전파 1.6배, 역전파 1.52배로 제시됐습니다.
• 제한된 TMEM에 스케일 팩터를 배치합니다. Blackwell의 TMEM은 512열 규모이며 기존 FlashAttention-4 커널에서도 대부분 사용됩니다. 새 구현은 앞선 단계에서 더 이상 필요하지 않은 영역을 스케일 팩터 저장에 겹쳐 사용하고, 비동기 MMA와의 쓰기 충돌을 막기 위해 필요한 배리어를 추가합니다.
• 양자화를 주변 연산에 융합합니다. RMSNorm과 양자화, GEMM과 양자화를 각각 하나의 커널로 결합해 후속 연산에 필요한 FP8 데이터와 스케일 팩터를 한 번에 생성합니다. 별도의 형식 변환 단계로 발생하는 오버헤드를 줄이는 방식입니다.
• softmax 결과를 온라인에서 변환합니다. 어텐션의 확률 행렬 P는 다음 P·V 연산을 위해 MXFP8으로 바뀌어야 합니다. 구현은 softmax 과정에서 이미 계산한 행별 최댓값을 재사용하고, Blackwell의 저수준 명령 최적화를 적용해 변환 비용을 낮춥니다.
• 가변 길이 입력에서 큰 Gather를 피합니다. jagged 모듈은 FP8 활성값을 패딩되지 않은 위치에 유지하고, 더 작은 스케일 팩터만 산재시키고 패딩한 뒤 재배열합니다. 그 결과 TMA 전송과 텐서 코어 연산에 맞는 주소 배치를 만들 수 있습니다.

의미와 주의점

Blackwell은 MXFP8, MXFP6, MXFP4, NVFP4 같은 마이크로스케일링 형식을 위한 블록 스케일 MMA 명령을 제공합니다. 하지만 양자화와 스케일 팩터 이동을 주 계산과 분리하면 변환과 동기화 비용이 하드웨어의 이점을 상쇄할 수 있습니다. 이번 작업의 핵심은 저정밀화를 하나의 행렬곱 커널에 국한하지 않고, 메모리 배치와 비동기 실행, softmax, 전후단 연산을 포함한 파이프라인 전체의 문제로 다뤘다는 점입니다.

MMA가 빨라질수록 softmax와 특수 함수 연산의 비용도 더 뚜렷해집니다. 팀은 KV 언롤 방식과 배리어 대기 위치를 조정해 타일 경계에서 생기는 파이프라인 공백을 줄였습니다. 이 구현은 Meta 내부 GEM 학습 작업에 사용되고 있으며 관련 코드도 공개됐습니다. 다만 제시된 성능은 특정 입력 형상과 커널 설정에 대한 결과이므로 모든 모델에서 동일한 가속을 보장하는 수치로 해석해서는 안 됩니다.

출처: PyTorch Blog