본문 바로가기
TYLER SONGBlog
← Blog index
LLM Research7 min read

Scaled Dot-Product Attention 구현: 직접 만든 attention이 F.scaled_dot_product_attention과 같은 숫자를 내는지 확인하기

Attention의 핵심 계산인 Scaled Dot-Product Attention을 PyTorch로 직접 구현하고, 라이브러리 함수 F.scaled_dot_product_attention과 결과를 비교합니다. batch·head 차원이 어떻게 붙는지, causal mask와 드롭아웃이 어디에 들어가는지 익힙니다. 시퀀스 길이를 늘리면 attention 메

By 송민성

오늘의 목표

직접 짠 Scaled Dot-Product Attention(SDPA) 코드가 PyTorch의 F.scaled_dot_product_attention과 같은 출력을 내는지 숫자로 확인한다. 그 과정에서 텐서의 shape이 어떻게 흘러가는지, causal mask가 미래 토큰을 정말 막는지 단위 테스트로 검증한다.

개념과 직관

Attention의 한 줄 요약: 각 토큰이 "나는 어떤 토큰의 정보를 얼마나 가져올까?"를 정해서, 다른 토큰들의 값(value)을 가중 평균해 오는 연산이다.

도서관 비유로 보자.

  • Query(Q, 질의): 내가 찾는 주제가 적힌 쪽지.
  • Key(K, 키): 각 책의 표지에 적힌 주제 라벨.
  • Value(V, 값): 각 책의 실제 내용.

쪽지(Q)와 라벨(K)이 얼마나 비슷한지 점수를 매기고, 점수가 높은 책의 내용(V)을 더 많이 섞어서 가져온다.

수식 (텍스트)

text
scores  = Q @ K^T / sqrt(d)
weights = softmax(scores)       # 마지막 축 기준, 각 행의 합 = 1
output  = weights @ V
  • @는 행렬 곱, K^T는 K의 전치(transpose, 행과 열을 뒤바꾼 것)다.
  • d는 Q, K 벡터의 길이(head_dim)다.
  • softmax는 숫자들을 합이 1인 확률처럼 바꿔 주는 함수다.
  • 왜 sqrt(d)로 나눌까? 벡터 길이 d가 커질수록 내적 값이 커지고, 값이 크면 softmax가 한 곳에 거의 100%를 몰아주는 극단적인 분포가 된다. 나눠서 값의 크기를 안정시키는 것이다.

Shape 직관: batch와 head 차원

실전에서 텐서는 이렇게 생겼다.

text
Q, K, V: (batch, heads, seq_len, head_dim)
scores : (batch, heads, seq_len, seq_len)
output : (batch, heads, seq_len, head_dim)
  • batch: 문장 여러 개를 한 번에 처리하는 묶음 크기.
  • head: attention을 여러 개의 독립된 시점으로 나눠 동시에 하는 것(멀티헤드). 오늘은 "head도 batch처럼 앞에 붙은 독립 차원"이라고만 이해하면 충분하다. 행렬 곱은 마지막 두 축에만 적용되고 앞쪽 축은 그대로 따라간다.
  • seq_len × seq_len: 모든 토큰 쌍의 점수표다. 오늘 실험의 핵심이다.

Causal mask (인과 마스크)

글을 생성하는 모델은 다음 단어를 맞혀야 하므로 미래 토큰을 보면 안 된다(정답을 훔쳐보는 셈). 점수표에서 "미래" 위치(대각선 위쪽 삼각형)를 -inf로 덮으면 softmax 후 그 가중치가 0이 된다.

text
허용(1) 여부, seq_len=4
토큰0: 1 0 0 0
토큰1: 1 1 0 0
토큰2: 1 1 1 0
토큰3: 1 1 1 1

어텐션 드롭아웃 (attention dropout): 학습 중 softmax로 얻은 가중치의 일부를 무작위로 0으로 만들어 특정 연결에 과하게 의존하지 않게 하는 규제 기법이다. 평가(eval) 때는 꺼야 한다.

코드로 직접 해보기

1단계: 라이브러리 없이 직접 구현

python
import math import torch import torch.nn as nn import torch.nn.functional as F class SDPA(nn.Module): def __init__(self, dropout_p=0.0): super().__init__() self.dropout = nn.Dropout(dropout_p) def forward(self, q, k, v, is_causal=False): # q, k, v: (batch, heads, seq_len, head_dim) d = q.size(-1) scores = q @ k.transpose(-2, -1) / math.sqrt(d) # (B, H, L, L) if is_causal: L, S = q.size(-2), k.size(-2) mask = torch.ones(L, S, dtype=torch.bool, device=q.device).tril() scores = scores.masked_fill(~mask, float("-inf")) weights = torch.softmax(scores, dim=-1) weights = self.dropout(weights) return weights @ v, weights

tril()은 아래쪽 삼각형(대각선 포함)만 남기는 함수다. ~mask는 True/False를 뒤집어 "가려야 할 위치"를 만든다.

2단계: 라이브러리와 수치 비교

python
torch.manual_seed(0) B, H, L, D = 2, 4, 8, 16 q = torch.randn(B, H, L, D) k = torch.randn(B, H, L, D) v = torch.randn(B, H, L, D) mine = SDPA(dropout_p=0.0).eval() for causal in [False, True]: out_mine, _ = mine(q, k, v, is_causal=causal) out_ref = F.scaled_dot_product_attention(q, k, v, is_causal=causal) print(causal, torch.allclose(out_mine, out_ref, atol=1e-5), (out_mine - out_ref).abs().max().item())

torch.allclose는 두 텐서가 허용 오차 안에서 같은지 검사한다. 부동소수점(float, 소수를 근사해서 저장하는 방식)은 계산 순서가 다르면 아주 작은 오차가 생기므로 == 대신 이 함수를 쓴다. True True가 나오면 성공이다.

3단계: causal mask 단위 테스트

python
def test_causal_mask(): torch.manual_seed(0) q = torch.randn(1, 1, 5, 8) k = torch.randn(1, 1, 5, 8) v = torch.randn(1, 1, 5, 8) _, w = SDPA().eval()(q, k, v, is_causal=True) # 1) 미래 위치의 가중치는 정확히 0 assert torch.all(w.triu(diagonal=1) == 0) # 2) 각 행의 합은 1 assert torch.allclose(w.sum(-1), torch.ones(1, 1, 5)) # 3) 미래 토큰을 바꿔도 앞쪽 출력은 변하지 않는다 k2, v2 = k.clone(), v.clone() k2[..., 3:, :] = torch.randn(1, 1, 2, 8) v2[..., 3:, :] = torch.randn(1, 1, 2, 8) o1, _ = SDPA().eval()(q, k, v, is_causal=True) o2, _ = SDPA().eval()(q, k2, v2, is_causal=True) assert torch.allclose(o1[..., :3, :], o2[..., :3, :]) test_causal_mask() print("causal mask OK")

3번 테스트가 가장 강력하다. 가중치 모양만 보는 것이 아니라 "미래를 바꿔도 과거 출력이 안 변한다"는 인과성 자체를 검증한다.

4단계: 드롭아웃 동작 확인

python
m = SDPA(dropout_p=0.5) m.train() _, w_train = m(q, k, v) m.eval() _, w_eval = m(q, k, v) print((w_train == 0).float().mean().item()) # 대략 0.5 근처 print((w_eval == 0).float().mean().item()) # 0.0

오늘의 실험

질문: 시퀀스 길이 L을 늘리면 attention의 메모리는 어떻게 커지는가?

점수표 scores는 (B, H, L, L) 모양이므로 원소 개수가 L의 제곱에 비례한다. 즉 L을 2배 하면 4배가 될 것으로 예상된다. 이를 직접 측정한다.

python
import time def measure(L, B=1, H=8, D=64, device="cpu"): q = torch.randn(B, H, L, D, device=device) k = torch.randn(B, H, L, D, device=device) v = torch.randn(B, H, L, D, device=device) theory_mb = B * H * L * L * 4 / 1024**2 # float32 = 4바이트 t = time.time() SDPA().eval()(q, k, v) elapsed = time.time() - t return theory_mb, elapsed for L in [128, 256, 512, 1024, 2048, 4096]: mb, sec = measure(L) print(f"L={L:5d} score matrix ≈ {mb:8.1f} MB time={sec:.3f}s")

GPU가 있다면 device="cuda"로 바꾸고, 실제 최대 메모리를 다음처럼 읽는다.

python
torch.cuda.reset_peak_memory_stats() # ... attention 실행 ... print(torch.cuda.max_memory_allocated() / 1024**2, "MB")

기록할 것 (표로 정리):

L이론 점수표 MB실측 시간(GPU) peak MB직전 L 대비 배율

관찰 포인트

  • L이 2배일 때 이론 메모리가 정확히 4배가 되는가?
  • 직접 구현과 F.scaled_dot_product_attention의 peak 메모리를 비교하면 어떤가? 환경(GPU 종류, PyTorch 버전)에 따라 라이브러리가 내부적으로 다른 커널을 고를 수 있어서 차이가 나기도 하고 나지 않기도 한다. 결과를 추측하지 말고 측정값을 그대로 적는다.
  • 어느 L에서 메모리 부족(OOM, out of memory)이 나는가?

왜 이걸 배우나 (LLM / Post-training 연결)

LLM의 모든 Transformer 층 안에는 오늘 만든 연산이 들어 있고, 긴 문맥(long context)이 비싼 이유가 바로 L² 증가다. 이후 Day에서 멀티헤드 Attention과 Transformer 블록을 조립할 때 오늘 만든 SDPA가 그대로 부품으로 쓰이며, SFT·선호 학습 같은 Post-training에서 시퀀스 길이와 배치 크기를 정할 때의 메모리 감각의 출발점이 된다. 라이브러리와 수치를 맞춰 보는 습관은 연구 코드의 버그를 잡는 기본 기술이다.

초보자가 막히는 지점

  • shape 혼동: k.transpose(-2, -1)을 빼먹거나 k.T를 쓰면 에러가 나거나 엉뚱한 결과가 나온다. 4차원 텐서에서 .T는 의도와 다르게 동작하므로 항상 마지막 두 축만 바꾸는 transpose(-2, -1)을 쓴다. 헷갈리면 중간마다 print(x.shape)를 찍는다.
  • softmax 축 실수: dim=-1(키 방향)이어야 각 행의 합이 1이다. weights.sum(-1)이 1인지 확인하면 바로 잡힌다.
  • mask 방향 반대: tril과 triu를 헷갈리면 과거를 가리고 미래를 보게 된다. 단위 테스트 3번(미래를 바꿔도 과거 출력 불변)이 이를 잡아준다.
  • `allclose`가 False: 먼저 (a - b).abs().max()로 오차 크기를 본다. 1e-6 수준이면 부동소수점 오차이므로 atol을 조정하고, 1e-1처럼 크면 논리 버그다. 또 비교할 때는 드롭아웃을 0으로 하거나 eval() 상태여야 한다. 드롭아웃은 무작위라 두 결과가 달라진다.
  • `dropout_p`는 학습 때만: F.scaled_dot_product_attention(..., dropout_p=0.1)은 training 상태를 자동으로 보지 않는다. 평가할 때는 직접 0.0을 넘겨야 한다.
  • `-inf`와 NaN: 어떤 행이 전부 -inf면 softmax가 NaN이 된다. causal mask에서는 대각선이 항상 열려 있어 생기지 않지만, 패딩 마스크를 추가하게 되면 만날 수 있다.
  • 메모리 측정 착시: CPU 시간은 측정마다 흔들린다. 같은 설정을 여러 번 돌려 중앙값을 적고, 첫 실행은 워밍업이라 버린다.

알고리즘 (병행 1시간)

LeetCode Medium 2문제(DP, 동적 계획법: 큰 문제를 작은 문제로 쪼개 답을 저장해 두고 재사용하는 방법)를 푼다. 예를 들어 1차원 DP 유형(계단 오르기 변형, 집 도둑 House Robber 같은 문제)에서 고르면 부담이 적다.

각 문제마다 다음을 기록한다.

  • 접근: 상태(dp[i]가 무엇을 뜻하는지)와 점화식(이전 값으로 현재 값을 구하는 식)을 한두 문장으로.
  • 시간·공간 복잡도: 예) 시간 O(n), 공간 O(n) → 변수 두 개로 줄여 O(1) 가능한지 생각.
  • 틀린 이유: 틀리거나 막혔다면 원인(초기값 오류, 인덱스 범위, 상태 정의 모호 등)을 한 줄로.

시간 제한은 문제당 약 30분이다. 막히면 30분 후 풀이를 보고, 덮은 뒤 처음부터 다시 직접 구현한다.