PyTorch 어텐션 프로파일링 Flash, cuDNN, Math 백엔드 완벽 비교

들어가며: 어텐션, 프로파일러로 들여다보기

트랜스포머 기반 모델을 다루다 보면 어텐션 연산의 성능이 전체 모델 속도를 좌우한다는 사실을 체감하게 됩니다. 하지만 막상 프로파일러를 열어보면 수많은 커널과 복잡한 타임라인에 압도되기 쉽죠. 이 글은 Hugging Face 블로그의 'Profiling in PyTorch' 시리즈 중 세 번째 편으로, 어텐션을 직접 구현하고 PyTorch의 SDPA(Scaled Dot Product Attention) 백엔드들과 비교하며 프로파일러 추적을 읽는 법을 다룹니다.

시리즈의 목표는 단순합니다. 프로파일러 추적을 '읽을 줄 아는' 개발자가 되는 것. Part 1에서는 기본 수학 연산을, Part 2에서는 선형 레이어와 MLP를 프로파일링했습니다. 이번 편에서는 어텐션을 다룹니다. 각 백엔드가 프로파일러에서 어떻게 다르게 보이는지 직접 확인해보세요.

참고: 이 글의 모든 실험은 NVIDIA A100-SXM4-80GB GPU에서 수행되었으며, 원본 스크립트는 Hugging Face 블로그에서 확인할 수 있습니다.

1. 순수 어텐션 구현과 프로파일링

먼저 어텐션을 기본 연산으로 직접 구현해봅시다. 어텐션은 다음과 같은 단계로 구성됩니다:

1. 점수 계산: matmul(q, k.T)
2. 스케일링: scores * scale
3. 인과 마스크 적용: scores.masked_fill(mask, "-inf")
4. 소프트맥스: softmax(scores)
5. 값 재가중치: matmul(attn, v)

import torch
import torch.nn as nn
import math

class NaiveCausalAttention(nn.Module):
    def __init__(self, head_dim):
        super().__init__()
        self.scale = 1.0 / math.sqrt(head_dim)

    def forward(self, q, k, v, mask):
        # q, k, v: [batch, heads, seq, head_dim]
        scores = torch.matmul(q, k.transpose(-2, -1))  # [batch, heads, seq, seq]
        scores = scores * self.scale
        scores = scores.masked_fill(mask, float("-inf"))
        attn = torch.softmax(scores, dim=-1)
        out = torch.matmul(attn, v)  # [batch, heads, seq, head_dim]
        return out

이 모듈을 프로파일링하면 CPU 레인에서 matmul, mul, masked_fill, softmax, matmul이 순서대로 나타납니다. GPU 레인을 펼쳐보면 예상치 못한 커널이 하나 보입니다. 바로 메모리 복사(Memcpy) 커널이죠.

숨겨진 Memcpy의 정체

masked_fill은 out-of-place 연산입니다. PyTorch는 기본적으로 텐서를 복사한 뒤 연산을 적용하고 결과를 반환합니다. 이 과정에서 Memcpy 커널이 발생하는 것이죠. 이를 in-place 연산(masked_fill_)으로 바꾸면 어떻게 될까요?

def forward(self, q, k, v, mask):
    scores = torch.matmul(q, k.transpose(-2, -1))
    scores = torch.mul(scores, self.scale)
    scores.masked_fill_(mask, float("-inf"))  # in-place로 변경
    attn = torch.softmax(scores, dim=-1)
    out = torch.matmul(attn, v)
    return out

GPU 레인에서 Memcpy 커널이 사라진 것을 확인할 수 있습니다. 단 한 줄의 변경으로 포워드 패스마다 커널 하나를 제거한 셈이죠. 레이어가 많고 시퀀스가 긴 LLM에서는 이 차이가 누적되어 큰 성능 향상으로 이어집니다.

주의: in-place 연산은 autograd가 역전파를 위해 저장한 텐서 값을 덮어쓸 수 있습니다. torch.no_grad() 컨텍스트 내에서만 안전하게 사용하세요.

2. SDPA 백엔드별 프로파일러 추적 비교

PyTorch는 어텐션 파이프라인 전체를 F.scaled_dot_product_attention 한 줄로 제공합니다. 이 함수는 내부적으로 여러 백엔드 중 최적의 것을 선택하는데, torch.nn.attention.sdpa_kernel 컨텍스트 매니저로 특정 백엔드를 강제할 수 있습니다.

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

with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    out = F.scaled_dot_product_attention(q, k, v, is_causal=True)

각 백엔드의 프로파일링 결과를 비교하면 다음과 같습니다:

math 백엔드: 레퍼런스 구현의 함정

가장 먼저 확인할 것은 math 백엔드입니다. 한 줄짜리 함수로 바꿨으니 빨라질 것 같지만, 실제로는 순수 구현보다 3.7배 느립니다. 이유는 명확합니다:

• Tensor Core 미사용: sgemm 커널은 FP32 CUDA 코어에서 실행됩니다. bf16 입력도 FP32로 업캐스트되어 데이터 이동량이 두 배가 됩니다.
• 마스크 매번 재생성: is_causal=True 플래그는 편의를 제공할 뿐, 내부적으로는 매 호출마다 [seq, seq] 삼각 행렬을 생성합니다.
• _safe_softmax: 완전히 마스킹된 행에서 NaN을 방지하는 안전 장치가 추가 커널을 발생시킵니다.

math 백엔드는 정확성이 보장되는 레퍼런스 구현입니다. 빠르지 않지만 항상 동작하죠. 다른 백엔드들은 이 20개의 커널을 하나의 융합 커널로 접어넣는 것을 목표로 합니다.

flash 백엔드: 낮은 점유율이 함정이 아닌 이유

FlashAttention-2는 가장 빠른 백엔드지만, 프로파일러는 **13%의 낮은 점유율(occupancy)**을 보고합니다. 이는 최적화가 덜 된 것이 아니라 의도된 설계입니다.

Flash는 타일 단위로 k와 v를 순회하며 온라인 소프트맥스 기법으로 전체 [seq, seq] 행렬을 HBM에 쓰지 않습니다. 대신 많은 레지스터와 공유 메모리를 사용해 데이터를 온칩에 유지하죠. 블록당 128스레드 × 255레지스터 = 32,640 레지스터가 필요해, 65,536 레지스터를 가진 SM에는 두 블록만 상주할 수 있습니다. 결과적으로 8 워프만 활성화되어 약 13% 점유율이 됩니다.

점유율이 낮다고 성능이 나쁜 것은 아닙니다. Flash는 레지스터와 공유 메모리를 의도적으로 많이 사용해 데이터 재사용을 극대화하고, 전역 메모리 접근을 최소화합니다.

cuDNN 백엔드: 생성된 커널의 양면성

cuDNN은 NVIDIA의 딥러닝 라이브러리로, 문제별로 커널을 생성하고 튜닝합니다. 커널 이름 cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0에서 알 수 있듯, 사전 컴파일된 바이너리가 아닙니다.

장점: 네이티브 [B, H, S, D] 레이아웃을 직접 소비하므로 aten::transpose가 전혀 발생하지 않습니다. flash와 efficient는 각각 4개의 메타데이터 전치를 삽입하죠.

단점: CPU 오버헤드가 큽니다. aten::scaled_dot_product_attention self time이 전체 실행의 26%를 차지하며, 이는 cuDNN 런타임이 매 호출마다 "knob" 검색을 수행하기 때문입니다. 프로파일러 추적이 깔끔해 보인다고 해서 작업이 사라진 것은 아닙니다. 작업이 프로파일러가 분해할 수 없는 곳으로 이동한 것이죠.

실무 팁: 어텐션 성능 최적화 시 먼저 F.scaled_dot_product_attention을 사용하세요. 특정 백엔드가 더 나은 성능을 보인다면 sdpa_kernel 컨텍스트로 고정할 수 있습니다. 하지만 항상 프로파일러로 실제 병목을 확인한 후 결정하세요.

결론: 추측하고, 확인하고, 불일치를 찾아라

이 시리즈 전체를 관통하는 하나의 습관이 있습니다. 먼저 추측하고, 그 다음에 확인하라는 것입니다.

프로파일러 추적을 열기 전에 "이 커널들이 보일 것이다"라고 말해보세요. 그리고 추적을 열어 예상과 다르면, 그 불일치를 화면에서 가장 흥미로운 것으로 취급하세요. 이 글의 모든 인사이트 — 숨겨진 Memcpy, math 백엔드의 20개 커널, flash의 낮은 점유율, cuDNN의 CPU 오버헤드 — 는 모두 추측과 추적의 불일치에서 나왔습니다.

프로파일링은 GPU 전문가만의 기술이 아닙니다. "잠깐, 왜 저러지?"라고 질문하고 답이 떠오를 때까지 들여다보는 습관일 뿐입니다. 이제 여러분은 자신의 모델에서 그 습관을 실천할 어휘와 반사신경을 갖추었습니다. 추적을 열고, 추측하고, 불일치를 찾아가세요.

한국 개발 생태계에서의 적용 맥락

국내에서 LLM이나 비전 모델을 파인튜닝하거나 서빙하는 팀이라면, 어텐션 최적화는 곧 비용 절감과 직결됩니다. 특히 A100/H100 GPU를 임대해 쓰는 환경에서는 F.scaled_dot_product_attention으로의 전환이 즉각적인 성능 향상을 가져올 수 있습니다. 다만 cuDNN 백엔드는 CPU 오버헤드가 크므로, 배치 크기가 작거나 호출 빈도가 높은 서빙 환경에서는 flash 백엔드가 더 유리할 수 있습니다. 실제 서비스에 적용하기 전에 반드시 자체 워크로드로 프로파일링하세요.

함께 보면 좋은 글

• Oracle에서 PostgreSQL로, AI가 마이그레이션을 혁신하는 방법 (feat. Azure)
• CSS 하이라이트 가상 요소 완벽 정리 search-text부터 target-text까지