TL;DR

  • Blackwell의 B200에서 Triton TLX 기반 FlashAttention을 평가한 결과, BF16 가변 배치 24개 사례의 기하평균 처리량이 기존 FA4보다 7.3% 향상됨.
  • 고정 시프트 밀집 소프트맥스로 반복 누산기 보정을 없애고, 저장한 역분모로 역전파의 행 정규화를 선형 전처리로 옮김.
  • 인과 경로에서는 첫 블록 앵커와 선택적 복구를 사용하고, 완전히 마스킹된 대각 수축 연산을 생략함.
  • 패킹된 BF16 지수 근사와 파이프라인형 dQ 게시로 원소당 연산량과 온칩 데이터 이동을 줄임.
  • 밀집 순전파와 역전파는 각각 13.0%, 10.4% 빨라졌으며, 인과 역전파는 8.0% 향상되고 인과 순전파는 1.5% 느려짐.

연구 배경과 구현

  • Blackwell에서는 행렬 곱셈이 이미 깊게 파이프라인화된 어텐션 커널도 정규화와 피연산자 이동이 병목이 될 수 있음.
  • Triton TLX에서 FA4 스타일 파이프라인을 구현하고, 역전파를 위해 고정 시프트 밀집 소프트맥스와 저장된 역분모를 사용함.
  • 고정 시프트는 반복적인 누산기 보정을 제거하고, 저장된 역수는 행 정규화를 이차 복잡도의 역전파 루프에서 선형 전처리 단계로 옮김.
  • 가드가 적용된 인과 경로는 첫 블록 앵커와 선택적 복구를 사용하며, 완전히 마스킹된 대각 수축 연산은 건너뜀.
  • 패킹된 BF16 지수 근사와 파이프라인형 dQ 게시를 적용해 원소당 연산량과 온칩 데이터 이동을 줄임.

성능 평가

  • 1,000W B200에서 준비된 커널 실행과 FA4의 기존 타이머를 사용해 측정함.
  • 헤드 차원 128, 시퀀스 길이 최대 32,768토큰인 FA4의 BF16 가변 배치 그리드 24개 사례에서 TLX가 기하평균 처리량을 7.3% 높임.
  • 밀집 순전파는 13.0%, 밀집 역전파는 10.4% 향상됨.
  • 인과 역전파는 8.0% 향상됐으며, 인과 순전파는 1.5% 느려짐.
  • TritonBench 비교에서는 순전파 기하평균 성능이 13.4%, 역전파가 24.4% 향상됨.

논문 정보

  • 제목: *Efficient FlashAttention on Blackwell via Fixed-Shift Softmax and Persistent Scheduling*.
  • 저자: Oleksandr Stashuk, Hongtao Yu, Jay Shah.
  • 제출일: 2026년 9월 24일.
  • arXiv 식별자: arXiv:2610.02229.
  • DOI: https://doi.org/10.48550/arXiv.2610.02229