안녕하세요, Shrijith Venkatramana입니다. 저는 모든 커밋에서 실행되는 AI 코드 리뷰어인 git-lrc를 구축하고 있습니다. 개발자들이 이 프로젝트를 발견할 수 있도록 Star Us를 눌러주세요. 꼭 한번 사용해 보시고 제품 개선을 위한 피드백을 공유해 주시기 바랍니다.

대규모 언어 모델 (Large language models)은 계속해서 커지고 있습니다.

컨텍스트 윈도우 (Context windows)는 수천 개의 토큰에서 수십만 개로 성장했으며, 일부 모델은 이제 수백만 개의 토큰으로 측정되는 컨텍스트 길이를 광고하고 있습니다.

하지만 수년 동안 Transformer의 한 부분이 병목 현상 (bottleneck)이 될 위협을 가해왔습니다:

Attention (어텐션).

수학적 계산이 너무 많이 필요해서가 아닙니다.

데이터를 너무 많이 이동시키기 때문입니다.

FlashAttention은 바로 그 문제를 해결하기 때문에 현대 AI 인프라에서 가장 중요한 최적화 중 하나입니다. 이는 Transformer 아키텍처를 변경하거나, 어텐션을 근사화(approximate attention)하거나, 새로운 모델 설계를 도입하지 않습니다. 대신, GPU에서 어텐션이 어떻게 실행되는지를 재고합니다.

그 결과로 메모리 사용량이 극적으로 감소하고, 더 빠른 학습 (training), 더 빠른 추론 (inference), 그리고 실용적인 긴 컨텍스트 (long-context) 모델이 가능해졌습니다.

어떻게 작동하는지 살펴보겠습니다.

1. 어텐션의 실제 비용

모든 Transformer 레이어는 세 가지 행렬을 사용하여 어텐션을 계산합니다:

• Query (Q)
• Key (K)
• Value (V)

표준 어텐션 방정식은 다음과 같습니다:

Attention(Q,K,V) = softmax(QK^T / sqrt(d))V

개념적으로, 모든 토큰은 자신을 다른 모든 토큰과 비교합니다.

N개의 토큰 시퀀스의 경우, 어텐션 스코어 행렬 (attention score matrix)은 다음 요소를 포함합니다:

N x N

16,384개의 토큰 시퀀스는 다음과 같은 어텐션 스코어를 필요로 합니다:

16,384 x 16,384 = 2억 6,800만 개

보통 첫 번째 본능은 다음과 같습니다:

"계산량이 정말 많구나."

하지만 현대 GPU에서는 메모리 트래픽 (memory traffic)이 더 큰 문제인 경우가 많습니다.

2. GPU는 연산 제한(Compute-Bound)이 아닌 메모리 제한(Memory-Bound)인 경우가 많다

현대 GPU는 초당 엄청난 수의 부동 소수점 연산 (floating-point operations)을 수행할 수 있습니다.

하지만 거의 효율적으로 수행하지 못하는 것은 메모리 계층 간에 방대한 양의 데이터를 이동시키는 일입니다.

단순화된 GPU 메모리 계층 구조는 다음과 같습니다:

HBM (GPU 메모리)
    |
L2 Cache
...

데이터가 연산 유닛(compute units)에서 멀어질수록, 데이터에 접근하는 비용은 더 커집니다.

기존의 어텐션(Attention) 방식은 전체 어텐션 행렬(attention matrix)을 HBM으로부터 반복적으로 쓰고 읽습니다.

워크플로우는 대략 다음과 같습니다:

Q 로드 (Load Q)
K 로드 (Load K)

...

유용한 연산을 수행하기보다는 데이터를 이동시키는 데 엄청난 양의 시간이 소비됩니다.

이것이 바로 FlashAttention이 해결하고자 하는 문제입니다.

3. 핵심 통찰: 어텐션 행렬을 실체화하지 마라 (Never Materialize the Attention Matrix)

FlashAttention의 이면에 있는 핵심 관찰은 놀라울 정도로 간단합니다:

전체 어텐션 행렬을 실제로 저장할 필요가 없다는 것입니다.

기존의 어텐션은 다음과 같이 명시적으로 생성합니다:

QK^T

그리고 이를 메모리에 기록합니다.

FlashAttention은 결코 그렇게 하지 않습니다.

대신, 어텐션을 작은 블록(blocks) 단위로 처리합니다.

개념적으로는 다음과 같습니다:

Q 블록 1 x K 블록 1
Q 블록 1 x K 블록 2
Q 블록 1 x K 블록 3
...

각 블록은 로드되어 처리되고, 최종 결과에 기여한 뒤 폐기됩니다.

거대한 어텐션 행렬은 메모리에 존재하지 않습니다.

이는 트랜스포머(Transformer) 실행 시 발생하는 가장 큰 메모리 병목 현상 중 하나를 즉각적으로 제거합니다.

4. 생각보다 어려운 이유

언뜻 보기에 블록 단위(block-wise) 처리는 불가능해 보입니다.

소프트맥스 (softmax) 연산은 행 전체의 정보를 필요로 하기 때문입니다.

softmax(x_i) = exp(x_i) / sum(exp(x_j))

단 하나의 확률을 계산하기 위해서도, 모든 점수(scores)에 의존하는 분모(denominator)가 필요합니다.

한 번에 하나의 블록만 볼 수 있다면, 어떻게 소프트맥스를 정확하게 계산할 수 있을까요?

여기서 FlashAttention의 영리함이 드러납니다.

모든 점수를 저장하는 대신, 블록을 처리하는 동안 실행 통계량(running statistics)을 유지합니다.

여기에는 다음 항목들이 포함됩니다:

• 실행 최댓값 (Running maximum)
• 실행 정규화 항 (Running normalization term)
• 실행 출력 누산기 (Running output accumulator)

각 블록이 도착할 때마다 이 값들은 업데이트됩니다.

모든 블록이 처리되면, 그 결과는 표준 어텐션과 수학적으로 동일합니다.

근사치가 아닙니다.

비슷한 수준도 아닙니다.

정확히 동일합니다.

5. 온라인 소프트맥스 (Online Softmax): 이를 가능하게 하는 기술

FlashAttention의 핵심적인 돌파구는 흔히 **온라인 소프트맥스 (Online Softmax)**라고 불립니다.

어텐션 점수(attention scores)가 다음과 같이 청크(chunks) 단위로 들어온다고 상상해 보세요:

Block A
Block B
Block C

단순한(naive) 구현 방식이라면 소프트맥스(softmax)를 계산하기 전에 모든 점수가 필요할 것입니다.

반면, 온라인 소프트맥스 (Online Softmax)는 결과를 점진적으로 업데이트할 수 있도록 충분한 정보를 유지합니다.

이 알고리즘은 다음 항목들을 추적합니다:

현재 최대 점수 (Current maximum score)
현재 정규화 계수 (Current normalization factor)
현재 가중치 출력 (Current weighted output)

새로운 블록이 도착하면:

최댓값 업데이트 (Update maximum)

이전 값들의 스케일 재조정 (Rescale previous values)
...

이를 통해 FlashAttention은 어텐션(attention)을 거대한 행렬이 아닌 하나의 스트림(stream)으로 처리할 수 있습니다.

메모리 절감 효과는 엄청납니다.

더 중요한 점은, 최종 출력이 표준 어텐션 (standard attention)이 생성했을 결과와 동일하다는 것입니다.

6. 타일링 (Tiling): GPU를 효율적으로 활용하기

FlashAttention은 흔히 IO-aware 알고리즘으로 설명됩니다.

IO-aware란 알고리즘이 순수하게 산술 복잡도(arithmetic complexity)가 아닌 메모리 이동(memory movement)을 중심으로 설계되었음을 의미합니다.

구현에는 타일링 (tiling) 기법이 사용됩니다.

FlashAttention은 거대한 행렬을 직접 연산하는 대신, 작은 청크들을 빠른 온칩 메모리 (on-chip memory)로 로드합니다:

타일을 공유 메모리 (shared memory)로 로드

어텐션 계산
...

공유 메모리는 HBM보다 훨씬 빠르기 때문에, GPU는 유용한 작업을 수행하는 데 훨씬 더 많은 시간을 할애할 수 있습니다.

이해를 돕기 위한 모델은 다음과 같습니다:

전통적인 어텐션 (Traditional Attention):

메모리 -> 연산 -> 메모리 -> 연산 -> 메모리

FlashAttention:

메모리 -> 연산 -> 연산 -> 연산

연산량은 대략 $O(N^2)$로 유지되지만, 메모리 트래픽 (memory traffic)은 극적으로 감소합니다.

이것이 대부분의 속도 향상이 발생하는 지점입니다.

7. FlashAttention-2, FlashAttention-3, 그리고 프로덕션 시스템

오리지널 FlashAttention 논문은 핵심 아이디어를 소개했습니다.

FlashAttention-2는 다음 사항들을 개선했습니다:

• GPU 활용도 (GPU utilization)
• 병렬성 (Parallelism)
• 작업 분할 (Work partitioning)
• 학습 처리량 (Training throughput)

그 결과 현대적인 가속기 (accelerators)에서 현저히 높은 성능을 보여주었습니다.

FlashAttention-3는 최신 NVIDIA Hopper GPU를 위해 성능을 더욱 끌어올렸으며, FP8과 같은 현대적인 저정밀도 형식 (low-precision formats)에 대한 지원을 도입했습니다.

오늘날 FlashAttention은 AI 생태계 전반에서 사용되고 있습니다.

다음과 같은 곳에서 찾아볼 수 있습니다:

• PyTorch
• Hugging Face Transformers
• vLLM
• TensorRT-LLM
• 많은 오픈 웨이트 (open-weight) LLM들

많은 개발자에게 FlashAttention을 활성화하는 것은 최적화된 어텐션 백엔드 (attention backend)를 선택하는 것만큼 간단할 수 있습니다.

모델 아키텍처 (model architecture)는 변경되지 않은 채 유지됩니다.

하지만 성능 특성은 극적으로 향상됩니다.

FlashAttention이 해결하지 못하는 것

FlashAttention은 강력하지만, 모든 스케일링 (scaling) 문제를 마법처럼 제거하지는 못합니다.

연산량은 시퀀스 길이 (sequence length)에 대해 여전히 대략 다음과 같습니다:

O(N^2)

매우 긴 컨텍스트 (contexts)는 여전히 상당한 연산량을 요구합니다.

FlashAttention은 주로 메모리 트래픽 (memory traffic)과 메모리 사용량 (memory footprint)을 줄여줍니다.

이는 어텐션을 훨씬 더 효율적으로 만들지만, 표준 어텐션의 근본적인 이차 함수적 상호작용 패턴 (quadratic interaction pattern)을 바꾸지는 않습니다.

그렇기 때문에 연구자들은 다음과 같은 대안들을 계속해서 탐구하고 있습니다:

• 슬라이딩 윈도우 어텐션 (Sliding-window attention)
• 선형 어텐션 (Linear attention)
• 상태 공간 모델 (State-space models)
• 하이브리드 아키텍처 (Hybrid architectures)

이러한 접근 방식들은 메모리 이동 (memory movement) 문제보다는 연산 스케일링 (computational scaling) 문제 자체를 해결하려고 시도합니다.

결론

FlashAttention은 거의 즉시 기초적인 기술이 된 보기 드문 돌파구 중 하나입니다.

이 기술은 Transformers를 대체하지 않았습니다.

새로운 어텐션 메커니즘을 발명하지도 않았습니다.

모델을 다시 학습시킬 필요도 없었습니다.

대신, 현대적인 GPU가 데이터를 이동시키는 데 엄청난 양의 시간을 소비한다는 점을 인식하고, 그 이동을 최소화하도록 어텐션을 재설계했습니다.

연산 (arithmetic)이 아닌 메모리 액세스 (memory access)를 병목 현상 (bottleneck)으로 취급함으로써, FlashAttention은 어텐션을 메모리 집약적인 작업에서 훨씬 더 하드웨어 효율적인 작업으로 변화시켰습니다.

오늘날의 많은 롱 컨텍스트 (long-context) LLM들은 FlashAttention 없이는 훨씬 더 느리거나, 더 비싸거나, 혹은 단순히 실용적이지 않았을 것입니다.

AI 시스템이 계속해서 확장됨에 따라, FlashAttention은 때때로 가장 큰 돌파구가 알고리즘 자체를 바꾸는 것이 아니라, 그 알고리즘이 실제 하드웨어와 어떻게 상호작용하는지를 이해하는 데서 온다는 중요한 사실을 상기시켜 줍니다.

지금까지 LLM을 위한 가장 중요한 인프라 기술적 돌파구는 무엇이었다고 생각하시나요: FlashAttention, 양자화 (quantization), KV 캐싱 (KV caching), 추측적 디코딩 (speculative decoding), 아니면 완전히 다른 무엇인가요?

*AI 에이전트는 코드를 빠르게 작성합니다. 하지만 사용자에게 알리지 않고 조용히 로직을 제거하거나, 동작을 변경하고, 버그를 유발하기도 합니다. 이는 종종 운영 환경(production)에 도달해서야 발견되곤 합니다.

git-lrc가 이를 해결합니다. git commit에 연결되어 모든 diff를 반영하기 전에 검토합니다. 60초면 설정이 완료됩니다. 완전히 무료입니다.*

어떠한 피드백이나 기여자도 환영합니다! 이 프로젝트는 온라인에 공개되어 있으며, 소스 코드를 확인할 수 있고(source-available), 누구나 사용할 준비가 되어 있습니다.

HexmosTech / git-lrc

커밋 시 실행되는 무료 마이크로 AI 코드 리뷰

| 🇩🇰 Dansk | 🇪🇸 Español | 🇮🇷 Farsi | 🇫🇮 Suomi | 🇯🇵 日本語 | 🇳🇴 Norsk | 🇵🇹 Português | 🇷🇺 Русский | 🇦🇱 Shqip | 🇨🇳 中文 | 🇮🇳 हिन्दी |

git-lrc

커밋에서 실행되는 무료, 초소형 AI 코드 리뷰