Memory Bottlenecks & FlashAttention

GPU가 빠르게 곱셈을 할 수 있어도 계산할 데이터를 기다리면 그 성능을 쓰지 못한다. Standard Attention에서는 큰 score와 probability matrices를 HBM에 쓰고 다시 읽는 이동이 반복된다. FlashAttention은 계산을 blocks로 나누어 빠른 SRAM 안에서 이어 처리함으로써 이 이동을 줄인다. Attention을 근사하는 방법이 아니라 같은 Attention을 실행하는 순서와 memory 사용을 바꾸는 방법이다.

1. $Q$/$K$/$V$ blocks
2. SRAM tile
3. score / online Softmax / weighted $V$
4. 최종 O만 HBM 기록

Two Memory Bottlenecks

Training/Prefill에서는 $n_{q}\times n_{k}$ intermediate, Decode에서는 request별 KV Cache가 서로 다른 memory 문제다.$32,768^{2}$ BF16 score matrix 하나는 약 2GiB/head/layer다. L32,H_KV8,d128,BF16의 cache는 token당 128KiB,32Kcontext에 4GiB다.

KV Cache는 recompute FLOPs↓/stored bytes↑, FlashAttention은 HBM bytes↓, GQA/MQA는 KV bytes↓와 일부 KV projection FLOPs↓다.

FlashAttention

Standard Attention은 $S=QK^\top$를 HBM에 쓰고 $P=\operatorname{Softmax}(S)$를 다시 읽고 쓴 뒤 PV를 계산한다. FlashAttention은 $Q$/$K$/$V$ blocks를 SRAM에서 tile로 처리하며 running max/sum의 online Softmax로 최종 output을 얻는다. S/P전체를 HBM에 물리적으로 남기지 않고 exact 결과를 유지한다(반올림 차이 가능).

FLOPs 차수는 $O(n^{2}d)$다. 강의의 2–4×speedup은 논문 hardware 설정의 보고값이다.

Why the Phase Matters

Prefill에서는 이미 주어진 prompt의 Queries를 한꺼번에 계산할 수 있어 $n_q\times n_k$ intermediate가 커진다. 반면 Decode에서는 새 token을 하나씩 만들며 이전 K/V를 재사용하므로, 요청별 cache를 저장하고 매 step 다시 읽는 부담이 생긴다. 같은 “memory 문제”여도 무엇을 저장하고 이동하는지가 다르다.

Cache 식에서 2는 K와 V 두 종류, $B$는 batch, $L$은 layers, $n$은 context positions, $H_{KV}$는 KV heads, $d_{\mathrm{head}}$는 head 폭이다.

위 128 KiB/token과 4 GiB 예시는 $B=1$이다. Batch와 context가 커지면 cache도 함께 커진다. 이 부분을 줄이는 GQA/MQA와 $n^2$ intermediate의 HBM 이동을 줄이는 FlashAttention은 같은 최적화가 아니다.

Online Softmax

Block마다 Softmax를 독립적으로 하고 결과를 단순히 더하면 전체 Keys에 대해 정규화한 Attention과 달라진다. 전체 분모가 필요하기 때문이다. Online Softmax는 지금까지 본 score의 최대값과 exponential 합을 유지하고, 새 block이 오면 이전 계산을 새 기준으로 rescale한다.

이전 최대값을 $m$, exponential 합을 $\ell$, 새 block의 scores를 $s_j$라 하자. 기준 최대값을 $m'$로 바꾸면 다음과 같이 갱신한다.

이전 항에 $e^{m-m'}$를 곱하는 것은 이전 scores를 새 최대값 기준으로 다시 표현하는 것과 같다. Value를 곱한 누적합에도 같은 rescaling을 적용하면, score matrix 전체를 보관하지 않고도 최종 가중합을 얻을 수 있다. 이것이 교수님이 설명한 “block으로 쪼개도 근사가 아닌 이유”다. 부동소수점의 반올림 차이는 있을 수 있다.

슬라이드의 HBM 80 GB, on-chip SRAM 합계 약 20 MB, 속도 차이 약 10배는 소개된 GPU 맥락의 수치다. SRAM이 작아도 tiles가 그 안에 들어가게 하면 score → Softmax → weighted V를 이어 처리할 수 있다. FLOPs 차수 $O(n^2d)$는 남지만, 큰 S/P를 HBM에 쓰고 읽는 이동을 피한다.

FLOPs와 memory bytes를 따로 세어야, 재계산을 줄이는 KV Cache와 이동을 줄이는 FlashAttention의 효과를 구분할 수 있다.

Exact Attention과 $O(n^{2}d)$를 유지하면서 HBM 이동을 줄인다.

FlashAttention · KV Cache · GQA/MQA

FLOPs와 bytes, Training/Prefill과 Decode의 문제를 구분한다.

Check

FlashAttention이 faster인 이유를 FLOPs와 bytes를 분리해 설명하라.

FlashAttention을 sparse Attention approximation 또는 KV Cache 감소 기법이라고 부르면 왜 틀리는가?

$n\times n$ intermediate 전체를 HBM에 쓰지 않고 tile/online Softmax로 exact 결과를 유지하는 기법은?

FlashAttention · KV Cache · GQA/MQA를 축별로 비교하고 그 차이가 필요한 이유를 설명하세요.

Related

• FLOPs
• KV Cache
• MHA · GQA · MQA · MLA