Flash Attention

표준 attention 의 수식 — O = softmax(QKT/√d) · V — 은 수학적으로 깔끔하다. 하지만 그 수식을 그대로 구현하면 중간 산출물 S = QKT 가 N×N 크기의 행렬이 되어 HBM 을 통째로 들었다 놨다 한다. seq_len 4K 에서 16M 원소, 32K 에서는 1G 원소.

강의가 풀려는 질문은 두 개다.

1. 같은 수학적 결과를 어떻게 N×N 행렬을 만들지 않고 계산하는가 — online softmax 와 tile streaming 의 조합.
2. 그 알고리즘을 CUDA 로 짜면 register / shared memory / Tensor Core 가 어떻게 분배되는가 — Thomas 가 직접 짠 flash_attention.cu (d=128, B_r=8, B_c=32) 의 해부.

표준 attention 의 비용을 메모리 시각으로 분해 — 같은 dimension d=128, seq_len N=4K 의 한 head 에 대해.

FlashAttention 의 가장 영리한 부분 — softmax 가 “전체 row 를 봐야 normalize 가능” 한 op 인데, tile 단위로 보면서 동시에 정확한 normalizer 를 누산해 나간다. 핵심 도구는 “running max” 와 그에 따른 보정 항.

Numerical stability 를 위해 softmax 는 softmax(x_i) = exp(x_i − max(x)) / Σ exp(x_j − max(x)) 로 계산한다. 문제는 — tile 별로 보면 그 tile 안의 max 만 알고, 다음 tile 에서 더 큰 값이 나오면 이미 누산해 둔 값이 틀린 normalizer 로 계산된 것이 된다. 그래서 max 가 갱신될 때마다 이전 누산을 “보정”한다.

Thomas 의 코드 (flash_attention.cu) 의 핵심 구조 — d=128, B_r=8 (Q tile rows), B_c=32 (K/V tile cols). 이 숫자들이 어떻게 정해지고 어떤 의미인지.

// shared memory: SRAM 에 올라가는 것들
__shared__ float Q_i[B_r][d];      // 8×128 = 4 KB
__shared__ float K_j[B_c][d];      // 32×128 = 16 KB
__shared__ float V_j[B_c][d];      // 32×128 = 16 KB
__shared__ float S[B_r][B_c];      // 8×32   = 1 KB

// register: thread 별 누산 상태
float l_i[B_r_over_bdy];        // running sum
float m_i[B_r_over_bdy];        // running max
float O_i[B_r_over_bdy][d_over_bdx];  // running output

// outer loop: Q tile 한 번 (per block)
for (int i = bix; i < bix + 1; i++) {
    // Q_i 를 한 번 shmem 으로
    load_Q(Q_i, Q + i * B_r);

    // inner loop: K/V tiles
    for (int j = 0; j < T_c; j++) {
        load_K_V(K_j, V_j, j);
        __syncthreads();

        // 1) S = scaling * Q @ K^T
        compute_S(S, Q_i, K_j, scaling);
        __syncthreads();

        // 2) online softmax 갱신 + O 누산
        update_running_state(m_i, l_i, O_i, S, V_j);
    }

    // 3) 최종 O = O_i / l_i, store l_i 도 (backward 위해)
    write_output(out, out_l, O_i, l_i, m_i, i);
}

Thomas 가 강의에서 가장 중요하게 보여준 진단 자리 — “register spilling from registers”. running 누산기 (m_i, l_i, O_i) 가 register 에 들어가는 게 FA 의 메모리 모델의 전제다. 만약 thread 당 register 가 부족하면 그 누산기들이 local memory (실제로는 HBM-backed) 로 spill 되고 — 이때 모든 IO advantage 가 사라진다.

FA 의 forward 가 N×N score 행렬을 안 만들고 했다면 — backward 도 같은 제약 안에서 풀려야 한다. 하지만 backward 에는 forward 의 attention probabilities P 가 다시 필요하다. 어떻게 푸는가.

// forward 에서 저장하는 한 줄
out_l[i] = m_i + logf(l_i);   // = logsumexp_i

// backward 에서 P 복원 (tile 단위)
float s_ij = scaling * Q_i[ii] @ K_j[jj];
float P_ij = expf(s_ij - out_l[i]);   // 정확

// 그 P_ij 로 dQ, dK, dV 누산
dV_j   += P_ij * dO_i;
dP_ij   = dO_i @ V_j;
dS_ij   = P_ij * (dP_ij - /* row reduction */);
dQ_i   += scaling * dS_ij @ K_j;
dK_j   += scaling * dS_ij^T @ Q_i;

FA 가 학계에 준 가장 큰 영향 중 하나 — algorithm 의 비용을 측정할 때 FLOPs 만 보는 게 아니라 “HBM ↔ SRAM 트래픽”도 별도 metric 으로 본다 는 시각.

FA 논문이 두 차례 나왔다. v1 (Dao 2022) 은 위에 깐 알고리즘. v2 (Dao 2023) 는 같은 알고리즘이지만 — “누가 어떤 일을 하느냐” 의 work partition 을 다시 짠다.

강의가 끝나면서 Thomas 가 짚는 미래 방향 — FlashAttention 의 알고리즘 패턴이 NVIDIA 의 다음 hardware 디자인을 끌었다. H100 의 새 instruction 들이 사실상 “FA 를 빠르게 돌리려면 무엇이 필요한가” 의 답으로 설계되어 있다.

attention 위에서 다시 일을 하게 됐을 때 — 이 핵심들이 빠르게 복원되어야 한다.

손에 새기기 — 실습 시퀀스

1. online softmax 한 페이지 짜기 — Python/numpy 로 chunk 단위 softmax 가 전체 softmax 와 같은지 verify. § 03 의 누산 표를 코드로 그대로.
2. flash_attention.cu 컴파일 + 실행 — Thomas 의 코드를 nvcc -Xptxas -v 와 함께 빌드. register 사용량 출력 확인. 작은 N 에서 표준 attention 과 결과 일치 확인.
3. spilling 비교 — flash_attention_spilling_from_registers.cu 와 정상 버전을 같은 입력으로 실행. NCU 로 “Local Memory Overhead” 비교.
4. B_r, B_c sweep — B_r ∈ {4, 8, 16}, B_c ∈ {16, 32, 64} 의 9 조합으로 시간 측정. spill 시작점과 SRAM 한계점 시각화.
5. F.scaled_dot_product_attention vs eager — sdp_kernel(enable_flash=True/False) 컨텍스트로 같은 입력의 forward+backward 시간 비교. 모델 크기·seq_len 별 cross-over 점 찾기.
6. backward 의 logsumexp 활용 — forward 에서 저장한 logsumexp 한 줄로 P 를 복원하는 mini-script. 정확한 backward 결과와 일치 검증.
7. IO complexity 측정 — NCU 의 “DRAM Read/Write Bytes” 메트릭으로 표준 vs FA 의 실측 트래픽 비교. 이론값과의 차이 분석.

FA 의 메커니즘이 이후 강의들의 핵심 패턴을 결정한다.

FA 는 빠르게 진화하는 자리라 — 강의 시점 (2024 March) 이후의 변화가 크다.

• FA-3 (Hopper-aware) 의 디테일 — TMA + WGMMA 가 어떻게 inner loop 와 맞물리는지. 강의는 v1/v2 까지. FA-3 논문 / repo 별도 추적 필요.
• FP8 attention 의 정확도 — § 09 에서 “정확도 손실 거의 없음” 으로 적었지만, 어떤 모델 / 어떤 길이에서 깨지는지. 실측 데이터 필요.
• strong-scaling 한계 — seq_len 이 100K 이상 갈 때 FA 의 한 GPU 한계는 어디인가. ring attention (L013) 으로 넘어가는 분기점.
• IO complexity 측정의 정확한 방법 — § 07 의 ~50 MB / ~12 MB 는 paraphrase. 실제 NCU 의 DRAM bytes 와 다를 수 있다. 자기 GPU 에서 직접 재기 필요.
• causal mask 의 효과 — autoregressive LM 에서 causal mask 가 적용되면 사실상 FA 의 절반만 일하면 된다. 그 절감이 algorithm 안에 어떻게 들어가는지 — 강의는 dense 에 집중.
• backward 의 numerical stability — recompute 가 forward 와 정확히 같은 P 를 복원하는지. 어떤 dtype 에서 깨지는지.
• multi-query / grouped-query attention — Q head 가 여러 개, K/V head 가 1개인 경우. inference 의 표준 패턴인데 — 이 때 FA 의 work partition 이 어떻게 변하는가.