Flash Attention

• Flash Attention은 GPU 메모리 계층 구조를 활용한 IO-aware 정확 어텐션 알고리즘 (by Tri Dao, NeurIPS 2022)
• N×N 어텐션 행렬을 HBM에 실체화하지 않고, 타일링(tiling)으로 SRAM에서 블록 단위 계산을 수행
• 근사가 아닌 정확한(exact) 어텐션이면서 최대 7.6배 빠르고 메모리 사용량을 O(N)으로 감소

해당 개념이 필요한 이유

• 표준 어텐션은 시퀀스 길이에 대해 O(N²) 시간/메모리 복잡도
• 근사 어텐션(Reformer, Performer 등)은 FLOP는 줄이지만 실제 wall-clock 속도 향상 미미
• 핵심 병목은 연산량이 아니라 메모리 접근(IO) — GPU HBM ↔ SRAM 간 데이터 이동

AS-IS: 표준 어텐션

sequenceDiagram
    autonumber
    participant HBM as GPU HBM
    participant SRAM as GPU SRAM

    HBM->>SRAM: Q, K 전체 로드
    SRAM->>SRAM: S = Q × K^T (N×N 행렬 계산)
    SRAM->>HBM: S를 HBM에 저장 (N×N)
    HBM->>SRAM: S 다시 로드
    SRAM->>SRAM: P = softmax(S)
    SRAM->>HBM: P를 HBM에 저장 (N×N)
    HBM->>SRAM: P, V 로드
    SRAM->>SRAM: O = P × V
    SRAM->>HBM: O 저장
    Note over HBM,SRAM: N×N 중간 행렬이 HBM을 왕복 → IO 병목

TO-BE: Flash Attention

sequenceDiagram
    autonumber
    participant HBM as GPU HBM (느림, 대용량)
    participant SRAM as GPU SRAM (빠름, 소용량)
    participant Compute as 연산 유닛

    HBM->>SRAM: K, V 블록 로드
    loop Q 블록 순회
        HBM->>SRAM: Q 블록 로드
        SRAM->>Compute: 블록 단위 어텐션 계산
        Compute->>SRAM: 부분 결과 + 정규화 팩터
    end
    SRAM->>HBM: 최종 결과 기록
    Note over HBM,Compute: N×N 행렬이 HBM에 실체화되지 않음

타일링(Tiling) 핵심 메커니즘

입력 Q, K, V를 블록으로 분할한 뒤, 느린 HBM에서 빠른 SRAM으로 블록 단위로 로드한다. 각 블록에서 어텐션 출력을 계산하고, 적절한 정규화 팩터로 스케일링하여 합산하면 정확한 결과를 얻는다.

타일링 덕분에 recomputation으로 인해 FLOP가 증가하더라도, HBM 접근이 대폭 줄어들어 실제 속도는 더 빠르다.

IO 복잡도 비교

방식	HBM 접근 횟수
표준 어텐션	Θ(Nd + N²)
Flash Attention	O(N²d²M⁻¹)

d: head 차원 (64-128), M: SRAM 크기 (100KB). 일반적으로 d² ≪ M이므로 수9배 적은 메모리 접근.

핵심 인사이트: FLOP ≠ 속도

근사 어텐션은 FLOP 절감에 집중했지만, FLOP 절감이 반드시 wall-clock 속도 향상으로 이어지지 않는다. Flash Attention은 오히려 recomputation으로 FLOP를 증가시키면서도, IO를 줄여 실제 속도를 향상시켰다. 이는 하드웨어 인식 알고리즘 설계의 중요성을 보여주는 대표 사례다.

참고 문서

• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (arXiv)
• FlashAttention — NeurIPS 2022 Paper
• What is Flash Attention? — Biao’s Blog
• Designing Hardware-Aware Algorithms — DigitalOcean