How FlashAttention 4 Works

FlashAttention 시리즈는 GPU 한 세대마다 다시 짜인다 — FA1 (V100/A100), FA2 (A100 최적화), FA3 (Hopper). 그리고 2025 년 Blackwell (B200) 이 등장하면서 FA4 가 등장. 이 강의는 Charles Frye 가 Modal 블로그에서 FA4 코드를 reverse-engineer 해 분석한 결과를 풀어낸 자리.

강의의 frame.

1. 새 칩의 새 명령 — Blackwell 의 tcgen05.mma.cta_group::1 같은 5세대 tensor core 명령. 이걸 쓰지 않으면 peak 의 한참 아래.
2. warp specialization 의 한 단계 위 — FA3 가 두 그룹의 ping-pong 이라면 FA4 는 Load / MMA / Softmax / Correction / Epilogue 의 5-way 분리.
3. SFU 우회 — softmax 의 exp 가 SFU bottleneck 에 부딪히는 자리에서 cubic polynomial approximation 으로 우회.
4. 결과 — cudnn 의 attention 대비 ~20% 빠르다. attention 이라는 op 가 얼마나 깊게 짜내질 수 있는지 의 새 표.

또 강조할 만한 변경 — softmax normalization update 의 lazy 화. FA3 까지는 row-max 가 갱신될 때마다 모든 partial output 을 rescale. FA4 는 "max 가 numerical stability 에 영향을 줄 만큼 변할 때만" rescale. "output rescaling 비용을 10배 줄였다" 는 Modal 블로그의 표현.

FA3 까지 TMA 는 warp 의 일부 thread 만 사용하는 모델이었다. FA4 는 한 step 더 — 한 warp 전체를 "Load warp" 로 specialize. async copy 발사만 하고 다른 일은 하지 않는다.

왜 의미가 있는가.

• register pressure 절감 — Load warp 는 multidimensional indexing 을 직접 안 한다. TMA descriptor 가 모든 stride / box shape 를 들고 있다. 그래서 register 가 매우 적게 든다.
• async copy 의 paralell 발사 — Load warp 가 여러 stage 의 copy 를 한 번에 발사할 수 있다. consumer warp 의 latency stall 이 가려진다.
• 다른 warp 의 register 자유 — load 를 안 하므로 MMA / Softmax warp 들이 가진 register 를 다른 데 쓸 수 있다.

Blackwell (B200) 이 H100 위에 추가한 핵심 자리.

• tcgen05.mma.cta_group::1 — 5세대 tensor core. cta_group::1 의 의미는 "이 명령은 한 CTA 안에서만 작동, cross-CTA op 없음" — 즉 tile scheduling 이 단순해진다.
• FP8 의 정식 first-class — H100 에서 실험적이던 FP8 이 Blackwell 에서는 standard. tensor core 가 FP8 ops 의 throughput 2× (BF16 대비).
• FP4 (NVFP4) — 새 dtype. 4-bit floating point. tensor core 가 4× throughput (BF16 대비). attention 의 attention-score / output 자리에 쓸 수 있다.
• tensor memory — 새로운 named address space. SMEM 보다 더 큰, register 보다 더 큰 자리. attention 의 KV cache 가 이 자리로 이동 (확인 필요).
• async copy 의 강화 — TMA 가 H100 보다 더 큰 box, 더 fine-grained barrier 모델.

FA4 의 핵심 architectural 결정 — warp 의 역할을 다섯으로 분리. 각자가 자기 stage 만 담당, 다른 stage 의 결정은 모른다.

이 5-way 모델의 핵심 이점 — each warp scheduler 가 매 cycle 다른 stage 로 switch 할 수 있다. SM 의 4 개 sub-partition 은 동시에 다른 일을 할 수 있고, 그것이 5 개 specialist 위에 mapping 되면 throughput 이 polynomial 로 올라간다.

attention 의 어느 자리가 low precision 으로 가도 안전한가의 답이 FA4 에서 더 명확해진다.

• Q · K^T (logits) — FP8 안전. tensor core 가 FP8 input → FP32 accumulator 라서 numerical headroom.
• softmax (max, exp, sum) — FP32 으로 유지. dynamic range 가 너무 커서 low-precision 위험.
• P · V (output) — FP8 안전. softmax 후 P 가 [0, 1] 범위 — narrow range 라서 안정.
• attention output store — BF16 또는 FP16. downstream 의 norm / GEMM 에서 처리.
• FP4 (실험) — Q·K^T 또는 P·V 자리에. accuracy 유지를 위해 추가 calibration 필요.

이 분포가 핵심 — tensor core 가 도는 자리는 low precision, scalar 작업 (softmax) 은 high precision. FP8 attention 이 BF16 대비 1.5-2× 속도, FP4 는 2-4× (확인 필요). 단 prefill 같은 long-context 에서는 numerical drift 가 누적할 risk.

Modal 이 reverse-engineer 한 디테일 중 단독 트릭들. 각자 5-15% effect, 곱하면 큰 차이.

• cubic polynomial approx for exp — 작은 head dim 에서 SFU bottleneck 을 우회. "BF16 값에 대해 SFU 의 출력과 일치". SFU 사용 0.
• lazy softmax rescale — row-max 변화가 numerical stability 에 영향을 줄 만큼이 아니면 partial output 의 rescale 을 skip. "output rescale 을 10× 절감".
• cta_group::1 사용 — single-CTA 만 돌게 해서 cross-CTA 통신 비용을 0 으로. tile 단위 스케줄링이 깔끔.
• multi-buffered pipeline — stage 갯수가 FA3 의 2-3 에서 더 깊어진다 (확인 필요 — 4-5 stage 가능성).
• register reuse 최적화 — 5-way specialization 으로 각 warp 의 register 가 할일에 집중. occupancy 유지.
• shared memory swizzling — SMEM 의 bank conflict 회피. tcgen05 의 layout 요구에 맞춰 직접 swizzle.
• scoreboard 최적화 — async barrier 의 mbarrier phase tracking. naïve 구현보다 register usage 작음.

• flash-attn 4.x — Tri Dao 의 공식 구현. ThunderKittens 위에 reference, 또는 직접 CUDA. (확인 필요)
• vLLM B200 백엔드 — Blackwell GPU 의 attention 핫패스에 FA4 를 default 로 채택.
• Together / Modal — 자체 inference 인프라에서 FA4 를 운영. Modal 이 reverse-engineer 한 동기.
• Anthropic / OpenAI — Blackwell cluster 에서의 추론 / 학습. 직접 fork 했을 가능성 (확인 필요).
• DeepSeek / Qwen 인프라 — large-scale 학습의 attention path. FA4 의 BF16 / FP8 path.

채택 곡선이 빠른 이유 — Blackwell 의 새 기능을 쓰지 않고 attention 을 짤 수 없기 때문. cudnn 은 NVIDIA 가 직접 짜지만 fine-grained 통제가 어려움. FA4 는 같은 hardware 를 더 직접 짜낸 결과.

• NVIDIA Rubin (2026) — 새 tensor core, 더 큰 SMEM, 새 async 모델. FA5 가 등장할 가능성. specialization 이 더 깊어질 것.
• FP4 의 mainstreaming — production 채택까지 calibration / mixed precision 의 표준화 필요.
• multi-GPU / distributed FA4 — Iris (L078) 같은 추상과 결합한 ring attention 의 fully-fused 버전.
• training 의 backward FA4 — forward 위주 FA4 의 backward 가 같은 5-way 추상 안에 들어가는지.
• kernel 합성으로의 흡수 — Mirage MPK 같은 super-optimizer 가 FA4 의 트릭을 자동 발견할 수 있는가.

손에 새기기 — 실습 시퀀스

1. FA4 코드 읽기 — flash-attention repo 의 sm_100 branch (Blackwell). 5 개 warp role 이 어떻게 분배되는지 line-by-line.
2. FA3 vs FA4 직접 측정 — 같은 attention shape 으로 H100 (FA3) 과 B200 (FA4) wallclock 비교.
3. cubic poly approx — exp 의 polynomial approximation 직접 짜보고 SFU 의 결과와 BF16 위에서 비교.
4. lazy rescale 실험 — naïve eager rescale vs lazy 의 wallclock 차이를 toy attention 에서 측정.
5. FP8 attention — flash-attn FP8 path 직접 사용. accuracy drift 측정.
6. nsys timeline — FA4 의 5-way specialization 이 SM 위에서 어떻게 동시 진행되는지 timeline 으로.

• FA4 의 정확한 stage 갯수 — multi-buffered pipeline 의 stage 가 4 인지 5 인지. Modal 블로그가 정확히 명시했는지 확인.
• Softmax warp 의 정확한 수 — Modal 의 8 warps 추정. 코드의 실제 갯수.
• FP4 attention 의 channels — Q·K^T 자리만, P·V 자리만, 또는 둘 다. accuracy 데이터.
• cubic poly 의 정확한 계수 — exp 의 polynomial form 의 정확한 식. fp16 도메인에서의 max error.
• lazy rescale 의 trigger 조건 — "stability 에 영향" 의 numerical 정의.
• backward pass 의 5-way specialization — forward 와 같은 추상이 backward 에 적용되는지.
• cudnn +20% 의 정확한 측정 — sequence length, head dim, dtype 별로 어떻게 다른지.