GPU에서 Transformer를 돌릴 때 병목이 어디서 오는지 묻는다면, 대부분은 곧바로 "attention의 O(N²) 복잡도"를 꼽는다. 맞는 말이긴 한데, 거기서 한 발 더 들어가야 Flash Attention의 설계가 보인다. 실제 지연은 O(N²) 연산 자체가 느려서가 아니라, 그 연산이 만들어내는 N×N 중간값을 GPU HBM에 쓰고 다시 읽는 왕복 비용에서 나온다.
Tri Dao et al.의 FlashAttention 원논문이 이걸 처음 수치로 증명했다. 연산량(FLOP)이 아니라 메모리 IO 횟수가 병목이라는 관점을 "IO-Aware Attention"이라 부른 것도 그래서다.
GPU 메모리 계층이 문제의 시작이다
A100의 메모리 계층은 세 단계로 나뉜다.
| 계층 | 대역폭 | 용량 |
|---|---|---|
| HBM (High Bandwidth Memory) | 2 TB/s | 80 GB |
| L2 캐시 | ~5 TB/s | 40 MB |
| SRAM (Shared Memory, SM당) | ~19 TB/s | 192 KB |
HBM은 용량이 크고 속도는 상대적으로 느리다. SRAM은 SM당 192KB에 불과하지만 대역폭이 HBM의 약 10배다. 문제는 Standard Attention이 SRAM에 올릴 수 없는 크기의 중간값을 HBM에 반복적으로 쓰고 읽는다는 것이다.
Standard Attention의 실제 병목: HBM 접근 횟수
Standard Attention을 PyTorch로 구현하면 세 단계로 쪼개진다.
# 개념적 흐름 (커널이 분리될 때 일어나는 일)
S = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d) # N×N → HBM에 기록
P = torch.softmax(S, dim=-1) # HBM에서 읽고 → 다시 기록
O = torch.matmul(P, V) # HBM에서 읽고 → 출력
S를 계산하면 N×N 크기의 행렬이 HBM에 기록된다. Softmax를 계산하기 위해 그걸 다시 읽는다. 결과 P를 HBM에 다시 쓰고, PV 곱셈을 위해 또 읽는다. S와 P의 HBM 왕복만 계산해도 쓰기 2회 + 읽기 2회, 원소 4N²개의 이동이 발생한다. fp16 기준이면 8N² 바이트다.
| 시퀀스 길이 N | Standard (S+P 중간값, fp16) | Flash (중간값 없음) |
|---|---|---|
| 512 | 2.1 MB | 0 |
| 1024 | 8.4 MB | 0 |
| 4096 | 134 MB | 0 |
| 16384 | 2.1 GB | 0 |
N이 두 배가 될 때마다 중간값 HBM 트래픽은 네 배로 불어난다. Horace He의 roofline 분석으로 보면 Standard Attention은 GPU 연산 유닛이 계산을 기다리는 compute-bound가 아니라 HBM 데이터가 돌아오기를 기다리는 memory-bandwidth-bound 연산이다. FLOP을 줄여봤자 속도가 개선되지 않는 이유가 여기 있다. 커널 퓨전 없이 matmul → softmax → matmul을 세 번으로 쪼개 실행하면 각 스텝마다 HBM 왕복이 강제된다.
Tiling: N² 행렬을 SRAM에서 처리하는 방법
Flash Attention의 핵심 아이디어는 N×N 행렬을 HBM에 기록하지 않는 것이다. Q, K, V를 작은 타일로 쪼개서 SRAM에 올린 뒤, 그 안에서 완전한 attention 출력을 계산해버린다.
타일 크기 Br=Bc=64, d=64, fp16 기준 SRAM 사용량을 계산하면:
Q 타일 : 64 × 64 × 2B = 8 KB
K 타일 : 64 × 64 × 2B = 8 KB
V 타일 : 64 × 64 × 2B = 8 KB
S 타일 : 64 × 64 × 2B = 8 KB (SRAM 내에서만 사용, HBM에 안 씀)
O 누산기 : 64 × 64 × 2B = 8 KB
──────────────────────────────
합계 : 40 KB
SM당 SRAM 192KB 중 40KB만 소비한다. N=1024에서 Q 방향으로 Tr=16개, K·V 방향으로 Tc=16개의 타일이 만들어진다. 각 타일 조합에서 S 블록을 SRAM 안에서 계산하고, softmax를 적용해 O를 갱신한 뒤, S는 그냥 버린다. HBM에 기록되는 건 최종 O뿐이다.
문제는 softmax가 전체 행의 최댓값을 알아야 수치 안정성이 보장된다는 점이다. 아직 처리하지 않은 타일의 값을 미리 알 수 없는 상황인데, Flash Attention은 이걸 online softmax로 해결한다. 타일을 처리할 때마다 현재까지의 최댓값 m과 분모 합계 l을 누적하고, 새 타일이 들어올 때 이전 O 누산기를 보정한다.
m_new = max(m_old, max(S_tile))
l_new = exp(m_old - m_new) × l_old + sum(exp(S_tile - m_new))
O_new = (exp(m_old - m_new) × l_old × O_old + exp(S_tile - m_new) × V_tile) / l_new
이 보정 덕분에 N×N 행렬 전체를 메모리에 올려두지 않아도 정확히 같은 softmax 결과가 나온다. 근사가 아니라 exact attention이다. 논문 제목에 "Exact Attention"이 붙은 이유가 여기 있다.
타일 크기 선택은 SRAM 용량과 GPU 아키텍처에 따라 튜닝한다. Br, Bc를 키우면 HBM 접근 횟수는 줄지만 SRAM 사용량이 늘어 다른 커널과의 shared memory 충돌이 생길 수 있다. flash-attn GitHub 저장소에 GPU 아키텍처별 권장 블록 크기가 정리되어 있다.
Recomputation: 메모리를 아끼기 위해 FLOP을 일부러 낭비하는 선택
학습에서 역전파는 forward에서 생성한 중간값을 필요로 한다. Standard Attention에서는 S와 P를 저장해두고 gradient를 계산하는데, 이게 N² 크기다. 시퀀스 길이가 늘면 이 중간값이 GPU 메모리를 잠식한다.
Flash Attention은 S와 P를 저장하는 대신 backward에서 재계산한다. 저장하는 건 online softmax의 통계량 m과 l뿐이고, 이건 O(N) 크기다. 명백히 FLOP 낭비인데, GPU에서 대역폭이 연산 자체보다 비싸기 때문에 합리적인 선택이 된다. N² 크기 텐서를 HBM에 저장했다가 다시 읽는 비용이 그 값을 SRAM에서 재계산하는 비용보다 크다. 메모리를 팔고 FLOP을 사는 트레이드오프다.
서빙(추론)에는 역전파가 없다. Recomputation은 gradient 계산을 위한 설계 결정이므로, 추론에서 Flash Attention의 이득은 순수하게 tiling이 만드는 HBM IO 감소에서 온다. 서빙 성능을 분석할 때 학습 메모리 절감 효과와 혼동하면 원인 파악이 틀린다.
FA2·FA3: 각 버전이 다음으로 겨냥한 병목
FA1이 HBM IO를 해결했다면, FA2와 FA3는 그 위에 남은 병목을 순서대로 다뤘다.
| FA1 (2022) | FA2 (2023) | FA3 (2024, H100) | |
|---|---|---|---|
| 핵심 변경 | Tiling + Recomputation | Q 기준 outer loop 전환 | Warp specialization + FP8 파이프라이닝 |
| Outer loop 방향 | K·V | Q | Q (유지) |
| 비matmul FLOP 최적화 | 없음 | 있음 | 있음 |
| FP8 지원 | 없음 | 없음 | 있음 |
| 달성 MFU (이론 대비) | 25–40% | 50–73% | ~75% (FP16) |
FA2 논문의 핵심은 outer loop 방향 전환이다. FA1에서 K·V를 바깥 루프로 돌리면 같은 Q 블록의 O를 갱신하는 작업이 여러 thread block에 분산되고 이를 합산하는 동기화가 필요하다. FA2는 Q를 바깥 루프로 바꿔 각 thread block이 독립적으로 O 블록을 완성하게 만들었다. occupancy가 올라가고, 동시에 softmax rescale 과정의 비matmul FLOP을 줄여 Tensor Core 활용률도 높였다. FA1이 이론 피크의 25–40%에 머물렀던 반면 FA2는 50–73%까지 끌어올렸다.
FA3는 Hopper 아키텍처(H100)에서만 가능한 기능을 활용한다. Warp specialization으로 matmul과 softmax를 서로 다른 warp에 분리 배정해 비동기로 겹쳐 실행한다. FP8 저정밀도 모드에서는 H100 기준 약 1.2 PFLOPs/s까지 올라간다.
서빙에서 Flash Attention 이득이 예상보다 작은 경우
이득의 크기가 N에 비례한다는 사실이 서빙 환경에서 자주 간과된다. Standard Attention의 N² 중간값 트래픽 대비 Flash Attention의 절약분은 N/d 비율로 커진다. d=64 고정이면 N=4096에서는 64배, N=512에서는 겨우 8배 수준이다.
짧은 시퀀스(128~512 토큰) 구간에서 FA2와 naive attention을 A100 batch=1로 비교하면 latency 차이가 미미하거나 경우에 따라 tiling setup overhead가 역전을 만든다. 절약할 N² 트래픽 자체가 적기 때문이다. FA1 원논문 실측에서도 N≤256에서는 speedup이 2× 미만으로 떨어진다.
batch=32처럼 배치가 커지면 GPU occupancy가 자연히 높아져 Flash Attention의 tiling 효과가 상대적으로 작아진다. 이 구간에서는 Flash Attention이 이득을 주는 게 맞지만, 이득의 주 원인이 IO 감소인지 occupancy 개선인지는 구분해서 봐야 한다.
GQA(Grouped Query Attention)나 MQA를 쓰는 모델에서는 헤드 수가 줄어 병렬화 기회도 감소한다. Llama 3 70B처럼 KV head가 8개인 모델에서 batch=1 decoding을 돌리면 Flash Attention이 처리하는 tile 수 자체가 작아지고, GPU occupancy 문제가 이득보다 크게 작용한다. 이 조건에서는 FlashDecoding이나 FlashInfer 같은 decoding 특화 커널이 더 적합한 선택이다.
Flash Attention이 극적인 이득을 주는 조건은 긴 시퀀스 + prefill 단계다. 짧은 시퀀스의 decoding에서 기계적으로 적용하면 이득이 거의 없고, 서빙 환경에서 이 조건을 먼저 확인해야 한다.
PagedAttention·Chunked Prefill과의 위치
vLLM과 SGLang은 Flash Attention을 커널 수준에서 쓰면서 그 위에 PagedAttention을 구현한다. 계층이 다르다. Flash Attention은 연속된 Q·K·V 텐서를 받아 tiled 방식으로 attention을 계산하는 커널이고, PagedAttention은 비연속 메모리 블록으로 KV cache를 관리하는 상위 구조다.
Flash Attention은 기본적으로 연속된 K·V 메모리 레이아웃을 가정한다. PagedAttention의 비연속 블록을 처리하려면 gather 연산으로 연속 버퍼에 복사하거나, flash-attn repo의 flash_attn_with_kvcache API처럼 paged 접근을 직접 지원하는 변형을 써야 한다. vLLM은 자체적으로 paged KV 레이아웃과 Flash Attention을 통합하는 커스텀 CUDA 커널을 추가로 구현했다.
Chunked Prefill도 마찬가지다. 긴 프롬프트를 청크로 나눠 prefill하는 방식인데, 각 청크 내 attention 계산에 Flash Attention 커널이 그대로 들어간다. 청크 간 KV cache 연결은 상위 레이어가 처리하고, Flash Attention은 어느 쪽에서도 청크 경계나 페이지 경계를 직접 알지 못한다—연속된 텐서 뷰만 받으면 된다.