출력을 누적해 메모리 이동을 줄이는 FlashAttention

지난 글에서는 점수를 조금씩 읽으며 최댓값과 지수합을 갱신하는 온라인 소프트맥스를 살펴봤습니다. 하지만 이 두 값만 구했다고 모든 위치의 확률이 출력되는 것은 아닙니다. 어텐션에서는 그 확률로 Value 벡터를 가중합한 최종 출력이 필요합니다. 이 출력까지 점수 묶음마다 누적할 수 있다면, 큰 점수·확률 행렬을 저장하지 않고 계산을 이어갈 수 있습니다.

이번 글에서는 먼저 Value 가중합을 구한 뒤 마지막에 지수합으로 나누어도 되는 이유를 알아봅니다. 이어서 새 점수를 읽을 때 지수합과 가중합을 같은 기준으로 보정하고 더하는 방법을 살펴봅니다. 이 규칙을 Q·K·V의 작은 타일에 적용해 출력이 완성되는 과정을 따라간 뒤, 큰 중간 행렬의 저장과 재읽기가 어떻게 줄어드는지 확인하겠습니다. 수학적인 누적 규칙을 실제 데이터 이동의 절약으로 연결하는 것이 FlashAttention의 핵심입니다.

이 글의 구체적인 계산 순서는 FlashAttention-2의 순전파 알고리즘을 기준으로 합니다. 타일 단위로 출력을 누적해 큰 중간 행렬의 저장을 피하는 원리는 첫 FlashAttention부터 사용되었고, 모든 타일을 처리한 뒤 마지막에 한 번 정규화하는 방식은 FlashAttention-2의 개선입니다.

Value 가중합을 먼저 계산하기

기본 어텐션은 Q와 K로 점수를 구하고, 소프트맥스로 확률을 만든 다음, 그 확률을 가중치로 V를 더합니다. 한 쿼리를 기준으로 보면 각 키의 점수가 하나씩 있고, 각 점수에는 같은 토큰의 Value 벡터가 대응합니다. 키는 얼마나 반영할지 정하는 데 쓰이고, Value는 출력에 반영할 내용을 담습니다.

안정적인 소프트맥스에서는 각 점수에서 같은 최댓값 m을 빼고 exp를 계산합니다. 그 exp 값을 모두 더한 지수합이 ℓ이며, 각 exp 값을 ℓ로 나누면 확률이 됩니다. 한 쿼리의 모든 확률은 같은 분모 ℓ을 사용하므로, 각각 나눈 뒤 Value를 가중합하는 대신, 먼저 가중합하고 마지막에 한 번 나눌 수 있습니다. 여기서 한 번 나눈다는 것은 출력 벡터의 모든 원소를 같은 ℓ로 나눈다는 뜻입니다.

그림 1에서 점수는 [1, 2], 대응하는 Value는 [2, 0]과 [0, 2]입니다. 최댓값 2를 빼면 exp 값은 약 [0.368, 1]이고 지수합은 약 1.368입니다. 먼저 확률을 구하면 약 [0.269, 0.731]이므로, 출력은 0.269 × [2, 0] + 0.731 × [0, 2] ≈ [0.538, 1.462]입니다. 이 글과 그림의 소수는 정확한 계산 후 반올림한 값입니다.

나눗셈을 마지막으로 미루면 exp 값으로 바로 Value를 가중합합니다. 0.368 × [2, 0] + 1 × [0, 2] ≈ [0.736, 2]를 구하고, 이 벡터를 지수합 약 1.368로 나누면 같은 출력 [0.538, 1.462]가 됩니다.

이처럼 아직 지수합으로 나누지 않은 Value 가중합을 a라고 하겠습니다. a는 스칼라 하나가 아니라 Value와 같은 길이의 벡터입니다. 앞으로는 exp 값들의 합 ℓ과, exp 값으로 Value를 가중합한 a를 함께 누적합니다. 전체 점수를 반영한 뒤 O = a ÷ ℓ을 계산하면 최종 출력 O를 얻습니다.

지수합과 가중합을 함께 갱신하기

점수를 나누어 처리하면 최댓값 m도 중간에 바뀔 수 있습니다. 이때 이전 지수합을 새 최댓값 기준으로 바꾸었던 것처럼, 이전 가중합 a도 같은 기준으로 바꾸어야 합니다. m이 커지면 이전 exp 값 모두에 같은 배율이 적용됩니다. 각각의 exp 값에 곱해지는 V는 그대로이므로, 그 곱들을 더해둔 a 전체에도 같은 배율을 곱할 수 있습니다.

그림 2는 한 쿼리의 점수 [1, 2, 0, 1, 3, 2]를 두 개씩 읽는 예시입니다. 첫 묶음 [1, 2]와 Value [2, 0], [0, 2]를 처리하면 m = 2, ℓ ≈ 1.368, a ≈ [0.736, 2]가 됩니다. 이 세 상태를 유지하고 다음 묶음을 읽습니다.

최댓값이 유지되면 새 기여 더하기

두 번째 점수 묶음은 [0, 1]이고, 대응하는 Value는 [1, 1], [2, 0]입니다. 두 점수 모두 기존 최댓값 2보다 작으므로 m은 2를 유지합니다. 이전 ℓ과 a도 이미 이 기준으로 계산되어 있어 그대로 사용할 수 있습니다.

새 점수의 exp 값은 약 [0.135, 0.368]입니다. 이 값들의 합 약 0.503을 ℓ에 더합니다. 가중합에는 0.135 × [1, 1] + 0.368 × [2, 0] ≈ [0.871, 0.135]를 더합니다. 두 묶음까지 반영하면 ℓ ≈ 1.871, a ≈ [1.607, 2.135]가 됩니다.

최댓값이 커지면 두 누적값 보정하기

세 번째 묶음 [3, 2]에서는 더 큰 점수 3을 만납니다. m을 2에서 3으로 바꾸면, 이전 점수의 exp 값에는 모두 exp(2 − 3) ≈ 0.368을 곱해야 합니다. 이전 점수나 Value를 다시 읽지 않고, 지금까지 더해둔 ℓ과 a에 이 배율을 곱하면 됩니다.

이전 지수합 약 1.871은 약 0.688로 보정됩니다. 이전 가중합 [1.607, 2.135]는 약 [0.591, 0.786]으로 보정됩니다. 보정 배율은 가중합 벡터의 모든 원소에 동일하게 적용됩니다.

이제 새 점수 [3, 2]도 새 최댓값 3을 기준으로 exp를 계산합니다. 그 값은 약 [1, 0.368]이고, 대응하는 Value는 [0, 2], [1, 1]입니다. 새 지수합 약 1.368과 새 가중합 약 [0.368, 2.368]을 각각 더하면, 최종적으로 ℓ ≈ 2.056, a ≈ [0.959, 3.153]가 됩니다. 모든 점수를 처리했으므로 a를 ℓ로 나누어 출력 약 [0.466, 1.534]를 얻습니다.

두 경우의 규칙은 같습니다. 새 최댓값을 정하고, 이전 지수합과 가중합을 같은 배율로 보정한 뒤, 새 점수와 Value의 기여를 더합니다. 최댓값이 그대로라면 보정 배율이 1일 뿐입니다. 각 타일에서 최종 확률을 완성할 필요 없이, 최종 출력을 구하는 데 필요한 정보를 함께 유지할 수 있습니다. FlashAttention-1의 알고리즘 1은 타일을 처리할 때마다 새 지수합으로 나눈 부분 출력 O를 유지합니다. 다음 타일에서는 이전 지수합을 곱해 가중합을 복원하고 보정한 뒤 새 기여를 더합니다. 반면 FlashAttention-2의 3.1.1절과 알고리즘 1은 정규화 전 가중합 a를 유지하고 마지막에 한 번만 나누어 반복 정규화를 줄입니다. exp 값으로 V를 가중합하는 계산 자체는 두 버전 모두에 있습니다.

타일을 순회하며 출력 완성하기

지금까지는 이미 주어진 점수에 Value를 연결했습니다. 실제 어텐션에서는 Q와 K를 곱해 점수 타일을 만들고, 그 자리에서 exp 계산과 Value 가중합까지 이어갑니다. 타일은 함께 처리하는 작은 행렬 영역입니다. Q 타일 하나를 유지한 채 K와 V를 같은 토큰 범위로 나누어 읽으면, 앞에서 배운 누적 규칙을 그대로 적용할 수 있습니다.

그림 3은 왼쪽에 쿼리 두 개, 위쪽에 키 네 개, 아래쪽에 Value 네 개, 오른쪽에 출력 누적값을 배치합니다. 왼쪽 Q 한 행과 오른쪽 출력 한 행은 서로 대응합니다. 위쪽 K와 아래쪽 V는 같은 토큰끼리 수직으로 정렬되어 있습니다. 토큰 위치를 맞추기 위해 K와 V는 그림에서 각각 Kᵀ·Vᵀ 형태로 펼쳤지만, 실제 Value 가중합 계산은 exp 가중치 타일과 V 타일의 곱입니다.

아래 그림에서 ‘다음’으로 단계를 넘겨보겠습니다. 숫자를 쉽게 따라가도록 스케일링과 마스크를 생략한 작은 예시입니다. 일반적인 어텐션에서는 Q·K의 내적을 키 벡터 차원 d의 제곱근으로 나누고, 필요하면 마스크를 적용한 점수를 사용합니다. 이 처리는 exp 계산 전에 이루어지며, 뒤에서 사용하는 누적 규칙은 같습니다.

그림 3을 별도 화면에서 단계별로 보기

Q를 유지하며 첫 K/V 타일 처리하기

첫 타일에서는 K₀·K₁과 V₀·V₁을 읽습니다. Q₀은 [1, 0], K₀은 [1, 2]이므로 두 벡터의 내적은 1 × 1 + 0 × 2 = 1입니다. 같은 방법으로 두 쿼리와 두 키를 곱하면 중앙에 2 × 2 점수 타일이 생깁니다. Q₀의 점수는 [1, 2], Q₁의 점수는 [2, 1]입니다.

각 행의 최댓값을 빼고 exp를 계산하면, Q₀은 약 [0.368, 1], Q₁은 약 [1, 0.368]이 됩니다. 두 행 모두 최댓값은 2이고 지수합은 약 1.368입니다. 중앙의 exp 값은 아직 최종 확률이 아닙니다. 뒤에 읽을 키의 점수도 지수합에 반영해야 합니다.

그렇지만 이 exp 값으로 Value를 가중합하는 일은 바로 할 수 있습니다. V₀은 [2, 0], V₁은 [0, 2]이므로, 첫 타일의 가중합은 Q₀에서 약 [0.736, 2], Q₁에서 약 [2, 0.736]입니다. 이 벡터들을 오른쪽에 유지합니다. 여러 쿼리를 함께 계산해도 최댓값·지수합·가중합은 쿼리 행마다 따로 관리합니다.

K/V를 옮기며 같은 출력에 누적하기

다음에는 위아래의 강조 영역이 K₂·K₃과 V₂·V₃으로 함께 이동합니다. 왼쪽 Q와 오른쪽 누적값의 위치는 그대로입니다. 같은 Q로 새 점수 타일을 계산하면 Q₀의 점수는 [3, 2], Q₁의 점수는 [0, 1]이 됩니다.

Q₀은 최댓값이 2에서 3으로 커졌으므로, 이전 ℓ과 a에 약 0.368을 곱합니다. Q₁은 기존 최댓값 2를 유지하므로 보정 배율은 1입니다. 이어서 새 기준으로 계산한 exp 값의 합을 각 ℓ에 더하고, 새 Value 가중합을 각 a에 더합니다. 같은 타일을 처리하는 중에도 행마다 필요한 보정은 다를 수 있습니다.

그림에서 첫 타일의 중앙 칸은 ‘소비됨’으로 바뀝니다. 점수와 exp 값이 지수합·가중합에 이미 반영되었으므로, 이 순전파 계산을 이어가기 위해 보관할 필요가 없다는 뜻입니다. 중앙 전체 격자는 계산 위치를 보여주는 지도이며, 전체 점수 행렬을 저장하는 공간이 아닙니다.

모든 타일을 반영한 뒤 최종 정규화하기

두 타일을 모두 처리하면 최종 지수합은 두 행 모두 약 1.871입니다. 가중합은 Q₀에서 약 [0.639, 3.104], Q₁에서 약 [2.368, 1.374]입니다. 각각을 같은 행의 지수합으로 나누면 최종 출력은 Q₀에서 약 [0.341, 1.659], Q₁에서 약 [1.266, 0.734]가 됩니다.

이 과정에서 수행한 두 행렬 곱은 분명합니다. 먼저 Q 타일과 K 타일로 점수를 계산하고, 이어서 exp 가중치 타일과 V 타일로 출력에 더할 벡터를 계산합니다. 그 사이를 온라인 소프트맥스의 보정·갱신이 연결합니다. 그림의 Q 타일처럼 다른 Q 타일도 각자의 출력과 상태를 유지하며 필요한 K/V 타일을 처리할 수 있습니다.

중간 행렬의 저장과 재읽기 줄이기

앞의 계산 순서가 중요한 이유는 GPU 안에서 이동하는 데이터가 달라지기 때문입니다. 그림 4의 기본 어텐션은 큰 점수 행렬 S를 전역 메모리에 저장하고, 소프트맥스 단계에서 다시 읽습니다. 그렇게 만든 큰 확률 행렬 P도 저장한 뒤, V와 곱하는 단계에서 다시 읽습니다. 시퀀스 길이가 N이라면 한 헤드의 점수·확률 행렬은 각각 N × N 크기이므로, 길이가 늘수록 중간 데이터의 부담이 빠르게 커집니다.

FlashAttention에서는 현재 점수 타일을 계산한 뒤 exp와 Value 가중합에 바로 사용합니다. Q 타일과 누적 상태를 온칩에 유지하면서 다음 K/V 타일을 읽고, 모든 타일을 반영하면 최종 출력을 기록합니다. 여기서 온칩 메모리는 계산에 가까운 공유 메모리와 레지스터 등을 가리킵니다. 큰 S·P 행렬을 통째로 전역 메모리에 저장하고 다음 단계에서 다시 읽는 과정을 피하는 것입니다. FlashAttention의 메모리 이동 최적화

그림은 HBM과 온칩 사이의 핵심 흐름을 단순화했습니다. 실제 전역 메모리 접근에는 캐시가 관여하므로 모든 접근을 실제 HBM 전송과 동일시하지는 않습니다. 또한 학습의 역전파를 위해 행별 정규화 통계를 저장하는 등의 처리는 이 그림에서 생략했습니다. 여기서 주목할 차이는 큰 N × N 중간 행렬을 저장하느냐, 작은 타일에서 계산하고 소비하느냐입니다.

앞서 배운 최적화들도 이 지점에서 연결됩니다. 타일링은 입력을 작은 영역으로 나누어 가져오고 재사용하게 합니다. 연산 결합은 점수 계산·소프트맥스·Value 가중합 사이의 중간 데이터 이동을 줄입니다. 온라인 갱신은 행 전체가 필요한 소프트맥스를 작은 타일로 나누어도 출력을 완성할 수 있게 합니다. 타일링과 연산 결합을 어텐션 전체에 적용할 수 있도록 해주는 연결 고리가 온라인 누적입니다.

FlashAttention은 일반적인 dense 어텐션의 이차 연산량을 선형으로 바꾸는 방법은 아닙니다. 근사 어텐션으로 바꾸지 않고도 데이터 이동과 중간 저장을 줄일 수 있다는 것이 핵심이며, 실제 부동소수점 결과에는 계산 순서에 따른 오차가 있을 수 있습니다. 이 글에서는 순전파의 원리에 집중했습니다. 버전별 스레드 배치나 파이프라인, 역전파의 재계산은 이 원리 위에 더해지는 별도의 최적화 주제입니다.