Attention 최적화가 어려운 이유

지난 글에서는 행렬 곱의 입력을 재사용하고, 데이터 이동을 계산과 동시에 진행해 시간을 줄이는 방법을 살펴봤습니다. 이번에는 Attention으로 넘어가겠습니다. Attention에서는 두 행렬 곱 사이에 큰 점수 행렬과 확률 행렬이 만들어집니다. 이 중간 결과를 전역 메모리에 저장했다가 다음 계산에서 다시 읽는 데에도 많은 시간이 들 수 있습니다. 따라서 Attention의 실행 시간을 줄이려면 행렬 곱의 계산 속도와 중간 결과의 데이터 이동을 함께 살펴봐야 합니다.

그렇다면 점수를 조금씩 계산하고, 곧바로 확률로 바꿔 다음 계산에 사용하면 어떨까요? 여기에는 일부 점수만으로는 그 점수의 최종 확률을 정할 수 없다는 어려움이 있습니다. Softmax에서 확률 하나를 구하려면 같은 행의 다른 점수도 반영해야 하기 때문입니다. 이번 글에서는 먼저 중간 행렬의 저장과 재읽기가 왜 부담이 되는지 살펴보고, 이 이동을 줄이려고 계산을 타일별로 이어갈 때 어디에서 어려움이 생기는지 확인하겠습니다.

Attention의 계산 흐름

Core Attention 글에서 살펴본 계산을 짧게 복습하겠습니다. 이번에는 토큰 T개를 함께 처리하는 한 헤드를 기준으로 합니다. 헤드 하나의 차원을 dₕ라고 하면 Q·K·V는 각각 T × dₕ 형태입니다.

Q의 한 행은 정보를 모을 토큰의 Query 벡터이고, K의 한 행은 참조할 토큰의 Key 벡터입니다. 두 벡터를 내적하면 토큰 사이의 점수 하나가 나옵니다. 이를 행렬 곱 QKᵀ로 계산하면 행은 Query 토큰, 열은 Key 토큰에 해당하는 T × T 점수 행렬 S가 됩니다.

그림 1의 두 번째 단계에서는 점수를 √dₕ로 나누어 크기를 조절하고, 참조할 수 없는 위치에 마스크를 적용합니다. 여기서는 자기 자신과 이전 토큰만 참조하는 causal mask를 사용합니다. 미래 토큰의 점수를 −∞로 두면 Softmax에서 그 위치의 확률은 0이 됩니다.

이어서 Softmax가 각 행의 점수를 합이 1인 확률로 바꿉니다. 이 확률은 각 토큰의 Value를 얼마나 반영할지 정하는 가중치입니다. 점수 하나마다 확률 하나가 대응하므로, 확률 행렬 P도 T × T입니다.

마지막으로 P와 V를 곱해 출력 O를 만듭니다. Query 토큰 하나에 대해 여러 Value 벡터를 확률에 따라 가중합하면 dₕ개의 성분이 남습니다. 전체 출력은 T × dₕ입니다. 따라서 Attention 안에서는 T × dₕ 입력으로부터 T × T 중간 행렬을 거쳐 다시 T × dₕ 출력을 만드는 과정이 일어납니다.

큰 중간 행렬의 저장과 재읽기

이 계산들을 나누어 실행하면서 중간 결과를 전역 메모리에 저장한다고 해보겠습니다. QKᵀ를 계산하는 커널은 점수 행렬 S를 저장하고 끝납니다. 다음 커널은 S를 읽어 스케일링·마스크와 Softmax를 수행한 뒤, 확률 행렬 P를 저장합니다. 마지막 행렬 곱은 P를 읽어 V와 곱합니다.

커널이 끝난 뒤에도 다음 커널에서 사용할 값은 남아 있어야 합니다. 앞선 커널이 사용하던 레지스터나 블록의 공유 메모리를 다음 커널에 그대로 넘길 수는 없으므로, 이 구현에서는 중간 텐서를 전역 메모리에 둡니다. 그림 2의 내려가는 화살표는 계산 결과의 저장, 올라가는 화살표는 다음 계산이 그 결과를 다시 읽는 과정입니다.

S와 P는 각각 T × T개의 값을 가집니다. 토큰 수가 두 배가 되면 행과 열이 모두 두 배가 되어, 각 행렬의 요소 수는 네 배가 됩니다. 예를 들어 T가 4,096이고 값 하나를 2 byte로 저장하면 행렬 하나가 32 MiB입니다. T가 8,192이면 같은 형식의 행렬 하나가 128 MiB가 됩니다. MiB는 1,048,576 byte입니다.

이 크기는 보관할 공간만의 문제가 아닙니다. 큰 행렬을 한 번 저장하고 한 번 다시 읽을 때마다 그만큼의 데이터를 전달해야 합니다. 그림에서는 S를 저장하고 읽은 뒤, 같은 크기의 P도 저장하고 읽습니다. 계산 도중 잠깐 필요한 중간 결과를 다음 단계에 넘기기 위해 큰 데이터가 반복해서 오가는 것입니다.

전역 메모리 접근에는 캐시가 관여하므로, 모든 읽기·쓰기가 실제 HBM 전송과 일대일로 대응하지는 않습니다. 하지만 중간 행렬이 커서 캐시만으로 전달을 처리하기 어려워지면 HBM과 주고받는 양도 커집니다. HBM 대역폭에는 한계가 있으므로 이동량이 많을수록 데이터를 저장하고 공급하는 데 필요한 시간도 길어집니다. 특히 스케일링·마스크·Softmax 구간은 큰 행렬을 읽고 처리해 다시 저장하므로, 데이터 이동이 실행 시간의 큰 부분을 차지할 수 있습니다. FlashAttention 원 논문의 기본 Attention 구현 분석

지난 글에서는 타일링을 통해 행렬 곱 안에서 입력을 재사용했습니다. 이번에는 두 행렬 곱 사이에서 S와 P를 저장하고 읽는 비용을 줄이기 위해, 앞의 계산이 만든 값을 가까운 메모리에 둔 채 다음 계산으로 이어갈 방법을 살펴보겠습니다.

점수 타일에서 확률로 넘어가기

앞서 원소별 연산에서는 여러 계산을 하나의 커널로 합치는 fusion을 살펴봤습니다. 같은 요소의 다음 계산으로 바로 이어가면, 중간값을 전역 메모리에 저장했다가 다시 읽는 일을 줄일 수 있었습니다. Attention에서도 작은 점수 타일을 계산하고 그 값을 곧바로 다음 단계에 사용한다면, 전체 점수 행렬을 전역 메모리에 저장했다가 다시 읽는 과정을 줄일 수 있습니다.

그림 3은 T가 4일 때 쿼리 토큰 2·3과 키 토큰 0·1의 점수를 먼저 계산한 경우입니다. 여기서 토큰 번호는 0부터 시작합니다. Q에서 두 행, Kᵀ에서 두 열을 골라 곱하면 2 × 2 점수 타일이 나옵니다. 타일 안의 점수 네 개는 각각 필요한 내적을 모두 마친 값입니다. 아직 계산하지 않은 것은 이 점수들의 나머지 성분이 아니라, 다른 키 토큰과의 점수입니다.

계산한 점수는 바로 √dₕ로 나눌 수 있습니다. 마스크도 해당 위치의 쿼리와 키 토큰 번호로 판단할 수 있습니다. 미래 토큰인지 확인하는 데 다른 점수의 값은 필요하지 않습니다. 스케일링과 마스크까지는 현재 타일에서 바로 처리할 수 있습니다.

Softmax에서는 상황이 달라집니다. 확률 하나는 해당 점수에 지수 함수 exp를 적용한 값을, 같은 행의 모든 유효한 점수에 exp를 적용해 더한 값으로 나누어 구합니다. 이 합을 지수합이라고 부르겠습니다. 한 행의 각 확률은 모두 같은 지수합을 분모로 사용합니다.

그림 3의 쿼리 토큰 2행을 보겠습니다. 키 0·1의 점수는 계산했지만 키 2의 점수는 아직 모릅니다. 키 3은 미래 토큰이므로 마스크로 제외됩니다. 키 0에 대한 확률을 구할 때에도 분모에는 키 0·1뿐 아니라 키 2 점수의 지수값까지 더해야 합니다. 키 2의 점수를 모르는 지금은 분모를 완성할 수 없습니다.

쿼리 토큰 3행에서는 키 0·1·2·3을 모두 참조할 수 있습니다. 현재 타일에는 키 0·1의 점수만 있으므로, 키 2·3의 점수를 더 반영해야 이 행의 지수합을 알 수 있습니다. 그림의 ‘미계산’ 칸을 0이나 제외된 값으로 취급해서는 안 됩니다. 이 점수들도 최종 확률을 구할 때 분모에 반영해야 합니다.

따라서 타일 안의 점수 계산이 끝났다는 사실만으로 그 타일의 최종 확률까지 정해지는 것은 아닙니다. 점수 하나의 내적에 필요한 입력은 모두 사용했어도, 그 점수를 확률로 바꾸는 데 필요한 같은 행의 정보는 아직 부족합니다. 다른 쿼리 행의 점수가 필요한 것은 아닙니다. 의존성이 생기는 범위는 해당 쿼리 행에서 참조할 수 있는 모든 키 위치입니다.

이 차이가 다음 행렬 곱 PV에도 이어집니다. P의 값은 V에 곱할 가중치이므로, 현재 점수 타일만으로 확정한 가중치를 만들어 바로 사용하려는 방식에는 문제가 생깁니다. 계산을 이어가려면 아직 모르는 점수가 나중에 반영될 때 가중치가 어떻게 달라지는지를 다룰 수 있어야 합니다.

나머지 점수가 바꾸는 확률

그림 4에서는 이 의존성을 숫자로 확인해보겠습니다. 네 위치가 모두 유효한 한 행에서, 스케일링을 마친 앞의 두 점수가 0과 0이라고 하겠습니다. 뒤의 두 점수는 아직 계산하지 않았습니다.

첫 점수의 지수값은 exp(0) = 1이므로 이미 알 수 있습니다. 하지만 확률의 분모는 네 점수의 지수값을 모두 더한 값입니다. 분자는 알아도 분모를 모르면 최종 확률을 정할 수 없습니다. 그림은 아직 모르는 두 점수에 서로 다른 값이 들어왔을 때를 비교합니다.

왼쪽처럼 나머지 점수도 0과 0이면 전체 점수는 [0, 0, 0, 0]입니다. 지수값이 모두 1이므로 지수합은 4이고, 첫 번째 확률은 1 ÷ 4 = 25%입니다.

오른쪽처럼 나머지 점수가 2와 2이면 전체 점수는 [0, 0, 2, 2]입니다. exp(2)는 약 7.39이므로 지수합은 1 + 1 + 7.39 + 7.39 ≈ 16.78입니다. 첫 번째 확률은 1 ÷ 16.78 ≈ 6%로 작아집니다. 앞의 점수 0은 바뀌지 않았지만, 뒤의 점수가 지수합을 키우면서 앞의 확률도 달라진 것입니다.

앞의 두 점수만으로 Softmax를 계산했다면 어떻게 될까요? [0, 0]의 지수합은 2이므로 각각 50%가 나옵니다. 하지만 이는 앞의 두 위치 사이에서만 합이 1이 되도록 계산한 확률입니다. 네 위치 전체를 대상으로 한 Attention의 확률은 아닙니다. 실제로 앞의 두 확률은 왼쪽에서는 각각 25%, 오른쪽에서는 각각 약 6%가 되어야 합니다.

그림 3으로 돌아가면, 점수 타일마다 Softmax를 따로 적용하고 그 결과를 그대로 이어 붙일 수 없는 이유가 드러납니다. 각 타일이 자기 안의 점수만으로 분모를 만들면, 같은 쿼리 행에 속한 값들이 서로 다른 기준으로 나누어집니다. 전체 행에서 각 Value를 얼마나 반영할지 정하려면, 타일 밖의 유효한 점수도 반영한 공통 분모가 필요합니다.

큰 중간 행렬의 저장과 재읽기를 줄이기 위해, 작은 묶음으로 계산한 결과를 다음 단계에서 바로 사용하는 방법을 생각해봤습니다. 그런데 일부 점수만으로 최종 확률을 정할 수는 없습니다. 여기서 필요한 것은 행 전체의 정보를 반영하는 방법입니다. 그 정보가 필요하다는 사실이, 모든 점수를 동시에 메모리에 보관해야 한다는 뜻일까요?

전체 점수를 저장하지 않고, 일부씩 읽으면서 정규화에 필요한 정보를 구할 수 있을까요? 다음 글에서는 이 질문에서 출발해 온라인 소프트맥스를 살펴보겠습니다.