A100 GPU에서 확인한 나이브 어텐션의 기본 실행 구조

PyTorch로 LLM이나 트랜스포머 모델을 구현할 때 라이브러리가 제공하는 기본 함수를 사용하면 실제 GPU 내부에서 연산이 어떤 순서로 실행되는지 알기 어렵다. 실행 경로를 알 수 없으면 최적화의 병목 지점을 찾기 어렵기 때문에, 최적화 기법을 적용하지 않은 나이브 어텐션 구현체의 커널 실행 구조를 분석해 하드웨어 수준의 동작을 확인했다.

분석은 NVIDIA A100-SXM4-80GB GPU 환경에서 진행했다. 어텐션 연산은 쿼리(Query, q), 키(Key, k), 밸류(Value, v) 텐서의 행렬 곱셈과 스케일링, 소프트맥스 연산의 연속적인 집합으로 구성된다. CPU와 GPU의 연산 레인을 동시에 분석할 수 있는 PyTorch Profiler를 사용해 시간 흐름에 따른 실행 경로를 추적했으며, 분석 코드는 [04_a_naive_attention.py](04_a_naive_attention.py)에 기록되어 있다.

프로파일러 분석 결과, CPU 레인에서는 `attn_fwd` 포워드 패스 과정에서 작성한 연산 순서가 그대로 나타났다. 실제 속도를 결정하는 GPU 레인에서는 총 5개의 커널이 런칭되는 것을 확인했다. 이 커널들은 행렬 곱셈, 스케일링, 소프트맥스 등 어텐션의 기본 단계에 각각 대응한다. 각 커널의 실행 시간과 하드웨어 자원 사용량은 향후 인플레이스 연산이나 통합 함수 적용 시 성능 개선 정도를 측정하는 정량적 기준점이 된다.

masked_fill 인플레이스 연산으로 Memcpy 커널 제거하기

PyTorch의 기본 연산은 데이터 무결성을 위해 결과값을 새로운 메모리 공간에 할당하며, 이 과정에서 오버헤드가 발생한다. 나이브 어텐션 구현체의 GPU 레인을 분석하면 연산 흐름과 무관한 Memcpy(메모리 복사) 커널이 실행되는 지점이 확인된다. 이는 `masked_fill` 함수가 기존 텐서를 수정하지 않고 새로운 복사본을 생성해 값을 채운 뒤 반환하기 때문이다.

메모리 복사 비용을 줄이려면 기존 메모리 주소의 값을 직접 덮어쓰는 인플레이스(in-place) 연산을 적용해야 한다. PyTorch에서는 함수 이름 끝에 언더바를 붙여 원본을 수정함을 명시하며, `masked_fill`을 `masked_fill_`로 교체하여 구현한다.

python

04_b_inplace_ops_attention.py

masked_fill을 masked_fill_로 변경하여 인플레이스 연산 적용

attn_weights.masked_fill_(mask == 0, float('-inf'))

이 변경을 적용하면 GPU 레인에서 Memcpy 커널이 완전히 사라진다. CPU의 디스패치 과정은 유지되지만, GPU 하드웨어에서 메모리를 복제하는 물리적 단계가 생략된다. 이러한 커널 제거 효과는 레이어 수가 수십 개에 달하는 LLM이나 디퓨전 모델에서 전체 추론 지연 시간을 낮추는 직접적인 요인이 된다. 또한 복사본을 생성하지 않아 로짓(logits)과 같은 대형 텐서를 다룰 때 GPU 메모리 점유율을 낮춰 더 큰 배치 크기를 확보할 수 있다.

단, 인플레이스 연산은 `torch.no_grad` 상태인 추론 환경에서만 안전하다. 학습 단계의 역전파(backward pass)에서는 그래디언트 계산을 위해 포워드 패스의 원본 텐서 값이 보존되어야 하는데, 인플레이스 연산은 이를 덮어써 계산 오류를 일으키기 때문이다. 따라서 프로파일러로 실제 커널 실행 경로를 확인해 불필요한 메모리 복사가 일어나는지 검증하는 과정이 필요하다.

SDPA math 백엔드가 나이브 구현보다 3.7배 느린 이유

PyTorch는 여러 어텐션 구현체를 하나로 통합한 SDPA(Scaled Dot Product Attention) 함수를 제공한다.

torch.nn.functional.scaled_dot_product_attention

이 함수는 입력 데이터 타입과 하드웨어 환경에 맞춰 최적의 백엔드를 자동 선택하며, `is_causal` 옵션으로 마스크 생성 과정을 생략할 수 있어 코드의 간결함과 유지보수 효율을 높인다. 하지만 자동 선택된 백엔드가 항상 최적의 성능을 보장하지는 않는다.

[04_c_sdpa_attention.py] 스크립트로 성능을 측정한 결과, SDPA의 `math` 백엔드는 직접 작성한 나이브 구현보다 3.7배 느린 실행 속도를 기록했다. 성능 저하의 원인은 GPU 커널 런칭 횟수의 급증에 있다. 나이브 구현이 포워드 패스당 5개의 커널을 실행할 때, SDPA `math` 백엔드는 동일 작업에 20개의 커널을 실행했다. 단 한 줄의 코드로 대체했음에도 GPU 내부에서는 4배 더 많은 작업 단위가 생성되어 오버헤드가 발생한 것이다. 이는 추상화된 함수가 하드웨어 자원을 비효율적으로 사용하는 사례다.

sgemm과 s16816 커널로 분석한 텐서 코어 활용 차이

성능 격차는 GPU 내부의 연산 유닛 활용 방식에서 발생한다. GPU의 SM(Streaming Multiprocessor)에는 범용 CUDA 코어와 행렬 연산 전용 가속기인 텐서 코어가 함께 들어있다. 텐서 코어는 작은 행렬 타일 전체를 단일 명령어로 처리하여 CUDA 코어보다 압도적인 속도를 낸다.

커널 이름은 사용된 하드웨어를 보여주는 지표가 된다. 나이브 구현의 `s16816` 커널은 bfloat16 텐서 코어 연산의 시그니처로, 하드웨어 가속 경로를 사용하고 있음을 의미한다. 반면 SDPA `math` 백엔드의 `sgemm` 커널은 CUDA 코어에서 실행되는 FP32 행렬 곱 커널이다. `math` 백엔드는 수치적 정확도를 위해 텐서를 FP32로 업캐스팅하며, 이 과정에서 입력값이 bf16일 때조차 데이터 이동량이 두 배로 늘어나고 텐서 코어를 배제한 채 느린 CUDA 코어에 의존하게 된다.

또한 `is_causal=True` 설정은 마스크 생성을 자동화하지만, `math` 백엔드 내부에서는 매 호출마다 마스크를 새로 생성한다. 이 과정은 CPU 레인에서 시작되어 GPU 레인의 `triu_tril_kernel`과 여러 개의 `where` 커널, `add_` 연산으로 이어진다. 마스크를 한 번만 생성해 재사용하는 나이브 방식과 달리, 이 옵션은 작업을 하위 레이어로 숨겨 매번 반복 실행하게 만든다.

SDPA 백엔드 제어 방법과 프로파일링 전략

SDPA는 dtype, 헤드 차원, 마스크 설정 및 하드웨어 사양을 고려해 백엔드를 자동 선택하지만, 특정 환경에서 성능을 검증하거나 강제로 백엔드를 지정해야 할 때는 `sdpa_kernel` 컨텍스트 매니저를 사용한다.

with torch.nn.attention.sdpa_kernel(torch.nn.attention.SDPBackend.MATH):

attention operation

사용 가능한 백엔드 목록은 `torch.nn.attention.SDPBackend` 열거형(enum)을 통해 참조할 수 있다. 이를 통해 개발자는 자동 선택 로직에 의존하지 않고 각 백엔드가 하드웨어에서 어떻게 작동하는지 개별적으로 프로파일링하며 최적의 조합을 찾을 수 있다. 상세한 제어 방법과 검증 과정은 [04_d_kernels_attention.py](04_d_kernels_attention.py) 파일에 구현되어 있다.

결국 라이브러리 함수의 편의성을 맹신하기보다 프로파일러를 통한 교차 검증이 중요하다. 실제 커널 런칭 횟수와 텐서 코어 같은 하드웨어 가속기가 의도한 대로 작동하는지 확인해야 병목 지점을 정확히 제거할 수 있다. 추상화 계층이 깊어질수록 하드웨어 수준의 실행 경로를 추적하고 실제 커널 이름을 분석하는 습관이 대규모 모델 최적화의 핵심이다.

PyTorch로 트랜스포머 모델을 구현하며 기본 함수를 사용하는 것은 효율적인 선택처럼 보인다. 그러나 SDPA의 math 백엔드가 텐서 코어라는 하드웨어 가속기를 활용하지 못해 나이브 구현보다 3.7배 느려지는 결과는 API의 편리함이 성능의 정답이 아님을 증명한다. `masked_fill_` 같은 인플레이스 연산으로 불필요한 메모리 복사 커널을 걷어내는 세밀한 조정이 실제 추론 속도를 가른다. 이제는 라이브러리를 맹신하는 대신 프로파일러를 통해 실제 실행 경로와 하드웨어 유닛 활용도를 대조하며 최적화 지점을 찾아야 한다.