FlashAttention CUDA 커널 한 줄씩 살펴보기
Flash 어텐션은 트랜스포머에서 어텐션을 계산하기 위한 메모리 효율적인 알고리즘입니다. CUDA 구현을 블록 단위로 하나씩 살펴봅시다. Flash 어텐션의 핵심(Core) 혁신은 최댓값과 부분 합을 세심하게 추적하여 수치적 안정성을 유지하면서 어텐션을 블록 단위로 처리한다는 점입니다. 이 알고리즘은 전체 어텐션 행렬을 한 번도 메모리에 실체화하지 않으면서도 표준 어텐션 메커니즘과 수학적으로 동등한 결과를 만들어 내어 메모리 효율성을 달성합니다.
먼저 몇 가지 작은 매크로가 필요합니다. CEIL_DIV는 두 수를 나눈 값의 올림을 계산하는 데 사용됩니다.
#include <cuda.h>
#include <cuda_runtime.h>
#define CEIL_DIV(x, y) ((x) >= 0 ? (((x) + (y) - 1) / (y)) : ((x) / (y)))
이제 커널 선언이 필요합니다:
template <const int Br, const int Bc>
__global__ void flash_attn_kernel(float* Q, float* K, float* V, int N, int d, int Tr, int Tc, float scale, float* l, float* m, float* O)
이 코드는 행과 열의 블록 크기를 나타내는 템플릿 매개변수 Br과 Bc를 가진 CUDA 커널을 선언합니다. 이 함수는 다음을 인자로 받습니다:
- Q, K, V: Query, Key, Value 행렬
- N: 시퀀스 길이
- d: 헤드 차원
- Tr, Tc: 행과 열 방향의 블록 개수
- scale: 어텐션 점수에 대한 스케일링 계수
- l, m: 수치적으로 안정적인 소프트맥스 계산을 위한 보조 배열
- O: 출력 행렬
스레드 및 블록 인덱스 설정
int tx = threadIdx.x;
int bx = blockIdx.x;
int by = blockIdx.y;
int qkv_off = (bx * gridDim.y * N * d) + (by * N * d);
int lm_off = (bx * gridDim.y * N) + (by * N);
이 부분은 어떤 스레드가 실행 중인지 식별하는 스레드 및 블록 인덱스(tx, bx, by)를 가져옵니다. 그런 다음 두 가지 중요한 오프셋 값을 계산합니다. Q, K, V 행렬에 접근하기 위한 qkv_off와 보조 배열 l과 m에 접근하기 위한 lm_off입니다. 이 오프셋들은 각 블록이 입력 데이터의 올바른 부분을 처리하도록 보장합니다.
공유 메모리 설정
extern __shared__ float smem[];
float* Qi = smem;
float* Kj = Qi + Br * d;
float* Vj = Kj + Bc * d;
float* Sij = Vj + Bc * d;
float* Oi = Sij + Br * Bc;
float* li = Oi + Br * d;
float* li_new = li + Br;
float* mi = li_new + Br;
float* mi_new = mi + Br;
float* mij_dash = mi_new + Br;
이 부분은 공유 메모리를 다음 용도의 여러 영역으로 분할합니다:
- Qi: Query 행렬의 현재 블록
- Kj: Key 행렬의 현재 블록
- Vj: Value 행렬의 현재 블록
- Sij: 어텐션 점수
- Oi: 출력 누산기
- li, li_new, mi, mi_new, mij_dash: 소프트맥스 계산 변수
주요 처리 루프
이 커널은 두 개의 중첩 루프를 사용합니다:
- Key/Value 블록(
j)에 대한 바깥쪽 루프 - Query 블록(
i)에 대한 내부 루프
Key-Value 로딩
for (int j = 0; j < Tc; j++) {
int loads_per_thread = CEIL_DIV(d, Br);
for (int e = 0; e < loads_per_thread; e++) {
int idx = e * (Br * Bc) + tx;
if (idx < Bc * d) {
int row = idx / d;
int col = idx % d;
if (j * Bc + row < N) {
Kj[row * d + col] = K[qkv_off + (j * Bc + row) * d + col];
Vj[row * d + col] = V[qkv_off + (j * Bc + row) * d + col];
}
}
}
__syncthreads();
이 부분은 K 행렬과 V 행렬의 현재 블록을 공유 메모리로 로드하며, 범위를 벗어난 접근을 막기 위해 경계 검사를 수행합니다.
각 쿼리 블록을 처리하는 주 계산 루프를 자세히 살펴봅시다.
Query 및 출력 데이터 로딩
for (int e = 0; e < loads_per_thread; e++) {
int idx = e * (Br * Bc) + tx;
if (idx < Br * d) {
int row = idx / d;
int col = idx % d;
if (i * Br + row < N) {
Qi[row * d + col] = Q[qkv_off + (i * Br + row) * d + col];
Oi[row * d + col] = O[qkv_off + (i * Br + row) * d + col];
}
}
}
이 섹션에서는 Query 행렬의 현재 블록과 출력 누산기를 공유 메모리로 로드합니다. 효율적인 메모리 접근 패턴을 보장하기 위해 각 스레드는 여러 요소를 로드할 수 있습니다. 경계 검사는 N이 블록 크기로 나누어떨어지지 않을 때 범위를 벗어난 메모리에 접근하지 않도록 보장합니다.
스레드 위치 계산 및 최댓값/합계 로딩
int s_row = tx / Bc;
int s_col = tx % Bc;
if (s_col == 0) {
mi[s_row] = m[lm_off + (i * Br) + s_row];
li[s_row] = l[lm_off + (i * Br) + s_row];
}
__syncthreads();
각 스레드는 공유 메모리 블록 내에서 자신의 위치를 계산합니다. 각 행의 첫 번째 스레드(s_col == 0)는 수치적으로 안정적인 소프트맥스 계산에 필요한 누적 최댓값과 누적 합계를 로드합니다.
어텐션 점수 계산
float acc = 0.f;
for (int k = 0; k < d; k++)
acc += Qi[s_row * d + k] * Kj[s_col * d + k];
acc *= scale;
Sij[s_row * Bc + s_col] = acc;
이 부분은 Query 행렬과 Key 행렬 사이의 스케일드 닷 프로덕트 어텐션 점수를 계산합니다. 각 스레드는 어텐션 점수 행렬 Sij의 요소 하나를 계산합니다.
수치적으로 안정적인 소프트맥스 계산
if (s_col == 0) {
float row_m = -INFINITY, row_l = 0.f;
// Find max for numerical stability
for (int c = 0; c < Bc; c++) {
float val = Sij[s_row * Bc + c];
if (val > row_m) {
row_m = val;
}
}
// Compute exponentials and sum
for (int c = 0; c < Bc; c++) {
float exp_val = expf(Sij[s_row * Bc + c] - row_m);
Sij[s_row * Bc + c] = exp_val;
row_l += exp_val;
}
// Update running max and sum
mij_dash[s_row] = row_m;
mi_new[s_row] = max(mi[s_row], row_m);
li_new[s_row] = expf(mi[s_row] - mi_new[s_row]) * li[s_row] +
expf(row_m - mi_new[s_row]) * row_l;
}
이 부분은 수치적으로 안정적인 소프트맥스 알고리즘을 구현합니다. 각 행에 대해 다음을 수행합니다.
- 수치적 안정성을 위해 최댓값을 찾습니다
- 최댓값으로 정규화한 지수를 계산합니다
- 정규화된 지수들을 합산합니다
- 온라인 소프트맥스 계산을 위해 누적 최댓값과 누적 합계를 갱신합니다
출력 계산 및 갱신
for (int col = s_col; col < d; col += Bc) {
float acc = 0.f;
for (int c = 0; c < Bc; c++)
acc += Sij[s_row * Bc + c] * Vj[c * d + col];
int global_row = (i * Br) + s_row;
if (global_row < N) {
Oi[s_row * d + col] = (1 / li_new[s_row]) *
((li[s_row] * expf(mi[s_row] - mi_new[s_row]) * Oi[s_row * d + col]) +
(expf(mij_dash[s_row] - mi_new[s_row]) * acc));
O[qkv_off + global_row * d + col] = Oi[s_row * d + col];
}
}
이 마지막 부분은 다음을 수행합니다.
- 소프트맥스 확률과 Value 행렬 사이의 행렬 곱셈을 계산합니다
- 온라인 소프트맥스 알고리즘을 사용하여 출력을 갱신합니다
- 각 스레드는 여러 출력 요소를 계산할 수 있습니다(d > Bc인 경우)
- 갱신 공식이 복잡한 이유는 부분 결과를 누적하는 동안 수치적 안정성을 보장하기 위해서입니다
다음 반복을 위한 상태 갱신
m[lm_off + (i * Br) + s_row] = mi_new[s_row];
l[lm_off + (i * Br) + s_row] = li_new[s_row];
마지막으로, 다음 반복을 위해 전역 메모리에 있는 누적 최댓값과 누적 합계를 갱신합니다. 이는 블록 전반에 걸친 온라인 소프트맥스 계산에 매우 중요합니다.