PyTorch profiling: nn.Linear에서 fused MLP까지 병목 읽기

한 줄 요약

Hugging Face의 Profiling in PyTorch 2편은 nn.Linear 하나에서 시작해 MLP 블록, torch.compile, fused Triton kernel, hand-tuned kernel까지 profiler trace로 따라가는 글이다. 핵심은 단순하다. 최적화는 감으로 하는 게 아니라 trace에서 CPU dispatch, GPU kernel, layout, launch overhead를 분리해서 본 뒤 해야 한다.

이 글은 원문을 한국어로 압축한 first draft다. 원문의 모든 수치를 재현한 것은 아니고, 개발자가 실제 모델 병목을 볼 때 어떤 순서로 의심해야 하는지에 초점을 맞춘다.

왜 이 주제가 실무적이냐

LLM/멀티모달 모델 운영에서 “더 작은 모델을 쓰자”, “quantization 하자”는 말은 자주 나온다. 그런데 막상 latency가 튀는 이유는 더 지루한 곳에 있는 경우가 많다.

• 작은 op가 너무 많이 쪼개져 GPU launch overhead가 쌓인다.
• Python/ATen dispatch 비용이 실제 compute보다 커진다.
• transpose/view 같은 metadata 연산을 실제 copy로 오해한다.
• torch.compile이 줄인 시간이 GPU 계산인지 CPU overhead인지 구분하지 못한다.
• fused kernel이 항상 이긴다고 생각하고 메모리 layout이나 shape 조건을 안 본다.

내 의견은 명확하다. serving 최적화의 첫 단계는 모델 교체가 아니라 profiler trace 읽기다. trace를 못 읽으면 compile, flash, fused, quantized 같은 단어에 끌려다니기 쉽다.

nn.Linear는 결국 matmul + bias다

PyTorch의 nn.Linear(in_dim, out_dim, bias=True)는 대략 아래 계산을 한다.

y = x @ w.T + b

여기서 중요한 디테일은 w.T다. 초보자는 transpose가 별도 GPU kernel이나 메모리 복사를 만든다고 생각하기 쉽지만, 많은 경우 PyTorch의 transpose는 실제 데이터를 옮기지 않는 view/stride metadata 변경이다.

원문이 보여주는 좋은 습관은 profiler에서 다음을 분리해서 보는 것이다.

• CPU lane: aten::linear, aten::t, aten::addmm 같은 dispatch/metadata 경로
• GPU lane: 실제 GEMM kernel이 몇 번 실행됐는지
• kernel name: _tn_, _nn_, dtype, tile shape 같은 layout 힌트

aten::t가 보인다고 해서 곧바로 “transpose가 느리다”라고 결론내리면 안 된다. GPU lane에 별도 transpose kernel이 없다면, 그건 대개 stride만 바꾼 view일 가능성이 높다.

torch.compile이 줄이는 것: 계산이 아니라 경로일 수 있다

torch.compile을 적용하면 trace에서 CPU 쪽 연산들이 줄어드는 경우가 있다. 이때 “GPU 계산이 빨라졌다”고 말하면 반은 맞고 반은 틀릴 수 있다.

nn.Linear 같은 단순 케이스에서는 compiler가 view chain과 dispatch 경로를 정리해 더 직접적인 호출로 만들 수 있다. 하지만 GPU lane에서 같은 GEMM kernel이 그대로 실행된다면, 실제 행렬곱 자체가 새 알고리즘으로 바뀐 것은 아니다.

실무적으로는 이렇게 읽는 편이 안전하다.

1. GPU kernel 수가 줄었나?
2. 같은 kernel인데 CPU dispatch 시간이 줄었나?
3. kernel 이름이 바뀌었나? dtype/layout/tile이 달라졌나?
4. compile warmup 비용을 제외하고 steady-state에서 이득이 있나?

특히 online serving에서는 compile warmup, dynamic shape, graph break가 중요하다. notebook에서 한 번 빠른 결과가 나왔다고 production path에 그대로 넣는 건 꽤 위험하다.

MLP로 가면 작은 op들이 문제를 만든다

Transformer 계열의 MLP/FFN 블록은 보통 선형층 여러 개와 activation으로 구성된다.

import torch.nn as nn
 
mlp = nn.Sequential(
    nn.Linear(hidden_dim, intermediate_dim),
    nn.GELU(),
    nn.Linear(intermediate_dim, hidden_dim),
)

게이트형 MLP라면 up_proj, gate_proj, down_proj처럼 선형층이 더 늘어난다. 여기서 profiler가 보여주는 병목은 단순히 “GEMM이 크다”가 아니다.

• Linear 사이에 activation kernel이 따로 뜬다.
• elementwise op가 memory bandwidth를 많이 먹는다.
• 작은 tensor shape에서는 kernel launch overhead가 상대적으로 커진다.
• GPU는 바쁜 것처럼 보이지만 실제로는 CPU가 다음 kernel을 계속 제출하는 구조가 된다.

그래서 fused MLP의 목표는 “수학을 바꾸는 것”이 아니라 중간 tensor materialization과 kernel launch 횟수를 줄이는 것에 가깝다.

fused Triton kernel을 볼 때 체크할 것

원문은 torch.compile을 통해 fused Triton kernel로 넘어가는 흐름을 보여준다. 여기서 중요한 질문은 “Triton이라서 빠른가?”가 아니다. 더 좋은 질문은 이것이다.

• fusion으로 사라진 op는 무엇인가?
• 중간 activation 결과를 global memory에 썼다가 다시 읽는 경로가 줄었나?
• shape가 compiler가 잘 다루는 정적/반정적 형태인가?
• batch size가 작을 때와 클 때 이득이 같은가?
• dtype이 bf16/fp16/fp8 중 무엇이고 tensor core 경로를 타는가?

fused kernel은 보통 아래 상황에서 유리하다.

• 같은 tensor를 elementwise op 여러 개가 반복해서 읽고 쓴다.
• 중간 결과를 굳이 메모리에 materialize할 필요가 없다.
• launch overhead가 총 latency에서 무시할 수 없을 만큼 크다.
• shape가 안정적이라 compile/cache 이득을 회수할 수 있다.

반대로 아래 상황에서는 기대보다 별로일 수 있다.

• 이미 큰 GEMM이 대부분의 시간을 먹고 있다.
• dynamic shape가 많아 compile graph가 자주 깨진다.
• fused kernel이 vendor-tuned GEMM보다 낮은 occupancy/throughput을 낸다.
• memory layout이 불리해서 fusion 이득을 상쇄한다.

hand-tuned kernel은 마지막 카드다

원문 후반부는 직접 튜닝된 kernel/library를 사용하는 방향도 다룬다. 이건 성능이 필요할 때 매력적이지만, 팀 운영 관점에서는 비용이 있다.

• 특정 GPU architecture에 강하게 묶일 수 있다.
• PyTorch 버전, CUDA 버전, driver 조합에 민감해진다.
• 일반 모델 코드보다 디버깅 난이도가 높다.
• fallback path와 accuracy regression test가 필요하다.

그래서 내 기준은 이렇다.

1. 먼저 profiler로 병목이 진짜 kernel/launch/memory 중 어디인지 확인한다.
2. torch.compile처럼 유지보수 비용이 낮은 선택지를 먼저 적용한다.
3. 그래도 SLA를 못 맞추고 hot path가 안정적이면 fused/custom kernel을 검토한다.
4. 도입 후에는 latency뿐 아니라 numerical drift, cold start, fallback까지 테스트한다.

custom kernel은 멋있지만, 모든 팀이 감당할 만한 추상화는 아니다. 성능 이득이 운영 복잡도를 이겨야 한다.

실무 체크리스트

PyTorch 모델이 느릴 때 아래 순서로 본다.

[trace 읽기]
- CPU lane과 GPU lane을 분리해서 본다.
- kernel launch 개수와 kernel duration을 같이 본다.
- 같은 op 이름보다 실제 kernel name/dtype/layout 변화를 본다.
 
[Linear/MLP]
- transpose가 실제 copy인지 view인지 확인한다.
- activation/elementwise op가 따로 kernel을 만들고 있는지 본다.
- 중간 tensor materialization이 많은지 본다.
 
[compile/fusion]
- warmup 제외 steady-state latency를 측정한다.
- dynamic shape와 graph break를 기록한다.
- fusion 후 kernel 수가 줄었는지, 아니면 CPU dispatch만 줄었는지 구분한다.
 
[운영]
- p50뿐 아니라 p95/p99를 본다.
- batch size별 trace를 따로 저장한다.
- fallback path와 accuracy regression test를 둔다.

작은 팁 하나를 더하면, 최적화 전후 trace 파일을 “성능 증거”로 저장해두는 게 좋다. PR 설명에 평균 latency 숫자만 쓰면 나중에 왜 빨라졌는지 잊어버린다. kernel 수, 주요 kernel 이름, CPU overhead 변화까지 남겨야 다음 최적화가 쉬워진다.

정리

이 주제의 가치는 nn.Linear 자체에 있지 않다. nn.Linear는 너무 기본적인 블록이라 오히려 좋다. 기본 블록의 profiler trace를 제대로 읽을 수 있으면 attention, MoE, KV cache, quantized inference도 덜 신비롭게 보인다.

최적화 글을 읽을 때 나는 이제 “몇 퍼센트 빨라졌다”보다 “어떤 trace evidence로 설명하는가”를 먼저 본다. 그 기준에서 Hugging Face의 PyTorch profiling 시리즈는 꽤 실용적이다. 과장된 벤치마크보다, trace를 읽는 습관을 만들어주는 글이 더 오래 간다.

참고 자료

• Hugging Face Blog — Profiling in PyTorch (Part 2): From nn.Linear to a Fused MLP
• Hugging Face Blog — Profiling in PyTorch (Part 1): A Beginner's Guide to torch.profiler
• PyTorch Docs — torch.profiler
• PyTorch Docs — torch.compile