tech · 2026-03-22
FlashAttention
1. Introduction
1.1 Attention Mechanism의 병목 구조
Transformer 아키텍처는 자연어 처리, 컴퓨터 비전 등 다양한 도메인에서 핵심 모델 구조로 자리잡은 구조임. 이 구조의 중심 연산인 Scaled Dot-Product Attention은 입력 시퀀스 길이에 대해 O(n²)의 연산 및 메모리 복잡도를 갖는 구조임. 특히, 시퀀스 길이가 길거나 배치 크기가 커질수록 연산량보다는 메모리 접근량(memory IO) 이 주요 병목으로 작용하는 경향임.
기존 구현에서는 attention 연산을 다음과 같은 3단계로 분리하여 수행하는 구조임:
Attention(Q, K, V) = softmax(QKᵀ) · V
1. Score = QKᵀ
2. Weights = softmax(Score)
3. Output = Weights · V
이 과정에서 attention score 행렬(QK^T) 및 normalized score, output 행렬이 모두 GPU global memory에 materialize됨. 이로 인해 연산량 대비 IO 양이 3~4배 이상 많아지는 비효율이 발생하는 구조임. 특히, softmax는 메모리 재접근이 필수적으로 요구되는 연산으로, register-level 최적화가 어렵다는 한계가 존재함.
1.2 FlashAttention의 설계 목적
FlashAttention은 이와 같은 attention 연산의 구조적 병목을 해결하기 위해 제안된 GPU 최적화 연산 구조임. 기존의 global memory 기반 처리 방식에서 벗어나, 다음과 같은 방식으로 attention 연산을 IO 중심으로 재설계한 구조임:
- softmax 연산을 online 방식으로 수행하여 intermediate tensor materialization 제거
- query chunking과 key streaming을 통해 attention score를 부분적으로 계산 및 누적
- 연산 결과를 register와 shared memory에서 처리하여 global memory 접근 최소화
이 방식은 연산량은 그대로 유지하면서, memory read/write 비용을 획기적으로 줄이는 효과를 가지는 구조임. 실제 논문에서는 기존 O(n²)의 IO complexity를 O(n) 수준으로 낮추는 것을 주요 성과로 제시함:
"FlashAttention reduces the memory reads/writes from O(n²·d) to O(n·d) using tiling and recomputation." — Dao et al., 2022 (NeurIPS)
1.3 버전별 흐름
FlashAttention은 버전별로 다음과 같은 진화를 거친 구조임:
- FlashAttention v1: Triton 기반의 forward 전용 memory-efficient 커널 구조. backward 연산에서는 memory 효율이 미흡했던 구조임.
- FlashAttention v2: native CUDA 기반의 full-fusion 구조 도입. forward-backward 전체를 IO-aware하게 최적화하며, softmax까지 모두 커널 내부에 통합한 구조임.
- FlashAttention v3: variable-length 시퀀스, multi-query attention, block-sparse attention에 대한 최적화 지원. 다양한 실전 모델과의 통합성이 강화된 구조임.
FlashAttention은 현재 GPT-3, PaLM, LLaMA, Mistral, Mixtral, Claude 등 대형 언어모델의 pretraining 및 serving 환경 모두에서 핵심 연산 커널로 활용되고 있음. 특히, TensorRT-LLM, vLLM, Hugging Face Transformers, Megatron-LM 등 다양한 최적화 프레임워크에서도 FlashAttention 기반 커널을 기본값으로 채택하고 있음.
2. Background
2.1 Scaled Dot-Product Attention의 연산 구조
Transformer에서 사용되는 Scaled Dot-Product Attention은 Query(Q), Key(K), Value(V)를 입력으로 받아 유사도 기반의 가중 합을 계산하는 형태임. 전체 연산은 다음과 같은 수식으로 정의됨:
Attention(Q, K, V) = softmax(QKᵀ / √dₖ) · V
여기서 Q는 시퀀스 길이 n과 임베딩 차원 d에 따라 크기가 (n × d)인 행렬이며, K와 V도 같은 형식으로 구성됨. Q와 K의 내적 결과인 attention score 행렬은 (n × n)의 크기를 가지며, 여기에 softmax 정규화를 적용한 뒤, V와 행렬 곱을 수행하여 최종 출력이 계산되는 방식임.
이 연산은 계산량 측면에서는 O(n²·d)의 복잡도를 갖지만, 실제 성능에는 메모리 접근 패턴이 더 큰 영향을 미치는 연산임. 특히, 중간 결과인 score 및 softmax 결과를 저장하고 재사용하는 과정에서 발생하는 memory IO 비용이 전체 latency의 주요 원인으로 작용하는 경향임.
2.2 GPU Memory Hierarchy와 IO 병목
GPU는 높은 병렬 연산 성능을 갖추고 있지만, 연산 자체보다 메모리 계층 간의 데이터 이동 속도가 전체 처리 시간의 병목으로 작용하는 경우가 많음. 특히 Transformer의 Attention 연산에서는 이러한 메모리 계층 간의 IO 부담이 성능 저하의 주요 원인으로 지적됨.
Register (Thread-Private Memory)
- 속도: 가장 빠름 (1-cycle access)
- 용도: 개별 스레드가 계산하는 스칼라 혹은 벡터 단위 결과값(c[i][j]) 저장
- 용량: 256KB per SM (H100)
- 예시: 각 스레드가 계산한 softmax 결과 또는 value vector 곱 결과를 저장
- 제한: 매우 작고, 너무 많은 변수 사용 시 register spilling 발생
Shared Memory (Thread Block-Shared Memory)
속도: 빠름 (shared L1 수준, low latency)
용도: 한 block 내 스레드들이 공동으로 사용하는 Q, K, V 블록(Tile) 저장
용량: 164KB per SM (H100)
예시:
- Q[i:i+64], K[j:j+64] tile을 loading 후, Q @ Kᵀ 연산
- softmax normalization을 위한 intermediate buffer로 활용
제한: 보통 48~96KB 정도로 제한적 → tile 단위 처리에 적합
L2 Cache (Device-Wide Cache)
속도: 중간 (100~200 cycles), 자동 관리
용도: 여러 block 간 재사용 가능성이 있는 Q, K, V 전체 블록 또는 matrix tile 캐싱
용량: 50MB (H100)
예시:
- Global memory에서 자주 참조되는 key matrix 조각들
- 이전 단계에서 생성된 중간 결과 (e.g., rescaled softmax sum)
특징: 사용자가 직접 관리하지 않고, 하드웨어가 자동으로 caching
L2 cache 데이터 유형 정리
| 데이터 유형 | 예시 | L2 Cache 저장 이유 |
|---|---|---|
| 재사용 K/V 타일 | K[j:j+64], V[j:j+64] |
Query chunk가 반복해서 동일한 K/V를 참조하므로 |
| Softmax 누적 통계량 | prev_max, prev_sum |
Online Softmax 연산에서 이전 chunk 통계 필요 |
| Position 정보 | Rotary Embedding, Position Bias | 모든 head에서 반복적으로 사용됨 |
| 전처리 중간 결과 | LayerNorm(x) | Q, K 생성을 위해 동일한 전처리 결과 참조 |
| Mask 관련 정보 | Dropout Mask, Causal Mask | 여러 block에서 반복적으로 사용될 수 있음 |
Global Memory (VRAM)
속도: 가장 느림 (400~800 cycles 이상)
용도: 전체 시퀀스의 Q, K, V 행렬, score 행렬, softmax 결과 등 모든 데이터를 저장
용량: 80GB (H100)
예시:
- Full Q[batch, head, seq, dim], K, V 행렬
- Attention output을 전 레이어로 전달하기 위한 저장
특징: 크고 범용적이지만, 지나친 read/write는 latency 급증의 원인
Attention 연산 내에서의 실제 데이터 흐름 예시
- Global Memory에서 Q, K, V 행렬의 일부분을 불러옴
- Shared Memory에 tile 단위로 적재 (예: Q[i:i+64], K[j:j+64])
- 각 스레드는 Register에 자신의 계산 결과 (c[i][j])를 임시 저장
- 필요 시 L2 Cache를 통해 tile 간 공유 데이터를 중간 유지
- 최종 연산 결과는 다시 Global Memory에 저장
병목의 핵심
기존 Attention 구현은 다음 두 단계 간에 매번 Global Memory read/write가 발생:
S = Q @ Kᵀ계산 후 → Global Memory 저장A = softmax(S)계산 후 → Global Memory 저장O = A @ V연산 시 → A를 다시 Global Memory에서 불러옴
이러한 불필요한 데이터 왕복이 GPU의 계산 성능 대비 낮은 효율을 유발하며, FlashAttention은 이러한 메모리 흐름을 재설계하여 이 문제를 해결하고자 함.
H100 GPU Memory Example
======================== H100 GPU Memory Hierarchy ========================
[ Global Memory ]
(80GB HBM3)
||
|| << high latency, high capacity >>
||
+------------------+
| L2 Cache |
| (50MB, shared) |
+------------------+
||
+----------------------------------+
| |
+---------------+ +---------------+
| SM #0 | ... | SM #N |
+---------------+ +---------------+
| Warp Scheduler | | Warp Scheduler |
| | | |
| Registers | | Registers |
| (per thread) | | (per thread) |
| | | |
| Shared Memory | | Shared Memory |
| (192 KB/SM) | | (192 KB/SM) |
+---------------+ +---------------+
|| ||
[ Tensor Cores / ALUs ] [ Tensor Cores / ALUs ]
=========================== Access Characteristics ===========================
- Registers : ultra-fast, per-thread, no access latency
- Shared Memory : very fast, low latency (per-SM, programmer-managed)
- L2 Cache : fast, shared by all SMs, hardware-managed
- Global Memory : slowest, large capacity, main data source/sink
===============================================================================
2.3 기존 Attention 구현의 한계
기존 Transformer 구현에서는 attention 연산을 세 단계로 나누어 처리하는 구조를 따름. 먼저 Q와 K의 내적을 계산하여 (n × n) 크기의 score 행렬을 만들고, 이 결과를 global memory에 저장함. 이후 이 score에 softmax를 적용하여 정규화한 뒤, 다시 그 결과를 V와 곱하여 최종 출력을 계산함.
이 구조는 다음과 같은 IO 중심의 병목을 유발함:
- Q × K^T 결과를 full matrix 형태로 저장 → 메모리 write 비용 발생
- softmax 결과를 저장하고 이후에 다시 불러옴 → 메모리 read 비용 발생
- softmax 연산은 수치적 안정성을 위해 log-sum-exp 기법을 사용하는데, 이를 위해 두 번에 걸쳐 score 행렬을 읽는 과정이 필요함
이로 인해 attention 연산은 연산량 자체는 적더라도, memory access가 집중되는 구간이 되어 latency의 주요 원인으로 작용함. BERT 모델 기준으로는 attention이 전체 연산량의 약 1020% 수준임에도 불구하고, 전체 실행 시간의 4060% 이상을 차지하는 현상이 관찰됨. 이는 attention이 계산 복잡도보다 IO 중심의 비효율이 큰 연산 패턴임을 보여주는 결과임.
주의사항: 위 연산량과 latency는 n < 2d 일 경우에만 적용됨. 즉, 입력 토큰수가 적었던 과거의 모델에 해당하는 이야기고 입력수가 임베딩 차원 수에 비해서 굉장히 큰 현대의 LLM은 attention이 차지하는 연산량 자체가 많음
2.4 FlashAttention의 설계 철학: IO-aware Attention
FlashAttention은 기존 attention 연산의 병목을 해소하기 위해, 연산 순서와 메모리 접근 방식을 근본적으로 재설계한 구조를 제안함. 해당 방식은 다음과 같은 핵심 전략을 기반으로 함:
- Online Softmax 처리: attention score 전체를 한 번에 계산하고 저장하지 않고, softmax의 분자와 분모를 블록 단위로 누적 계산하여 register나 shared memory에서 직접 처리하는 방식
- Query Chunking 및 Key Streaming: Q와 K를 작은 단위로 나누어 순차적으로 처리함으로써 full attention score 행렬을 만들지 않고도 연산을 수행하는 방식
- Recomputation 기반 메모리 최적화: backward 과정에서 softmax 결과를 저장하지 않고, forward 연산을 재계산함으로써 메모리 사용량을 줄이는 전략
FlashAttention은 이러한 전략을 통해 attention의 수학적 정의는 그대로 유지하면서도, 메모리 접근량을 최소화하는 최적화된 커널 구조를 구현함. 이 구조는 단순한 속도 개선을 넘어서, Transformer 모델의 시퀀스 확장성과 추론 속도를 동시에 확보하는 데 중요한 기법으로 평가받는 구조임.
3. FlashAttention Core Algorithm
3.1 FlashAttention 전체 구조 개요
1. Motivation: GPU 메모리 계층과 Attention 구조의 충돌 기존 Scaled Dot-Product Attention은 QKᵀ score 계산 → Softmax 정규화 → Weighted Sum 계산의 세 단계를 순차적으로 수행함. score matrix(QKᵀ)의 크기는 시퀀스 길이 n일 때 n × n으로 구성되며, FP16 기준 시퀀스 길이 1024만 되어도 score matrix만 약 2MB 메모리를 차지함. Q, K, V 행렬 및 중간 결과까지 포함하면 최소 10MB 이상이 필요하므로, GPU의 register, shared memory에는 적재가 불가능하고 global memory(HBM)를 사용할 수밖에 없음. 이로 인해 연산보다 메모리 read/write가 전체 latency의 병목이 되는 구조적 한계 발생.
2. Approach: Chunk 단위 연산과 Streaming 방식의 설계 FlashAttention은 연산 흐름을 tile 단위로 분해하여, Q는 chunk 단위로 register에 고정, K/V는 stream-in 방식으로 block 단위 shared memory에 임시 적재. score 계산 이후 softmax 및 weighted sum까지 register/shared memory 내부에서 통합 수행하여, 중간 결과 저장 없이 global memory 접근 최소화. 기존의 O(n²) 메모리 접근 패턴을 O(n) 수준으로 줄이도록 설계.
3. Problem: Softmax 연산의 전역 의존성과 분할 불가능성 QKᵀ와 V 연산은 tile 단위 분할이 가능하나, softmax는 한 query에 대한 모든 score를 알아야 정규화가 가능함. softmax(x) = exp(xᵢ) / ∑ⱼ exp(xⱼ)의 특성상, 분모 ∑ⱼ exp(xⱼ)를 계산하려면 전체 tile을 메모리에 올려야 하는 구조적 병목 발생. shared memory에는 모든 score vector를 수용할 수 없어 tile 기반 streaming이 softmax에서 적용 불가능. 결과적으로 softmax 단계가 FlashAttention의 병목 해소에 가장 큰 장애물이 됨.
4. Solution: Online Softmax 및 연산 통합 커널 설계 이 문제를 해결하기 위해 FlashAttention은 Online Softmax 기법을 도입. tile마다 최대값을 유지하면서 overflow를 방지하고, partial sum을 누적하여 정규화를 실시간으로 계산. softmax score가 확정되는 순간 V와의 가중합까지 동시에 계산하여 softmax 결과를 따로 저장하지 않음. 이는 softmax와 weighted sum을 하나의 fused kernel로 처리하게 하며, memory-intensive한 구조를 compute-intensive한 구조로 전환. 결과적으로 전체 연산을 메모리 접근 최소화 상태에서 고속으로 처리 가능하게 함.
5. Result: 메모리 사용량 및 latency의 획기적 절감
- Global memory 접근 횟수: O(n²) → O(n)
- 중간 결과 저장 불필요 (QKᵀ, softmax, weighted sum 모두 memory-free 처리)
- register/shared memory 기반 연산으로 latency 단축
- 학습 시 forward-backward 통합으로 추가 메모리 절감 달성
3.2 Query Chunking & Key Streaming
Query Chunking의 동기와 구조
FlashAttention은 시퀀스 길이와 임베딩 차원이 커짐에 따라 발생하는 메모리 병목 문제를 해결하기 위해 Query 행렬을 일정한 블록 단위로 분할하여 처리함. 이 방식을 Query Chunking이라 부르며, 각 Chunk는 register 또는 shared memory에 적재 가능한 수준으로 크기를 제한함.
예를 들어, 시퀀스 길이 4096, 임베딩 차원 1024, FP16 기준으로 Query 전체 텐서는 약 8MB의 메모리를 요구함. 이는 H100의 shared memory (164kb)나 register 용량 (256kb)을 초과하기 때문에, 전체 Query를 한 번에 처리하는 방식은 불가능함. 이에 따라 FlashAttention은 일반적으로 64~128개의 Query Tokens를 하나의 Chunk로 구성하여 반복 처리함. (즉, 시퀀스를 잘라서 chunk를 만듦 → hidden_dim을 자르는 것은 아님)
Key/Value Streaming의 원리
Query를 Chunk 단위로 처리하기 위해, Key와 Value는 스트리밍 방식으로 반복 접근함. Query Chunk가 고정된 상태에서, Key/Value는 슬라이딩 윈도우처럼 Block 단위로 shared memory에 적재되며, 각 Block은 연산 후 삭제됨. 이 방식은 전체 K/V 행렬을 메모리에 유지하지 않고도 QKᵀ 연산을 반복 수행할 수 있도록 설계됨.
각 Query Chunk는 전체 Key/Value Block을 순차적으로 순회하며 attention score를 계산함. 이때 Q는 register 혹은 shared memory에 고정 유지되며, K/V는 순차적으로 로딩 및 삭제되는 구조로, memory read/write 횟수를 최소화함.
for q_chunk in all_query_sequences:
load q_chunk to register
p_i, m_i, s_i = init values
for kv_stream in all_kv_sequences:
load kv_stream to shared memory
compute q_chunk @ kv_stream^T
update softmax numerator/denominator/V sum
kv_memory_access = O(n x b)
*b: # of chunks
*kv_stream은 b번 만큼 재사용 되므로 L2 cache에 적재해서 재사용
Memory 계층 기반 최적화 전략
FlashAttention은 GPU의 메모리 계층 구조에 맞추어 연산 단위를 배치함.
- Query Block: register 혹은 shared memory에 고정 유지
- Key/Value Block: streaming 방식으로 shared memory에 순차적으로 적재
- 중간 score 결과: softmax 계산을 위한 구조 없이 바로 다음 연산으로 전달
- 최종 Output: global memory로 direct write 수행
이러한 구조는 global memory IO를 줄이고, 대부분의 연산을 register-level에서 수행할 수 있도록 하여 memory bandwidth 병목을 해소함.
| 구성 요소 | 저장 위치 | 설명 |
|---|---|---|
| Q Chunk | Register / Shared | 반복 동안 고정 유지 |
| K/V Block | Shared | Streaming 방식 순차 로딩 및 삭제 |
| Score Block | Register | Block 단위 score 계산 결과 |
| Final Output | Global Memory | Accumulated Output 저장 |
연산 흐름 내 반복 패턴
Query Chunk 하나당 수행되는 반복 연산은 다음과 같음:
- Key/Value Block을 shared memory에 로드
- Query와 Key의 내적을 통한 score 계산
- 결과 score를 기반으로 후속 연산 수행 (online softmax 구조는 다음 섹션에서 설명)
- Key/Value Block 삭제, 다음 Block 로드
이러한 반복 구조는 고정된 Query Chunk를 기준으로 sliding 형태의 Key/Value streaming을 통해 메모리 효율성을 극대화하는 방식임.
# Pseudocode illustrating query chunking and KV streaming
for q_chunk in query_chunks:
load q_chunk into shared_memory
for kv_block in key_value_blocks:
load kv_block into shared_memory
score_block = matmul(q_chunk, kv_block.T)
# (Optional) Apply softmax or pass to next stage
accumulate(score_block)
┌──────────────┐
│ Q Chunk │ ← 고정 유지 (shared/register)
└──────────────┘
│
▼
┌──────────────┐ ┌──────────────┐ ┌──────────────┐
│ K/V Block1 │ → │ K/V Block2 │ → │ K/V Block3 │ → ...
└──────────────┘ └──────────────┘ └──────────────┘
↓ ↓ ↓
Compute 1 Compute 2 Compute 3
↓ ↓ ↓
Accumulate partial results → 최종 Attention Output
3.3 Online Softmax
Softmax 연산의 메모리 병목 문제
Query Chunking 및 Key/Value Streaming 기법을 통해 QKᵀ 행렬을 블록 단위로 순차적으로 계산하더라도, softmax 연산 자체는 전체 attention score를 기반으로 정규화를 수행해야 하는 제약이 존재함. 일반적인 softmax는 각 row에서 모든 QKᵀ 값이 준비된 후, 해당 row의 최대값 max와 지수합 sum(exp(...))을 계산해야 하기 때문에, 전체 메모리 상에 score 행렬이 존재하지 않는 streaming 구조에서는 적용 불가능함.
온라인 처리 기반 Softmax 누적 구조
FlashAttention에서는 Online Softmax라는 개념을 도입하여 위 문제를 해결함. 이 구조는 softmax의 분자 exp(x)와 분모 ∑exp(x) 를 한번에 계산하지 않고, 블록 단위로 나누어 계산한 후 누적하는 방식임.
이를 위해 다음 세 가지 값을 누적 관리함:
m_i: 현재까지 관측된 attention score의 최대값s_i: 정규화를 위한 누적된 지수합p_i: softmax 결과에 곱해질 V 벡터와의 중간 곱셈 결과
m_i = max(m_i, max(block_scores))
s_i = s_i * exp(m_i_prev - m_i) + sum(exp(block_scores - m_i))
p_i = p_i * exp(m_i_prev - m_i) + exp(block_scores - m_i) @ V_block
Online Softmax와 일반 Softmax의 동일성 검증 Example
x = [1.0, 2.0, 0.5, 3.0, 2.5, 1.5]
이를 3개의 블록으로 분할:
- Block 1: [1.0, 2.0]
- Block 2: [0.5, 3.0]
- Block 3: [2.5, 1.5]
일반 Softmax 계산
Global max 계산: max(x) = 3.0
분자 계산: exp(x_i - max(x)) for each x_i
| Index | xᵢ | xᵢ - max | exp(xᵢ - max) |
|---|---|---|---|
| 0 | 1.0 | -2.0 | 0.1353 |
| 1 | 2.0 | -1.0 | 0.3679 |
| 2 | 0.5 | -2.5 | 0.0821 |
| 3 | 3.0 | 0.0 | 1.0000 |
| 4 | 2.5 | -0.5 | 0.6065 |
| 5 | 1.5 | -1.5 | 0.2231 |
분모 합: 총합 ≈ 2.4149
Softmax 결과: [0.056, 0.152, 0.034, 0.414, 0.251, 0.092]
Online Softmax 계산
각 단계별 누적 max/sum 재정규화 및 softmax 계산 과정:
Input vector x = [1.0, 2.0, 0.5, 3.0, 2.5, 1.5]
Assume chunking into 3 blocks:
Block 1: [1.0, 2.0]
Block 2: [0.5, 3.0]
Block 3: [2.5, 1.5]
# Block 1
max_1 = 2.0
sum_1 = exp(1.0 - 2.0) + exp(2.0 - 2.0)
= exp(-1.0) + exp(0) ≈ 0.3679 + 1.0 = 1.3679
# Block 2
max_2 = max(2.0, 3.0) = 3.0
sum_1_rescaled = sum_1 × exp(2.0 - 3.0) = 1.3679 × exp(-1.0) ≈ 0.5034
sum_2 = exp(0.5 - 3.0) + exp(3.0 - 3.0) = exp(-2.5) + 1.0 ≈ 0.0821 + 1.0 = 1.0821
acc_sum_2 = sum_1_rescaled + sum_2 = 0.5034 + 1.0821 = 1.5855
# Block 3
sum_3 = exp(2.5 - 3.0) + exp(1.5 - 3.0)
= exp(-0.5) + exp(-1.5) ≈ 0.6065 + 0.2231 = 0.8296
final_sum = acc_sum_2 + sum_3 = 1.5855 + 0.8296 ≈ 2.4151
# Comparison with Standard Softmax
Standard softmax uses:
max = 3.0
sum = ∑exp(xᵢ - 3.0) for all xᵢ
= exp(-2.0) + exp(-1.0) + exp(-2.5) + exp(0) + exp(-0.5) + exp(-1.5)
≈ 0.1353 + 0.3679 + 0.0821 + 1.0 + 0.6065 + 0.2231 = 2.4151
최종적으로 Online Softmax 결과는 일반 Softmax와 완전히 동일:
Softmax_online(x_i) = exp(x_i - max) / final_sum
final_sum ≈ 2.4149, max = 3.0
결과적으로 Softmax_online(x) = Softmax(x)와 동일함
log-sum-exp 기반 수치 안정성 확보
지수 함수의 underflow/overflow 문제를 방지하기 위해, FlashAttention은 log-sum-exp 기반의 정규화 구조를 사용함. 각 블록의 softmax 계산은 다음과 같은 방식으로 안정화됨:
softmax(x) = exp(x - max(x)) / ∑exp(x - max(x))
이 계산을 블록 단위에서 수행하기 위해, 블록별 최대값을 추적하며 그에 따라 이전 sum과 numerator들을 재정규화함. 모든 값은 상대적인 max 기준으로 처리되기 때문에 수치적인 안정성이 크게 향상됨.
for q_chunk in query_chunks:
m_curr = max(m_prev, max(qk_chunk))
s_curr = s_prev * exp(m_prev - m_curr) + sum(exp(qk_chunk - m_curr))
p_curr = p_prev * exp(m_prev - m_curr) + exp(qk_chunk - m_curr) @ v_chunk
효율성 및 정확도 요약
- 전체 attention score를 저장하지 않고 softmax를 수행 가능
- register/shared memory 기반 누적 연산으로 메모리 IO 감소
- softmax의 memory-bound 성질을 compute-bound 방식으로 전환
3.4 IO Complexity 분석: O(n²) → O(n)으로의 개선
IO 복잡도 비교
| 항목 | Naive Attention | FlashAttention |
|---|---|---|
| QKᵀ 메모리 접근량 | O(n²) | O(n × b) |
| Softmax + QKV 접근량 | O(n²) | O(n) |
| K/V 읽기 횟수 | O(n²) | O(n) |
b: # of query chunk
기존 Attention에서는 각 query마다 전체 key와 value에 접근해야 했으므로, K/V는 n번 반복 접근됨. 반면 FlashAttention은 K, V를 한 번 stream 하며 모든 Q chunk에 사용함으로써, 메모리 접근 횟수를 단일화함.
IO-Bound에서 Compute-Bound로의 전환
기존 구현은 연산량보다 memory traffic이 병목인 memory-bound 구조였음. FlashAttention은 이를 compute-bound 구조로 전환함. 이는 연산량이 약간 증가한다고 하더라도 메모리 이동 비용이 낮아지고, GPU의 연산 유닛이 더 효율적으로 사용됨을 의미함.
예를 들어, A100 기준으로 global memory bandwidth는 1.6TB/s, SM당 peak 연산량은 수십 TFLOPs에 달하지만, global memory 접근 병목으로 인해 이를 활용하지 못하는 경우가 많음. FlashAttention은 이런 병목을 해소하고 SM의 연산 자원을 적극 활용함.
3.5 CUDA 관점에서 본 Memory-Efficient Attention 구조
Shared Memory 활용
FlashAttention은 고속의 메모리 접근이 가능한 Shared Memory를 적극적으로 활용함. 기존 Attention은 Q, K, V 행렬 전체를 Global Memory에 상주시켜 접근할 때마다 고비용의 메모리 I/O를 발생시켰음. 이에 반해 FlashAttention은 Query 단위로 행렬을 분할하여 Shared Memory에 로딩한 후, 해당 chunk에 대해 Key와 Value 행렬을 Streaming 방식으로 블록 단위로 순차 접근함으로써, 높은 Locality를 유지하고 반복적인 Global Memory 접근을 방지함.
Warp 병렬화 구조
CUDA Warp 단위(32 threads)로 Softmax 및 MatMul 연산을 병렬화하여 연산 효율을 높임. 각 Query token은 한 개의 Warp가 처리하며, 해당 Warp 내에서 병렬 연산을 수행함으로써 Latency를 감소시킴. Warp 간의 통신은 __shfl_sync를 통해 Register 간 직접 교환 방식으로 이루어지며, 이를 통해 Shared Memory 접근을 최소화함. Softmax 연산 시 필요한 최대값 및 합산 연산 또한 Warp 단위 reduction을 통해 처리됨.
Register Pressure 및 Spilling 최소화
Softmax 계산에서 Log-Sum-Exp 기법을 적용하며, 중간 계산값을 float32 precision으로 Register에 유지함. 이로 인해 연산 정확도를 보장하면서도 메모리 접근을 줄일 수 있음. 그러나 Query chunk 수가 많아지고, Head 수가 커질수록 Register 사용량이 증가하게 되어 Register spilling이 발생할 수 있음. 이를 완화하기 위해 병렬도를 일정 수준으로 제한하거나, Fusion 전략을 통해 계산 그래프를 통합함으로써 Register 재사용을 유도함.
Tiling 및 Double Buffering 전략
FlashAttention은 Query-Key score 계산 시 Tiling 기법을 활용하여, 연산을 작은 블록 단위로 분할하고 각 블록을 순차적으로 처리함. 이때 Softmax normalization의 정합성을 유지하기 위해 block-wise Softmax 결과를 Online Softmax로 누적 처리함. Double Buffering 기법을 통해 Shared Memory에서 다음 Key block을 Prefetch하는 동안 현재 블록 연산을 동시에 수행함으로써 메모리-계산 간 Overlap을 최대화함. 이는 실제 실행 시점에서 메모리 지연 없이 연산을 지속할 수 있게 하는 중요한 기법임.
Mixed Precision 연산 및 수치 안정성 확보
Attention score의 누적 계산은 float32로 수행하여 Softmax의 수치적 안정성을 유지함. 반면 MatMul 연산은 FP16이나 BF16으로 수행함으로써 연산량 감소와 메모리 사용 최적화를 동시에 달성함. 이러한 Mixed Precision 구조는 Tensor Core를 효과적으로 활용하면서도 underflow/overflow 없이 안정적인 출력 결과를 생성하게 함.
L2 Cache 활용
Key와 Value 행렬의 Streaming 시, L2 Cache Locality를 활용하여 메모리 대역폭을 줄임. 동일한 Key block이 여러 Query chunk에서 반복적으로 사용되므로, Cache hit rate가 높아지는 구조이며, 이를 통해 전체 메모리 대역폭 요구량을 줄이고 Throughput을 개선함. 이러한 접근은 L2 cache-aware tiling strategy와도 결합될 수 있음.