---
title: "FlashAttention: 어텐션을 빠르게 만든 건 계산량보다 데이터 이동이었다"
author: "Dne Cna"
published_at: "2026-09-30T22:00:04.271Z"
updated_at: "2026-09-30T08:00:35.577Z"
tags: ["ai","머신러닝","논문리딩","flashattention"]
canonical_url: "https://ariser.io/@writer_8706b598dfee/flashattention%EC%9D%84-%EC%9D%BD%EA%B3%A0-%EC%96%B4%ED%85%90%EC%85%98%EC%9D%98-%EB%B3%91%EB%AA%A9%EC%9D%80-%EB%A9%94%EB%AA%A8%EB%A6%AC-%EC%9D%B4%EB%8F%99%EC%97%90-%EC%9E%88%EC%97%88%EB%8B%A4"
ai_assisted: false
---

# FlashAttention: 어텐션을 빠르게 만든 건 계산량보다 데이터 이동이었다

[원문](https://ariser.io/@writer_8706b598dfee/flashattention%EC%9D%84-%EC%9D%BD%EA%B3%A0-%EC%96%B4%ED%85%90%EC%85%98%EC%9D%98-%EB%B3%91%EB%AA%A9%EC%9D%80-%EB%A9%94%EB%AA%A8%EB%A6%AC-%EC%9D%B4%EB%8F%99%EC%97%90-%EC%9E%88%EC%97%88%EB%8B%A4)

GPU가 바쁘게 계산하는데 학습은 좀처럼 빨라지지 않는 순간이 있다. 수식에서 연산량을 줄이면 실행 시간도 줄어들 것 같지만, 실제 장치는 숫자를 계산하기 전에 그 숫자를 가져와야 한다. 어텐션에서는 이 이동이 뜻밖에 큰 비용이 된다. FlashAttention을 읽을 때 재미있는 지점도 여기다. 익숙한 어텐션 수식을 그대로 두고, 수식을 실행하는 순서와 중간 결과의 거처를 바꿨더니 속도와 메모리 사용량이 함께 달라진다.

## N×N 행렬은 꼭 저장해야 할까

표준 어텐션은 대략 `O = softmax(QKᵀ / √d)V`로 쓸 수 있다. 길이가 N인 시퀀스라면 Q와 K를 곱한 점수 행렬은 N×N이다. 구현을 수식 순서대로 나누면 점수 행렬을 만들고, 소프트맥스 결과를 만들고, 다시 V와 곱한다. 각 단계의 결과를 GPU의 큰 메모리에 저장했다가 다음 단계에서 읽어 오는 식이다.

여기서 큰 메모리와 빠른 메모리는 다르다. HBM은 용량이 크지만 연산 장치 가까이 있는 SRAM보다 데이터 접근에 더 많은 비용이 든다. 창고에서 작업대로 재료를 가져오는 장면을 생각하면 이해하기 쉽다. 작업대에서 덧셈 몇 번을 더 하는 것보다, 이미 만든 커다란 표를 창고에 넣었다가 꺼내는 일이 더 비쌀 수 있다.

저자들의 핵심 선택은 Q, K, V를 작은 블록으로 나누어 SRAM에 가져온 뒤, 점수 계산과 정규화와 가중합을 이어서 처리하는 것이다. 전체 N×N 점수 행렬을 HBM에 쓰지 않는다. 이때 정확한 어텐션을 유지하려면 블록마다 계산한 소프트맥스를 올바르게 합쳐야 한다. 그 방법이 이 논문의 알고리즘적인 중심이다. [FlashAttention 원문](https://arxiv.org/abs/2205.14135)

## 소프트맥스는 나누어 계산할 수 있다

소프트맥스의 분모는 한 행의 모든 점수를 필요로 한다. 따라서 첫 블록만 보고 확률을 확정할 수는 없다. 대신 지금까지 본 점수의 최댓값, 그 최댓값을 기준으로 계산한 지수합, 정규화에 필요한 출력 누적값을 보관한다. 새 블록에 더 큰 점수가 나타나면 이전 누적값을 새 기준에 맞춰 재조정한다. 마지막 블록까지 보면 전체 행에 대해 계산한 결과와 같은 수학적 결과가 나온다.

아주 작은 예를 직접 만들어 보자. 한 쿼리가 네 키에 준 점수가 1, 2, 3, 4라고 하자. 앞의 두 점수와 뒤의 두 점수를 각각 소프트맥스한 다음 결과를 이어 붙이면 틀린다. 각 묶음에서 확률의 합이 1이므로 전체 합은 2가 되기 때문이다. 필요한 것은 묶음별 확률을 붙이는 일이 아니라, 두 묶음의 지수합을 같은 기준에서 합치는 일이다.

이 예에서 첫 묶음은 최댓값 2를 기준으로 계산할 수 있다. 두 번째 묶음을 만나면 전체 최댓값이 4로 바뀐다. 그러면 첫 묶음의 지수합과 가중합을 `exp(2−4)`만큼 줄이고, 최댓값 4를 기준으로 계산한 두 번째 묶음을 더한다. 분모 전체를 저장하지 않아도 분모에 필요한 정보는 남는다. 큰 행렬을 없애는 비결이 정보를 버리는 데 있는 것이 아니라, 나중에 합칠 수 있는 형태로 정보를 압축해 유지하는 데 있는 셈이다.

여기서 '정확하다'는 표현은 어텐션을 희소화하거나 근사 커널로 바꾸지 않는다는 뜻이다. 부동소수점 연산은 순서에 따라 작은 오차가 달라질 수 있으므로, 기존 구현과 모든 비트가 같다는 주장으로 읽으면 곤란하다. 공식 구현도 기준 구현과의 수치적 차이를 검사한다. [공식 구현과 테스트 안내](https://github.com/Dao-AILab/flash-attention)

## 계산을 다시 하는 편이 싸다는 역설

학습에서는 역전파도 생각해야 한다. 순전파에서 중간 행렬을 모두 저장해 두면 역전파 때 꺼내 쓸 수 있지만, 그 보관 비용이 크다. FlashAttention은 필요한 정규화 정보를 남기고, 역전파 중 작은 블록의 어텐션 값을 다시 계산한다. 그래서 연산 횟수가 늘어나는 부분이 있어도 HBM 왕복을 줄여 실제 실행 시간을 줄일 수 있다.

이 대목은 최적화를 볼 때 우리가 무엇을 세는지 묻는다. FLOPs는 계산량을 비교하는 유용한 지표지만, 메모리 계층이나 커널 호출 비용을 대신 설명해 주지는 않는다. 종이에 적힌 알고리즘에서 중간 결과 하나를 재사용한다고 해도, 장치가 그 결과를 멀리서 가져와야 한다면 재사용이 공짜가 아니다. 반대로 재계산은 항상 낭비라는 직관도 성립하지 않는다.

나는 이 선택을 모델 연구와 시스템 연구가 만나면서 생기는 좋은 긴장으로 읽었다. 학습 알고리즘은 '무엇을 계산해야 하는가'를 정하고, 실행 알고리즘은 '어떤 값이 어느 순간 어디에 있어야 하는가'를 정한다. 둘을 함께 보아야 같은 모델을 더 싸게 학습할 길이 열린다. 이 문장은 논문의 새로운 실험 결과라기보다, 논문에서 얻은 나의 해석이다.

## 속도 그래프에서 놓치기 쉬운 것

원문은 어텐션 커널의 실행 시간뿐 아니라 Transformer의 전체 학습 시간도 비교한다. BERT, GPT-2, 긴 시퀀스 벤치마크에서 개선을 보고하고, 더 긴 입력으로 학습할 수 있게 되었을 때 모델 품질이 달라지는 사례도 제시한다. 다만 커널이 빨라진 비율과 전체 학습이 빨라진 비율은 같지 않다. 어텐션 밖의 연산, 데이터 로딩, 통신이 차지하는 비중이 있기 때문이다. [원문 실험과 한계](https://arxiv.org/html/2205.14135v2#S4)

긴 문맥에서 품질이 좋아졌다는 결과도 구분해서 읽어야 한다. 같은 입력과 가중치에 FlashAttention을 적용했다고 지능이 추가되는 것은 아니다. 메모리 비용을 줄여 더 긴 입력으로 학습하는 선택이 가능해졌고, 그 선택이 일부 과제에서 도움이 된 것이다. 실행 최적화가 모델을 직접 바꾸지 않더라도, 실험 가능한 범위를 바꾸면 모델 연구의 결과에 영향을 줄 수 있다.

반면 연산량의 제곱 증가 자체가 사라진 것은 아니다. N이 늘면 모든 쿼리와 키 사이의 상호작용을 여전히 계산한다. 어텐션 중간 행렬을 저장하지 않는 것과, 긴 입력의 계산 비용이 선형이 되는 것은 다른 주장이다. 전체 학습 메모리에는 가중치, 옵티마이저 상태, 다른 층의 활성값도 남아 있다.

## 실제로 적용한다면 먼저 재볼 것

내가 이 아이디어를 적용한다면 '지원되는가'와 '우리 작업에서 이득인가'를 따로 확인하겠다. GPU 종류, 데이터 타입, 헤드 차원, 마스크 형태가 구현의 지원 범위에 맞는지 확인한 다음, 실제 시퀀스 길이와 배치 크기로 측정한다. 공식 저장소는 후속 버전도 함께 다루므로 2022년 원문의 결과와 현재 라이브러리의 기능을 같은 것으로 취급하지 않는 편이 좋다.

측정표에는 커널 시간보다 전체 스텝 시간, 최대 메모리, 처리한 토큰 수를 넣고 싶다. 같은 배치에서 빨라졌는지, 남는 메모리로 배치를 늘려 총 처리량이 개선됐는지도 나누어 본다. 두 변화는 운영 방식에 서로 다른 의미를 가진다. 메모리 부족으로 실험을 접었던 사람에게는 최대 입력 길이 자체가 더 중요한 성과일 수 있다.

FlashAttention이 남기는 질문은 간단하고 오래간다. 지금 반드시 저장한다고 생각하는 중간 결과가, 사실은 최종 결과를 얻기 위해 잠깐만 필요한 값일 수 있지 않을까. 그 질문을 제대로 던지는 순간, 수식은 그대로여도 프로그램은 상당히 달라질 수 있다.

읽은 논문: Tri Dao 외, *FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness*, NeurIPS 2022. [논문](https://arxiv.org/abs/2205.14135) · [공식 코드](https://github.com/Dao-AILab/flash-attention)