SDPA란 무엇인가? — 수동으로 작성한 Attention을 한 줄로 대체하기 전에 알아야 할 3가지 포인트 - Insights

PyTorch에서
softmax(Q @ K.T / sqrt(d)) @ V

를 직접 작성해 본 사람이 있나요?

이 전제 지식은 이 한 줄뿐입니다. Attention의 메커니즘 자체도 본문에서 1부터 다시 정의합니다.

글을 다 읽고 나면, 손으로 작성한 Attention을 F.scaled_dot_product_attention

로 대체해도 되는지 판단할 수 있고, 만약 대체한다면 마스크(mask), 드롭아웃(dropout), 인과 마스크(causal mask)의 세 가지 인수를 어떻게 수정해야 할지 결정할 수 있게 됩니다.

대체 여부를 판단하는 자료: 수동 버전과의 수치 차이, 메모리 및 속도 실측값
대체 시 조용히 깨지는 3가지 지점과 각각을 막는 방법
그대로 붙여넣기 가능한 대체 완료된 Attention 모듈
FlashAttention 커널 내부의 상세 내용. 아이디어만 언급하고 자세한 내용은 관련 기사로 넘깁니다.
GPU에서의 속도 비교. 이번 측정은 모두 CPU입니다.
FlexAttention 등 SDPA 외부의 새로운 API

소형 GPT를 스크래치(scratch)로 구축하는 교재를 만들었을 때, Attention은 교과서대로 손으로 작성했습니다. 교재로서는 그것이 정답이었습니다. 그런데 Hugging Face의 Transformers 라이브러리를 읽다 보니 attn_implementation="sdpa"

가 당연하다는 듯 나열되어 있어서, 제가 직접 작성한 버전을 이 하나의 함수로 대체해도 되는지, 대체한다면 무엇이 달라지는지를 설명할 수 없었습니다. 그것이 분해서 공식 문서를 처음부터 끝까지 읽고, 제 손으로 측정해 보았습니다.

독자들에게 생각해 주었으면 하는 것은 속도에 대한 이야기가 아니라, 같은 식을 계산하는 함수라도 인수의 의미가 다르면 조용히 깨진다는 점입니다. 대체하기 전에 등가성을 측정하는 습관을 자신의 코드베이스에서 어떻게 돌릴 수 있을지. 그 자료로 읽어주세요.

측정 환경은 GPU가 없는 클라우드 작업 환경입니다. Xeon 2코어, 메모리 약 8GB로 빈약하지만, 그만큼 수동 버전이 8GiB 확보에 실패하여 다운되는 경계가 명확하게 보였습니다.

항목	값
PyTorch	2.14.1
...
이 섹션에서 알 수 있는 것: SDPA가 어떤 함수로 계산하는지, 그리고 수동 버전과 결과가 어디까지 일치하는지

본문에서는 Attention을 도서관의 레퍼런스 사서에 비유합니다. 이용자가 질문 카드를 들고 오면, 사서는 장서의 책등(spine)과 한 권씩 대조하며 관련성이 깊어 보이는 책 내용들을 비율을 정해 섞어 답으로 돌려줍니다. 이 작업 1회 분량이 Attention입니다.

기호는 사용하기 직전에 하나씩 도입합니다.

$Q$는 Query(질문 카드 다발)이며, 타입은 4차원 텐서이고, 형태는 $(N, H_q, L, E)$입니다.
$L$은 질문 카드의 장수이며, 정수입니다. 공식 문서에서는 target sequence length라고 불립니다.
$E$는 카드 한 장에 적을 수 있는 항목 수이며, 정수입니다. 임베딩 차원(embedding dimension)이라고도 부릅니다.
$K$는 Key(책등 다발)이며, 형태는 $(N, H, S, E)$입니다.
$S$는 장서의 권수이며, 정수입니다. source sequence length라고 불립니다.
$V$는 Value(책 내용 다발)이며, 형태는 $(N, H, S, E_v)$입니다.
$B$는 Bias(책에 붙이는 열람표)이며, 형태는 $(L, S)$의 행렬이고, 값은 열람 가능하면 $0$, 열람 금지면 $-\infty$입니다.

도구가 갖춰졌으니 식을 작성합니다.

\mathrm{Attention}(Q, K, V) = \mathrm{softmax}\!\left(\frac{QK^{\top}}{\sqrt{E}} + B\right) V

이 식이 말하는 것은 다음 4단계입니다.

$QK^{\top}$로 모든 질문 카드 $\times$ 모든 책등의 적합성 표를 만듭니다. 크기는 $L \times S$입니다.
$\sqrt{E}$로 나누어, 항목 수가 많을수록 적합 점이 부풀어 오르는 편향을 상쇄합니다.
$B$를 더하여, 열람 금지된 책의 적합 점을 $-\infty$로 만듭니다.
softmax로 적합 점을 '어떤 책을 몇 할 읽을지'의 배분으로 바꾸고, 그 비율로 $V$를 섞습니다.

기호	읽는 법	도서관에서의 역할	타입
$Q$	Query (쿼리)	질문 카드 다발	텐서
...
이 식을 그대로 하나의 함수로 만든 것이 `torch.nn.functional.scaled_dot_product_attention`

, 줄여서 SDPA입니다. 시그니처는 다음과 같으며, scale만 키워드 인수로만 전달할 수 있습니다.

F.scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0,
is_causal=False, scale=None, enable_gqa=False)

공식 문서는 이 함수가 위의 수식과 등가인 효율적인 구현이라고 설명하고 있습니다. 말만으로는 대체 여부를 판단할 근거가 되지 않기 때문에, 직접 작성한 버전(handwritten version)과 비교했습니다.

import math
import torch
import torch.nn.functional as F
...

dtype	인과 마스크	직접 작성한 버전과의 최대 절대 오차
float32	없음	4.47e-07
...
float32에서 $10^{-7}$, float64에서 $10^{-15}$의 차이는 둘 다 해당 타입의 반올림 오차 수준입니다. SDPA는 직접 작성한 버전과 동일한 답을 반환한다고 말할 수 있는 결과였습니다.

이 섹션에서 알 수 있는 것: SDPA로 대체했을 때 메모리와 시간 변화량

앞선 절에서, SDPA의 출력은 직접 작성한 버전과 반올림 오차 범위 내에서 일치함을 확인했습니다. 답이 같다면 차이는 계산 과정에만 있을 뿐입니다. 그래서 중간에 사용되는 메모리와 시간을 측정합니다.

직접 작성한 버전(handwritten version)은 유사도 행렬 $QK^{ op}$ 전체를 메모리에 올립니다. 도서관으로 비유하자면, 모든 질문 카드 $ imes$ 모든 장서의 표를 책상 가득 펼쳐놓고 작업을 시작하는 방식입니다. 이 표의 크기는 다음과 같이 표현할 수 있습니다. 새로운 기호는 없습니다.

\text{유사도 행렬 바이트 수} = H \times L \times S \times 4

이 공식이 의미하는 것은, 사서 $H$ 명이 각각 $L imes S$ 크기의 표를 가지고 있으며, 한 칸이 float32의 4바이트를 차지한다는 것입니다. $L$과 $S$를 2배로 늘리면 표는 4배가 됩니다.

반면, SDPA가 선택하는 융합 커널(여러 연산을 하나로 통합하고 중간 결과를 메모리에 쓰지 않는 구현)은 이 표 전체를 펼치지 않습니다. 선반 구획별로 책등을 보며 다니면서, 그 시점의 최대 유사도 점과 분배된 합계만 손에 들고 적어 나갑니다. 마지막에 메모에서 계산을 맞추기 때문에, 책상은 작은 상태를 유지할 수 있습니다. 이 생각의 자세한 내용은 Flash Attention 기사에 작성했습니다.

import subprocess
import sys
CHILD = """
...

프로세스를 하나씩 별도로 실행하며 최대 상주 메모리 증가량을 측정하고 있습니다. $H=8$, $E=64$, float32, 추론 모드입니다.

L	직접 작성한 버전	SDPA・MATH 강제	SDPA・지정 없음
1024	116MB / 0.03s	99MB / 0.04s	19MB / 0.01s
...
읽어낼 수 있는 것은 세 가지입니다.

직접 작성한 버전의 $L=8192$는 4117MB였습니다. 공식으로 계산된 유사도 행렬은 $8 \times 8192^2 \times 4$ 바이트로 2048MiB이므로, 유사도 행렬과 softmax 이후의 행렬 두 장 분량이 딱 들어갑니다.
$L=16384$의 직접 작성한 버전은,
you tried to allocate 8589934592 bytes
에서 실패했습니다. $8 \times 16384^2 \times 4 = 8{,}589{,}934{,}592$로, 위의 공식과 바이트 단위까지 일치합니다. 지정 없음(None)의 SDPA는, $L$을 16배로 늘려도 19MB에서 49MB로만 증가했고, 속도 역시 직접 작성한 버전의 약 4배였습니다.

예상 밖이었던 것은 가운데 열입니다. 같은 SDPA임에도 MATH 백엔드를 강제하면 직접 작성한 버전보다 무거워졌습니다. 성능을 결정하는 것은 SDPA라는 함수 이름이 아니라, 내부적으로 누가 계산하고 있느냐입니다.

GPU에서도 상황은 같아서, 직접 작성한 버전의 유사도 행렬은 VRAM을 $L^2$에 비례하여 사용합니다.

이 섹션에서 알 수 있는 것: SDPA 내부를 선택하는 메커니즘과 자신의 환경에서 무엇이 선택되었는지 확인하는 방법

앞선 절에서, 같은 SDPA임에도 MATH를 강제하면 직접 작성한 버전보다 무거워진 것을 보았습니다. 즉 대체 효과는 SDPA가 어떤 구현을 선택하느냐에 따라 결정됩니다. 여기서는 그 선택 방식을 확인합니다.

SDPA는 호출될 때마다 입력의 형태, 타입, 장치(device)를 보고 조건을 충족하는 구현체 중 최적의 것을 자동으로 선택합니다. 이 내부적인 구현체를 백엔드(backend)라고 부릅니다. 도서관에 비유하자면, 같은 창구 안쪽에 작업 스타일이 다른 담당자 팀들이 대기하고 있으며, 질문 종류에 따라 담당자가 배정되는 시스템과 같습니다.

백엔드	내용물	도서관에서의 작업 스타일
`FLASH_ATTENTION`	FlashAttention-2 계열의 융합 커널	구역별로 순회하며 메모만 사용하는 방식
`EFFICIENT_ATTENTION`	xFormers 기반 메모리 효율 커널	역시 표를 펼치지 않는 별도의 팀
`CUDNN_ATTENTION`	NVIDIA cuDNN 구현	외부에서 온 전문 팀
`MATH`	수식을 그대로 PyTorch 연산으로 구성한 것	표를 책상 가득 펼쳐 놓는 방식

후보군을 좁히고 싶을 때는 torch.nn.attention.sdpa_kernel

컨텍스트 매니저(context manager)를 사용합니다. CPU에서 하나씩 강제로 테스트해 보았습니다.

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
...

백엔드	CPU에서의 결과
MATH	작동함
...
여기서 두 번째 놀라운 점이 있었습니다. 공식 문서의 Note에는 최적화 커널을 호출할 때는 CUDA 백엔드를 사용하고, 그 외 경우에는 PyTorch 구현을 사용한다는 취지의 설명이 있습니다. 그런데 CPU에서도 `FLASH_ATTENTION`
가 작동했고, 이전 절에서 언급된 49MB라는 수치도 MATH로는 설명이 되지 않습니다. 그래서 프로파일러(Profiler)를 이용해 실제로 실행된 연산 이름을 확인했습니다.

import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel
...

호출 방식	실제로 실행된 주요 연산
지정 없음	`aten::_scaled_dot_product_flash_attention_for_cpu`
MATH 강제
`aten::_scaled_dot_product_attention_math`와 `aten::_safe_softmax`

PyTorch 2.14.1의 CPU에서는 지정 없이도 CPU용 Flash 계열 커널이 선택되었습니다. 문서의 설명보다 구현이 앞서 나간 것으로 보입니다. 자신의 환경에서 무엇이 선택될지는 문서를 읽는 것보다 프로파일러로 보는 것이 더 확실합니다. MATH 쪽에서 나오는 _safe_softmax
는 제4절에서 중요하게 다뤄집니다.

선택 과정을 정리하면 다음과 같습니다.

GPU에서 특정 융합 커널을 사용하고 싶다면, sdpa_kernel
로 MATH를 후보군에서 제외해 두면, 조건이 맞지 않아도 조용히 MATH로 떨어지는 것을 방지할 수 있습니다. 만약 떨어진 경우 그 이유는 경고(warning) 형태로 표시됩니다.

절	알게 된 점
1	SDPA는 손으로 작성한 버전과 반올림 오차 범위 내에서 같은 답을 반환한다
...
여기까지로도 대체할 가치가 있다고 할 수 있습니다. 다만 제1절에서 일치 여부를 확인한 것은 마스크가 없는 경우와 단순 인과 마스크의 두 가지 경우뿐이었습니다. 실제 모델에는 패딩 마스크(padding mask), 드롭아웃(dropout), 추론 시의 KV 캐시(KV cache) 등이 관련됩니다. 다음 절에서는 이 부분을 측정할 것입니다.

이 섹션에서 알 수 있는 것: 손으로 작성한 버전의 감각 그대로 인수를 전달했을 때 결과가 달라지는 3곳과 그 수정 방법

이전 절까지로, 수식은 같고 빠르며 최적화된 구현체가 자동으로 선택된다는 것을 확인했습니다. 하지만 수식이 같더라도, 인수의 의미가 손으로 작성한 버전이나 nn.MultiheadAttention
의 관행과 다른 부분이 3군데 있습니다. 이들 모두 오류를 발생시키지는 않지만, 숫자만 조용히 틀어집니다.

SDPA의 attn_mask에 부울 마스크(boolean mask)를 전달하면, True는 '이 쌍을 Attention에 사용한다'는 의미입니다. 반면, nn.MultiheadAttention
의 key_padding_mask

이 경우, True는 "패딩이므로 버린다"를 의미합니다. 공식 문서의 Note에서도 MHA(Multi-Head Attention)에서 전환할 때는 부울 마스크(Boolean mask)를 반전하라고 주의하고 있습니다. 도서관에 비유하자면, 같은 빨간 표지판이 한 건물에서는 열람 가능을, 다른 건물에서는 열람 금지를 의미하는 상태와 같습니다.

import torch
import torch.nn.functional as F
torch.manual_seed(0)
...

전달 방식	정답과의 최대 절대 오차
MHA의 방식을 그대로 전달	3.29e+00
반전하여 전달	5.96e-08

그대로 전달하면, 패딩만 읽고 실제 토큰을 무시하는 계산이 됩니다. 반전해야 일치합니다.

추가 결과도 적어두겠습니다. 모든 책이 열람 금지인 행, 즉 $B$의 1행이 모두 $-\infty$인 경우, 수기로 작성한 softmax는 NaN(Not a Number)을 반환합니다. SDPA는 CPU 지정 여부와 관계없이 MATH 강제 모드에서도 NaN 대신 0을 반환했습니다. MATH 경로에서는 제3절에서 본 _safe_softmax가 이러한 안전한 동작을 담당한다고 생각됩니다. 왼쪽 정렬 패딩으로 시작 행 전체가 마스킹되어도, SDPA라면 NaN이 전파되기 어렵다는 것입니다.

nn.Dropout은 model.eval()로 자동적으로 멈춥니다. 하지만 SDPA의 dropout_p는 멈추지 않습니다. 공식 문서 경고문에는 이렇게 쓰여 있습니다.

This function always applies dropout according to the specified dropout_p argument.

번역: 이 함수는, 지정된 dropout_p 인자(argument)에 따라 항상 dropout을 적용합니다.

도서관에 비유하자면, 교육용으로 "무작위로 책을 숨겨 사서를 훈련시키는 훈련"을 본 서비스 창구에서도 계속 해버리는 상태와 같습니다.

import torch
import torch.nn as nn
import torch.nn.functional as F
...

구현	eval 후, 동일 입력으로 2회 추론했을 때의 차이
`dropout_p=self.p` 고정	7.07e-01
`self.training`으로 분기	0.00e+00

torch.no_grad()를 붙여도 멈추지 않습니다. 추론할 때마다 출력이 바뀌는데, 오류나 경고가 나오지 않는 것이 이 함정의 무서운 점입니다.

문장 생성의 추론 과정에서는 과거 토큰의 $K$와 $V$를 캐시(cache)에 저장하고, 새로운 1개 토큰 분량의 $Q$만 계산합니다. 이때 $L=1$, $S$는 과거 길이이며, $L \ne S$가 됩니다.

is_causal=True의 인과 마스크(미래 토큰을 보지 못하게 하는 하삼각 마스크)는, $L \ne S$일 때 좌상단 정렬로 생성됩니다. $L=1, S=6$이라면, 열람 가능한 것은 첫 번째 책 한 권뿐입니다. 정말 필요한 것은 우하단 정렬, 즉 최신 질문 카드가 전 장서를 읽을 수 있는 형태입니다.

import torch
import torch.nn.functional as F
from torch.nn.attention.bias import causal_lower_right
...

전달 방식	정답과의 최대 절대 오차
`is_causal=True`	3.15e+00
`is_causal=False`	1.19e-07
`attn_mask=causal_lower_right(1, S)`	1.19e-07

마지막의 print는 [[True, False, False, False, False, False]]를 출력합니다. 최신 토큰이 문두의 첫 번째 토큰만 보고 있습니다. 토큰을 하나씩 생성한다면 is_causal=False로 충분합니다. 여러 토큰을 모아서 처리하는 청크 추론(chunk inference)에서는 우하단 정렬인 causal_lower_right(L, S)를 사용해야 합니다.

이 함정은 학습 시에는 절대 발생하지 않습니다. 학습 때는 $L = S$이기 때문에, 좌상단 정렬과 우하단 정렬이 일치하기 때문입니다. 추론에서만 품질이 떨어진다는 형태로 나중에 발견하게 됩니다.

이 섹션에서 알 수 있는 것: 수기 Attention(handwritten Attention)으로부터의 구체적인 대체 방법과 바로 사용할 수 있는 모듈

앞서 발견한 3가지 함정은 모두 인자 하나로 막을 수 있습니다. 따라서, 제1절의 수기 버전을 SDPA에 대체할 때의 차이점부터 살펴보겠습니다.

scores = q @ k.transpose(-2, -1) / math.sqrt(q.size(-1))
scores = scores.masked_fill(is_pad[:, None, None, :], float("-inf"))
scores = scores.masked_fill(~causal_tri, float("-inf"))
...
패딩과 인과 마스크를 모두 사용할 경우, 두 가지를 AND 연산으로 합성하여 attn_mask에 전달합니다. 이는 attn_mask와 is_causal=True를 동시에 전달하면 오류가 발생하기 때문입니다.

이를 헤드 분할(head splitting)과 출력 투영(output projection)까지 포함한 모듈로 만들었습니다. Query 헤드 수와 Key/Value 헤드 수를 별도로 지정하면, GQA (Grouped Query Attention: 여러 Query 헤드가 한 쌍의 Key/Value 헤드를 공유하는 방식)가 됩니다. 도서관에 비유하자면, 32명의 사서가 8개의 책장을 각각 4명씩 공유하는 형태입니다. SDPA에서는 enable_gqa=True를 전달하면, 책장 공유를 SDPA 측에서 처리해 줍니다.

import torch
import torch.nn as nn
import torch.nn.functional as F
...

이 모듈은 학습(training)과 전체 추론(full inference)을 가정하고 설계되었습니다. KV 캐시를 이용해 토큰 단위로 생성하는 경로를 추가할 때는, 트릭 3에서와 같이 is_causal=False 또는 causal_lower_right로 전환해야 합니다.

작동 확인을 위해 왼쪽 패딩(left-padding)이 적용된 입력과 패딩이 처음부터 없는 입력을 비교했습니다.

import torch
from sdpa_attention import CausalSelfAttention
torch.manual_seed(0)
...

KV 헤드 수	마스크 있음	마스크 누락	eval 후 재현성
8	0.0e+00	8.4e-01	0.0e+00
2 (GQA)	0.0e+00	7.1e-01	0.0e+00

마스크를 올바르게 전달하면 패딩의 유무와 관계없이 출력이 달라지지 않으며, eval 후에는 매번 동일한 출력이 반환됩니다. GQA에서도 마찬가지입니다.

기존 모델을 사용하는 입장이라면, 직접 모듈을 작성할 필요가 없습니다. 로드 시 지정하기만 하면 됩니다.

from transformers import AutoModelForCausalLM
model = AutoModelForCausalLM.from_pretrained("모델명", attn_implementation="sdpa")

Hugging Face 문서에서는 sdpa를 PyTorch 내장 SDPA 구현을 사용할 수 있는 선택지로 설명하고 있습니다. 3절의 sdpa_kernel로 감싸면, 생성 시 사용되는 커널(kernel)까지 제한할 수 있습니다.

실제로 발생했던 오류와 출력이 미묘하게 어긋나는 증상을 정리했습니다.

⚠️ [IMG:N] 형식 토큰은 이미지 placeholder 입니다. 번역하지 말고 원래 위치에 그대로 유지하세요.

SDPA 사용 시 발생하는 오류 및 해결 방법

증상/오류	원인	대처 방안
`Explicit attn_mask should not be set when is_causal=True`	마스크와 `is_causal=True`를 함께 사용함	인과 마스크(Causal Mask)를 직접 만들고, AND 연산으로 합성하여 `attn_mask`만 전달한다.
`takes from 3 to 6 positional arguments but 7 were given`	`scale`을 위치 인수(positional argument)로 전달함	키워드(`scale=0.125`와 같이)로 전달한다.
`Expected attn_mask dtype to be bool or float or to match query dtype`	int64 또는 float64 타입의 마스크를 사용함	`bool` 타입으로 하거나, `query`와 동일한 부동소수점(floating-point) 타입을 사용한다.
`The size of tensor a (32) must match the size of tensor b (8)`	Q와 KV의 헤드 개수가 다른데 GQA를 활성화하지 않음	`enable_gqa=True`를 전달한다.
`No viable backend for scaled_dot_product_attention was found`	`sdpa_kernel`에서 해당 환경에서 작동하지 않는 백엔드만 지정함	후보에 MATH를 포함하거나, 조건을 충족하는 입력으로 만든다.
파딩(Padding) 처리 후 출력이 이상함	부울 마스크의 True 의미가 반대임	`~mask`로 반전시킨다.
평가(eval) 이후에도 추론할 때마다 출력이 달라짐	`dropout_p`를 고정값으로 전달하고 있음	`self.p if self.training else 0.0`을 사용한다.
학습은 되는데 생성만 품질이 낮음	$L
e S$일 때 `is_causal=True`로 설정함	`is_causal=False` 또는 `causal_lower_right`를 사용한다.
긴 입력에서 메모리가 부족함	MATH가 선택되어 표(matrix)를 넓히고 있음	프로파일러(Profiler)로 연산명을 확인하고, 융합 커널의 조건을 충족시킨다.

주요 용어 정리

용어	의미	도서관에서의 대응 비유
SDPA	Scaled Dot Product Attention. PyTorch에서는 `F.scaled_dot_product_attention`을 사용한다.	레퍼런스 계수(reference coefficient)의 창구 그 자체
융합 커널 (Fusion Kernel)	여러 연산을 하나의 구현에 모아 중간 결과를 메모리에 쓰지 않는 계산 방식	표를 넓히지 않고 순회하며 메모만 기록하는 작업
백엔드 (Backend)	SDPA 내부에서 실제로 계산을 수행하는 구현체	창구 깊숙한 곳에 대기하는 담당자 팀
부울 마스크 (Boolean Mask)	True/False로, 어떤 쌍(pair)의 값을 사용할지 지정하는 마스크	열람표(閲覧札)

SDPA란 무엇인가? — 수동으로 작성한 Attention을 한 줄로 대체하기 전에 알아야 할 3가지 포인트

요약

핵심 포인트

SDPA 사용 시 발생하는 오류 및 해결 방법

주요 용어 정리

관련 개념 및 참고 자료

댓글