Profiling in PyTorch (Part 3): Attention is all you profile

🖼️ 인포그래픽

🖼️ 4컷 인포그래픽

💡 한 줄 요약

파이토치 프로파일러로 어텐션 구현을 비교한 결과, 제자리 마스킹은 불필요한 메모리 복사를 제거했지만 SDPA 수학 백엔드는 텐서 코어 미사용과 매 호출 마스크 생성으로 인해 단순 구현보다 3.7배 느렸다.

📌 핵심 요약

이 글은 파이토치 프로파일링 연재의 세 번째 편으로, 트랜스포머의 핵심 연산인 어텐션이 구현 방식에 따라 프로파일러에서 어떻게 다르게 나타나는지를 분석한다.
단순 인과 어텐션은 쿼리·키 행렬곱, 크기 조정, 인과 마스킹, 소프트맥스, 어텐션 가중치·값 행렬곱이라는 연속된 기본 연산으로 구성된다.
일반 masked_fill을 제자리 연산인 masked_fill_로 바꾸자 GPU의 메모리 복사 커널이 사라져, 순전파마다 실행되는 커널 하나와 추가 메모리 복사를 줄일 수 있었다.
파이토치의 SDPA는 입력 자료형, 헤드 차원, 마스크, 하드웨어에 따라 여러 백엔드 중 하나를 선택하지만, 실험의 수학 백엔드는 단순 제자리 구현보다 평균 CUDA 실행 시간이 3.7배 길었다.
수학 백엔드의 성능 저하는 입력을 FP32로 승격해 텐서 코어를 사용하지 않은 점, GPU 커널 수가 5개에서 20개로 늘어난 점, 인과 마스크를 호출할 때마다 새로 만든 점으로 설명된다.

🧩 주요 포인트

이 글은 파이토치 프로파일링 연재의 세 번째 편으로, 트랜스포머의 핵심 연산인 어텐션이 구현 방식에 따라 프로파일러에서 어떻게 다르게 나타나는지를 분석한다.
단순 인과 어텐션은 쿼리·키 행렬곱, 크기 조정, 인과 마스킹, 소프트맥스, 어텐션 가중치·값 행렬곱이라는 연속된 기본 연산으로 구성된다.
일반 masked_fill을 제자리 연산인 masked_fill_로 바꾸자 GPU의 메모리 복사 커널이 사라져, 순전파마다 실행되는 커널 하나와 추가 메모리 복사를 줄일 수 있었다.
파이토치의 SDPA는 입력 자료형, 헤드 차원, 마스크, 하드웨어에 따라 여러 백엔드 중 하나를 선택하지만, 실험의 수학 백엔드는 단순 제자리 구현보다 평균 CUDA 실행 시간이 3.7배 길었다.
수학 백엔드의 성능 저하는 입력을 FP32로 승격해 텐서 코어를 사용하지 않은 점, GPU 커널 수가 5개에서 20개로 늘어난 점, 인과 마스크를 호출할 때마다 새로 만든 점으로 설명된다.

🧠 상세 정리

1. 프로파일링 연재와 이번 글의 목표

이 글은 프로파일러의 표와 실행 추적을 읽는 능력을 단계적으로 익히는 파이토치 프로파일링 연재의 세 번째 편이다. 첫 번째 편에서는 덧셈과 곱셈 같은 기본 수학 연산을 측정하면서 프로파일러 표로 병목을 찾고, 실행 추적으로 알고리즘의 시간상 실행 순서를 확인했다. 두 번째 편에서는 이 연산들을 선형 계층과 다층 퍼셉트론으로 확장하고, 융합 커널과 수작업으로 조정한 커널까지 비교했다. 이번 편은 트랜스포머 구조에서 다음으로 살펴볼 핵심 알고리즘인 어텐션을 대상으로 삼으며, 이차 시간 복잡도를 줄이는 모든 기법을 설명하기보다 각 구현이 프로파일러에 어떻게 다르게 보이는지를 밝히는 데 초점을 둔다. 실험 스크립트는 단순 구현, 제자리 연산 구현, SDPA 구현, 커널 구현으로 구분되며 측정에는 엔비디아 A100-SXM4-80GB GPU가 사용됐다.

2. 단순 인과 어텐션을 이루는 다섯 단계

어텐션은 쿼리, 키, 값 사이의 상호작용을 여러 기본 연산으로 연결한 알고리즘이다. 먼저 쿼리와 전치한 키를 행렬곱해 어텐션 점수를 만들고, 헤드 차원의 제곱근에 기반한 크기 조정값을 곱한다. 이어 미래 위치를 참조하지 못하도록 인과 마스크가 적용되는 지점을 음의 무한대로 채운 뒤, 소프트맥스로 점수를 정규화해 어텐션 가중치를 얻는다. 마지막으로 이 가중치와 값을 다시 행렬곱해 출력 텐서를 계산하며, 코드에서는 matmul, 곱셈, masked_fill, softmax, matmul의 순서로 그대로 표현된다. 따라서 프로파일링 전에도 두 번의 행렬곱과 크기 조정, 마스킹, 소프트맥스에 대응하는 연산과 커널이 실행 추적에 나타날 것이라고 구체적으로 예상할 수 있다.

3. CPU와 GPU 실행 추적에서 발견된 추가 복사

단순 어텐션의 CPU 실행 구간에서는 예측한 대로 행렬곱, 크기 조정용 곱셈, 인과 마스킹, 소프트맥스, 두 번째 행렬곱이 순서대로 확인됐다. GPU 구간을 펼쳐 한 번의 순전파를 확대하면 쿼리·키 행렬곱, 크기 조정, 메모리 복사, 인과 마스킹, 소프트맥스, 어텐션 가중치·값 행렬곱이라는 여섯 개의 커널이 나타난다. 어텐션 계산 자체에서 예상한 다섯 커널 외에 메모리 복사가 하나 더 포함된 것이 핵심적인 차이다. 원문은 이 복사의 원인을 일반적인 masked_fill이 입력 텐서를 직접 수정하지 않는 비제자리 연산이라는 점에서 찾는다. 파이토치는 이러한 연산에서 흔히 텐서 복사본을 만들고 그 복사본에 연산을 적용해 반환하므로, 코드 한 줄의 의미만 볼 때 드러나지 않던 복사 비용이 GPU 실행 추적에서는 별도 커널로 분명하게 노출된다.

4. 제자리 마스킹으로 제거한 메모리 복사 커널

두 번째 구현은 masked_fill을 파이토치의 제자리 연산 표기인 masked_fill_로 바꾸는 한 줄 수정만 적용한다. CPU 실행 추적을 비교하면 제자리 버전의 마스킹 단계 안에 포함된 연산 수가 비제자리 버전보다 훨씬 적어져, 내부 작업이 감소했다는 신호를 먼저 확인할 수 있다. GPU 실행 구간에서는 차이가 더 직접적으로 드러나며, 기존 구현에 있던 메모리 복사 커널이 완전히 사라진다. 그 결과 한 번의 어텐션 순전파에서 전체 커널 하나가 줄어들고, 큰 점수 텐서의 추가 복사도 피하게 된다. 단일 호출만 보면 변화가 작아 보일 수 있지만, 대규모 언어 모델이나 확산 모델처럼 어텐션이 여러 계층에서 반복되는 트랜스포머에서는 같은 절감이 계층마다 누적된다는 것이 글의 설명이다.

5. 제자리 연산의 적용 조건과 자동 미분 주의점

비제자리 연산이 파이토치의 기본값인 데에는 자동 미분과 관련된 이유가 있다. 역전파 공식은 순전파 때 계산된 텐서 값을 다시 사용하는 경우가 많으므로, 자동 미분 시스템은 필요한 값을 보존해야 한다. 제자리 연산으로 그 값을 메모리에서 덮어쓰면 역전파가 원래 값이 아닌 변경된 값을 읽어 잘못된 계산을 할 수 있다. 이 글의 실험은 torch.no_grad 아래에서 순전파만 실행하므로 뒤이어 수행할 역전파가 없고, 그 조건에서는 제자리 마스킹을 안전하게 사용할 수 있다고 설명한다. 또한 제자리 연산은 별도의 복사본을 만들지 않기 때문에 실행 시간을 줄일 뿐 아니라 메모리도 절약하며, 로짓처럼 크기가 큰 텐서에서는 이 메모리 절감도 중요한 이점이 된다.

6. 한 번의 호출로 감싼 SDPA와 백엔드 선택

파이토치는 손으로 작성한 어텐션의 전체 절차를 scaled_dot_product_attention이라는 단일 함수로 제공한다. 쿼리, 키, 값을 전달하고 is_causal을 참으로 지정하면 직접 인과 마스크를 만들 필요도 없어, 여러 줄의 연산이 함수 호출 한 줄로 대체된다. 그러나 SDPA는 내부 구현이 하나로 고정된 함수가 아니라 입력 자료형, 헤드 차원, 마스크, 하드웨어 등의 조건에 따라 지원 가능한 여러 백엔드 중 적합한 구현을 선택하는 디스패치 계층이다. 글에서 열거한 선택지는 수학, 플래시 어텐션, 효율적 어텐션, cuDNN 어텐션 백엔드다. 일반적으로는 SDPA가 백엔드를 자동 선택하지만, 실험에서는 sdpa_kernel 문맥 관리자로 특정 백엔드를 고정해 각 구현의 실행 추적과 성능 차이를 개별적으로 읽을 수 있게 했다.

7. 예상을 뒤집은 SDPA 수학 백엔드의 측정 결과

여러 기본 연산을 한 줄짜리 SDPA 호출로 대체했기 때문에 저자들은 실행 추적이 단순해지고, CPU 디스패치와 GPU 커널 수가 줄며, 융합 커널을 통해 더 빨라질 가능성을 예상했다. 하지만 프로파일러 표에서 단순 제자리 어텐션의 순전파 평균 CUDA 시간은 1.955밀리초였고, SDPA 수학 백엔드는 7.239밀리초로 측정됐다. 전체 자체 CUDA 시간도 각각 7.194밀리초와 27.279밀리초였으며, 원문은 이를 수학 백엔드가 3.7배 느린 결과라고 정리한다. 실행 추적에서는 단순 제자리 구현이 순전파 한 번에 GPU 커널 5개를 실행한 반면, 수학 백엔드는 20개를 실행했다. 즉, 사용자 코드가 한 줄로 짧아졌다는 사실은 내부 실행이 융합되거나 단순해졌다는 뜻이 아니며, 실제 성능은 프로파일러 표와 커널 추적으로 확인해야 했다.

8. 텐서 코어 미사용과 반복적인 인과 마스크 생성

첫 번째 원인은 두 구현의 행렬곱 커널 이름에서 확인된다. 단순 구현의 커널 이름에 나타난 s16816은 bfloat16 텐서 코어의 16×8×16 행렬곱 명령을 사용한다는 표식인 반면, 수학 백엔드의 sgemm은 일반 CUDA 코어에서 실행되는 FP32 행렬곱을 뜻한다. A100의 텐서 코어는 작은 행렬 타일의 곱셈과 누산을 한 명령으로 처리하지만, 수학 백엔드는 수치 정확도를 위해 bf16 입력도 FP32로 승격해 이동 데이터량을 늘리고 텐서 코어의 빠른 경로를 사용하지 않았다. 두 번째 원인은 인과 마스크 생성으로, 단순 구현은 미리 만든 마스크를 재사용했지만 수학 백엔드는 is_causal 설정에 따라 호출할 때마다 하삼각 행렬과 음의 무한대 채움값, 가산 편향을 새로 만들었다. CPU에서는 ones, tril, scalar_tensor, fill_, where 연산으로 나타나고 GPU에서는 삼각 행렬 커널과 여러 where 커널, add_ 연산으로 나타나므로, 편리한 옵션이 마스크 생성 작업을 없앤 것이 아니라 내부 계층으로 옮겨 매 순전파마다 반복한 셈이다.

🧾 핵심 주장 / 시사점

프로파일링에서는 상위 수준 함수의 코드 길이나 이름보다 실제 GPU 커널의 수, 종류, 자료형과 실행 시간을 함께 확인해야 한다. 한 줄짜리 SDPA 수학 백엔드가 다섯 단계의 직접 구현보다 더 많은 커널을 실행하고 더 느렸다는 측정이 이를 보여준다.
작은 제자리 연산 변경도 반복 구조에서는 누적 효과를 만든다. masked_fill_는 한 번의 순전파에서 메모리 복사 커널 하나를 제거했으며, 어텐션이 여러 계층에 반복되는 모델에서는 같은 절감이 계층마다 이어진다.
편의 기능이 계산 자체를 제거한다고 볼 수는 없다. is_causal은 사용자가 마스크를 직접 만드는 코드를 없앴지만 수학 백엔드는 마스크를 호출마다 생성했으므로, 추상화 내부의 실제 작업은 실행 추적으로 별도 검증해야 한다.

✅ 액션 아이템

단순 인과 어텐션에서 masked_fill을 masked_fill_로 교체해 순전파당 메모리 복사 커널을 제거하고 커널 실행 수를 줄인다.
SDPA는 입력 자료형·헤드 차원·마스크·하드웨어 조합으로 백엔드를 바꾸므로, 수학 백엔드의 3.7배 지연 구간을 환경별로 검증해 사용 영역을 정한다.
수학 백엔드의 FP32 승격, 마스크 재생성, 커널 수 급증 비용을 근거로 텐서 코어 경로와 대체 구현을 비교해 기본 경로를 정한다.