AI & 기술

셀프 어텐션, CUDA 커널, 그리고 FlashAttention

2026-11-15

마이크로컨트롤러 메모리 블록을 보여 주는 예시 이미지이며, 기사에서 다루는 GPU가 아닙니다.

현대의 트랜스포머는 셀프 어텐션 행렬 연산, Triton 같은 고수준 GPU 컴파일러, FlashAttention 같은 메모리 인식 CUDA 커널에 의존해 방대한 토큰 시퀀스를 하드웨어 최고 속도로 처리한다.

1. 비전 트랜스포머에서 셀프 어텐션은 수학적으로 어떻게 동작하는가

비전 트랜스포머(ViT)에서 이미지(예: 224×224 픽셀)는 16×16 패치 격자로 잘려 N = 196개의 평탄화된 패치가 된다. 각 패치는 차원 D(예: 768)의 벡터로 선형 투영되고 위치 임베딩과 결합되어 입력 행렬 X ∈ R^(N × D)를 이룬다.

쿼리, 키, 값 만들기. 모델은 패치 행렬 X에 학습된 가중치 행렬 세 개(W_Q, W_K, W_V ∈ R^(D × d_k))를 곱해 모든 이미지 패치에 대해 서로 다른 세 가지 표현을 만든다:

Q=XW_(Q), K=XW_(K), V=XW_(V)

쿼리(Q): 특정 이미지 패치가 무엇을 찾고 있는지를 나타낸다(예: "나는 털 질감이고, 근처에서 눈이나 귀를 찾고 있다").
키(K): 패치가 어떤 정보를 담고 있는지를 알린다(예: "나는 날카롭게 휜 에지를 포함하고 있다").
값(V): 일치가 일어날 경우 앞으로 전달될 실제 시각적 특징 페이로드를 담고 있다.

어텐션 점수 계산. 모든 패치가 다른 모든 패치와 얼마나 강하게 관련되는지 측정하기 위해, 모델은 모든 쿼리와 키 사이의 내적을 계산한다:

Scores=(QK^(T))/(sqrt(d_(k)))

sqrt(d_(k))로 나누면 큰 내적 값이 이후의 소프트맥스 함수를 그래디언트가 아주 작은 영역으로 밀어 넣는 것을 막는다. 그 결과 N × N 어텐션 행렬이 만들어지며, 항목 (i, j)는 패치 i가 패치 j에 얼마나 주목해야 하는지를 점수로 나타낸다.

소프트맥스 가중치와 집계. 행 단위 소프트맥스를 적용하면 원시 점수가 합이 1인 확률로 바뀐다. 이 확률에 값 행렬 V를 곱하면 각 패치에 대한 혼합된 표현이 만들어진다:

Attention(Q,K,V)=softmax((QK^(T))/(sqrt(d_(k))))V

멀티헤드 셀프 어텐션에서는 이 과정이 여러 헤드(예: 12개)에서 병렬로 그대로 실행되어, 서로 다른 헤드가 다양한 공간적 거리, 색상, 객체 부위에 동시에 집중할 수 있다.

2. 실습 코드 비교: 순수 CUDA C++ 대 OpenAI Triton

GPU 프로그래밍이 어떻게 진화했는지 보기 위해, N개 원소에 대한 간단한 원소별 벡터 덧셈(C = A + B)을 생각해 보자.

순수 CUDA C++(스레드 수준 관리)

CUDA C++에서는 단일 스칼라 스레드의 관점에서 프로그래밍한다. 스레드 ID를 직접 계산하고, 배열 경계를 확인하고, 호스트 측 실행 구성을 위한 상용구 코드를 작성해야 한다.

#include <cuda_runtime.h>

// Device kernel: executed by thousands of individual threads
__global__ void vector_add_cuda(const float* A, const float* B, float* C, int N) {
    // Manually compute global thread index from block and thread IDs
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    
    // Guard against out-of-bounds memory access
    if (idx < N) {
        C[idx] = A[idx] + B[idx];
    }
}

// Host launch setup
void launch_cuda(const float* d_A, const float* d_B, float* d_C, int N) {
    int threads_per_block = 256;
    int blocks_per_grid = (N + threads_per_block - 1) / threads_per_block;
    vector_add_cuda<<<blocks_per_grid, threads_per_block>>>(d_A, d_B, d_C, N);
}

OpenAI Triton(블록 수준 파이썬 프로그래밍)

Triton에서는 파이썬으로 데이터 블록(타일)의 관점에서 직접 프로그래밍한다. Triton 컴파일러가 스레드 동기화, 메모리 코얼레싱, 벡터화를 내부적으로 자동 처리한다.

import triton
import triton.language as tl

@triton.jit
def vector_add_triton(A_ptr, B_ptr, C_ptr, N, BLOCK_SIZE: tl.constexpr):
    # Identify which block of data this program instance handles
    pid = tl.program_id(axis=0)
    
    # Create a vector of offsets for the entire block at once
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N
    
    # Load chunks of memory into fast SRAM, add them, and store back
    a = tl.load(A_ptr + offsets, mask=mask)
    b = tl.load(B_ptr + offsets, mask=mask)
    tl.store(C_ptr + offsets, a + b, mask=mask)

# Launch directly from PyTorch tensors
def launch_triton(A, B, C):
    N = A.numel()
    grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),)
    vector_add_triton[grid](A, B, C, N, BLOCK_SIZE=1024)

3. FlashAttention: 트랜스포머의 메모리 병목 해결

표준 셀프 어텐션은 고해상도 이미지나 긴 텍스트 시퀀스를 처리할 때, 어텐션 행렬 S = QK^T가 시퀀스 길이 N에 따라 제곱(N × N)으로 커지기 때문에 심각한 한계에 부딪힌다.

메모리 대역폭 병목(HBM 대 SRAM). 현대 NVIDIA GPU에는 두 가지 주요 메모리 계층이 있다:

HBM(고대역폭 메모리 / VRAM): 용량이 크지만(예: 80GB) 읽기와 쓰기가 상대적으로 느리다.
SRAM(온칩 공유 메모리): 매우 작지만(모든 스트리밍 멀티프로세서를 합쳐 약 20MB) HBM보다 대략 10배에서 20배 빠르다.

표준 어텐션에서 GPU는 S = QK^T를 계산하고, 거대한 N × N 행렬을 느린 HBM에 써 넣은 뒤, 소프트맥스를 적용하기 위해 다시 읽고, 확률을 HBM에 다시 쓰고, V와 곱하기 위해 또 한 번 읽는다. 텐서 코어는 대부분의 시간을 메모리 전송을 기다리며 유휴 상태로 보낸다.

타일링(N × N 행렬을 절대 구체화하지 않기). Tri Dao 등이 2022년에 발표한 FlashAttention은 커널 융합을 사용해 어텐션 계산 전체를 빠른 온칩 SRAM 안에 유지한다. Q, K, V를 SRAM에 들어가는 작은 블록으로 잘라 타일 단위로 어텐션을 계산하고, 최종 출력 행렬만 HBM에 다시 쓴다. 거대한 N × N 중간 행렬은 VRAM에 저장되지 않는다.

온라인 소프트맥스 기법. 보통 소프트맥스는 어느 한 원소를 정규화하기 전에 행 전체의 최댓값과 지수 합을 알아야 한다. FlashAttention은 블록 단위로 실행 중인 최댓값과 정규화 계수를 추적하는 온라인 소프트맥스 알고리즘을 사용하며, 새 타일을 처리할 때마다 부분 출력을 수학적으로 다시 스케일링한다.

역전파 중 재계산. FlashAttention은 순전파 중에 N × N 어텐션 행렬을 저장하지 않기 때문에, 학습 중에는 영리한 절충을 사용한다. 작은 정규화 통계만 HBM에 저장하고, 역전파 중에 빠른 SRAM 안에서 어텐션 타일을 즉석에서 다시 계산한다. SRAM에서 추가 연산을 하는 편이 HBM에서 큰 행렬을 읽는 것보다 훨씬 빠르며, 실제 경과 시간 기준 2배에서 4배 속도 향상을 얻는 동시에 메모리 사용량을 제곱 O(N²)에서 선형 O(N)으로 줄인다.

더 많은 뉴스: AI & 기술