Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP

🖼️ 인포그래픽

Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP 내용을 설명하는 본문 이미지

🖼️ 4컷 인포그래픽

💡 한 줄 요약

이 글은 프로파일러 추적을 통해 nn.Linear의 전치·행렬곱·편향 처리가 실제로 어떻게 실행되는지 분석하고, 단일 선형 계층에서 GeGLU MLP의 다중 커널 구조로 관찰 범위를 확장한다.

📌 핵심 요약

nn.Linear는 입력과 전치된 가중치의 행렬곱에 편향을 더하는 연산이며, 내부적으로 aten::linear를 거쳐 aten::addmm를 호출한다.
가중치 전치에 사용되는 aten::t는 데이터를 복사하거나 GPU 커널을 실행하지 않고 텐서의 모양과 스트라이드 메타데이터만 바꾼다.
편향 덧셈은 GEMM 커널의 에필로그에 포함되므로 별도의 aten::add 또는 GPU 덧셈 커널로 나타나지 않으며, 추가적인 고대역폭 메모리 접근도 피한다.
단일 nn.Linear를 torch.compile로 컴파일해도 GPU에서는 같은 cuBLAS GEMM 커널이 실행되며, 주요 차이는 전치 뷰를 처리하던 일부 CPU 디스패치 오버헤드가 사라지는 데 있다.
세 개의 선형 계층과 GeLU·원소별 곱으로 구성한 GeGLU MLP에서는 순전파마다 총 다섯 개의 GPU 커널이 실행되며, 저자는 추적을 보기 전에 예상 연산과 커널 수를 세우는 습관을 강조한다.

🧩 주요 포인트

nn.Linear는 입력과 전치된 가중치의 행렬곱에 편향을 더하는 연산이며, 내부적으로 aten::linear를 거쳐 aten::addmm를 호출한다.
가중치 전치에 사용되는 aten::t는 데이터를 복사하거나 GPU 커널을 실행하지 않고 텐서의 모양과 스트라이드 메타데이터만 바꾼다.
편향 덧셈은 GEMM 커널의 에필로그에 포함되므로 별도의 aten::add 또는 GPU 덧셈 커널로 나타나지 않으며, 추가적인 고대역폭 메모리 접근도 피한다.
단일 nn.Linear를 torch.compile로 컴파일해도 GPU에서는 같은 cuBLAS GEMM 커널이 실행되며, 주요 차이는 전치 뷰를 처리하던 일부 CPU 디스패치 오버헤드가 사라지는 데 있다.
세 개의 선형 계층과 GeLU·원소별 곱으로 구성한 GeGLU MLP에서는 순전파마다 총 다섯 개의 GPU 커널이 실행되며, 저자는 추적을 보기 전에 예상 연산과 커널 수를 세우는 습관을 강조한다.

🧠 상세 정리

1. 프로파일링 학습의 두 번째 단계

이 글은 PyTorch 프로파일러 추적을 단계적으로 읽고 그 결과를 최적화 판단에 활용하는 연재의 두 번째 편이다. 첫 번째 편에서는 torch.add와 torch.matmul로 구성한 직접적인 행렬곱·덧셈을 대상으로 CPU 디스패치 체인, 커널 실행 오버헤드, 오버헤드 지배 구간과 연산 지배 구간의 차이, torch.compile의 일부 내부 동작을 살펴봤다. 이번에는 같은 계산을 실제 딥러닝 모델에서 널리 쓰이는 nn.Linear로 바꾸고, 이어서 세 개의 선형 계층과 활성화 연산을 쌓은 MLP 블록까지 분석 범위를 넓힌다. 실험 스크립트는 02_linear.py, 03_simple_mlp.py, 03_kernels_mlp.py이며, 제시된 실행 환경은 NVIDIA A100-SXM4-80GB GPU다. 글 전반에서는 GPU 커널이 많은 GPU 스레드에서 병렬로 실행되는 프로그램이고, CPU가 그 커널을 예약하고 시작하며, 프로파일러에서 관찰되는 PyTorch 오버헤드의 상당 부분이 이 CPU 측 작업이라는 두 가지 전제를 반복해서 사용한다.

2. 행렬곱과 덧셈에서 nn.Linear로

nn.Linear는 앞선 글에서 분석한 행렬곱과 덧셈을 모듈 형태로 감싼 구성 요소로, 가중치와 편향을 매개변수로 소유하고 익숙한 forward 인터페이스를 제공한다. 편향을 사용하는 선형 계층의 계산은 y = x @ w.T + b로 표현되며, 입력 x에 전치된 가중치 w를 곱한 뒤 편향 b를 더한다. 글에서는 배치 1024, 입력 차원 32, 출력 차원 64로 02_linear.py를 실행하고, 대기 1회·준비 1회·활성 기록 3회로 구성된 프로파일러 일정을 사용한다. 이 설정 때문에 CPU와 GPU 레인에는 세 개의 프로파일 단계가 나타나며, 각 단계에서 선형 계층의 순전파가 어떤 CPU 연산과 GPU 커널로 이어지는지 비교할 수 있다. 분석의 핵심은 소스 코드에 적힌 수식만 보는 것이 아니라, 실제 디스패치 체인과 커널 실행이 그 수식을 어떤 방식으로 구현하는지 추적에서 확인하는 데 있다.

3. aten::t가 수행하는 전치의 실체

nn.Linear의 추적을 확대하면 행렬곱과 편향 덧셈을 담당하는 aten::addmm 앞에 aten::t가 나타나며, 이는 선형 계층이 가중치를 전치된 형태로 사용하기 때문이다. 그러나 aten::t는 GPU에서 데이터를 다시 배열하거나 새로운 텐서 내용을 복사하는 연산이 아니라, CPU에서 텐서의 모양과 스트라이드 메타데이터를 바꾸어 전치된 뷰를 만드는 연산이다. 따라서 이 전치에는 별도의 GPU 커널이 실행되지 않으며, GPU 레인에 전치 커널이 없는지 확인하거나 프로파일러 표에서 aten::t의 CUDA 시간이 없는지 살펴보면 이를 검증할 수 있다. 같은 원시 메모리를 유지한 채 데이터를 읽는 순서만 바꾸므로, 추적에 CPU 연산이 보인다는 사실만으로 실제 데이터 이동이나 GPU 계산이 발생했다고 해석해서는 안 된다. 이 구분은 프로파일러의 연산 이름과 실제 하드웨어 작업을 일대일로 동일시하지 않기 위해 필요한 기초다.

4. 편향 덧셈을 흡수하는 GEMM 에필로그

선형 계층의 디스패치 체인에는 편향을 더하는 별도의 aten::add가 나타나지 않는데, 이는 편향 덧셈이 행렬곱 커널의 에필로그에 포함되기 때문이다. 에필로그는 GEMM 커널이 계산 결과를 GPU의 고대역폭 메모리에 기록하기 직전에 수행하는 작은 후처리로, 편향 덧셈뿐 아니라 활성화 적용이나 상수 배율 계산도 대표적인 사례다. nn.Linear는 torch.nn.functional.linear를 호출하고, 이는 다시 aten::linear로 이어지며, 편향이 전달된 경우 aten::linear는 행렬곱과 덧셈을 따로 실행하는 대신 aten::addmm를 선택한다. GPU에서 실행되는 cuBLAS GEMM 커널에는 편향 덧셈 변형이 이미 마련되어 있으므로, 덧셈은 행렬곱 결과를 기록하는 과정의 일부가 되고 별도의 커널 실행과 추가 메모리 왕복이 필요하지 않다. 이 때문에 추적에서 독립적인 덧셈 커널이 보이지 않는 것은 연산이 누락된 것이 아니라 더 낮은 수준에서 이미 결합되어 실행된 결과다.

5. 단일 nn.Linear에서 torch.compile의 한계

단일 nn.Linear의 forward를 torch.compile로 컴파일한 추적과 즉시 실행 방식의 추적을 비교하면, GPU에서는 두 경우 모두 같은 cuBLAS GEMM 커널이 실행되고 CPU에서도 핵심 연산은 같은 aten::addmm로 나타난다. 컴파일된 추적에는 컴파일 경로에 고유한 일부 CPU 행이 추가되지만, 행렬곱과 편향 덧셈 자체의 GPU 계산은 바뀌지 않는다. 이는 즉시 실행 방식의 nn.Linear가 이미 편향을 GEMM 에필로그에 결합한 addmm 커널을 사용하기 때문에, 컴파일러가 단일 연산에서 새롭게 융합할 여지가 거의 없기 때문이다. 따라서 모델이 느리다는 이유만으로 torch.compile을 적용하면 항상 GPU 커널이 줄거나 계산이 빨라질 것이라고 기대해서는 안 된다. 글은 이를 컴파일러의 결함이 아니라, 융합 가능성을 만들려면 서로 결합할 수 있는 연산이 둘 이상 존재해야 한다는 구조적 특성으로 설명한다.

6. 스트라이드와 뷰를 이용한 CPU 전처리 제거

즉시 실행 방식에서는 aten::linear 내부에서 aten::t 다음에 aten::addmm가 호출되지만, 컴파일된 경로에서는 전치 호출 없이 aten::addmm가 직접 나타난다. 텐서는 실제 데이터를 연속된 메모리에 저장하고, 모양과 스트라이드 메타데이터를 이용해 행과 열을 따라 이동할 간격을 정의하므로, 스트라이드를 서로 바꾸면 데이터를 옮기지 않고도 전치된 뷰를 표현할 수 있다. Inductor는 컴파일 시점에 이 뷰 연산 체인을 추적하고 결과 스트라이드를 미리 계산한 뒤, 해당 스트라이드를 고정한 aten::addmm 호출을 생성한다. 그 결과 사라지는 것은 GPU 전치 커널이 아니라 전치 뷰를 디스패치하던 몇 마이크로초 규모의 CPU 작업이며, GPU는 전과 동일한 수학 연산을 수행한다. 다만 실제 입력 데이터의 스트라이드가 컴파일러가 미리 계산한 조건을 위반하면 오류가 발생할 수 있으므로, 이 최적화는 예상된 레이아웃을 전제로 한다.

7. 커널 이름으로 읽는 데이터 레이아웃

두 추적의 GPU 레인에는 순전파마다 정확히 하나의 커널이 있으며, 제시된 이름도 cutlass_80_wmma_tensorop_bf16_s161616gemm_bf16_32x32_32x1_tn_align8로 동일하다. 별도의 전치 커널 없이 가중치를 전치된 순서로 읽을 수 있는 이유는 커널 이름의 tn과 같은 레이아웃 설명자에 있으며, t는 전치된 입력을, n은 전치되지 않은 입력을 뜻한다. cuBLAS와 CUTLASS는 입력 레이아웃 조합별로 서로 다른 커널 바이너리를 미리 준비하고, 디스패처는 입력 스트라이드를 조사해 그 조건에 맞는 커널을 선택한다. 따라서 두 실행에서 커널 이름이 같다면 GPU가 같은 작업을 수행한다고 볼 수 있고, tn과 nn, bf16과 fp16, 서로 다른 GEMM 형식처럼 이름이 달라지면 디스패처가 다른 실행 경로를 선택한 것이다. 글은 CPU 연산 행만 비교하는 데 그치지 않고 커널 이름의 레이아웃·자료형·연산 형식을 읽는 습관이 추적 비교에서 매우 유용하다고 강조한다.

8. 세 개의 선형 계층으로 구성한 GeGLU MLP

다음 분석 대상은 gate_proj, up_proj, down_proj라는 세 개의 편향 없는 nn.Linear와 GeGLU 활성화 흐름으로 구성된 SimpleGeGLUMLP다. 순전파에서는 입력을 gate_proj와 up_proj에 각각 통과시키고, gate 결과에 tanh 근사를 사용하는 GeLU를 적용한 뒤 up 결과와 원소별로 곱하며, 마지막으로 down_proj를 실행한다. 추적을 열기 전에 예상하면 세 번의 aten::linear 디스패치와 GeLU·곱셈을 위한 두 번의 점별 커널 실행이 있어야 하며, 실제 프로파일러에서도 순전파 한 번당 정확히 다섯 개의 GPU 커널이 관찰된다. 선형 투영의 CPU 레인에는 점유율 조회와 cudaLaunchKernel 호출이 보이고, GeLU와 원소별 곱에는 각각 벡터화된 점별 커널 실행이 대응한다. 저자들은 추적을 처음 본 뒤 이야기를 만드는 대신, 코드에서 예상 연산과 커널 수를 먼저 세운 다음 추적이 그 예상을 확인하는지 또는 깨뜨리는지 검사하는 방식을 프로파일링에서 가장 유용한 습관으로 제시한다.

🧾 핵심 주장 / 시사점

프로파일러에서 CPU 디스패치 행이 줄었다고 해서 GPU 계산까지 바뀐 것은 아니며, 두 실행의 GPU 커널 이름과 개수를 함께 비교해야 실제 변화의 범위를 판단할 수 있다.
nn.Linear는 편향 덧셈을 GEMM 에필로그에 이미 포함하므로 단일 계층에는 torch.compile이 새로 융합할 연산이 거의 없고, 컴파일의 효과는 연산 구조에 따라 달라진다.
코드를 바탕으로 예상되는 디스패치와 커널 수를 먼저 정리한 뒤 추적과 대조하면, 누락된 커널·추가 오버헤드·다른 레이아웃 선택을 더 체계적으로 발견할 수 있다.

✅ 액션 아이템

nn.Linear의 실행 경로를 aten::linear→aten::addmm로 정리하고, 프로파일 로그 해석 시 동일한 기준으로 수렴한다.
가중치 전치에서 aten::t가 뷰 변경임을 반영해 복사·GPU 커널 가설을 배제하고 커널 원인 분석을 재정의한다.
GeGLU MLP(3개 선형·GeLU·원소별 곱) 기준으로 예상 연산과 커널 수를 사전 정의해 torch.compile 전후 커널 수 일치 여부를 점검한다.