8.1 FlashAttention-1 - IO 인지를 통한 정확한 어텐션의 가속

8.1 FlashAttention-1 - IO 인지를 통한 정확한 어텐션의 가속

2025-12-21, G30DR

1. 서론: 트랜스포머의 확장과 메모리 장벽의 도래

2017년 구글이 ‘Attention Is All You Need’ 논문을 통해 트랜스포머(Transformer) 아키텍처를 세상에 내놓은 이래, 인공지능 분야는 거대한 패러다임의 전환을 맞이했다. 자연어 처리(NLP)에서 시작된 이 혁명은 컴퓨터 비전(Computer Vision), 생물학(Biology), 강화 학습(Reinforcement Learning) 등 데이터의 시퀀스를 다루는 모든 영역으로 확장되었다. 모델의 파라미터 수는 수억 개를 넘어 수천억, 조 단위로 치솟았고, 이에 비례하여 모델이 처리해야 할 입력 데이터의 길이, 즉 문맥(Context)의 크기 또한 기하급수적으로 증가했다. 바야흐로 ’트랜스포머 싱귤래리티(Transformer Singularity)’가 도래한 것이다.

하지만 이러한 확장의 이면에는 언제나 하드웨어라는 물리적 제약이 존재했다. 트랜스포머의 핵심인 셀프 어텐션(Self-Attention) 메커니즘은 입력 시퀀스 길이 N에 대해 시간 복잡도와 공간 복잡도가 모두 O(N^2)로 증가하는 2차 함수적(Quadratic) 특성을 가진다.1 이는 짧은 문장을 처리할 때는 큰 문제가 되지 않지만, 시퀀스 길이가 수천, 수만 토큰으로 늘어나면 연산량과 메모리 사용량이 감당할 수 없을 정도로 폭증하게 됨을 의미한다.

이러한 문제를 해결하기 위해 수많은 연구가 진행되었다. Reformer, Linformer, Performer와 같은 ‘효율적 어텐션(Efficient Attention)’ 기법들은 어텐션 행렬을 희소(Sparse)하게 만들거나 저랭크(Low-rank) 근사를 통해 복잡도를 낮추려 시도했다.3 그러나 이러한 접근법은 이론적인 연산량(FLOPs)을 줄이는 데는 성공했을지 몰라도, 실제 하드웨어 상에서의 실행 시간(Wall-clock time)을 단축하는 데는 실패하는 경우가 많았다. 심지어 근사 과정에서 모델의 정확도가 떨어지는 치명적인 단점을 안고 있었다.1

2022년 스탠포드 대학의 Tri Dao 연구진이 제안한 FlashAttention은 기존의 접근 방식을 완전히 뒤집는 혁신이었다. 그들은 문제의 본질이 연산량(FLOPs)에 있는 것이 아니라, 메모리 접근(Memory Access), 즉 **IO(Input/Output)**에 있음을 간파했다.4 현대의 GPU는 연산 속도가 메모리 대역폭보다 훨씬 빠르게 발전해 왔기 때문에, 대부분의 어텐션 연산은 데이터를 메모리에서 읽어오는 시간을 기다리느라 멈춰 있는 ‘메모리 병목(Memory Bound)’ 상태에 빠져 있었던 것이다.

FlashAttention은 ’IO 인지(IO-Awareness)’라는 새로운 설계 철학을 도입하여, GPU의 메모리 계층 구조를 효율적으로 활용하는 정확한(Exact) 어텐션 알고리즘을 제시했다.6 이를 통해 어텐션 연산의 메모리 복잡도를 O(N^2)에서 선형 복잡도인 O(N)으로 획기적으로 낮추었으며, 근사를 통한 정확도 손실 없이 기존 베이스라인 대비 2~4배의 속도 향상을 달성했다.7 본 장에서는 트랜스포머 아키텍처의 효율화 역사에서 가장 중요한 분기점이 된 FlashAttention-1의 원리와 구현, 그리고 그 파급 효과를 심층적으로 분석한다.

graph TD
    A["트랜스포머(Transformer) 등장"] --> B["모델 파라미터 증가 (수천억 개)"]
    A --> C["입력 시퀀스(Context) 길이 증가"]

    B --> D{"하드웨어 제약 발생"}
    C --> D

    D --> E["셀프 어텐션의 2차 함수적 복잡도 O(N^2)"]

    E --> F["연산량(FLOPs) 폭증"]
    E --> G["메모리 사용량 폭증"]

    F --> H["기존 해결책: 효율적 어텐션 (희소/저랭크)"]
    G --> H

    H -- "문제점" --> I["정확도 하락 & 실제 속도 개선 미비"]

    I --> J["FlashAttention 등장: IO 인지(IO-Awareness)"]

2. 하드웨어 배경과 문제의 재정의

FlashAttention의 혁신성을 온전히 이해하기 위해서는 먼저 딥러닝 모델이 구동되는 하드웨어, 즉 GPU 가속기의 구조적 특성을 깊이 파악해야 한다. 소프트웨어 알고리즘은 하드웨어 위에서 실행되며, 하드웨어의 특성을 무시한 최적화는 공허할 뿐이다.

2.1 GPU 메모리 계층 구조의 이해

현대의 GPU(Graphics Processing Unit)는 대규모 병렬 연산을 처리하기 위해 특화된 장치이다. NVIDIA의 A100이나 H100과 같은 최신 데이터센터용 GPU는 크게 두 가지의 메모리 영역을 가지고 있다.8

  1. 고대역폭 메모리 (HBM, High Bandwidth Memory):
  • 정의: GPU 보드에 장착된 메인 메모리다. 일반적으로 ’GPU 메모리’라고 부르는 40GB, 80GB의 용량이 바로 HBM이다.
  • 특성: 용량은 크지만 속도는 상대적으로 느리다. A100 GPU 기준으로 대역폭은 약 1.5~2.0 TB/s 수준이다.8
  • 역할: 모델의 파라미터, 입력 데이터, 그리고 연산 중간에 생성되는 활성화(Activation) 값들이 모두 이곳에 저장된다.
  1. 온칩 SRAM (On-chip SRAM):
  • 정의: GPU 연산 유닛인 스트리밍 멀티프로세서(SM, Streaming Multiprocessor) 내부에 위치한 초고속 메모리다. L1 캐시(Cache) 또는 공유 메모리(Shared Memory)라고도 불린다.
  • 특성: 속도는 매우 빠르지만 용량은 극도로 작다. A100의 경우 SM당 192KB에 불과하며, 전체 108개 SM을 합쳐도 수십 MB 수준이다. 대역폭은 19 TB/s 이상으로 HBM보다 10배 이상 빠르다.8
  • 역할: 연산 코어(Tensor Core 등)가 데이터를 처리하기 직전에 데이터를 임시 보관하는 역할을 한다.
메모리 종류위치용량 (A100 기준)대역폭 (Bandwidth)비유
SRAM칩 내부 (On-chip)192 KB / SM (총 ~20MB)~19 TB/s책상 위 (작지만 바로 작업 가능)
HBM칩 외부 (Off-chip)40GB ~ 80GB1.5 ~ 2.0 TB/s책장 (크지만 가져오는데 시간 걸림)

표 8.1.1 GPU 메모리 계층 비교: 속도와 용량의 트레이드오프 8

하드웨어 배경: GPU 메모리 계층 구조

graph TD
    subgraph "GPU 칩 외부 (Off-chip)"
        HBM[("HBM (고대역폭 메모리)\n용량: 40~80GB\n속도: 1.5~2.0 TB/s")]
    end

    subgraph "GPU 칩 내부 (On-chip)"
        SRAM[("SRAM (L1 캐시/공유 메모리)\n용량: ~20MB (총합)\n속도: ~19 TB/s")]
        Cores["GPU 연산 코어 (Tensor Cores)"]
    end

    HBM -- "데이터 로드 (느림/병목 구간)" --> SRAM
    SRAM -- "데이터 공급 (초고속)" --> Cores
    Cores -- "연산 결과" --> SRAM
    SRAM -- "결과 저장 (느림/병목 구간)" --> HBM

    style HBM fill:#f9f,stroke:#333,stroke-width:2px
    style SRAM fill:#bbf,stroke:#333,stroke-width:2px
    style Cores fill:#ff9,stroke:#333,stroke-width:2px

2.2 메모리 병목 현상 (Memory Bottleneck)

딥러닝 연산의 성능을 결정짓는 요소는 크게 두 가지다. 하나는 연산 자체의 속도(Compute)이고, 다른 하나는 메모리 접근 속도(Memory Access)다.

  • 컴퓨트 바운드 (Compute Bound): 행렬 곱셈(Matrix Multiplication)과 같이 연산량이 매우 많고 메모리 접근은 상대적으로 적은 작업. 이 경우 GPU의 FLOPS(초당 부동소수점 연산 횟수) 성능이 전체 속도를 결정한다.
  • 메모리 바운드 (Memory Bound): 요소별 연산(Element-wise operation)인 Dropout, Softmax, Masking 등과 같이 연산량은 적은데 데이터의 모든 요소를 읽고 써야 하는 작업. 이 경우 GPU의 메모리 대역폭이 전체 속도를 결정한다.4

트랜스포머의 어텐션 메커니즘은 전통적으로 이 두 가지 특성이 혼재되어 있으나, 시퀀스가 길어질수록 중간 결과물의 크기가 커지면서 심각한 메모리 바운드 상태에 빠진다. 특히 표준 어텐션 구현에서는 N \times N 크기의 거대한 어텐션 점수 행렬을 HBM에 썼다가 다시 읽어들이는 과정이 반복되면서, 연산 코어는 데이터가 도착하기를 기다리며 유휴(Idle) 상태에 머무르게 된다. 이것이 바로 하드웨어 성능을 100% 활용하지 못하게 만드는 ’메모리 장벽’의 실체다.

2.3 표준 어텐션의 비효율성 분석

표준 어텐션(Standard Attention)의 수식은 다음과 같다.
\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right)V
PyTorch나 TensorFlow와 같은 딥러닝 프레임워크에서 이 수식은 일련의 연산 커널(Kernel)들로 나뉘어 실행된다. 이 과정에서 발생하는 HBM 접근 패턴을 분석해보면 비효율성이 명확히 드러난다.9

  1. MatMul 1 (S = QK^T): 쿼리(Q)와 키(K)를 HBM에서 읽어와 행렬 곱을 수행하고, N \times N 크기의 점수 행렬 S를 HBM에 **저장(Write)**한다.
  2. Masking: HBM에서 S를 읽어와(Read) 마스킹 연산을 수행하고 다시 HBM에 저장한다.
  3. Softmax (P = \text{softmax}(S)): HBM에서 S를 읽어와 소프트맥스를 계산하고, N \times N 크기의 확률 행렬 P를 HBM에 저장한다.
  4. Dropout: (훈련 시) HBM에서 P를 읽어와 드롭아웃을 적용하고 다시 저장한다.
  5. MatMul 2 (O = PV): HBM에서 P와 밸류(V)를 읽어와 행렬 곱을 수행하고 최종 결과 O를 HBM에 저장한다.

이 과정에서 N \times N 크기의 행렬 S와 P가 HBM에 **구체화(Materialization)**된다는 점이 치명적이다.5 시퀀스 길이 N이 1,000일 때는 10^6개 요소지만, N이 10,000이 되면 10^8개, 즉 1억 개의 요소가 된다. 이를 HBM에 쓰고 읽는 비용은 N^2에 비례하여 증가하며, 이는 전체 연산 시간의 대부분을 차지하게 된다.

표준 어텐션(Standard Attention)의 IO 흐름

graph TD
    Input[("입력 Q, K, V (HBM)")]

    op1["MatMul 1 (Q x K^T)"]
    S_mat[("S 행렬 (N x N) HBM 저장")]
    op2["Masking 연산"]
    S_masked[("Masked S (HBM 저장)")]
    op3["Softmax 연산"]
    P_mat[("P 행렬 (N x N) HBM 저장")]
    op4["Dropout 연산"]
    P_drop[("Dropped P (HBM 저장)")]
    op5["MatMul 2 (P x V)"]
    Output[("최종 결과 O (HBM 저장)")]

    Input --> op1
    op1 -- "쓰기(Write)" --> S_mat
    S_mat -- "읽기(Read)" --> op2
    op2 -- "쓰기(Write)" --> S_masked
    S_masked -- "읽기(Read)" --> op3
    op3 -- "쓰기(Write)" --> P_mat
    P_mat -- "읽기(Read)" --> op4
    op4 -- "쓰기(Write)" --> P_drop
    P_drop -- "읽기(Read)" --> op5
    op5 --> Output

    style S_mat fill:#f99,stroke:#333,stroke-width:2px,stroke-dasharray: 5 5
    style P_mat fill:#f99,stroke:#333,stroke-width:2px,stroke-dasharray: 5 5
    style P_drop fill:#f99,stroke:#333,stroke-width:2px,stroke-dasharray: 5 5

3. FlashAttention의 핵심: IO 인지(IO-Awareness)

FlashAttention의 저자들은 기존의 연구들이 간과했던 ’IO 비용’에 주목했다. 그들은 연산 횟수(FLOPs)를 줄이는 대신, HBM 접근 횟수를 줄이는 것이 실제 속도 향상의 열쇠라고 판단했다.1 이를 위해 알고리즘을 하드웨어 친화적으로 재설계하는 ‘IO 인지’ 접근법을 제안했다.

FlashAttention의 핵심 전략은 다음과 같이 요약할 수 있다.

  • HBM 접근 최소화: N \times N 크기의 중간 행렬(S, P)을 HBM에 절대 저장하지 않는다.
  • SRAM 활용 극대화: 빠른 SRAM 내부에서 블록 단위로 연산을 수행하고, 중간 결과들을 융합(Fusion)한다.
  • 재연산(Recomputation): 역전파 시 필요한 데이터를 저장하는 대신, 빠른 연산력을 이용해 다시 계산한다.

이러한 전략을 구현하기 위해 FlashAttention은 **타일링(Tiling)**과 **재연산(Recomputation)**이라는 두 가지 핵심 기술을 도입했다.

3.1 기술적 도약 1: 타일링(Tiling)과 커널 퓨전

타일링은 큰 행렬을 작은 블록(Tile)으로 쪼개어 처리하는 기법으로, 컴퓨터 구조론에서는 오래된 최적화 방식이다. FlashAttention은 이를 어텐션 메커니즘에 적용하여, 전체 Q, K, V 행렬을 한 번에 처리하는 대신, SRAM 용량(예: 100KB)에 맞는 작은 블록 단위로 나누어 로드한다.11

3.1.1 타일링의 작동 원리

  1. 입력 행렬 Q, K, V를 블록 크기 B_r \times d, B_c \times d로 분할한다.
  2. HBM에서 K, V의 블록을 하나씩 SRAM으로 로드한다.
  3. Q의 블록을 SRAM으로 로드하여, SRAM 내부에서 QK^T를 계산한다.
  4. 계산된 부분 어텐션 점수에 대해 소프트맥스를 적용하고, 바로 V 블록과 곱한다.
  5. 이 모든 과정이 SRAM 내부에서 이루어지므로, 중간 결과인 S나 P 행렬을 HBM에 기록할 필요가 없다.

이 방식을 통해 기존에 분리되어 있던 MatMul, Masking, Softmax, Dropout, MatMul 연산들이 하나의 거대한 CUDA 커널(Kernel)로 통합된다. 이를 **커널 퓨전(Kernel Fusion)**이라 한다.6 퓨전은 커널 실행 오버헤드를 줄일 뿐만 아니라, 데이터를 레지스터와 SRAM에 유지한 채로 연쇄적인 연산을 수행하게 하여 메모리 접근을 획기적으로 감소시킨다.

FlashAttention: 타일링 및 커널 퓨전

graph TD
    HBM_In[("HBM: 전체 Q, K, V")]
    HBM_Out[("HBM: 최종 결과 O")]

    subgraph "SRAM (고속 메모리) 내부 처리"
        Block_Load["블록 단위 로드 (Q_i, K_j, V_j)"]
        Compute1["S_ij = Q_i * K_j^T (연산)"]
        Softmax["Online Softmax (통계량 갱신)"]
        Compute2["O_i = P_ij * V_j (연산)"]

        Block_Load --> Compute1
        Compute1 -- "HBM 기록 없음" --> Softmax
        Softmax -- "HBM 기록 없음" --> Compute2
    end

    HBM_In -- "타일링(Tiling)" --> Block_Load
    Compute2 -- "최종 결과만 기록" --> HBM_Out

    style Compute1 fill:#cfc,stroke:#333
    style Compute2 fill:#cfc,stroke:#333
    style Softmax fill:#cfc,stroke:#333

3.2 기술적 도약 2: 온라인 소프트맥스 (Online Softmax)

타일링을 어텐션에 적용하는 데 있어 가장 큰 난관은 소프트맥스 함수였다. 소프트맥스는 벡터의 모든 요소 x_i에 대해 e^{x_i} / \sum e^{x_j}를 계산해야 하므로, 분모인 전체 합(\sum)을 구하기 위해서는 벡터 전체를 봐야 한다. 즉, 데이터를 블록으로 쪼개서 처리할 때, 국소적인 정보만으로는 전체에 대한 정확한 소프트맥스 값을 구할 수 없는 것처럼 보였다.3

FlashAttention은 온라인 소프트맥스(Online Softmax) 알고리즘을 도입하여 이 문제를 해결했다. 이 알고리즘은 2018년 Milakov와 Gimelshein에 의해 제안된 것으로, 소프트맥스 연산을 스트리밍 방식으로 처리할 수 있게 해준다.

3.2.1 온라인 소프트맥스의 수학적 원리

전체 벡터를 두 개의 블록 x^{(1)}과 x^{(2)}로 나누었다고 가정하자.

각 블록에 대해 로컬 최대값 m과 로컬 지수 합 l을 계산한다.

  • 블록 1: m_1 = \max(x^{(1)}), l_1 = \sum \exp(x^{(1)} - m_1)
  • 블록 2: m_2 = \max(x^{(2)}), l_2 = \sum \exp(x^{(2)} - m_2)

이제 두 블록을 합친 전체 벡터의 글로벌 최대값 m과 글로벌 지수 합 l은 다음과 같이 유도할 수 있다.
m = \max(m_1, m_2)

l = e^{m_1 - m} l_1 + e^{m_2 - m} l_2

이 공식을 이용하면, 새로운 블록이 들어올 때마다 기존의 통계량(m, l)과 결과값(O)을 갱신(Update)하는 방식으로 전체 소프트맥스를 계산할 수 있다.4 즉, 이전 블록의 계산 결과에 적절한 보정 계수(Scaling Factor)를 곱해주기만 하면, 전체 데이터를 한 번에 보지 않고도 수학적으로 완벽하게 동일한(Exact) 결과를 얻을 수 있다.

FlashAttention 알고리즘은 HBM에 최종 결과 행렬 O와 함께, 정규화를 위한 통계량 L과 m만을 저장하며 반복적으로 갱신한다. 이를 통해 O(N^2) 메모리를 사용할 필요가 없어진다.

온라인 소프트맥스 (Online Softmax) 작동 원리

graph TD
    Block1["블록 1 데이터"] -- "계산" --> Local1["로컬 최대값 m1\n로컬 합 l1"]
    Block2["블록 2 데이터"] -- "계산" --> Local2["로컬 최대값 m2\n로컬 합 l2"]

    Local1 --> Merge{"정보 결합 및 보정"}
    Local2 --> Merge

    Merge -- "Scaling Factor 적용" --> Global["글로벌 통계량 갱신\n(Global m, Global l)"]

    Global --> Result["정확한 Softmax 결과 도출"]

3.3 기술적 도약 3: 역전파를 위한 재연산 (Recomputation)

딥러닝의 훈련 과정에서 역전파(Backpropagation)는 필수적이다. 역전파를 수행하려면 순전파(Forward pass) 때 계산했던 N \times N 어텐션 확률 행렬 P와 점수 행렬 S의 값이 필요하다. 표준 어텐션은 이를 위해 순전파 시 생성된 거대한 행렬들을 HBM에 저장해두는데, 이것이 메모리 부족(OOM)의 주된 원인이 된다.

FlashAttention은 여기서 과감한 결단을 내린다. “저장하지 말고, 다시 계산하자.”

이를 재연산(Recomputation) 또는 Gradient Checkpointing 기법이라 한다.6

3.3.1 재연산 메커니즘

  1. 순전파: N \times N 행렬을 저장하지 않는다. 대신, 블록별 정규화 통계량인 m (최대값)과 l (지수 합)만을 HBM에 저장한다. 이들은 시퀀스 길이 N에 비례하는 O(N) 크기이므로 메모리 부담이 거의 없다.
  2. 역전파: 그라디언트를 계산할 때, 저장해둔 Q, K, V와 통계량 m, l을 HBM에서 다시 읽어온다. 그리고 SRAM 내부에서 순전파 때 수행했던 연산을 다시 수행하여 S와 P 값을 순간적으로 복원하고, 이를 이용해 그라디언트를 계산한다.

언뜻 보기에 연산을 두 번 하는 것은 비효율적으로 느껴질 수 있다. 그러나 앞서 설명했듯 GPU는 연산 속도가 메모리 접근 속도보다 훨씬 빠르다. HBM에 거대한 행렬을 썼다가 다시 읽어오는 시간보다, 빠른 SRAM과 텐서 코어를 이용해 데이터를 다시 계산하는 시간이 훨씬 짧게 걸린다.4 결과적으로 재연산은 메모리 사용량을 획기적으로 줄일 뿐만 아니라, 전체 실행 속도까지 향상시키는 역설적인 효과를 가져온다.

역전파를 위한 재연산(Recomputation) 전략

graph TD
    subgraph "순전파 (Forward Pass)"
        FW_Calc["어텐션 연산 수행"]
        Save_Stats[("통계량 m, l 저장 (크기: O(N))")]
        No_Save["N x N 행렬 저장 안함 (메모리 절약)"]

        FW_Calc --> Save_Stats
        FW_Calc -.-> No_Save
    end

    subgraph "역전파 (Backward Pass)"
        Read_Stats[("통계량 m, l 로드")]
        Read_QKV[("Q, K, V 로드")]
        Recompute["SRAM에서 어텐션 값 재계산 (Recomputation)"]
        Grad_Calc["그라디언트(Gradient) 계산"]

        Save_Stats --> Read_Stats
        Read_Stats --> Recompute
        Read_QKV --> Recompute
        Recompute --> Grad_Calc
    end

4. 이론적 분석: IO 복잡도의 승리

FlashAttention의 우수성은 단순히 구현의 묘미에 있는 것이 아니라, 알고리즘의 **IO 복잡도(IO Complexity)**가 개선되었다는 이론적 근거에 있다. 논문에서는 표준 어텐션과 FlashAttention의 HBM 접근 횟수를 수학적으로 분석하여 비교했다.1

4.1 표준 어텐션의 IO 복잡도

입력 시퀀스 길이 N, 헤드 차원 d라고 할 때, 표준 어텐션은 N \times N 크기의 행렬을 HBM에 쓰고 읽어야 한다.
\text{IO}_{Standard} = \Theta(N^2)
이는 시퀀스 길이가 늘어날수록 메모리 접근량이 제곱으로 증가함을 의미한다.

4.2 FlashAttention의 IO 복잡도

FlashAttention은 타일링을 통해 SRAM 크기 M을 활용한다. 분석 결과 IO 복잡도는 다음과 같다.
\text{IO}_{Flash} = \Theta\left(\frac{N^2 d^2}{M}\right)
여기서 중요한 점은 분모에 있는 M이다. SRAM의 크기 M이 클수록 HBM 접근 횟수가 반비례하여 줄어든다.1 이는 FlashAttention이 단순한 O(N^2) 알고리즘이 아니라, 가용한 빠른 메모리 자원을 최대한 활용하여 느린 메모리 접근을 줄이는 최적의(Optimal) 알고리즘임을 시사한다. 논문은 FlashAttention이 SRAM 크기에 대해 점근적으로 하한(Lower bound)에 도달하는 최적 알고리즘임을 증명했다.

5. 비교 분석: 기존 효율적 어텐션과의 차이

FlashAttention 등장 이전에도 O(N^2) 문제를 해결하려는 수많은 시도가 있었다. 이들과 FlashAttention의 결정적인 차이는 무엇일까?

graph LR
    Target["목표:<br>트랜스포머 가속 및<br>긴 문맥 처리"]

    Target --> Standard["표준 어텐션"]
    Target --> Approx["효율적 어텐션<br>(Reformer 등)"]
    Target --> Flash["FlashAttention"]

    Standard -- "단점" --> S_Cons["메모리 O(N^2)<br>느린 속도 (IO 병목)"]
    Approx -- "단점" --> A_Cons["정확도 손실 (근사)<br>하드웨어 비효율적"]

    Flash -- "장점 1" --> F_Adv1["정확함<br>(Exact Attention)"]
    Flash -- "장점 2" --> F_Adv2["메모리<br>O(N) (선형)"]
    Flash -- "장점 3" --> F_Adv3["속도 2~4배 향상<br>(IO 최적화)"]
    Flash -- "장점 4" --> F_Adv4["초장문 처리 가능<br>(16K~64K+)"]

    F_Adv1 --> Result["BERT, GPT 등<br>n모델 성능 유지/향상"]
    F_Adv4 --> Result

5.1 희소 어텐션 (Sparse Attention)

Reformer나 Sparse Transformer와 같은 모델들은 어텐션 행렬의 일부만 계산하는 방식을 택했다.

  • 문제점: 희소한 데이터 접근 패턴(Random Access)은 GPU의 병렬 처리 효율을 떨어뜨린다. 이론적인 FLOPs는 줄었지만, 메모리 접근이 불규칙하여 실제 속도는 빨라지지 않는 경우가 많았다.3 또한, 전체 문맥을 보지 못하므로 정확도가 하락한다.
  • FlashAttention: 밀집(Dense) 연산을 유지하여 GPU의 행렬 연산 유닛(Tensor Core) 효율을 극대화하면서도 IO를 줄여 속도를 높인다. 정확도 손실이 전혀 없다.

5.2 저랭크 근사 (Low-Rank Approximation)

Linformer와 같은 모델은 어텐션 행렬을 저차원으로 투영하여 근사한다.

  • 문제점: 근사 과정에서 정보 손실이 발생하며, 모델의 품질이 저하된다. 특정 작업에서는 성능이 급격히 떨어진다.1
  • FlashAttention: 수학적으로 표준 어텐션과 동일한 결과값을 보장한다(Exact Attention). 연구자들은 모델 성능 저하에 대한 우려 없이 FlashAttention을 즉시 도입할 수 있다.
특성표준 어텐션희소/근사 어텐션 (Reformer 등)FlashAttention
정확도정확함 (Exact)근사 (Approximate)정확함 (Exact)
메모리 복잡도O(N^2)O(N \log N) 또는 O(N)O(N)
실행 속도느림이론상 빠름, 실제론 느릴 수 있음매우 빠름 (2~4배 가속)
IO 인지 여부아니오아니오예 (IO-Aware)
긴 문맥 지원제한적 (1K~2K)가능하지만 품질 저하우수함 (16K~64K+)

표 8.1.2 다양한 어텐션 메커니즘의 비교 분석 1

6. 실험 결과 및 성능 검증

FlashAttention의 성능은 다양한 벤치마크를 통해 입증되었다. 단순히 “빠르다“는 것을 넘어, 기존에는 불가능했던 모델링을 가능하게 했다는 점에서 그 의의가 크다.

6.1 훈련 속도 (Training Speedup)

논문에서는 BERT-large와 GPT-2 모델에 대해 훈련 속도를 측정했다.1

  • BERT-large (Seq Len 512): NVIDIA의 MLPerf 1.1 훈련 기록보다 15% 더 빠른 속도를 달성했다. 이는 이미 극도로 최적화된 산업 표준 베이스라인을 넘어섰다는 점에서 의미가 크다.
  • GPT-2 (Seq Len 1K): HuggingFace 구현 및 Megatron-LM 구현 대비 **3배(3x)**의 속도 향상을 보였다.
  • Long-Range Arena (Seq Len 1K-4K): 긴 시퀀스를 다루는 벤치마크에서 베이스라인 대비 2.4배 빨랐다.

이러한 가속 효과는 시퀀스 길이가 길어질수록 더욱 극적으로 나타났다. 이는 FlashAttention이 긴 문맥 처리에 특화되어 있음을 보여준다.

6.2 메모리 효율성과 문맥 길이의 확장

FlashAttention은 메모리 사용량을 선형(O(N))으로 줄임으로써, 트랜스포머가 처리할 수 있는 문맥의 길이를 획기적으로 늘렸다.2

  • 메모리 절감: 시퀀스 길이에 따라 최대 20배의 메모리 절감 효과를 보였다.7
  • 초장문 처리: 표준 어텐션이 메모리 초과(OOM)로 실패하는 16K(Path-X), 64K(Path-256) 길이의 시퀀스에서도 FlashAttention은 문제없이 작동했다.
  • 성능 향상: 더 긴 문맥을 볼 수 있게 되면서 모델의 성능도 향상되었다. GPT-2 모델의 경우 퍼플렉서티(Perplexity)가 0.7 포인트 개선되었으며, 긴 문서 분류 작업에서도 성능 향상이 관찰되었다.1

6.3 블록 희소 어텐션으로의 확장 (Block-Sparse FlashAttention)

FlashAttention은 그 자체로도 강력하지만, 희소 어텐션 기법과 결합되었을 때 더욱 강력한 성능을 발휘한다.10 저자들은 FlashAttention 알고리즘을 블록 희소 어텐션에 적용한 Block-Sparse FlashAttention을 함께 제안했다.

이는 어텐션 행렬의 특정 블록만을 계산하도록 커널을 수정한 것으로, IO 효율성을 유지하면서 연산량(FLOPs)까지 줄여준다. 그 결과 기존의 근사 어텐션 방법들보다 더 빠른 속도를 달성했으며, 시퀀스 길이를 64K까지 늘리는 데 성공했다. 이는 트랜스포머가 책 한 권 분량의 텍스트를 한 번에 처리할 수 있는 길을 열어준 것이다.10

7. 결론: 트랜스포머 싱귤래리티의 가속화

제8장 1절에서 살펴본 FlashAttention-1은 단순한 알고리즘 개선이 아니다. 이것은 AI 모델링의 무게중심을 ’순수 알고리즘’에서 ’시스템과 하드웨어의 조화’로 옮겨놓은 전환점이다. Tri Dao 연구진은 무한히 늘어나는 연산량과 메모리 요구량 앞에서, 하드웨어의 물리적 특성(IO 병목)을 정확히 파악하고 이를 소프트웨어적으로 해결하는 우아한 해법을 제시했다.

FlashAttention-1이 가져온 변화는 다음과 같이 요약할 수 있다.

  1. O(N^2) 메모리 장벽의 붕괴: 긴 문맥 처리를 가로막던 가장 큰 장애물을 제거했다.
  2. 정확한 어텐션의 고수: 성능 타협 없는 최적화가 가능함을 증명하여 산업계 전반의 즉각적인 채택을 이끌어냈다.
  3. LLM 시대의 기반 마련: 이후 등장하는 GPT-4, Llama 등 거대 언어 모델들이 32K, 128K의 긴 문맥 윈도우를 가질 수 있었던 것은 FlashAttention과 같은 메모리 효율화 기술이 뒷받침되었기 때문이다.11

FlashAttention-1은 이후 등장할 FlashAttention-2와 3, 그리고 분산 처리를 위한 Ring Attention 등으로 이어지는 ‘효율적 어텐션’ 계보의 시발점이 되었다. 이제 트랜스포머는 단순한 문장 처리기를 넘어, 방대한 지식을 한 번에 담고 처리할 수 있는 진정한 의미의 지능형 에이전트로 진화할 준비를 마쳤다. 다음 절에서는 이러한 FlashAttention-1의 한계를 보완하고 병렬 처리 효율을 더욱 극대화한 FlashAttention-2에 대해 다룰 것이다.

8. 참고 자료

  1. FlashAttention: Fast and Memory-Efficient Exact Attention … - arXiv, https://arxiv.org/abs/2205.14135
  2. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness - arXiv, https://arxiv.org/pdf/2205.14135
  3. FLASHATTENTION: Fast and Memory-Efficient Exact Attention with IO-Awareness, https://proceedings.neurips.cc/paper_files/paper/2022/file/67d57c32e20fd0a7a302cb81d36e40d5-Supplemental-Conference.pdf
  4. How FlashAttention Eliminates Transformer Memory Bottlenecks - Galileo AI, https://galileo.ai/blog/stanford-flashattention-algorithm
  5. Basic idea behind flash attention (V1) - Damek Davis, https://damek.github.io/random/basic-idea-behind-flash-attention/
  6. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness | OpenReview, https://openreview.net/forum?id=H4DqfPSibmx
  7. FlashAttention-2: Faster Attention with Better Parallelism … - Tri Dao, https://tridao.me/publications/flash2/flash2.pdf
  8. ELI5: FlashAttention. Step by step explanation of how one of… | by …, https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad
  9. Designing Hardware-Aware Algorithms: FlashAttention | DigitalOcean, https://www.digitalocean.com/community/tutorials/flashattention
  10. Sparse Transformer Algorithms (FlashAttention) - Emergent Mind, https://www.emergentmind.com/topics/sparse-transformer-algorithms-flash-attention
  11. FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-precision | Tri Dao, https://tridao.me/blog/2024/flash3/
  12. A Deep Dive into FlashAttention Algorithm-part 3 | by Sachin Kalsi | Towards AI, https://pub.towardsai.net/a-deep-dive-into-flashattention-v1-part-3-0f58153f431b
  13. Fast Attention Over Long Sequences With Dynamic Sparse Flash Attention - François Fleuret, https://fleuret.org/papers/pagliardini-et-al-neurips2023.pdf
  14. Long Sequences Transformers: a review of the SOTA | by Achraff Adjileye - Medium, https://medium.com/@adjileyeb/long-sequences-transformers-a-review-of-the-sota-734ef9e07446
  15. Attention Types in LLM. Attention mechanisms in generative AI… | by Sulbha Jain - Medium, https://medium.com/@sulbha.jindal/attention-types-in-llm-self-multi-head-cross-multi-head-latent-e0bff7572515
  16. Block Sparse Flash Attention - arXiv, https://arxiv.org/html/2512.07011v1
  17. Dao-AILab/flash-attention: Fast and memory-efficient exact attention - GitHub, https://github.com/Dao-AILab/flash-attention