최종 업데이트: 2026-06-22

어텐션 메모리

표준 트랜스포머 어텐션은 모든 토큰을 다른 모든 토큰과 비교해, 소프트맥스 이전에 점수 행렬을 만듭니다. 이 행렬은 시퀀스 길이에 대해 정사각이므로 메모리가 컨텍스트의 제곱으로 커집니다. 이것이 그 유명한 어텐션의 제곱 비용입니다. 이 계산기는 그 실체화된 행렬의 크기를 시퀀스 길이, 어텐션 헤드 수, 배치 크기, 점수 하나당 바이트로부터 추정합니다. 이는 정확히 FlashAttention이 저장을 피하는 메모리입니다.

제곱 항

$s$개의 토큰으로 이루어진 시퀀스에 대해 어텐션은 $s \times s$ 점수 행렬을 만듭니다. $i$행 $j$열에는 토큰 $i$가 토큰 $j$를 얼마나 어텐션하는지가 담깁니다. 이를 저장하는 데에는 $s^2$에 비례하는 메모리가 듭니다. 컨텍스트에 선형으로 커지는 모델 가중치나 키-값 캐시와 달리, 점수 행렬은 그 제곱으로 커지므로 긴 컨텍스트에서는 가장 큰 중간 버퍼가 되며, 이것이 단순한 어텐션이 메모리를 소진하는 이유입니다.

공식

각 헤드는 자기 점수 행렬을 만들고 배치의 각 시퀀스는 자기 사본을 들고 있으므로, 바이트 단위 메모리는 다음과 같습니다.

여기서 $B$는 배치 크기, $H$는 어텐션 헤드 수, $s$는 시퀀스 길이, $e$는 저장 점수당 바이트입니다. $10^9$로 나누면 기가바이트가 됩니다. 긴 컨텍스트를 비싸게 만드는 것은 $s$에 붙은 제곱입니다. 헤드 수와 배치는 선형으로만 곱해집니다.

이 메모리가 무엇인가

이 값은 표준 어텐션이 쿼리-키 곱과 값에 대한 소프트맥스 가중합 사이에 메모리에 써넣는, 실체화된 점수 행렬을 측정합니다. 이는 단순한 구현이 반드시 들고 있어야 하는 활성화 버퍼입니다. 이것이 지속되는지는 상황에 따라 다릅니다. 평범한 순전파에서는 일시적이어서 레이어 사이에서 재사용할 수 있지만, 학습 프레임워크는 체크포인팅이나 융합 커널이 개입하지 않는 한 역전파를 위해 레이어별 사본을 유지할 수 있습니다.

FlashAttention

FlashAttention은 전체 $s \times s$ 행렬을 한 번도 써넣지 않고 동일한 출력을 계산합니다. 키와 값을 작은 블록 단위로 훑으며 실행 중인 소프트맥스 통계를 유지하므로, 한 번에 한 블록만 메모리에 두면 됩니다. 이로써 저장 공간이 시퀀스 길이에 대해 제곱에서 선형으로 바뀝니다. 이 계산기가 크기를 재는 행렬이 바로 FlashAttention이 저장을 거부하는 대상이므로, 여기서 나온 결과는 융합 어텐션 커널이 주어진 컨텍스트 길이에서 절약하는 메모리의 좋은 근사치입니다.

계산 예시

4,096 토큰짜리 시퀀스 하나를 어텐션 헤드 32개, 16비트 정밀도로 처리한다고 합시다.

한 레이어의 점수 행렬이 약 1.07 GB입니다. 컨텍스트를 8,192 토큰으로 2배 늘리면 제곱 항이 위력을 발휘해, 같은 식이 대략 4.29 GB가 됩니다. 길이가 2배인데 메모리는 4배입니다. 같은 어텐션 예산의 키-값 측면은 KV 캐시 크기 계산기에서, 모델 가중치 부분은 LLM 추론 VRAM 계산기에서 다룹니다.