FlashAttention-4는 비원인적(noncausal) 추론과 인과적(causal) 학습을 위한 새로운 경로를 도입하여 Blackwell의 4비트 부동소수점(FP4) 텐서 코어의 성능 한계를 해결합니다. Direct-P라고 불리는 이 방법은 행렬 곱이 축소됨에 따라 지배적이 되는 소프트맥스 변환 오버헤드를 우회하기 위해 점수를 직접 FP4 확률로 매핑합니다.

  • Direct-P는 NVIDIA GB200에서 비원인적 추론을 위해 bfloat16(BF16) 순방향 처리량 대비 최대 2.13배의 성능을 달성합니다.
  • 인과적 경로는 순방향 양자화를 역방향으로 직접 전달하며, 저장된 양자화된 쿼리와 키를 사용하여 8비트 부동소수점(FP8) 그라디언트 연산자를 통해 확률을 복원합니다.
  • 이 인과적 접근법은 단일 GPU에서 80억 파라미터 업데이트 전체를 최대 1.14배 가속화합니다.
  • 일치된 분산 학습은 FP8 확률과 값을 유지하며, 테스트된 모든 MXFP4 확률/값 학습 궤적이 발산하기 때문입니다.

저자들은 이것이 Blackwell 아키텍처에서 효율적인 하드웨어 인식 FP4 어텐션을 가능하게 하여 FP4 텐서 코어로부터 자동 속도 향상을 방해하는 의존성을 극복하므로 중요하다고 봅니다.