PyTorch 프로파일링 (파트 3): Attention이 전부다

PyTorch 프로파일링 (파트 3): Attention이 전부다

이 블로그 포스트는 NVIDIA A100-SXM4-80GB GPU를 사용하여 PyTorch의 어텐션 메커니즘을 프로파일링하고, 각 변형이 프로파일러 트레이스에 어떻게 나타나는지 보여줍니다.

단순 어텐션

원시 연산 (matmul, mul, masked_fill, softmax, matmul)으로 구성된 단순 어텐션 구현은 순전파당 6개의 GPU 커널을 실행하며, 여기에는 out‑of‑place masked_fill에 의해 발생하는 예상치 못한 메모리 복사가 포함됩니다.

인플레이스 인과 마스킹을 적용한 단순 어텐션

masked_fill을 인플레이스 버전인 masked_fill_으로 교체하면 메모리 복사 커널이 사라져 순전파당 GPU 커널 수가 6개에서 5개로 감소합니다.

스케일드 닷 프로덕트 어텐션

PyTorch의 F.scaled_dot_product_attention은 단일 라인 인터페이스를 제공하며 여러 백엔드로 디스패치됩니다; 각 백엔드는 고유한 프로파일링 특성을 보여줍니다.

Math 백엔드

Math 백엔드에 고정하면 SDPA는 순전파당 약 20개의 GPU 커널을 실행하고, CUDA 코어에서 FP32로 동작하며, 매 호출마다 인과 마스크를 재구성하고, NaN을 방지하기 위해 _safe_softmax를 사용합니다; 이는 단순 인플레이스 버전보다 대략 3.7배 느립니다.

Efficient 백엔드

Efficient 백엔드(xformers)는 순전파당 단일 융합 fmha_cutlassF_bf16_aligned_64x64_rf_sm80 커널을 실행하여 연산을 Tensor 코어에서 bfloat16으로 유지하고, 중간 매트릭스 쓰기를 HBM에 기록하지 않습니다.

Flash 백엔드

Flash 백엔드는 순전파당 단일 융합 pytorch_flash 커널(FlashAttention‑2)을 실행합니다; 프로파일러에서 추정 점유율이 낮게(약 13%) 표시되지만, 레지스터와 공유 메모리를 사용해 어텐션 타일을 칩 내에 유지하기 때문에 가장 빠른 백엔드입니다.

cuDNN 백엔드

cuDNN 백엔드는 순전파당 단일 생성 커널을 실행합니다(예: cudnn_generated_fort_native_sdpa_sm80_flash_fprop_wmma_f16_knob_6_128x64x64_4x1x1_cga1x1x1_kernel0_0); CPU에서 전치 연산을 피하지만 런타임 플랜 선택으로 인해 CPU 시간이 더 높게(~214 µs) 소요되며, GPU 시간은 Efficient와 Flash 백엔드 사이에 위치합니다.

전체 내용 한눈에 보기

변형	변경 내용	순전파당 커널 수	트레이스에서 밝혀진 내용
단순 어텐션	원시 연산(matmul, mul, mask, softmax, matmul)으로 직접 구현한 어텐션	6	out‑of‑place masked_fill에 의해 발생하는 숨겨진 Memcpy.
인플레이스 단순	masked_fill → masked_fill_	5	Memcpy 커널이 완전히 사라집니다.
SDPA Math	F.scaled_dot_product_attention을 Math 백엔드에 고정	20	참조: CUDA 코어에서 FP32, 매 호출마다 마스크 재구성, _safe_softmax. 정확하지만 약 3.7배 느림.
SDPA Efficient	Efficient (xformers) 백엔드	1	단일 융합 fmha_cutlassF 커널, Tensor 코어에서 bf16 유지.
SDPA Flash	Flash 백엔드	1	단일 융합 pytorch_flash 커널(FlashAttention‑2). 가장 빠르지만 “잘못된” 13% 점유율을 보임.
SDPA cuDNN	cuDNN 백엔드	1	문제별 생성 커널: 전치 없음, cuLaunchKernelEx, 하지만 비용이 큰 CPU 부분으로 이동.

시리즈 마무리

핵심 요점은 프로파일러 트레이스가 어떻게 보일지 추측을 만든 뒤 트레이스를 검토하고, 불일치를 가장 흥미로운 인사이트로 다루는 것입니다; 이 습관을 통해 숨겨진 Memcpy, Math 백엔드의 20개 커널, Flash의 낮은 점유율, cuDNN의 CPU 측 비용을 발견했습니다.

이제 이 추측‑검증 방식을 적용해 자신의 모델을 프로파일링할 수 있습니다.