어텐션은 왜 느린가, 그리고 왜 '메모리'가 문제인가
트랜스포머 모델의 심장은 어텐션(attention)이에요. GPT든 클로드든 라마든, 결국 이 연산을 수십 층 쌓아서 만든 거죠. 그런데 이 어텐션에는 태생적인 약점이 있어요. 입력 토큰이 N개면 계산량과 메모리가 N의 제곱으로 늘어난다는 거예요. 토큰 1,000개면 100만 개짜리 행렬, 토큰 32,000개면 10억 개짜리 행렬을 다뤄야 해요. 그래서 2022년 이전만 해도 '긴 문맥'은 사치였어요.
이 벽을 무너뜨린 게 2022년 스탠퍼드의 트리 다오(Tri Dao) 팀이 발표한 FlashAttention이에요. 오늘 소개할 글은 한 개발자가 FlashAttention 논문을 읽으면서 정리한 개인 노트인데요. 논문 자체가 GPU 아키텍처 지식을 꽤 요구해서 진입장벽이 높은 편인데, 이 노트는 그걸 차근차근 풀어놓았어요. 저도 그 흐름을 따라 한국어로 다시 정리해 볼게요.
어텐션이 뭐냐면
이게 뭐냐면, 문장 속 각 단어가 다른 단어들을 '얼마나 참고할지' 점수를 매기는 과정이에요. 예를 들어 「그 고양이는 배가 고파서 밥을 먹었다」에서 '먹었다'는 '고양이'와 '밥'을 많이 참고해야 하고 '그'는 덜 참고해도 되겠죠. 이 점수를 계산하려고 Q(질문), K(열쇠), V(값)라는 세 행렬을 만들고, Q와 K를 곱해서 점수표 S를 만든 다음, 소프트맥스로 확률로 바꾸고(P), 그걸 V에 곱해서 결과를 얻어요.
문제는 중간 산물인 S와 P예요. 크기가 N×N이라서, 토큰이 길어지면 이 행렬을 저장하는 것만으로 GPU 메모리가 터져요. 게다가 표준 구현은 S를 계산해서 메모리에 쓰고, 다시 읽어서 소프트맥스하고, 또 쓰고, 다시 읽어서 V랑 곱해요. 계산 자체보다 메모리를 왔다 갔다 하는 데 시간이 더 걸리는 거죠.
GPU 메모리 계층: 진짜 병목은 여기
FlashAttention의 통찰은 「어텐션은 계산이 느린 게 아니라 메모리 읽기·쓰기가 느리다」는 거예요. GPU에는 두 종류의 메모리가 있어요. 하나는 HBM이라고 부르는 큰 메모리인데, A100 기준 40~80GB짜리고 초당 약 1.5~2TB를 읽을 수 있어요. 다른 하나는 SRAM인데, 각 연산 유닛 바로 옆에 붙어 있는 아주 작은 메모리예요. 크기는 A100 전체를 합쳐도 20MB 정도로 아주 작지만 속도는 HBM보다 10배 이상 빨라요.
비유하자면 HBM은 창고고 SRAM은 작업대예요. 표준 어텐션은 재료를 창고에서 꺼내 작업대에서 한 단계 가공하고, 다시 창고에 넣고, 또 꺼내서 다음 단계를 하는 식이었어요. FlashAttention은 「작업대에 올릴 수 있는 만큼만 재료를 잘라서 가져오고, 작업대 위에서 끝까지 가공한 다음 완성품만 창고에 넣자」는 전략이에요. 이걸 'IO-aware', 그러니까 입출력을 의식한 알고리즘이라고 불러요.
핵심 기법 1: 타일링(Tiling)
N×N 행렬을 한 번에 만들지 않고, Q와 K, V를 작은 블록으로 잘라요. 예를 들어 64행씩요. 그리고 Q 블록 하나와 K 블록 하나를 SRAM에 올려서 그 부분의 점수만 계산하고, 바로 V 블록과 곱해서 결과에 누적해요. 이렇게 하면 전체 S 행렬이 HBM에 존재한 적이 없어요. 메모리 사용량이 N의 제곱에서 N에 비례하는 수준으로 떨어지는 거죠.
핵심 기법 2: 온라인 소프트맥스(Online Softmax)
그런데 여기서 걸리는 게 있어요. 소프트맥스는 한 행 전체의 최댓값과 합계를 알아야 계산할 수 있거든요. 블록 단위로 잘라서 보면 전체를 모르잖아요? 이걸 해결하는 트릭이 온라인 소프트맥스예요. 블록을 하나씩 처리하면서 '지금까지의 최댓값'과 '지금까지의 합계'를 계속 갱신하고, 새 블록에서 더 큰 값이 나오면 이전까지 누적한 결과를 보정 계수로 다시 스케일링해요. 수학적으로는 exp(a-m1)을 exp(a-m2)로 바꾸려면 exp(m1-m2)만 곱하면 된다는 성질을 이용한 거예요. 스트리밍으로 평균을 구하는 것과 비슷한 발상이죠.
핵심 기법 3: 역전파에서 재계산(Recomputation)
학습할 때는 역전파를 위해 순전파의 중간값 P가 필요해요. 표준 방식은 이걸 저장해 두는데, 그게 또 N×N이에요. FlashAttention은 저장하는 대신 역전파 때 Q, K와 저장해 둔 소프트맥스 통계값(최댓값, 합계)만 가지고 P를 다시 계산해요. 계산은 조금 늘지만 HBM 접근이 훨씬 줄어서 결과적으로 더 빨라요. '계산은 싸고 메모리는 비싸다'는 현대 GPU의 특성을 정확히 이용한 거예요.
업계 맥락: FA2, FA3, 그리고 이후
FlashAttention 1은 GPT-2 학습을 3배 빠르게 만들었고, 이후 FlashAttention-2(2023)는 병렬화 구조를 바꿔서 A100에서 이론 최대치의 70% 수준까지 끌어올렸어요. FlashAttention-3(2024)는 H100의 비동기 연산과 FP8을 활용해서 다시 1.5~2배를 더 냈고요. 지금은 PyTorch에 scaled_dot_product_attention이라는 이름으로 내장되어 있어서, 여러분이 모르는 사이에 이미 쓰고 있을 가능성이 높아요.
비슷한 시기에 나온 xFormers의 메모리 효율 어텐션도 같은 아이디어를 구현했고, vLLM의 PagedAttention은 추론 시 KV 캐시를 운영체제의 페이징처럼 관리해서 서빙 효율을 높였어요. 방향은 조금씩 다르지만, 「어텐션의 병목은 메모리다」라는 같은 깨달음에서 출발한 기술들이에요.
한국 개발자에게 주는 시사점
LLM을 파인튜닝하거나 서빙하는 분이라면 FlashAttention을 켜는 것만으로 메모리와 속도가 크게 달라져요. Hugging Face에서 attn_implementation 옵션 하나로 켤 수 있고, vLLM이나 SGLang은 기본으로 사용해요. 그리고 더 중요한 건 사고방식이에요. 「알고리즘 복잡도보다 메모리 접근 패턴이 성능을 좌우한다」는 건 GPU뿐 아니라 CPU 캐시, 데이터베이스, 분산 시스템에서도 똑같이 적용되는 원칙이거든요. FlashAttention은 그 원칙을 가장 극적으로 보여준 사례예요.
마무리
한줄 정리: FlashAttention은 어텐션 연산을 새로 발명한 게 아니라, GPU 메모리 계층에 맞게 '계산 순서'를 재배치해서 N×N 행렬을 메모리에 쓰지 않고도 같은 결과를 얻는 기법이에요.
여러분은 성능 최적화를 하다가 '알고리즘이 아니라 메모리가 문제였다'는 걸 깨달은 경험이 있으신가요? 어떤 상황이었는지 댓글로 나눠주세요.
🔗 출처: Hacker News