이 기사에서는 프로파일러 트레이스를 사용하여 PyTorch의 어텐션 메커니즘을 프로파일하고 성능 병목 지점과 최적화 기회를 식별하는 방법을 시연합니다. 커널 동작을 설명하기 위해 단순 어텐션 구현과 PyTorch의 내장 Scaled Dot Product Attention (SDPA)를 비교합니다.
- 단순 인과적 어텐션은 matmul, 스케일링, 마스크 처리, softmax 및 아웃오브플레이스 연산으로 인한 예상치 못한 메모리 복사를 포함하는 여러 커널을 사용합니다.
- masked_fill을 인플레이스 masked_fill_로 교체하면 GPU 트레이스에서 불필요한 메모리 복사 커널이 제거됩니다.
- PyTorch의 SDPA 함수는 입력 유형과 하드웨어에 따라 Flash Attention이나 cuDNN과 같은 가장 빠른 사용 가능한 백엔드로 자동으로 디스패치합니다.
- 개별 SDPA 백엔드를 프로파일하면 고유한 커널 시그니처가 드러나 개발자가 활성화된 구현을 확인할 수 있습니다.
이러한 프로파일러 트레이스를 이해하면 사용자는 불필요한 메모리 작업을 제거하고 PyTorch의 효율적인 어텐션 백엔드를 활용하여 트랜스포머 모델을 최적화할 수 있습니다.