119편 · 약 16분

FlashAttention-4: Blackwell를 위한 알고리즘-커널 공동 설계

요약

FlashAttention-4(FA4)는 2026년 3월 Tri Dao와 Together AI가 발표한 어텐션 커널로, NVIDIA Blackwell 아키텍처의 비대칭 확장 문제를 해결하기 위해 알고리즘과 하드웨어 커널을 함께 재설계했습니다. 텐서 코어 처리량은 Hopper 대비 2배 증가했지만 소프트맥스·공유 메모리는 동일한 속도에 머물러 기존 FlashAttention-3 커널은 실효 대역폭의 25~60%를 낭비하고 있었습니다. FA4는 소프트웨어 기반 exp() 에뮬레이션, 2-CTA MMA, TMEM(Tensor Memory) 직접 접근, 더 깊은 비동기 파이프라인을 통해 이 병목을 제거합니다.

배경

Blackwell의 비대칭 확장

NVIDIA Hopper(H100)에서 Blackwell(B200, SM100)로 넘어갈 때 모든 유닛이 같은 비율로 빨라지지 않았습니다.

유닛	Hopper → Blackwell 배율
WGMMA/UMMA 텐서 코어 (BF16)	×2
TMEM 용량	신규 (256 KB/SM)
exp() 특수 함수 유닛	×1 (동일)
Shared Memory 대역폭	×1 (동일)
L2 캐시	×1~1.5

FlashAttention-3가 Hopper 기준으로 설계된 이유는 분명합니다. Hopper에서는 WGMMA 타일이 64×64×16이고 소프트맥스 오버헤드가 전체 사이클의 15~20%에 불과했습니다. Blackwell에서 UMMA 타일은 128×256×16(BF16)으로 두 배 커졌는데, 소프트맥스가 처리할 수 있는 원소 수는 그대로여서 비율이 역전됩니다.

병목 계산 예시 (시퀀스 길이 8K, head\_dim 128):

UMMA 처리량: H100 대비 약 2×
exp() 오버헤드: 기존 비율 기준 소요 시간 2× 증가 상당
결과: 소프트맥스가 전체 어텐션 연산의 40~60%를 차지

그림 1. FlashAttention-3와 FA4의 파이프라인 비교

Blackwell 하드웨어 신기능

TMEM: 텐서 메모리

TMEM(Tensor Memory)은 Blackwell에서 새롭게 추가된 256 KB/SM 크기의 온칩 스크래치패드입니다. 공유 메모리와 달리 텐서 코어 유닛에 물리적으로 직접 연결되어 있어, 행렬 연산 결과를 공유 메모리를 경유하지 않고 곧바로 TMEM에 누적할 수 있습니다.

Hopper:  [텐서 코어] → [WGMMA 결과 레지스터] → [공유 메모리] → [다음 연산]
Blackwell: [텐서 코어] ←→ [TMEM 직접] → [다음 연산 또는 HBM 기록]

FA4는 부분 소프트맥스 통계(행별 최댓값 m, 로그 합산값 ℓ)를 TMEM에 유지하고, 조건부 재조정 계산을 소프트웨어로 직접 구현합니다.

UMMA: 통합 MMA

Hopper의 WGMMA(Warpgroup MMA)가 64×64×16 타일을 처리했다면, Blackwell의 UMMA(Unified MMA)는 기본 128×256×16 타일을 지원합니다. 타일 크기가 4배 커졌으므로 텐서 코어 처리량도 해당 비율로 증가합니다.

2-CTA MMA: 두 개의 CTA(Cooperative Thread Array)가 각자의 TMEM을 합산하여 256×256×16 타일을 형성하는 협력 모드입니다. 대형 시퀀스에서 한 CTA가 담당하는 연산 블록이 클수록 소프트맥스 오버헤드 비중이 줄어드는 원리를 극대화합니다.

TMA: 텐서 메모리 어드레서

TMA(Tensor Memory Accelerator)는 global HBM ↔ 공유 메모리/TMEM 사이의 비동기 데이터 이동을 전담하는 하드웨어 유닛입니다. CUDA 스레드가 메모리 복사를 기다리지 않고 다음 연산을 즉시 발행할 수 있어, 데이터 적재와 행렬 곱셈이 타임라인 상에서 겹칩니다.

FA4의 핵심 알고리즘 변경

소프트웨어 기반 exp() 에뮬레이션

Hopper까지는 하드웨어 특수 함수 유닛(SFU)이 exp()를 실행했습니다. Blackwell에서는 텐서 코어 처리량이 2× 증가했으나 SFU 처리량은 그대로이기 때문에, exp()가 파이프라인 병목이 됩니다. FA4는 exp()를 다음 방식으로 소프트웨어 에뮬레이션합니다.

exp2 변환: exp(x) = exp2(x / ln2) — 이진 지수는 비트 시프트와 LUT로 빠르게 계산 가능
다항식 보정: 2~3차 미니맥스 다항식으로 잔차를 보정해 BF16 정밀도 유지
TMEM 누적: 부분 통계를 레지스터 대신 TMEM에 누적해 레지스터 압력 감소

에뮬레이션 코드는 Dao-AILab/flash-attention 저장소의 flash_attn/cute/ 디렉터리에서 cute:: 커널 추상화를 통해 구현됩니다.

조건부 소프트맥스 재조정

FlashAttention 시리즈는 블록 단위로 어텐션을 계산하고 각 블록마다 최댓값 m과 합산값 ℓ을 갱신합니다. FA4는 블록 간 최댓값 변화량이 임계값 이하일 때 재조정 과정을 건너뜀으로써 연산 횟수를 줄입니다.

if |m_new - m_old| < ε:
    # 재조정 생략 → TMEM 누적만 수행
else:
    # 완전 재조정 후 누적

이 조건부 처리 덕분에 길이가 긴 시퀀스에서 소프트맥스 실효 연산이 약 30% 감소합니다(논문 측정치).

더 깊은 비동기 파이프라인

FA3의 파이프라인은 2단계(TMA 적재 → WGMMA)였지만, FA4는 4단계(TMA 적재 → UMMA 1 → SW-exp → UMMA 2)로 깊어졌습니다. 각 단계가 독립적으로 실행될 수 있어 SM 활용률이 높아집니다. 구체적으로는 CUDA cooperative groups와 TMEM 장벽 명령어(TMEM_ARRIVE, TMEM_WAIT)를 조합해 단계 간 동기화를 구현합니다.

성능 측정

아래 수치는 arXiv:2603.05451(논문 원문)과 vLLM 0.27.0 릴리스 노트(2026-08-09)에서 인용한 것입니다. 환경에 따라 차이가 있을 수 있습니다.

설정	FA3 (H100)	FA4 (B200)	비고
seq=2K, head\_dim=128, BF16	380 TFLOPS	720 TFLOPS	B200 기준
seq=8K, head\_dim=128, BF16	350 TFLOPS	890 TFLOPS	긴 시퀀스 이득 큼
seq=32K, head\_dim=128, BF16	310 TFLOPS	940 TFLOPS	소프트맥스 비중 감소
이론 peak (BF16)	1,979 TFLOPS	4,500 TFLOPS	—
실효 peak 달성율	~18%	~21%	—

Open question: B300 Blackwell Ultra(SM103) 기준 공식 벤치마크는 논문 게재 이후 발표 예정이며 현재 독립 측정값이 없습니다.

운영 가이드

vLLM에서 FA4 활성화 (0.27.0+)

FlashAttention-4는 vLLM 0.27.0(2026-08-09 릴리스)부터 Blackwell GPU 환경에서 자동 감지됩니다.

# 설치 (Blackwell CUDA 환경 필요)
pip install vllm==0.27.0

# B200 서버에서 실행 — FA4 자동 선택
vllm serve meta-llama/Llama-3.1-70B-Instruct \
  --tensor-parallel-size 4 \
  --max-model-len 131072

# 어텐션 백엔드 명시 (선택)
VLLM_ATTENTION_BACKEND=FLASH_ATTN_4 vllm serve ...

수동으로 백엔드를 확인하려면 로그에서 Using FlashAttention-4 backend 문자열을 확인합니다.

주의 사항

조건	내용
GPU 요구 사항	Blackwell (SM100 이상) 필수; Hopper(SM90)에서는 FA3 유지
CUDA 버전	CUDA 12.8 이상 권장 (TMA 비동기 API)
배치 크기	2-CTA MMA는 배치 크기가 작으면 오버헤드가 있어 seq≥4K에서 효과가 두드러짐
정밀도	BF16·FP16 지원; FP8 지원은 별도 커널 (개발 중, Open question)
causal mask	지원 (FlashAttention 시리즈 동일)
MQA/GQA	지원

FlashAttention 세대 선택 기준

GPU 세대 감지
  ├─ SM100+ (Blackwell)   → FA4 (vLLM 0.27.0+)
  ├─ SM90  (Hopper)       → FA3 (vLLM 0.7.x+)
  └─ SM80  (Ampere)       → FA2

요점 정리

Blackwell에서 텐서 코어가 2× 빨라졌지만 소프트맥스·공유 메모리는 그대로여서 기존 어텐션 커널이 병목 구간을 만든다.
TMEM(256 KB/SM)은 텐서 코어에 직접 연결된 온칩 메모리로, FA4가 부분 소프트맥스 통계를 여기에 유지해 레지스터 압력을 낮춘다.
UMMA 128×256×16 타일과 2-CTA 256×256×16 협력 모드로 대형 타일 연산을 활용한다.
exp()를 exp2 변환 + 다항식 보정으로 소프트웨어 에뮬레이션하여 SFU 병목을 우회한다.
4단계 비동기 파이프라인으로 TMA 적재, UMMA, SW-exp가 겹쳐 실행된다.
긴 시퀀스(32K+)에서 B200 기준 약 940 TFLOPS, H100 FA3 대비 약 3×의 실효 처리량을 달성한다.
vLLM 0.27.0(2026-08-09)부터 Blackwell GPU에서 FA4 백엔드가 자동 선택된다.

References

Tri Dao, Daniel Y. Fu et al., "FlashAttention-4: Co-Designing Algorithm and Kernel for Blackwell," arXiv:2603.05451 (2026-03-05). https://arxiv.org/abs/2603.05451
Dao-AILab/flash-attention GitHub: https://github.com/Dao-AILab/flash-attention
vLLM 0.27.0 Release Notes (2026-08-09): https://github.com/vllm-project/vllm/releases/tag/v0.27.0
NVIDIA Blackwell Architecture Technical Brief: https://resources.nvidia.com/en-us-blackwell-architecture
NVIDIA CUDA Programming Guide — Tensor Memory Accelerator (TMA): https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#tensor-memory-accelerator
FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision (arXiv:2407.08608): https://arxiv.org/abs/2407.08608