Flash Attention

Flash Attention(FlashAttention)은 Transformer 모델의 정확한 어텐션을 계산하면서 전체 어텐션 행렬을 GPU의 주 메모리에 기록하지 않는 알고리즘이자 GPU 커널 계열입니다. N개의 토큰에 대한 표준 셀프 어텐션은 N × N 점수 행렬을 만들고 이를 고대역폭 메모리(HBM)에 기록한 다음, 소프트맥스 및 값과의 곱셈을 위해 다시 읽습니다. 따라서 어텐션 실행 시간의 대부분은 데이터 이동에 소요됩니다. 스탠퍼드의 Tri Dao와 동료들은 2022년 논문 FlashAttention: IO 인식 기반의 빠르고 메모리 효율적인 정확한 어텐션에서 이 방법을 소개했습니다. 이 방법은 작업을 빠른 온칩 SRAM에 들어가는 타일로 분할하므로 메모리 사용량은 시퀀스 길이에 선형적으로 증가하고, 출력은 부동소수점 반올림 오차 범위 내에서 표준 어텐션과 일치합니다. 이 커널은 PyTorch 내부와 vLLM 같은 대규모 언어 모델(LLM) 추론 엔진에서 실행됩니다.

Flash Attention의 작동 방식#

GPU에는 연산 유닛 옆에 소량의 초고속 메모리가 있고, 칩 외부에는 용량이 크지만 속도가 느린 메모리가 있습니다. FlashAttention 논문은 NVIDIA A100의 HBM 40~80 GB, 대역폭 1.5~2.0 TB/s, 그리고 108개의 스트리밍 멀티프로세서 각각에 있는 SRAM 192 KB와 추정 대역폭 19 TB/s를 설명합니다. 어텐션은 읽는 바이트당 연산량이 적으므로 속도는 데이터가 두 메모리 계층 사이를 오가는 빈도에 좌우됩니다. 이 방법은 이러한 전송을 계산하며, 이것이 바로 "IO 인식"의 의미입니다.

FlashAttention은 쿼리 타일을 불러온 다음 키와 값 타일을 SRAM으로 스트리밍합니다. 각 타일에 대해 점수를 계산하고, 각 쿼리 행의 실행 중인 최댓값과 합계를 갱신하며, 최댓값이 커질 때마다 부분 출력을 재조정하고 새 타일의 기여분을 더합니다. 이 온라인 소프트맥스는 행 전체를 한꺼번에 보지 않고도 정확한 소프트맥스 결과를 제공하며, 최종 출력과 행마다 두 개의 숫자만 HBM으로 돌아갑니다. 역전파 단계에서는 N × N 행렬을 저장하는 대신 Q, K, V에서 점수 타일을 다시 계산합니다. 이로 인해 연산량은 늘어나지만 메모리를 훨씬 적게 읽으므로 여전히 더 빠르게 실행됩니다. 한 번에 하나의 키-값 타일에 대해 모든 쿼리를 처리하는 PyTorch의 동일한 루프는 다음과 같습니다:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

가장 큰 점수 타일은 4,096 × 256으로, 4,096 × 4,096 행렬의 16분의 1입니다. 결과는 표준 어텐션과 최대 4.3e-07만큼 차이가 나며, 이는 FP32 반올림 오차의 크기입니다. 실제 커널은 쿼리도 타일로 나누어 각 타일을 GPU 스레드 블록 하나에 할당하고, 전체 루프를 단일 커널로 융합하여 타일이 SRAM을 벗어나지 않도록 합니다.

FlashAttention-2, -3, -4#

각 버전은 더 새로운 GPU 세대를 대상으로 합니다. FlashAttention-2(2023)는 스레드 간 작업 분할 방식을 재구성해 기존보다 약 두 배 빠르게 실행되었으며, A100 최대 처리량의 50~73%에 도달했습니다. FlashAttention-3(2024)는 H100의 비동기 Tensor Core를 사용하고 FP8을 지원하여 FP16에서 740 TFLOPs/s, FP8에서 약 1.2 PFLOPs/s에 도달했습니다. FlashAttention-4(2026)는 Blackwell GPU를 대상으로 하며 B200에서 BF16 기준 최대 1,613 TFLOPs/s, 즉 최대 처리량의 71%를 보고합니다. 이 커널은 flash-attention 오픈 소스 라이브러리에 포함되어 있으며 FP16 및 BF16 입력에서 작동합니다. PyTorch의 scaled_dot_product_attention은 GPU와 입력이 지원하는 경우 FlashAttention-2 커널을 자동으로 선택하고, 그렇지 않으면 메모리 효율형 또는 기본 구현으로 대체합니다. Ultralytics 패키지에서 실행되는 Meta의 SAM 2는 Hiera 이미지 인코더와 메모리 어텐션에서 해당 함수를 통해 어텐션을 계산하므로, 하드웨어와 입력 정밀도가 허용하는 경우 PyTorch는 해당 계층을 FlashAttention으로 라우팅합니다.

Flash Attention과 관련 용어 비교#

• 표준 어텐션: 동일한 수학적 연산이지만 전체 점수 행렬과 확률 행렬을 HBM에 기록하여 계산합니다.
• 희소 어텐션: 선택된 쿼리-키 쌍만 점수를 계산합니다. FlashAttention 논문에서는 비어 있는 타일을 건너뛰는 블록 희소 변형도 소개했습니다.
• 선형 어텐션: 소프트맥스를 특징 맵으로 대체해 비용이 시퀀스 길이에 선형적으로 증가하도록 합니다. 이 방식은 모델 출력을 바꿉니다.
• PagedAttention: 서빙 중 KV 캐시를 고정 크기 페이지에 저장하는 메모리 관리자입니다. FlashAttention을 포함한 어텐션 커널은 해당 페이지에서 키와 값을 읽습니다.
• 링 어텐션: 매우 긴 시퀀스를 여러 기기에 분산하고, 각 기기는 자신의 샤드에서 블록 단위 어텐션을 계산하며 키-값 블록은 링을 따라 전달합니다.

학습 및 긴 컨텍스트에 미치는 영향#

메모리가 시퀀스 길이에 선형적으로 증가하기 때문에 FlashAttention은 긴 컨텍스트 윈도우를 실용적으로 학습할 수 있게 했습니다. 원 논문은 MLPerf 1.1 학습 기록에서 시퀀스 길이 512인 BERT-large의 종단 간 속도가 15% 향상되었고, 시퀀스 길이 1,024인 GPT-2는 3배, Long Range Arena 벤치마크에서는 2.4배 빨라졌다고 보고했습니다. Vision Transformer(ViT)나 디퓨전 Transformer가 고해상도 이미지와 비디오를 긴 토큰 시퀀스로 변환하는 비전 분야에도 동일한 절감 효과가 적용됩니다. 더 긴 시퀀스는 정확도도 높였습니다. 논문은 128 × 128 이미지의 픽셀 16,384개를 일렬로 펼쳤을 때 두 점이 경로로 연결되어 있는지 묻는 작업인 Path-X에서 우연 수준을 넘은 최초의 Transformer 결과를 보고했습니다.