플래시 어텐션(FlashAttention) — 어텐션을 메모리로 다시 설계하다 | 기술노트

플래시 어텐션(FlashAttention)은 트랜스포머의 셀프 어텐션을, 결과는 똑같이 내면서 GPU 메모리를 훨씬 덜 오가도록 다시 짠 알고리즘이다. 핵심은 거대한 어텐션 행렬을 통째로 메모리에 쓰지 않는다는 것. 수학은 그대로지만 계산 순서를 바꿔서, 느린 메모리(HBM) 왕복을 줄이고 빠른 온칩 메모리(SRAM) 안에서 최대한 처리한다. 그래서 근사가 아니라 정확히 같은 값을 더 빠르고 더 적은 메모리로 얻는다.

이걸 처음 접했을 때 나는 "어텐션은 원래 O(n²)인데 어떻게 메모리를 줄이지?"가 이해가 안 됐다. 답은 의외로 단순했다. 연산량은 여전히 n²이지만, n×n짜리 중간 행렬을 저장하지 않으면 메모리는 선형으로 줄일 수 있다는 것이다.

진짜 병목은 연산이 아니라 메모리 이동

사람들은 GPU가 곱셈을 못 따라가서 느리다고 오해한다. 어텐션에서는 반대다. 최신 GPU의 연산 유닛은 남아도는데, 데이터를 HBM에서 실어 오고 다시 쓰는 왕복이 병목이다. 이런 걸 메모리 바운드(memory-bound) 연산이라고 부른다.

표준 어텐션은 이렇게 흐른다. 우선 S = QKᵀ 라는 n×n 점수 행렬을 만들어 HBM에 쓴다. 그걸 다시 읽어 softmax를 하고 또 쓴다. 또 읽어서 V와 곱한다. 시퀀스가 길면 이 n×n 행렬 자체가 어마어마하고, 그걸 읽고 쓰는 왕복만으로 시간이 다 간다. 실제로 나는 긴 시퀀스에서 GPU 사용률(util)은 100%인데 체감 속도가 안 나오는 걸 보고 이 개념을 뒤늦게 이해했다.

타일링과 online softmax

플래시 어텐션의 무기는 두 가지다. 첫째는 타일링(tiling). Q·K·V를 작은 블록으로 잘라, 한 블록씩 SRAM으로 불러와 그 안에서 점수 계산부터 V 가중합까지 끝내고 결과 조각만 밖으로 내보낸다. n×n 행렬 전체가 메모리에 앉아 있을 필요가 없어진다.

문제는 softmax다. softmax는 원래 한 행 전체를 봐야 정규화가 되는데, 블록을 쪼개면 전체를 한 번에 못 본다. 여기서 online softmax가 등장한다. 블록을 하나씩 보면서 지금까지 본 최댓값과 누적 합을 들고 다니다가, 새 블록이 더 큰 값을 가져오면 이전 누적치를 비율로 보정한다. 이 재조정 트릭 덕분에 조각내서 계산해도 정확히 같은 결과가 나온다.

# online softmax의 감각 (의사코드)
m = -inf          # 지금까지 본 최댓값
l = 0             # 지수합 누적
acc = 0           # 출력 누적
for block in blocks:
    s = q @ block.k.T                 # 이 블록 점수
    m_new = max(m, s.max())
    p = exp(s - m_new)                # 새 기준으로 지수화
    scale = exp(m - m_new)            # 이전 누적치 보정 비율
    l = l * scale + p.sum()
    acc = acc * scale + p @ block.v   # V 가중합도 함께 보정
    m = m_new
out = acc / l

항목	표준 어텐션	플래시 어텐션
중간 행렬 저장	n×n 전체를 HBM에	저장 안 함(블록 처리)
메모리 사용	O(n²)	O(n)
결과	기준	수치적으로 동일

왜 실무에서 중요한가

긴 컨텍스트가 현실이 된 데에는 이 계열 기법의 공이 크다. n×n 행렬을 안 만드니 메모리 벽에 덜 부딪히고, HBM 왕복이 줄어 학습·추론 둘 다 빨라진다. 요즘은 대부분의 학습 프레임워크와 추론 엔진이 기본으로 이걸 켠다. 우리가 특별히 뭘 안 해도 이미 그 위에서 돌고 있는 셈이다.

플래시 어텐션은 "더 똑똑한 근사"가 아니라 "같은 답을 메모리에 덜 쓰며 얻는 재배열"이다. 정확도를 포기하지 않는다는 점이 이 기법의 미덕이다.

자주 묻는 질문

어텐션의 연산량 O(n²)이 O(n)으로 줄어드는 건가요?

아니다. 흔한 오해다. 곱셈 횟수 자체(연산 복잡도)는 여전히 O(n²)이다. 줄어드는 건 메모리 사용량으로, n×n 중간 행렬을 저장하지 않기 때문에 O(n)이 된다. 속도가 빨라지는 이유는 연산이 줄어서가 아니라 느린 메모리 왕복이 줄어서다.

결과가 근사값인가요, 정확히 같은 값인가요?

정확히 같은 값이다. 타일링과 online softmax는 계산 순서만 바꿀 뿐 수학적으로 동일한 결과를 낸다. 부동소수점 반올림 차이 정도만 있고, 희소 어텐션이나 근사 어텐션처럼 정보를 버리는 방식과는 성격이 다르다.

따로 설치하거나 켜야 하나요?

대개는 아니다. 최신 GPU를 쓰는 주요 학습·추론 스택에서는 기본 경로로 자동 사용되는 경우가 많다. 다만 특정 GPU 아키텍처나 헤드 차원, 데이터 타입 조합에서는 지원 여부가 갈리므로, 성능이 기대만큼 안 나오면 해당 커널이 실제로 활성화됐는지 확인해 볼 가치가 있다.

플래시 어텐션을 쓰면 KV 캐시 문제도 해결되나요?

별개의 문제다. 플래시 어텐션은 어텐션 연산의 메모리·속도를 개선하지만, 생성 단계에서 쌓이는 KV 캐시의 총량 자체를 줄여 주지는 않는다. 캐시 용량은 GQA, 캐시 양자화, PagedAttention 같은 다른 기법으로 다뤄야 한다.