플래시 어텐션 원리부터 트라이톤 GPU 커널 구현까지: 어텐션 메모리 최적화 완전 정리

핵심 메시지

• 플래시 어텐션은 계산 결과를 바꾸지 않으면서, 느린 GPU 글로벌 메모리(HBM) 접근을 줄여 어텐션 속도를 끌어올리는 기법이다.
• 어텐션 연산은 계산량보다 메모리 입출력이 병목인 'IO 바운드' 문제이며, 작고 빠른 공유 메모리(SRAM) 안에서 블록 단위로 계산하는 것이 핵심이다.
• 소프트맥스를 수치적으로 안정화하고, 블록을 이어 붙이며 정규화 값을 보정하는 '온라인 소프트맥스'가 이 알고리즘을 가능하게 한다.
• 트라이톤(Triton)은 파이썬 코드를 GPU 커널로 컴파일해 주므로, C++이나 CUDA를 직접 쓰지 않고도 커널을 작성할 수 있다.
• 순전파뿐 아니라 역전파까지 구현하려면 도함수·그래디언트·야코비안 같은 미분 개념으로 소프트맥스와 행렬곱의 기울기를 직접 유도해야 한다.

쉽게 이해하기

이 강의는 '제1원리(first principles)' 방식으로 플래시 어텐션을 다룬다. 강사는 플래시 어텐션 논문이 아예 존재하지 않는다고 가정하고, 표준 어텐션 계산을 들여다보며 그 안에 어떤 문제가 있는지 찾아낸 뒤 단계적으로 해결해 나간다. 목표는 원리를 깊이 이해하는 동시에 실제로 코드까지 작성해 보는 것이다.

먼저 멀티헤드 어텐션을 복습한다. 핵심 공식은 소프트맥스(Q와 K의 전치를 곱하고 헤드 차원의 제곱근으로 나눈 값)에 V를 곱하는 것이다. 입력 토큰들을 WQ·WK·WV 세 개의 선형 투영으로 쿼리·키·값으로 바꾸고, 이를 여러 헤드로 나눠 각 헤드에서 독립적으로 어텐션을 계산한 뒤 다시 이어 붙인다.

문제의 핵심은 GPU 메모리 구조다. HBM(글로벌 메모리)은 A100 기준 40GB처럼 크지만 접근이 매우 느리고, 공유 메모리(SRAM)는 훨씬 작지만 아주 빠르다. 플래시 어텐션 논문의 주장은 어텐션이 'IO 바운드' 연산이라는 것이다. 즉 느린 이유는 계산 자체가 아니라 느린 글로벌 메모리를 계속 드나들기 때문이며, 해결책은 어텐션을 빠른 공유 메모리 안에서 블록 단위로 계산하는 것이다.

이를 가능하게 하는 두 가지 도구가 수치적으로 안정한 소프트맥스와 온라인 소프트맥스다. 소프트맥스는 분자와 분모에 같은 상수를 곱해도 값이 변하지 않는 성질을 이용해 지수 값이 폭발하거나 사라지지 않도록 안정화한다. 온라인 소프트맥스는 전체 행렬을 한꺼번에 올리지 않고, 블록을 하나씩 처리하면서 이전 블록의 정규화 값을 지수 트릭으로 보정해 정확한 소프트맥스를 점진적으로 계산한다.

이후 강의는 GPU와 CPU의 차이, 커널, 텐서의 메모리 배치(행/열 우선, 스트라이드), 블록 행렬 곱셈, 트라이톤의 파이프라인 최적화를 거쳐 순전파를 구현한다. 여기서 멈추지 않고 역전파까지 구현하는데, 커스텀 연산이라 자동미분에 기댈 수 없으므로 도함수·그래디언트·야코비안을 이해하고 소프트맥스와 행렬곱의 기울기를 직접 유도한다.

주요 인사이트

• 성능 병목을 다시 정의한다: 어텐션이 느린 이유는 O(n^2) 계산 자체보다, 커다란 중간 행렬을 느린 메모리에 반복해서 읽고 쓰기 때문이다. 이 관점 전환이 플래시 어텐션의 출발점이다.
• 플래시 어텐션은 근사가 아니라 정확한 재구성이다. 계산 순서와 메모리 사용 방식만 바꿀 뿐, 표준 어텐션과 같은 결과를 얻는다.
• 온라인 소프트맥스 덕분에 전체 점수 행렬을 메모리에 올리지 않고도 블록 단위로 소프트맥스를 정확히 계산할 수 있다. 이것이 공유 메모리 안 계산을 실현하는 열쇠다.
• 트라이톤은 개별 스레드가 아니라 스레드 그룹(블록) 단위로 파이썬에서 로직을 기술하게 해, GPU 커널 작성의 진입장벽을 크게 낮춘다.
• 역전파는 파이토치의 자동미분이 모르는 커스텀 연산이므로, 소프트맥스와 행렬곱의 기울기를 손으로 유도해야 한다. 이 과정이 딥러닝 내부 동작을 깊이 이해하게 해 준다.

자주 묻는 질문

관련 AI 소식