왜 표준 어텐션이 병목인가
Transformer의 self-attention은 시퀀스 길이 N에 대해 N×N 크기의 어텐션 스코어 행렬을 만든다. 문제는 연산량(FLOPs)보다 메모리다. 표준 구현은 S = QKᵀ, P = softmax(S), O = PV를 순서대로 계산하며 중간 행렬 S와 P를 GPU의 HBM(글로벌 메모리)에 통째로 저장한다. N이 4K, 8K로 늘면 이 행렬이 O(N²)로 커져 메모리를 잡아먹고, 무엇보다 HBM 읽기/쓰기(I/O)가 폭증한다.
최신 GPU에서 어텐션은 대부분 연산이 아니라 메모리 대역폭에 묶여 있다(memory-bound). 연산 유닛은 놀고 있는데 HBM에서 큰 행렬을 오가는 데 시간이 다 쓰인다는 뜻이다.
FlashAttention의 핵심 아이디어
FlashAttention은 결과를 근사하지 않는다. 수학적으로 동일한 값을 계산하되, N×N 행렬을 HBM에 절대 쓰지 않는 방식으로 재구성한다. 두 가지 기법이 핵심이다.
- 타일링(tiling): Q, K, V를 작은 블록으로 쪼개 SRAM(온칩 공유 메모리)에 올리고, 블록 단위로 스코어와 출력을 누적한다.
- 온라인 softmax: 전체 행을 보지 않고도 블록을 하나씩 처리하며 running max와 running sum을 갱신해 softmax를 정확히 계산한다.
결과적으로 HBM I/O가 O(N²)에서 O(N²/M)(M은 SRAM 블록 크기 관련 상수) 수준으로 줄어 실제 벽시계 시간이 크게 단축된다. 메모리 사용량도 시퀀스 길이에 대해 선형(O(N))으로 떨어진다.
비교: 무엇이 달라지나
| 항목 | 표준 어텐션 | FlashAttention |
|---|---|---|
| 중간 행렬 저장 | N×N을 HBM에 저장 | 저장 안 함(SRAM에서 처리) |
| 메모리 복잡도 | O(N²) | O(N) |
| 주 병목 | HBM I/O | 연산으로 이동 |
| 결과값 | 정확 | 정확(동일) |
PyTorch에서 바로 켜기
가장 손쉬운 경로는 PyTorch 2.x의 scaled_dot_product_attention(SDPA)이다. 조건이 맞으면 내부적으로 FlashAttention 커널을 자동 선택한다.
import torch
import torch.nn.functional as F
from torch.nn.attention import sdpa_kernel, SDPBackend
q = torch.randn(2, 16, 4096, 64, device="cuda", dtype=torch.float16)
k = torch.randn_like(q)
v = torch.randn_like(q)
# FlashAttention 백엔드를 명시적으로 강제
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
out = F.scaled_dot_product_attention(q, k, v, is_causal=True)
print(out.shape) # torch.Size([2, 16, 4096, 64])
is_causal=True는 마스크 행렬을 메모리에 만들지 않고 커널 내부에서 처리하므로 별도 마스크 텐서보다 효율적이다.
flash-attn 패키지 직접 사용
가변 길이 배치나 최신 기능이 필요하면 Dao-AILab의 flash-attn을 직접 쓴다. 설치와 최소 호출 예시는 다음과 같다.
# 설치 (컴파일에 시간이 걸림, CUDA 툴킷 필요)
pip install flash-attn --no-build-isolationfrom flash_attn import flash_attn_func
# 입력 레이아웃: (batch, seqlen, nheads, head_dim)
q = torch.randn(2, 4096, 16, 64, device="cuda", dtype=torch.bfloat16)
k = torch.randn_like(q)
v = torch.randn_like(q)
out = flash_attn_func(q, k, v, causal=True, dropout_p=0.0)
패딩 없는 가변 길이 시퀀스는 flash_attn_varlen_func와 누적 시퀀스 길이(cu_seqlens)를 쓰면 패딩 토큰 연산을 완전히 제거할 수 있다.
실무 주의점
- dtype 제약: FlashAttention 커널은 fp16/bf16에서 동작한다. fp32 입력은 커널이 선택되지 않고 표준 경로로 폴백된다.
- head_dim 제한: 버전에 따라 지원하는 head 차원 상한(예: 128, 256)이 다르다. 초과하면 조용히 폴백되니 실제 커널이 쓰였는지 프로파일러로 확인하라.
- 커스텀 마스크: 임의의 additive bias 마스크는 지원이 제한적이다. causal이나 sliding window처럼 커널이 아는 패턴으로 표현하는 편이 빠르다.
- 하드웨어 의존: Ampere 이상에서 이득이 크고, 아키텍처별로 최적 커널이 다르다. GPU가 바뀌면 다시 벤치마크한다.
- 수치 재현성: 값은 수학적으로 동일하지만 누적 순서 차이로 마지막 비트 수준의 차이가 날 수 있다. 엄격한 비트 단위 재현이 필요한 테스트라면 감안하라.
정리
FlashAttention은 정확도를 희생하지 않고 HBM I/O를 줄여 긴 컨텍스트에서 속도와 메모리를 동시에 개선한다. 대부분은 PyTorch SDPA로 충분하고, 가변 길이·세밀한 제어가 필요할 때 flash-attn을 직접 붙인다. 켜기 전후로 실제 커널 선택 여부와 지연 시간을 반드시 프로파일링해 폴백을 놓치지 않는 것이 핵심이다.