FlashAttention: GPU 메모리 계층을 정복하여 어텐션 효율을 새롭게 정의하다

FlashAttention: GPU 메모리 계층을 정복한 어텐션 효율의 혁신

트랜스포머 모델이 인공지능, 특히 자연어 처리(NLP) 분야의 혁신을 이끌면서, 그 심장과도 같은 '어텐션(Attention)' 메커니즘의 중요성은 날로 커지고 있습니다. 어텐션은 모델이 입력 시퀀스 내 다양한 정보 조각들에 선택적으로 집중하여 정교한 문맥 이해를 가능하게 하는 핵심 기술입니다. 하지만 이 강력한 능력의 이면에는 시퀀스 길이가 길어질수록 계산량과 메모리 사용량이 눈덩이처럼 불어나는 문제가 숨어 있었습니다. 이는 모델 학습 및 추론 속도를 저해하고, 더 긴 문맥을 처리하는 데 있어 명확한 한계점으로 작용했습니다. 이러한 난제를 해결하기 위해 혜성처럼 등장한 'FlashAttention'은 마치 숙련된 외과 의사가 막힌 혈관을 정교하게 뚫어내듯, 어텐션 연산의 효율성을 극적으로 끌어올린 최적화 알고리즘입니다. 기존 어텐션 방식의 고질적인 병목 현상을 GPU 메모리 계층 구조의 특성을 날카롭게 파고들어 해결함으로써, 더 빠르고, 더 적은 메모리로, 정확한 어텐션 연산을 수행하는 시대를 열었습니다.

---

기존 어텐션의 숨 막히는 비효율성: 느린 메모리와의 사투

FlashAttention의 혁신성을 제대로 이해하려면, 기존 표준 어텐션 메커니즘이 왜 비효율적이었는지 살펴볼 필요가 있습니다.

개념: 표준 어텐션은 쿼리(Query, Q), 키(Key, K), 밸류(Value, V)라는 세 개의 행렬을 활용하여 입력 시퀀스 내 요소들 간의 상호 연관성을 계산합니다. 이 과정의 핵심은 Q와 K 행렬을 곱하여 생성되는 '어텐션 스코어(Attention Score)' 행렬인데, 이 행렬의 크기는 시퀀스 길이에 따라 매우 커질 수 있습니다. 예를 들어 시퀀스 길이가 N이라면, 이 행렬은 N x N 크기가 됩니다.

문제점:

• 메모리 병목 현상: 가장 큰 문제는 이 거대한 어텐션 스코어 행렬을 GPU의 주 메모리인 HBM(High Bandwidth Memory)에 쓰고 다시 읽어오는 과정에서 발생합니다. 왼쪽 '메모리 계층(Memory Hierarchy)'을 보면, GPU 내부의 초고속 캐시 메모리인 SRAM은 초당 19 테라바이트(TB/s)라는 엄청난 속도를 자랑하지만 용량은 20 메가바이트(MB)로 매우 작습니다. 반면, HBM은 초당 1.5 테라바이트의 속도에 40 기가바이트(GB)의 비교적 큰 용량을 가집니다. 표준 어텐션은 마치 비좁은 작업대(SRAM)에서 감당하기 어려운 거대한 부품(어텐션 스코어 행렬)을 상대적으로 느린 창고(HBM)에 계속 옮겨가며 작업하는 것과 같습니다. 시퀀스가 길어질수록 이 중간 행렬은 HBM의 대역폭을 한계까지 몰아붙여 전체 연산 속도를 현저히 저하시킵니다.
• 제곱에 비례하는 복잡도: 어텐션 스코어 행렬의 크기와 계산량은 시퀀스 길이에 제곱으로 비례합니다. 쉽게 말해, 시퀀스 길이가 두 배 늘어나면 메모리 사용량과 계산량은 네 배로 폭증하는 것입니다. 이는 긴 문맥 정보를 처리해야 하는 과제에서 치명적인 약점으로 작용했습니다.

비유: 기존 어텐션 방식은 마치 거대한 도서관(HBM)에서 아주 두꺼운 참고서(어텐션 스코어 행렬) 전체를 매번 꺼내와, 작은 독서대(SRAM) 위에서 한 페이지씩 넘겨 내용을 확인하고(계산), 다시 도서관 서가에 꽂아 넣는 작업을 반복하는 사서와 같습니다. 참고서가 두꺼울수록(시퀀스 길이가 길수록) 도서관을 오가는 시간과 노력이 기하급수적으로 늘어나는 것이죠.

---

FlashAttention의 마법: GPU 메모리 계층을 지배하다

FlashAttention은 이러한 비효율을 해결하기 위해 GPU 메모리 계층 구조를 정교하게 활용하는 IO-인식(IO-Awareness) 알고리즘을 제시합니다. 핵심 아이디어는 "느린 HBM으로의 불필요한 데이터 이동을 최소화하고, 대부분의 계산을 번개처럼 빠른 SRAM 내에서 완료하자"입니다.

1. 타일링 (Tiling): 잘게 쪼개어 빠르게 정복한다

• 개념: 입력 행렬 Q, K, V를 GPU의 훨씬 빠른 온칩 메모리(SRAM)에 효율적으로 적재할 수 있는 작은 블록(타일) 단위로 분할합니다.
• 작동 방식: 'FlashAttention'을 보면, K와 V 행렬의 블록들이 먼저 SRAM으로 복사되고(Copy Block to SRAM), 이후 Q 행렬의 블록이 SRAM으로 들어와 어텐션 계산의 핵심 부분(Q와 K 블록 간의 곱셈으로 스코어 계산, 소프트맥스 적용, 결과적으로 V 블록과의 가중 합)이 전적으로 SRAM 내에서 수행됩니다(Compute Block on SRAM). 이 과정은 Q, K, V 행렬의 모든 블록에 대해 반복되며, 최종 결과 블록만이 HBM으로 출력됩니다(Output to HBM). 덕분에 거대한 전체 어텐션 스코어 행렬을 느린 HBM에 기록하거나 다시 읽어올 필요가 없어집니다. 소프트맥스 계산 시 수치적 안정성을 확보하기 위해 블록 단위로 영리하게 처리하는 기법이 사용됩니다.
• 효과: HBM으로의 읽기/쓰기 횟수를 획기적으로 줄여 메모리 병목 현상을 크게 완화합니다. 이는 마치 두꺼운 참고서 전체를 옮기는 대신, 필요한 핵심 내용이 담긴 몇 쪽(블록)만 작은 메모장(SRAM)에 빠르게 복사해와 내용을 파악하고, 그 결과만 간결하게 정리하여 기록하는 것과 같습니다.
• 예시: 만약 시퀀스 길이가 8이고 각 단어가 64차원으로 표현된다고 할 때, FlashAttention은 이 데이터들을 SRAM에 들어갈 수 있는 작은 덩어리(예: 4x64 크기의 블록)로 나누어 처리합니다. 이렇게 하면 전체 8x8 크기의 거대한 중간 계산 결과를 HBM에 저장하지 않고도 효율적인 계산이 가능해집니다.

2. 커널 퓨전 (Kernel Fusion): 흩어진 작업을 하나로 모아 효율 극대화

• 개념: 표준 어텐션에서는 행렬 곱셈, 드롭아웃, 소프트맥스, 마스킹 등의 연산이 각각 별도의 GPU 작업 단위(커널)로 실행되어 여러 번의 HBM 접근을 유발합니다. FlashAttention은 이러한 여러 개별 연산들을 하나의 최적화된 '융합 커널(Fused Kernel)'로 통합하여 실행합니다.
• 효과: 여러 번에 걸쳐 HBM을 읽고 쓰는 대신, 한 번의 데이터 로딩으로 여러 연산을 SRAM 내에서 연속적으로 처리함으로써 메모리 접근 횟수를 줄이고 연산 속도를 높입니다. 'GPT-2 어텐션 성능 비교(Attention on GPT-2)'를 보면, PyTorch의 표준 어텐션이 여러 단계(Matmul, Dropout, Softmax, Mask 등)로 나뉘어 총 17밀리초(ms) 정도 소요되는 반면, FlashAttention은 이를 단일 '융합 커널(Fused Kernel)'로 통합하여 처리 시간을 약 2.5밀리초로 대폭 단축시키는 것을 확인할 수 있습니다.
• 비유: 여러 작은 가게에서 각각 다른 물건을 사기 위해 여러 번 줄을 서고 계산하는 번거로움 대신, 모든 물건을 한 곳에 모아둔 대형 할인점에서 한 번의 동선으로 쇼핑을 마치고 한 번에 계산하는 것과 유사합니다. 작업의 효율성이 극대화되는 것이죠.

3. 역방향 전파 최적화 (재계산, Recomputation): 똑똑하게 버리고 필요할 때 다시 만든다

• 개념: 트랜스포머 모델을 훈련시킬 때는 순방향 계산뿐만 아니라 역방향 전파(그래디언트 계산) 과정도 필수적입니다. 기존 방식에서는 이 역방향 계산을 위해 순방향에서 생성된 거대한 어텐션 스코어 행렬을 HBM에 고스란히 저장해두어야 했습니다. 이는 엄청난 메모리를 소모하는 주된 원인이었습니다. FlashAttention은 이 큰 어텐션 행렬을 저장하는 대신, 역방향 전파에 필요한 부분을 SRAM 내에서 즉석으로 다시 계산하는(recompute) 방식을 채택합니다.
• 효과: 약간의 추가 계산 비용을 감수하는 대신, 메모리 사용량을 획기적으로 줄입니다. 이로 인해 시퀀스 길이에 대한 메모리 복잡도가 기존의 '제곱으로 증가하는 방식'에서 '선형적으로, 즉 시퀀스 길이에 정비례하여 증가하는 방식'으로 극적으로 감소합니다.
• 비유: 중요한 시험을 준비할 때, 교과서의 모든 내용을 처음부터 끝까지 노트에 빼곡히 옮겨 적어 엄청난 양의 필기 노트를 만드는 것(메모리 많이 차지)이 아니라, 핵심 요약과 키워드(작은 데이터)만 간추려 정리해두고, 실제 문제를 풀 때 그 요약을 바탕으로 세부 내용을 머릿속에서 논리적으로 재구성하여(재계산) 답을 찾는 스마트한 학습 전략과 같습니다.

---

FlashAttention이 가져온 놀라운 변화

FlashAttention의 등장은 어텐션 연산 효율성에 다음과 같은 혁신적인 개선을 가져왔습니다.

• 경이로운 속도 향상: 기존 방식 대비 최대 서너 배의 속도 향상을 기록했으며, 특정 모델에서는 전체 학습 속도도 눈에 띄게 빨라졌습니다. 최신 버전은 하드웨어의 이론적 최대 연산 능력에 근접하는 높은 활용률을 보이기도 합니다.
• 획기적인 메모리 절감: 시퀀스 길이에 대한 메모리 사용량이 크게 줄어들어, 긴 시퀀스에서도 약 10배에서 20배에 달하는 놀라운 메모리 절감 효과를 보였습니다. 이는 훨씬 긴 시퀀스도 제한된 GPU 메모리 내에서 효율적으로 처리할 수 있게 되었음을 의미합니다.
• 더 긴 컨텍스트 처리 능력의 확장: 메모리 제약이 크게 완화되면서 모델이 이전에는 상상하기 어려웠던 훨씬 긴 시퀀스를 처리할 수 있게 되었습니다. 이는 장문의 문서를 이해하고 요약하거나, 여러 번 오고 간 긴 대화를 기억하고 일관성 있는 답변을 생성하는 등, 더욱 복잡하고 긴 문맥 정보를 요구하는 작업에서 모델 성능을 비약적으로 향상시키는 결정적인 계기가 되었습니다. 실제로 대규모 언어 모델(LLM)의 처리 가능 문맥 길이는 지난 몇 년간 크게 늘어났는데, FlashAttention이 이러한 발전에 지대한 공헌을 했습니다.
• 새로운 모델 아키텍처 및 기능 탐색의 문 활짝: 어텐션 연산의 효율성 향상은 더 복잡하고 정교한 모델 아키텍처를 실험하거나, 이미지나 음성과 같은 다른 유형의 매우 긴 데이터를 사용하는 것을 현실적으로 가능하게 만들었습니다.

---

마무리하며: 어텐션의 미래를 밝히다

FlashAttention은 단순한 최적화 기법을 넘어, GPU 하드웨어의 메모리 계층 구조라는 본질을 깊이 이해하고 이를 알고리즘 설계에 창의적으로 반영함으로써 어텐션 연산의 패러다임을 성공적으로 전환시킨 기념비적인 혁신입니다. 타일링, 커널 퓨전, 재계산과 같은 핵심 기법들을 통해 HBM으로의 느린 접근을 최소화하고 SRAM의 초고속 처리 능력을 극대화함으로써, 어텐션의 정확성은 그대로 유지하면서도 속도와 메모리 효율성을 전례 없는 수준으로 끌어올렸습니다. 이러한 눈부신 발전은 대규모 언어 모델(LLM)이 더 광범위한 문맥을 이해하고, 더 복잡한 문제를 해결하며, 이전에는 불가능했던 새로운 기능을 선보일 수 있는 탄탄한 기반을 마련해주었습니다. FlashAttention은 현대 AI 모델, 특히 트랜스포머 기반 아키텍처의 지속적인 발전에 있어 없어서는 안 될 핵심 동력으로 확고히 자리매김했으며, 앞으로도 인공지능 기술의 지평을 넓히는 데 중추적인 역할을 수행할 것으로 기대됩니다.

---

https://arxiv.org/abs/2205.14135

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness