FlashAttention 원 논문 해부: attention 행렬을 HBM에 쓰지 않고도 정확히 계산하는 법

2단계 열네 번째 논문. 이 글은 Tri Dao 외의 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv v2 PDF, 2022-06-24)을 직접 읽고 썼다. FlashAttention은 self-attention(자기-어텐션)의 결과를 바꾸지 않으면서, 큰 N×N score·probability 행렬을 GPU HBM(High Bandwidth Memory, 고대역폭 메모리)에 materialize(중간 결과 전체를 실제 메모리에 저장)하지 않고, 작은 타일과 online softmax(순차 결합 softmax) 상태로 계산하는 IO-aware(입출력 인식) 알고리즘이다.

이 글에서 얻을 답과 범위

앞 글 LoRA는 “큰 model 전체를 task마다 다시 학습·저장하지 않으려면?”을 다뤘다. FlashAttention은 다른 층의 질문이다.

같은 Transformer, 같은 Q·K·V, 같은 softmax, 같은 output인데

왜 어떤 attention은 긴 입력에서 GPU memory를 터뜨리고 느린가?
왜 FlashAttention은 FLOP(부동소수점 연산) 수를 줄이지 않아도 빨라질 수 있는가?

정답은 “GPU가 행렬 곱을 빨리해서”만은 아니다. attention 구현이 큰 N×N intermediate(중간값)를 HBM에 썼다가 다시 읽는지, 작은 on-chip SRAM(Static Random Access Memory, 칩 위 정적 메모리) 안에서 바로 다음 계산으로 넘기는지가 wall-clock time(실제 경과 시간)을 바꾼다.

이 글을 다 읽으면 다음 질문에 답할 수 있어야 한다.

1. standard attention의 S=QKᵀ, P=softmax(S), O=PV에서 어떤 tensor가 왜 N²인가?
2. FlashAttention이 정확한 softmax를 유지하면서 한 tile씩 처리하려면 왜 row maximum m, normalizer l, unnormalized output accumulator(정규화 전 출력 누산기)만 남기면 되는가?
3. “linear memory”와 “quadratic compute”가 동시에 참인 이유는 무엇인가?
4. HBM·SRAM·kernel·tiling·recomputation(재계산)이 forward와 backward에서 각각 무엇을 읽고 쓰는가?
5. 원 논문 Theorem 2의 IO complexity와 Figure 2·Table 1–4 숫자는 어떤 GPU·batch·sequence length 조건의 결과인가?
6. 원 논문 시점의 공식 CUDA source는 Python call에서 어떤 shape·dtype·device 검사를 거쳐 kernel을 launch하는가?
7. “FlashAttention을 쓰면 context가 무한히 길어진다”, “N²이 사라진다”, “모든 GPU에서 자동으로 빠르다”가 왜 틀린가?

범위: 원 논문 v2의 §1–4, Appendix B–C·E와, 원 논문 v2보다 열흘 앞선 2022-06-14의 공식 Dao-AILab/flash-attention commit 765741c를 다룬다. 2022 구현의 Python interface, C++ dispatcher, CUDA kernel entry를 읽었다. FlashAttention-2, FlashAttention-3, 이후 PyTorch SDPA(Scaled Dot Product Attention, 스케일된 내적 어텐션) dispatch, AMD·Hopper의 최신 kernel은 원 논문과 다른 후속 구현이므로 이 글에서 성능 수치를 섞지 않는다.

먼저 외울 한 문장

FlashAttention은 N²번의 “누가 누구를 볼지” 계산을 없애지 않는다. 대신 그 거대한 표를 HBM에 저장했다가 다시 읽지 않는다. 작은 표 조각을 SRAM에서 계산하고, 각 query 행의 최대값·분모·출력 합만 갱신해 마지막에 정확히 같은 softmax 결과를 만든다.

왜 필요한가: standard attention은 큰 표를 세 번 HBM에 오가게 만든다

12살 비유: 반 친구 전원의 점수표를 벽에 붙이는 문제

교실에 N명의 학생이 있다. 각 학생은 “내가 다른 친구 말을 얼마나 들을까?”를 정한다.

학생 i가 학생 j를 보는 점수     → S[i,j]
점수를 0~1 비율로 바꾼 결과      → P[i,j]
그 비율로 친구들의 정보를 섞은 답 → O[i]

N=5면 점수표는 25칸이다. N=16,384면 한 head(헤드)만 해도 268,435,456칸이다. standard attention은 보통 다음처럼 큰 점수표 S와 확률표 P를 HBM에 쓴다.

Q, K를 읽어 S = QKᵀ 계산 → S[N×N]를 HBM에 씀
S를 다시 읽어 P = softmax(S) → P[N×N]를 HBM에 씀
P와 V를 다시 읽어 O = PV → O[N×d]를 HBM에 씀

마치 교실 칠판보다 훨씬 느린 창고 벽(HBM)에 전교생 점수표를 붙이고, 확률표로 다시 바꾸고, 또 읽어 답을 만드는 셈이다. FlashAttention은 점수표 전체를 벽에 붙이지 않는다. 작은 구역의 친구들만 잠깐 빠른 책상(SRAM)에 올리고, 현재까지의 합계를 작은 메모에 갱신한다.

N²가 실제로 어느 정도인지 먼저 계산한다

한 attention head에서 sequence length(시퀀스 길이) N=4,096, head dimension(헤드 차원) d=64, FP16(16-bit 부동소수점, 원소 2 byte)을 생각하자.

[ N^2 = 4096^2=16{,}777{,}216 ]

이는 batch 1, head 1의 forward 중간값 감각을 위한 산술이다. 실제 training은 batch·head 수, dtype, mask, dropout, autograd(자동 미분)가 곱해진다. 핵심은 N을 두 배로 하면 S,P는 네 배가 된다는 점이다. base model parameter가 고정되어 있어도 long-context training에서 attention activation이 먼저 병목이 될 수 있다.

HBM과 SRAM은 “메모리 있음/없음”이 아니라 속도·크기가 다른 작업대다

원 논문 §2.1은 A100 예로 HBM은 40–80 GB, 대역폭 1.5–2.0 TB/s이고, 각 streaming multiprocessor(SM, GPU에서 thread block을 실행하는 연산 단위)의 on-chip SRAM은 192 KB, 추정 대역폭은 약 19 TB/s라고 설명한다. 수치는 A100이라는 특정 GPU의 예시다. 모든 GPU의 상수가 아니다.

HBM  : 크지만 SRAM보다 멀고, 읽기·쓰기가 상대적으로 비싸다.
SRAM : 매우 작지만 연산 장치와 가깝고 빠르다.

kernel: HBM → registers/SRAM에 필요한 조각을 읽음 → 계산 → HBM에 결과를 씀

여기서 IO(Input/Output, 이 글에서는 주로 HBM과 on-chip SRAM 사이의 읽기·쓰기)는 disk I/O가 아니다. GPU memory hierarchy(메모리 계층) 안의 데이터 이동이다. FLOP가 충분히 많아도, 큰 tensor를 여러 번 HBM에서 왕복하면 attention은 memory-bound(계산기보다 메모리 이동이 병목인 상태)가 될 수 있다.

kernel fusion만으로는 training의 N² 저장을 완전히 없애기 어렵다

여러 elementwise operation(원소별 연산)을 한 CUDA kernel로 합치면 S를 한 번 읽고 mask·softmax·dropout을 이어 처리할 수 있다. 이를 kernel fusion(커널 융합)이라고 한다. 하지만 backward(역전파)가 P 또는 dropout mask를 필요로 하면 forward 뒤에 이를 HBM에 보존해야 할 수 있다.

원 논문의 출발점은 더 강하다.

1. 큰 S와 P를 HBM에 materialize하지 않는다.
2. softmax가 전체 row를 한 번에 보지 않아도 정확히 결합되게 한다.
3. backward에는 S/P를 저장하는 대신 Q/K/V와 작은 통계를 사용해 SRAM에서 다시 계산한다.

따라서 FlashAttention은 단순히 “softmax와 mask를 fuse한 kernel”이 아니다. memory access schedule(언제 어떤 block을 어디에 읽고 쓰는지의 계획) 자체를 바꾼 exact attention algorithm(근사하지 않는 어텐션 알고리즘)이다.

무엇인가: output은 같고, N×N HBM intermediate를 없앤 IO-aware exact attention

원 논문의 정의

한 attention head에서 Q,K,V∈ℝ^{N×d}라 하자. N은 token 수, d는 한 head의 vector 길이다. τ는 보통 1/√d인 softmax scale이다.

[ S=\tau QK^\top\in\mathbb{R}^{N\times N},\qquad P=\operatorname{softmax}(S),\qquad O=PV\in\mathbb{R}^{N\times d} ]

FlashAttention의 최종 O는 이 식의 O와 같다. 원 논문 Theorem 1은 forward가 O(N²d) FLOP를 수행하고 input·output 외 추가 memory가 O(N)임을 보인다. 여기서 “linear memory”는 Q/K/V/O가 사라진다는 뜻이 아니다. S,P 같은 N×N 전역 intermediate를 저장하지 않는다는 뜻이다.

정확히 하지 않는 일

LoRA가 “어떤 weight update를 학습할까”의 선택이라면, FlashAttention은 “같은 attention을 GPU memory에서 어떤 순서로 실행할까”의 선택이다. 두 기술은 서로 대체하지 않는다.

선행 개념: Q·K·V, stable softmax, HBM IO, tile, causal mask

Q·K·V와 한 행의 의미

각 token은 같은 hidden state에서 query(q_i), key(k_j), value(v_j) vector를 만든다. token i의 점수는 모든 token j에 대해 다음이다.

[ s_{ij}=\tau q_i\cdot k_j ]

s_{ij}가 클수록 token i는 j의 value를 더 많이 섞는다. row-wise softmax(행마다 하는 softmax)는 한 행의 모든 점수를 양수 비율로 바꾸고 합을 1로 만든다.

[ p_{ij}=\frac{e^{s_{ij}}}{\sum_{t=1}^{N}e^{s_{it}}},\qquad o_i=\sum_{j=1}^{N}p_{ij}v_j ]

causal mask(인과 마스크)는 autoregressive decoder에서 미래 token을 못 보게 한다. j>i인 score를 -∞로 두면 e^{-∞}=0이 되어 그 value의 비율은 0이다. FlashAttention도 이 수학적 mask 결과를 유지해야 한다.

stable softmax는 overflow를 막는 필수 변형이다

e^{1000}은 floating-point에서 overflow(표현 가능한 최대값을 넘음)할 수 있다. 한 row의 최대값 m=max_j s_{ij}를 빼도 softmax는 바뀌지 않는다.

[ \operatorname{softmax}(s)_j= \frac{e^{s_j-m}}{\sum_t e^{s_t-m}} ]

왜냐하면 분자와 분모에 같은 e^{-m}이 곱해지기 때문이다. FlashAttention의 핵심 state m은 단지 속도용 숫자가 아니라, tile을 나눠도 numerical stability(수치 안정성)를 유지하는 기준점이다.

tile은 큰 표를 SRAM에 들어가는 작은 직사각형으로 자르는 일이다

Q_i block의 행 수를 B_r, K_j,V_j block의 행 수를 B_c라고 하자. 한 번의 score tile은 S_{ij}=Q_iK_jᵀ이며 shape는 B_r×B_c다.

전체 score S: N × N            → HBM에 전부 쓰지 않는다.
현재 score tile Sᵢⱼ: Bᵣ × B𝚌 → SRAM에 잠깐 만들고 바로 소비한다.

B_r,B_c는 크게 잡을수록 반복 HBM IO가 줄 수 있지만, Q/K/V/score/output/statistics가 동시에 SRAM에 들어가야 한다. 너무 크게 잡으면 kernel이 launch조차 못 하거나 occupancy(동시에 실행 가능한 block 수)가 낮아질 수 있다. 그래서 tile size는 model shape와 GPU architecture에 묶인 tuning parameter다.

m, l, u만 있으면 앞 tile을 버려도 된다

각 query row에 대해 지금까지 본 key tile의 상태를 세 숫자 묶음으로 요약한다.

최종 output은 o=u/l이다. 이 세 상태가 있으면 이미 지나간 score/probability tile을 저장하지 않아도 다음 tile과 정확히 합칠 수 있다. 이 아이디어를 online softmax 또는 algebraic aggregation(대수적 결합)이라고 부른다.

밑바닥 원리와 알고리즘: 새 tile의 softmax를 옛 상태에 정확히 접는 법

두 tile을 한 row로 합치는 식

이미 본 key tile들에 대한 상태가 (m,l,u)라고 하자. 새 tile의 score들을 안정적으로 계산해 (\tilde m,\tilde l,\tilde u)를 얻는다.

[ \tilde m=\max(\text{new scores}),\quad \tilde l=\sum_j e^{s_j-\tilde m},\quad \tilde u=\sum_j e^{s_j-\tilde m}v_j ]

둘의 기준점이 다르므로 먼저 큰 기준점으로 맞춘다.

[ m'=\max(m,\tilde m),\qquad a=e^{m-m'},\qquad b=e^{\tilde m-m'} ]

그 뒤 분모와 분자를 같은 비율로 합친다.

[ l'=a l+b\tilde l,\qquad u'=a u+b\tilde u,\qquad o'=u'/l' ]

처음에는 m=-∞, l=0, u=0이다. 따라서 첫 tile에서 a=0이 되고 새 tile만 들어간다. 다음 tile에서 더 큰 score가 나오면 과거의 l,u에 e^{m-m'}를 곱해 새 기준점으로 낮춘다. 분자와 분모를 똑같이 재스케일하기 때문에 최종 비율은 전체 row softmax와 같다.

12살 숫자 예: 앞 표를 버려도 합계가 안 틀리는 이유

앞 tile의 상태가 m=2, l=1.5, u=6이라고 하자. 따라서 앞 tile만의 output은 u/l=4다. 새 tile의 score가 [3,1], value가 [10,20]이면:

[ \tilde m=3,\quad \tilde l=1+e^{-2}\approx1.1353,\quad \tilde u=1\times10+e^{-2}\times20\approx12.7067 ]

새 기준점은 m'=3이다. 과거 scale은 a=e^{2-3}=e^{-1}\approx0.3679, 새 tile scale은 b=1이다.

[ l'\approx0.3679\times1.5+1.1353=1.6871 ]

[ u'\approx0.3679\times6+12.7067=14.9140,\qquad o'\approx8.840 ]

앞 tile의 원래 score 전체를 다시 보지 않았다. 그러나 m,l,u가 앞 tile의 필요한 정보를 보존했기 때문에, 새 tile과 합친 o'는 네 score를 한 번에 softmax한 결과와 같다.

원 논문 Algorithm 1의 block 실행 순서

원 논문은 K,V block을 outer loop로, Q block을 inner loop로 둔다. 단순화한 흐름은 다음과 같다.

HBM에 Q, K, V, O, m, l이 있다.

for K/V tile j:
  Kⱼ, Vⱼ를 HBM → SRAM으로 한 번 읽는다.

  for Q/output tile i:
    Qᵢ, 현재 Oᵢ, mᵢ, lᵢ를 HBM → SRAM으로 읽는다.
    Sᵢⱼ = τ Qᵢ Kⱼᵀ 를 SRAM에서 만든다.
    mask를 적용하고 tile의 m̃, l̃, ũ를 계산한다.
    위 결합식으로 Oᵢ, mᵢ, lᵢ를 갱신해 HBM에 쓴다.

return O

중요한 경계는 두 개다.

1. Sᵢⱼ는 존재한다. 다만 B_r×B_c만 SRAM에 잠깐 있고 전체 N×N으로 HBM에 저장되지 않는다.
2. 출력 O, 통계 m,l은 tile마다 HBM에 갱신된다. FlashAttention이 HBM IO를 0으로 만드는 것은 아니다. 큰 quadratic intermediate의 read/write를 피하는 것이다.

forward가 끝난 뒤 backward는 무엇을 저장하고 무엇을 다시 계산하는가

standard training은 backward에서 P를 써서 dV=PᵀdO, dP=dOVᵀ, softmax gradient, dQ,dK를 계산한다. P가 N×N이라 저장 비용이 크다.

FlashAttention Appendix B.3–B.4는 forward에서 O, row statistics m,l, 그리고 dropout을 위한 pseudo-random number generator(PRNG, 의사 난수 생성기) state를 저장한다. backward에서는 tile별로 S와 P를 SRAM에서 다시 만든다. causal/padding mask와 dropout mask도 같은 PRNG state에서 재생성해야 forward와 같은 경로가 된다.

저장: Q, K, V, O, m, l, RNG state              → O(Nd) + O(N)
저장하지 않음: 전체 S, P, dropout mask          → O(N²)을 피함
재계산: 현재 Sᵢⱼ, Pᵢⱼ, dropout mask Zᵢⱼ       → SRAM tile에서 즉시 소비

재계산은 FLOP를 조금 더 쓸 수 있다. 그러나 논문의 요지는 “FLOP가 약간 늘어도 HBM 왕복이 크게 줄면 total runtime이 더 짧을 수 있다”다. 이것이 memory–compute trade-off(메모리와 연산의 교환)의 정확한 방향이다.

내부 구조와 실제 실행 흐름: 2022 공식 source는 Python에서 CUDA tile kernel까지 어떻게 내려가는가

조사한 source의 시점과 경계

공식 repository의 전체 history를 확인해, 원 논문 v2와 가까운 commit 765741c1eeb86c96ee71a3291ad6968cfbf4e4a1(2022-06-14, “More explanation”)을 고정했다. 이때 README는 alpha release 0.1로 CUDA 11·NVCC·Turing/Ampere GPU, FP16, head dimension 16/32/64/128을 적었다. 2026년의 repository main은 FlashAttention-2/3/4 계열을 포함하므로, 아래 line은 2022 원 논문 구현의 사실이지 오늘 pip install flash-attn의 지원 범위 선언이 아니다.

호출 경로를 한 장으로 그리면

PyTorch model
  qkv: [batch, sequence, 3, heads, head_dim]
  ↓ FlashAttention.forward
padding이 있으면 unpad, 없으면 [batch×sequence, 3, heads, head_dim]로 reshape
  ↓ flash_attn_func
FlashAttnFun.forward (autograd.Function)
  ↓ flash_attn_cuda.fwd(...)
C++ mha_fwd: CUDA/contiguous/shape/head-size 검사, output·softmax_lse buffer 생성
  ↓ run_fmha_fp16_sm80
CUDA kernel: K/V·Q/O tile을 shared memory에 두고 score, mask, softmax, V 곱을 처리
  ↓
context output + softmax_lse를 Python autograd context에 저장
  ↓ backward
QKV/O/softmax_lse/RNG state로 tile score·probability·dropout mask를 재생성해 dQ,dK,dV 계산

1. Python interface는 scale·RNG state·작은 통계를 autograd에 넘긴다

원본 flash_attn_interface.py를 직접 읽었다.

• FlashAttnFun.forward는 dropout이 있을 때 CUDA RNG state를 저장한다. softmax_scale=None이면 qkv.shape[-1] ** (-0.5), 즉 1/√d를 쓴다.
• 같은 함수는 flash_attn_cuda.fwd에서 context, softmax_lse, 선택적 S_dmask를 받고, qkv, context, S_dmask, softmax_lse, cu_seqlens, rng_state를 ctx.save_for_backward에 보관한다.
• backward는 저장한 RNG state로 CUDA RNG를 되돌린 뒤 extension의 bwd를 호출하고 이전 RNG state를 복구한다. 이것이 dropout mask 재생성의 deterministic(결정적) 계약이다.
• normal path의 flash_attn_func는 return_attn_probs=False로 호출된다. 전체 attention probability를 debugging/testing 때문에 요청하는 별도 class만 return_softmax=True를 선택한다. 확률표 전체를 반환하라고 요청하면 N² 저장 회피라는 이점이 약해진다.

2. module 경계는 모든 PyTorch attention을 받아들이지 않는다

원본 FlashAttention.forward는 다음을 assert한다.

assert not need_weights
assert attn_mask is None
assert qkv.dtype == torch.float16
assert qkv.is_cuda

즉 2022 alpha module은 arbitrary attn_mask, attention weight 반환, CPU, BF16/FP32를 일반 지원하지 않았다. padding mask가 있으면 원본 lines 55–66처럼 유효 token만 unpad하고 cu_seqlens(각 batch sequence의 누적 token offset)와 max_s를 kernel에 준다. padding 0을 계산하지 않는 것은 tile 알고리즘과 별개로 실제 throughput에 중요하다.

3. C++ dispatcher는 N² buffer를 언제 만들지 명시적으로 고른다

원본 mha_fwd는 다음 실행 경로를 확인하게 해 준다.

1. qkv.is_cuda(), cu_seqlens.is_cuda(), contiguous(연속 배치), 1D cu_seqlens, 4D QKV, sizes[THREE_DIM] == 3을 검사한다.
2. head size가 16, 32, 64, 128 중 하나인지 검사하고, 당시 SM75/SM80 device property를 확인한다.
3. ctx output은 {total, num_heads, head_size}로, softmax_lse는 {batch_size, num_heads, seq_len} FP32로 할당한다.
4. return_softmax일 때만 s = {batch_size, num_heads, seq_len, seq_len}을 할당한다(lines 166–181). normal output은 이 N×N tensor를 만들지 않는다.
5. Q/K/V pointer, strides, scale, dropout keep probability, causal flag를 params에 채우고 kernel launcher를 호출한다.

이 source 근거는 “FlashAttention은 score를 전혀 만들지 않는다”가 틀렸음을 다시 보여 준다. debug 반환용 s를 명시적으로 만들 수 있고, normal path만 그 할당을 피한다.

4. CUDA entry는 실행 조건별 template을 고르고 shared memory 크기를 계산한다

원본 fmha_block_fprop_fp16_kernel.sm80.cu는 Is_dropout, Is_causal, Return_softmax template 조합을 선택한다. loop_steps와 softmax statistics tile을 포함해 dynamic shared memory 크기를 정한 뒤, grid(heads, batch)로 kernel을 launch한다. head_dim에 따라 FMHA_kernel_traits<256, 16|32|64, ...>도 바뀐다.

여기서 얻을 실무 교훈은 “FlashAttention”이 하나의 Python 함수 이름이 아니라는 점이다. 실제 성능은 dtype × head_dim × sequence length × causal/dropout × padding × GPU architecture × kernel variant 조합이 결정한다. package version 하나만 기록해서 재현 가능하다고 말하기 어렵다.

직접 검증과 재현: CPU에서 online softmax tile 결과가 standard attention과 같은지 확인했다

확인 범위와 실행 환경

직접 실행 확인: macOS, Python 3.9.6, 표준 library만 사용했다. 이 machine에는 nvidia-smi와 PyTorch가 없었다. 따라서 CUDA extension build, A100 kernel runtime, 원 논문의 GPU wall-clock benchmark는 실행하지 않았다. 대신 다음 두 범위를 직접 검증했다.

1. 원 논문 수식의 online-softmax tile recurrence가 standard attention과 같은 output인지 pure Python으로 계산했다.
2. 2022 공식 source checkout의 flash_attn_interface.py, flash_attention.py를 python3 -m py_compile로 syntax compile했다. 이것은 CUDA runtime 성공 검증이 아니다.

원 논문 PDF SHA-256은 ca7f9fda10b90fc05dd291a3accc85e9c1a4a860b99b31928dab03ed3fcb14e4, 직접 실행 script SHA-256은 9282ea9b45ce76e1db56b4342062dfad16332941d0784a392eb50e9dd40663f9다.

복사 가능한 최소 재현 코드

아래는 GPU 성능 측정 코드가 아니다. N=5, value dimension 2인 고정 Q/K/V로 standard attention과 tile 크기 1·2·3의 online recurrence를 비교하는 수학 검산이다. causal=False와 causal=True 모두 assert한다.

from math import exp, sqrt


def dot(a, b):
    return sum(x * y for x, y in zip(a, b))


def softmax(scores):
    m = max(scores)
    e = [0.0 if x == float("-inf") else exp(x - m) for x in scores]
    z = sum(e)
    return [x / z for x in e]


def standard(q, k, v, causal):
    scale, result = 1 / sqrt(len(q[0])), []
    for i, qi in enumerate(q):
        scores = [scale * dot(qi, kj) if not causal or j <= i else float("-inf")
                  for j, kj in enumerate(k)]
        p = softmax(scores)
        result.append([sum(weight * value[d] for weight, value in zip(p, v))
                       for d in range(len(v[0]))])
    return result


def flash_tiled(q, k, v, block_columns, causal):
    scale, result = 1 / sqrt(len(q[0])), []
    for i, qi in enumerate(q):
        m, l, u = float("-inf"), 0.0, [0.0] * len(v[0])
        for start in range(0, len(k), block_columns):
            end = min(start + block_columns, len(k))
            scores = [scale * dot(qi, k[j]) if not causal or j <= i else float("-inf")
                      for j in range(start, end)]
            tile_m = max(scores)
            new_m = max(m, tile_m)
            old_scale = 0.0 if m == float("-inf") else exp(m - new_m)
            weights = [0.0 if x == float("-inf") else exp(x - new_m) for x in scores]
            l = old_scale * l + sum(weights)
            u = [old_scale * u[d] + sum(w * v[start + offset][d]
                 for offset, w in enumerate(weights)) for d in range(len(v[0]))]
            m = new_m
        result.append([x / l for x in u])
    return result


q = [[3, -2, 1], [0.5, 4, -1], [-3, 1, 2], [2, 2, -4], [1, -5, 3]]
k = [[2, 1, -1], [-1, 3, 2], [4, -2, 0], [0, 1, 5], [-3, 2, 1]]
v = [[1, 0], [0, 2], [3, -1], [-2, 4], [5, 1]]

for causal in (False, True):
    expected = standard(q, k, v, causal)
    for block_columns in (1, 2, 3):
        actual = flash_tiled(q, k, v, block_columns, causal)
        assert all(abs(a - b) <= 1e-12
                   for left, right in zip(expected, actual) for a, b in zip(left, right))

직접 실행 결과의 핵심은 다음과 같았다.

causal=False; block_columns=1, 2, 3 모두 standard output과 오차 1e-12 이내 일치
causal=True;  block_columns=1, 2, 3 모두 standard output과 오차 1e-12 이내 일치

causal=True 첫 행 output=[1.0, 0.0]
→ 첫 token은 mask 때문에 첫 value만 볼 수 있다.

row 4, block_columns=2 상태:
(0, 2, m=-3.464102, l=1.099321)
(2, 4, m= 8.082904, l=1.099331)
(4, 5, m= 8.082904, l=1.099332)
→ 더 큰 score tile을 만나도 기존 l/u를 새 m 기준으로 재스케일해 합친다.

동일 입력에서 materialized score와 probability는 각각 5×5=25 entries다. streamed implementation이 HBM에 오래 보존하는 것은 output 5×2=10 entries와 row statistics m,l 각 5개, 합계 10 entries다. 실제 GPU tile의 SRAM·register 사용량까지 이 CPU script가 측정한 것은 아니다. 검산 범위는 정확한 output을 만드는 recurrence다.

값을 하나 바꾸는 경계 실험과 오류 해석

block_columns를 1, 2, 3으로 바꿔도 assert가 통과했다. tile partition이 달라도 수학적 결과는 같아야 한다. 반대로 stable softmax에서 m을 빼지 않고 큰 score에 exp를 적용하면 overflow 또는 nan이 나올 수 있다. causal mask에서 한 row의 모든 score를 -∞로 만드는 잘못된 mask는 max=-∞, 분모 0으로 이어진다. production kernel 오류를 볼 때는 “FlashAttention이 틀렸다”보다 먼저 다음을 확인한다.

1. 모든 query row에 적어도 하나의 unmasked key가 있는가?
2. Q/K/V shape, stride, head dimension, dtype, device가 kernel contract와 맞는가?
3. train/eval dropout policy와 RNG state 재생성 경로가 일치하는가?
4. causal alignment와 padding/unpadding offset이 맞는가?
5. baseline과 비교할 때 dtype·mask·scale·tolerance를 같게 했는가?

FP16/BF16 GPU reduction은 tile order가 달라 bitwise equality(비트 단위 동일성)를 보장하지 않을 수 있다. 정확성 검증은 같은 model mode와 fixed input에서 allclose tolerance, downstream loss, task metric을 함께 본다. “exact”는 standard attention 수학식과 같은 알고리즘적 결과를 뜻하지, 모든 floating-point implementation이 마지막 bit까지 같다는 뜻은 아니다.

성능, 복잡도와 트레이드오프: FLOP보다 HBM access가 병목이면 더 많은 계산도 빨라질 수 있다

세 종류의 비용을 분리한다

Theorem 2를 “FlashAttention은 항상 asymptotically better”라고 짧게 말하면 틀린다. M이 더 큰 SRAM이거나 d가 작을수록 N²d²/M이 practical하게 작아진다는 것이 논문의 조건부 주장이다. tile을 너무 키우면 SRAM에 안 들어가고, 충분히 큰 뒤에는 arithmetic 또는 다른 kernel overhead가 병목이 된다. 논문 Figure 2 가운데도 A100에서 block size를 키우면 runtime이 감소하다가 다른 요인에 막힘을 보인다.

원 논문 benchmark는 조건을 떼지 않고 읽는다

논문은 benchmark runtime을 A100 40 GB 한 장, 8 heads·head dimension 64·batch 16, dropout 0.1·정해진 masking 조건, attention call 100회 평균으로도 별도 측정했다(Appendix E). speedup은 sequence length, masking, dropout, head dimension, GPU HBM bandwidth, baseline kernel에 따라 바뀐다. 따라서 내 환경에서는 “FlashAttention 3×”를 목표치로 쓰지 말고 같은 input distribution·batch·precision에서 baseline과 프로파일해야 한다.

실패, 한계, 장애와 운영: long context의 모든 비용과 모든 kernel 호환성을 해결하지는 않는다

실패 1: O(N) memory를 O(N) attention compute로 잘못 읽는다

증상: context를 4배 늘리면 compute time이 크게 늘었는데 “FlashAttention인데 왜 느리지?”라고 판단한다.

원인: original FlashAttention은 exact dense attention의 QKᵀ와 PV를 계산하므로 forward FLOP가 O(N²d)다. 피하는 것은 S,P의 HBM materialization과 그에 따른 O(N²) 추가 activation memory다.

대응: memory와 compute budget을 별개로 산정한다. context length별로 prefill latency, attention FLOPs, peak VRAM, batch size, output length를 측정한다. sequence가 너무 길어 compute 자체가 SLA(Service Level Agreement, 서비스 수준 협약)를 넘으면 block-sparse/approximate attention, retrieval, chunking, architecture 변경을 별도로 검토한다.

실패 2: attention map을 반환하면서 memory 절약을 기대한다

증상: need_weights=True 또는 debug 옵션을 켠 run에서 memory가 다시 N²로 증가한다.

원인: 2022 source의 normal forward는 return_softmax=False라서 큰 s를 만들지 않지만, 반환 요청 경로는 {batch, heads, sequence, sequence} tensor를 할당한다. “attention weight를 직접 inspect하고 싶다”는 요구는 “N² tensor를 보관하지 않는다”는 설계와 충돌한다.

대응: production은 summary statistic, sampled row, small sequence debug fixture를 별도로 기록한다. 전 길이 attention map은 짧은 test input에서만 수집한다. observability(관측 가능성)를 위해 전체 P를 항상 저장하는 설계는 long-context memory 목표를 무너뜨린다.

실패 3: mask·dropout·tile order를 가벼운 옵션으로 취급한다

증상: baseline과 output이 다르거나, training loss가 재현되지 않거나, 특정 padded batch만 nan이 된다.

원인: causal mask의 방향, left/right padding offset, 1/√d scale, dropout probability, PRNG state, train/eval mode는 모두 attention 수식의 일부다. backward에서 다른 dropout mask를 만들면 같은 loss graph의 gradient가 아니다. 모든 key가 mask된 row도 softmax 분모 0이라는 정의되지 않은 상태를 만든다.

대응: shape test만 하지 말고 causal/non-causal, unequal sequence length, padding, dropout on/off, extreme score를 포함한 golden test를 둔다. same seed·same input의 baseline과 tolerance 비교를 CI(Continuous Integration, 지속적 통합)에 넣는다. nan은 먼저 mask coverage, scale, dtype range, invalid padding offset을 확인한다.

실패 4: 논문 수식의 M과 GPU의 “총 VRAM”을 혼동한다

증상: 80 GB GPU면 M=80 GB라고 넣어 tile·IO speedup을 추정한다.

원인: Theorem 2의 M은 algorithm이 fast memory로 모델링한 on-chip SRAM capacity다. HBM 용량은 큰 base model과 activation을 담는 공간이고, SRAM은 한 tile이 동시에 들어갈 작은 작업대다. 또한 asymptotic formula의 원소·상수·alignment·register·occupancy를 byte 단위 제품 계산으로 바로 바꾸면 안 된다.

대응: 논문 식은 병목 방향을 판단하는 데 쓰고, 실제 block size와 kernel selection은 framework profiler와 target GPU benchmark로 결정한다. device memory total, allocated/reserved peak, HBM bandwidth, kernel duration, achieved occupancy를 같은 dashboard에서 본다.

실패 5: old source의 지원 범위를 current package의 계약으로 선언한다

증상: CPU·FP32·지원하지 않는 head dimension에서 import/build/forward 실패가 난다. 또는 기존 code가 attn_mask·attention weight 반환을 요구한다.

원인: 이 글이 확인한 commit은 alpha 0.1이다. source에 FP16 CUDA와 제한된 head dimension assert가 있다. 반대로 최신 package는 후속 세대·새 API·다른 hardware를 포함할 수 있다. 논문 코드와 현재 dependency의 version을 섞어 지원 여부를 추정하면 안 된다.

대응: 도입 전에는 현재 runtime의 공식 compatibility matrix와 installed build log를 별도로 확인한다. torch.__version__, CUDA runtime/driver, GPU compute capability, flash-attn wheel/source commit, model dtype/head dimension, target call signature을 artifact로 남긴다. native CUDA extension은 package hash·빌드 image·compiler도 공급망 검토 대상이다.

실패 6: prefill의 attention 최적화를 decode·KV cache 문제의 해답으로 착각한다

증상: 긴 prompt prefill은 빨라졌는데 token-by-token generation throughput이나 concurrency가 기대만큼 늘지 않는다.

원인: prefill은 prompt 전체 Q/K/V로 attention을 계산해 N² score path가 중요하다. decode는 새 query 한두 행이 과거 KV cache를 읽으며 next token을 만드는 다른 access pattern이다. weight read, KV cache capacity/fragmentation, continuous batching, scheduler가 병목일 수 있다. FlashAttention은 KV cache eviction, tenant isolation, paged allocation을 설계하지 않는다.

대응: latency를 request 전체 하나로 뭉치지 말고 tokenizer, queue, prefill, first-token latency, decode token/s, cache hit/miss, p95 VRAM으로 나눈다. FlashAttention 적용 전후에도 이 지표를 따로 비교한다.

운영 runbook: 이상한 output·OOM·느린 kernel을 만났을 때

1. 계약 확인
   model SHA, framework/extension version, CUDA/driver, GPU, dtype, head_dim,
   causal/dropout, sequence length 분포를 request/experiment ID에 묶는다.

2. 정확성 확인
   작은 fixed Q/K/V에서 baseline과 Flash path를 같은 mask·scale·eval mode로 allclose 비교한다.

3. NaN 확인
   all-masked row, padding offset, extreme score, scale, loss-scaling, RNG state를 먼저 본다.

4. 메모리 확인
   attention map 반환 여부, max sequence, batch, heads, activation checkpointing,
   KV cache와 동시 요청 수를 분리해 peak VRAM을 측정한다.

5. 성능 확인
   prefill/decode 분리, kernel timeline, HBM traffic, achieved occupancy, p50/p95를 비교한다.

6. 완화/rollback
   검증한 standard SDPA 또는 이전 extension image로 canary route를 되돌린다.
   정확성 실패 중에는 tool write·업무 side effect가 있는 agent traffic을 보수적으로 차단한다.

마지막 줄은 FlashAttention 논문 자체의 기능이 아니라 AI system 운영 해석이다. backend·DBA 경험자는 model kernel을 믿고 끝내지 않고 version, input contract, metric, profiler trace, rollback 가능한 deployment를 묶어야 한다.

대안, 비교와 선택 기준: “같은 attention을 빠르게”와 “attention 자체를 덜 하게”를 분리한다

| 상황 | 먼저 검토할 선택 | 왜 | FlashAttention만으로 부족한 부분 | | --- | --- | --- | | standard softmax attention 품질을 보존하며 long prefill training memory를 줄이고 싶은가? | exact FlashAttention 계열 | S,P global materialization을 피하고 output 수학을 보존 | N²d compute와 base/activation/KV memory는 남음 | | 짧은 sequence·CPU·unsupported shape에서 correctness를 먼저 보고 싶은가? | framework standard attention baseline | 단순하고 attention map debug가 쉬움 | HBM IO·activation memory가 더 클 수 있음 | | sequence가 매우 길어 exact N² compute도 감당 못 하는가? | block-sparse·linear/low-rank approximate attention | 계산량과 IO를 함께 줄일 수 있음 | 원 attention과 output이 달라질 수 있어 task 품질 검증 필요 | | 최신 외부 문서·정책을 읽어야 하는가? | RAG(Retrieval-Augmented Generation, 검색 증강 생성)·DB/API | context를 무한히 키우기보다 필요한 근거만 선별 | retrieval recall·권한·freshness·citation을 별도 설계 | | 생성 decode가 느리거나 동시 요청이 많아지는가? | KV cache 관리·continuous batching·scheduler profiling | decode 병목은 prefill attention과 다를 수 있음 | cache allocation·isolation·eviction은 FlashAttention의 범위 밖 | | 오늘의 package API를 도입하려는가? | 현재 framework가 선택한 optimized SDPA/Flash backend 검증 | post-2022 kernel·hardware contract가 다름 | 논문 Table 숫자만으로 compatibility·speedup을 보장 못 함 |

원 논문 §3.3의 block-sparse FlashAttention은 approximate attention의 한 예다. 사전에 정한 block mask에서 zero block을 건너뛰어 IO를 sparsity ratio에 비례해 더 줄인다. 하지만 원 FlashAttention은 exact, block-sparse extension은 approximate라는 층위를 섞지 말아야 한다.

내 데이터+AI 시스템 학습에서의 적용 순서

FlashAttention을 “CUDA를 잘해야 쓸 수 있는 논문”으로 멀리 둘 필요는 없다. 데이터·AI pipeline을 만들 때 다음 질문 순서로 연결하면 된다.

1. 문제 크기 측정
   실제 prompt/token 길이 histogram, prefill/decode 비율, p95 latency, OOM 사례를 수집한다.

2. 정확성 baseline
   작은 fixed input에서 standard attention과 optimized path의 output/loss/gradient를 비교한다.

3. 비용 실험
   같은 model·dtype·batch·mask에서 VRAM peak, HBM traffic, kernel time, tokens/s를 기록한다.

4. 운영 release
   model SHA + tokenizer + framework + CUDA image + extension build + kernel policy를 묶는다.
   canary metric, fallback backend, rollback 기준을 저장한다.

이 흐름은 AI career에서 큰 무기가 된다. 논문 수식을 읽을 줄 아는 것과, 실제 workload가 memory-bound인지 측정하고, data length 분포를 파악하고, native dependency를 재현 가능하게 배포하는 것은 다르다. backend·DBA 출신자는 이 사이의 계약·관측·복구를 설계할 수 있다.

흔한 오해와 최초 질문의 답: FlashAttention을 어디까지 이해하고 언제 써야 하는가

오해 1. “FlashAttention은 N² attention을 N으로 바꾼다.”

아니다. 원 논문 forward FLOP는 O(N²d)다. S,P를 HBM에 저장하는 추가 memory가 O(N²)에서 O(N)으로 줄고, HBM access schedule이 바뀌는 것이다. 근사/희소 attention과 같은 말이 아니다.

오해 2. “score matrix는 아예 계산하지 않는다.”

아니다. S_ij=Q_iK_jᵀ score tile은 SRAM에서 계산한다. 전체 S[N×N]를 HBM에 materialize하지 않을 뿐이다. local tile이 없다면 softmax weight도 만들 수 없다.

오해 3. “FLOP가 더 많으면 반드시 더 느리다.”

아니다. paper Figure 2의 특정 A100 setup에서 FlashAttention은 standard attention보다 75.2 대 66.6 GFLOPs로 더 많았지만, HBM R/W가 4.4 대 40.3 GB로 줄어 runtime은 7.3 대 41.7 ms였다. memory-bound operation에서는 데이터 이동이 더 큰 비용일 수 있다.

오해 4. “모든 GPU·dtype·head size에서 그대로 실행된다.”

아니다. 원 논문 시점 source의 alpha contract는 FP16 CUDA와 제한된 head size다. 후속 버전의 지원 범위는 별도로 확인해야 한다. package 이름과 논문 이름이 같아도 source commit이 다르면 기능 계약도 다르다.

오해 5. “FlashAttention을 켜면 context length 제한이 사라진다.”

아니다. model weights, Q/K/V/O, MLP activation, optimizer state, batch, KV cache, quadratic compute가 남는다. FlashAttention은 큰 attention intermediate를 없애 더 긴 context를 가능하게 할 수 있지만, 무한한 context와 무료 latency를 주지 않는다.

오해 6. “attention map을 항상 저장해 설명 가능성을 얻어도 memory 이점은 같다.”

아니다. N×N attention map을 반환·저장하면 원래 피하려던 allocation이 다시 생긴다. 큰 production trace는 sampled diagnostics와 aggregate metric으로 설계해야 한다.

처음의 질문: 이 논문을 지금 어떻게 읽을까?

top-down 학습에서는 “긴 context가 왜 OOM·느려짐을 만드는가?”라는 system 문제부터 시작한다. 그 다음 QKᵀ→softmax→PV를 복습하고, m,l,u online recurrence까지 bottom-up으로 내려간다. 권장 순서는 다음이다.

20분: Abstract → Figure 1 → §2.2 standard attention → Algorithm 1 → Theorem 1·2
40분: §3.1의 stable softmax 결합식 → Appendix B.3/B.4 forward/backward
40분: 2022 Python interface → C++ mha_fwd → CUDA kernel entry → 작은 CPU recurrence 검산
프로젝트: 내 서비스의 length histogram → baseline/optimized profile → versioned release/rollback

체크리스트를 통과하면 “FlashAttention이 빠르다”가 아니라 “왜 이 workload에서 HBM IO가 병목이고, exact output을 어떻게 유지하며, 어떤 계약과 측정으로 안전히 도입할지”를 설명할 수 있다.

• Q,K,V,S,P,O의 shape를 N,d로 손으로 적고 N² intermediate를 계산한다.
• stable softmax에서 max를 빼는 이유와 m,l,u tile 결합식을 설명한다.
• “additional O(N) memory”와 “O(N²d) FLOPs”를 한 문장에 함께 설명한다.
• original source의 normal path가 N×N s를 만들지 않는 조건과 debug 반환 경계를 확인한다.
• causal·padding·dropout·RNG state를 포함해 baseline과 allclose test를 만든다.
• prefill/decode, p50/p95, peak VRAM, kernel time, extension version을 분리해 관측한다.
• GPU·CUDA·framework·extension image와 fallback backend를 포함한 rollback release를 만든다.

다음 논문 Scaling LLM Test-Time Compute는 “model weight나 attention kernel을 바꾸지 않고, 답을 만들 때 더 많은 계산을 쓰면 어떤 품질·비용 교환이 생기는가?”를 시스템·추론 전략 층에서 잇는다.

출처

1차 자료

1. Dao, Fu, Ermon, Rudra, Ré, FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness, arXiv:2205.14135v2, 2022-06-24. 이 글의 standard/Flash algorithm, Theorem 1·2·5, Algorithm 0–4, block-sparse extension, A100 memory hierarchy와 Table 1–4·Figure 2 benchmark의 원문이다. §2–4, Appendix B–C·E를 직접 대조했다.
2. Dao-AILab, 공식 flash-attention repository, commit 765741c1eeb86c96ee71a3291ad6968cfbf4e4a1, 2022-06-14, 확인일 2026-07-24. Python autograd interface, module input contract, C++ mha_fwd, CUDA forward entry를 직접 읽었다.

직접 검증 기록

1. /tmp/flashattention-paper-research/verify_flashattention_online_softmax.py를 Python 3.9.6으로 직접 실행했다. stable standard attention과 tiled online recurrence를 causal/non-causal 각각 block size 1·2·3으로 비교했고 모든 원소가 1e-12 이내에 일치했다. script SHA-256: 9282ea9b45ce76e1db56b4342062dfad16332941d0784a392eb50e9dd40663f9.
2. 2022 고정 source의 flash_attn_interface.py, flash_attention.py를 python3 -m py_compile로 compile했다. 이 machine에는 PyTorch·NVIDIA driver/GPU가 없어 CUDA extension build, GPU forward/backward, 원 논문 A100 benchmark는 실행하지 않았다. compile 성공은 GPU runtime 검증이 아니며, 그 한계를 본문에 명시했다.

보조 자료

이 글의 핵심 동작·수치에는 보조 기술 블로그를 사용하지 않았다. 12살 비유, 운영 runbook, data/AI pipeline 선택 기준은 원 논문·공식 source를 바탕으로 한 이 글의 해석이며, 논문 benchmark 사실과 분리해 표기했다.