FlashAttention-4, Blackwell 병목에 맞춘 overlap 중심 커널 설계 공개

r/LocalLLaMA에서 공유된 FlashAttention-4는 NVIDIA Blackwell 세대에서 attention 커널을 다시 설계한 사례다. 핵심 메시지는 단순하다. 최신 GPU는 tensor core 처리량은 크게 늘지만 shared memory bandwidth와 SFU 처리량은 같은 속도로 늘지 않기 때문에, GEMM 단독 최적화보다 연산 겹침(overlap) 전략이 더 중요해졌다는 것이다.

작성팀은 H100에서 B200으로 갈 때 BF16 tensor throughput이 약 1.0에서 2.25 PFLOPs로 증가하는 반면, SFU 수와 shared memory bandwidth는 크게 늘지 않았다고 설명한다. 이 전제를 바탕으로 FlashAttention-4는 다음 두 구간을 집중 공략한다.

• Forward: softmax exponential 비용을 MMA와 최대한 겹치고, hardware/software 혼합 exp 경로를 사용.
• Backward: TMEM 배치와 Blackwell 2-CTA MMA를 통해 shared-memory 트래픽과 atomic 부담을 축소.

원문 기술 글에는 ping-pong 타일 스케줄, conditional online softmax rescaling, TMEM 재사용, DSMEM 교환 기반 dQ 분해가 상세히 나온다. 또한 backward에서 reduction 순서를 고정하는 deterministic mode를 제공하며, benchmark 기준으로 nondeterministic 대비 약 85-90% 처리량을 제시한다.

공개 수치도 강하다. B200 BF16에서 최대 1605 TFLOPs/s(약 71% utilization), forward 기준 cuDNN 9.13 대비 최대 1.1-1.3x, Triton 대비 최대 2.1-2.7x 개선을 보고했다. 문서에는 cuDNN 팀과의 협업 및 최신 버전 비교도 포함되어 있다.

구현 측면에서 실무자가 주목할 지점은 CuTe-DSL 사용이다. 팀은 CUTLASS Python DSL 기반 구현으로 C++ template 중심 워크플로 대비 compile time을 약 20-30x 단축했다고 주장한다.

물론 성능 수치는 작성팀이 공개한 benchmark이므로 실제 배치 크기, mask 패턴, sequence 길이에 맞춘 재검증이 필요하다. 그럼에도 attention이 비용 중심인 LLM 학습·추론 스택에는 매우 실질적인 업데이트다.

커뮤니티 출처: r/LocalLLaMA 스레드
원문: Together AI FlashAttention-4