How to profile CUDA kernels in PyTorch

대부분의 사람이 CUDA 를 시작할 때 부딪히는 문제는 문법이 아니라 경로의 문제다 — 어디서 시작해서 무엇으로 검증하고 어디까지 내려갈지 모른다는 것. 이 강의는 그 경로를 도구 단위로 깐다.

강의가 풀려고 하는 질문은 두 개로 압축된다.

1. PyTorch 프로그램 안에 커스텀 CUDA 커널을 어떻게 끼워넣는가 — C/C++ 빌드 시스템 다루지 않고.
2. 그 커널이 실제로 빠른지 어떻게 검증하는가 — 그리고 더 빠르게 만들 여지가 있는지 어떻게 보는가.

Mark 의 입장은 명시적으로 실용주의다. “GPU 가 어떻게 도는지 모든 디테일을 이해하지 않아도, 프로파일러로 black box 를 시각적으로 보면서 유용한 일을 할 수 있다.” 이 시각이 강의 전체를 끌고 간다 — 교과서식 bottom-up (메모리 모델 → warp → SM → kernel) 대신, top-down 으로 이미 도는 코드를 어디서 어떻게 보는지부터 깐다.

그래서 이 강의의 끝에 도착했을 때 손에 잡혀 있어야 하는 건 4개 도구와 1개의 의사결정 사다리다 — autograd profiler, PyTorch profiler(Chrome trace), Triton interpret 모드, NCU. 그리고 “언제 torch 에서 Triton 으로, Triton 에서 CUDA 로 내려갈지” 의 판단 기준.

강의에서 Mark 가 “CUDA 에 대해 가장 먼저 알아야 할 것 하나만 고르라면” 이라고 못 박은 사실 — CUDA 는 비동기다. 호스트가 kernel<<<...>>> 을 호출해도 GPU 가 끝나기를 기다리지 않는다. python 의 time.time() 으로 측정하는 건 launch overhead 일 뿐, 커널이 실제로 도는 시간이 아니다.

# Mark 의 time_pytorch_function — pytorch_square.py 그대로
def time_pytorch_function(func, input):
    # CUDA IS ASYNC — python time 모듈 쓰면 안 된다
    start = torch.cuda.Event(enable_timing=True)
    end   = torch.cuda.Event(enable_timing=True)

    # Warmup — 첫 호출에 context init 비용이 섞임
    for _ in range(5):
        func(input)

    start.record()
    func(input)
    end.record()
    torch.cuda.synchronize()
    return start.elapsed_time(end)

강의에서 Mark 가 보여준 첫 번째 작은 실험 — torch.square(b) vs b * b vs b ** 2 — 도 이 세팅 위에서만 의미가 있다. 결과: torch.square 와 b ** 2 는 똑같이 aten::pow 를 부르고, b * b 는 aten::mul 을 부르며 미세하게 더 빠르다. 같은 수학 연산이 백엔드에서 다른 op 로 디스패치된다는 사실이 이 첫 측정에서 이미 드러난다.

첫 프로파일러는 autograd profiler. context manager 안에 코드를 넣으면, GPU 위에서 어떤 op 들이 호출됐고 각자 얼마나 걸렸는지 표로 떨어진다. 코드 한 줄 더 추가하지 않아도 “내 함수는 aten::pow 를 부르는구나”, “aten::mul 이 더 빠르네” 같은 사실이 드러난다.

# torch.square / a*a / a**2 가 각자 어떤 op 를 부르는지 본다
with torch.profiler.profile() as prof:
    torch.square(b)

print(prof.key_averages().table(
    sort_by="cuda_time_total",
    row_limit=10,
))

--- 출력 (핵심 줄만) ---
Self CPU time total: 211.000us
Self CUDA time total: 162.001us

aten::pow                    7.66%   16.000us  ...   2  CUDA  pow_tensor_scalar
aten::result_type            0.00%    0.000us  ...   2
vectorized_elementwise_kernel
                           100.00%  162.001us  ...   2  CUDA

표 위에서는 op 들이 같은 시간에 어떻게 겹치는지, memcpy 가 커널과 어떤 순서로 도는지가 안 보인다. PyTorch profiler 의 Chrome trace 가 그 자리를 채운다 — JSON 으로 export 해서 chrome://tracing 에 드래그하면 바로 timeline 뷰.

from torch.profiler import profile, ProfilerActivity

with profile(activities=[
    ProfilerActivity.CPU,
    ProfilerActivity.CUDA,
]) as prof:
    for _ in range(10):
        a = torch.square(torch.randn(10000, 10000).cuda())

prof.export_chrome_trace("trace.json")
# 그 다음 chrome://tracing 에 trace.json 을 드래그

이 뷰에서 처음 잡히는 새로운 정보가 두 가지 더 있다.

• memcpy host-to-device 가 커널 앞에 항상 박힌다 — 입력을 .cuda() 로 보낼 때 발생. 이게 길면 커널을 빠르게 짜봐야 의미가 없다. 강의의 큰 메시지 중 하나가 여기서 시작된다.
• 커널 이름이 vectorized_elementwise_kernel<4> — 4 는 block 수다. 같은 사이즈의 입력에 대해 PyTorch 가 자동으로 4 블록을 띄운다는 사실. 큰 입력에서 작은 grid 면 GPU 자원이 비어 있을 가능성이 높다 (§ 08 NCU 가 이 hint 를 직접 준다).

커스텀 CUDA 커널을 PyTorch 에 넣을 때 전통적으로 만나는 두 벽 — pybind11 보일러플레이트와 setup.py / CMake / make. torch.utils.cpp_extension.load_inline 은 둘 다 우회한다. C++ 소스와 CUDA 소스를 문자열로 그냥 넣으면, 안에서 ninja 가 build 해서 .so 로 import 가능한 모듈을 돌려준다.

# load_inline.py — 강의 repo 그대로
cuda_source = '''
__global__ void square_matrix_kernel(
    const float* matrix, float* result,
    int width, int height) {
  int row = blockIdx.y * blockDim.y + threadIdx.y;
  int col = blockIdx.x * blockDim.x + threadIdx.x;
  if (row < height && col < width) {
    int idx = row * width + col;
    result[idx] = matrix[idx] * matrix[idx];
  }
}

torch::Tensor square_matrix(torch::Tensor m) {
  auto h = m.size(0); auto w = m.size(1);
  auto out = torch::empty_like(m);
  dim3 t(16,16);
  dim3 b((w+15)/16, (h+15)/16);
  square_matrix_kernel<<<b,t>>>(
    m.data_ptr<float>(),
    out.data_ptr<float>(), w, h);
  return out;
}
'''
cpp_source = "torch::Tensor square_matrix(torch::Tensor);"

from torch.utils.cpp_extension import load_inline

ext = load_inline(
    name='square_matrix_extension',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['square_matrix'],
    with_cuda=True,
    extra_cuda_cflags=["-O2"],
    build_directory='./load_inline_cuda',
)

# 이제 그냥 파이썬 함수처럼 부른다
a = torch.tensor([[1.,2.,3.],
                  [4.,5.,6.]],
                 device='cuda')
print(ext.square_matrix(a))

Triton 은 파이썬 안 DSL 이지만 CUDA 코드를 만들지 않는다 — 직접 PTX 까지 내려간다. 통합은 trivial: 그냥 파이썬 함수를 부른다. 그런데 강의에서 Mark 가 보여준 첫 실험에서 torch.square 보다 자기 Triton 커널이 더 느렸다. 원인이 뭐였나 — 그게 이 섹션의 본론.

@triton.jit(interpret=True)
def square_kernel(out_ptr, in_ptr, n,
                  BLOCK: tl.constexpr):
    pid    = tl.program_id(0)
    offs   = pid * BLOCK + tl.arange(0, BLOCK)
    mask   = offs < n
    x      = tl.load(in_ptr + offs, mask=mask)
    breakpoint()                # ← 여기서 멈춘다
    tl.store(out_ptr + offs, x * x, mask=mask)

# pdb 안에서 print(x.tensor) 같이 wrapped tensor 의 내부를 본다

강의 후반부에서 Mark 가 가장 여러 번 추천한 트릭. 새 커널을 처음부터 짜지 말고, torch.compile 이 같은 연산에 대해 만들어내는 Triton 커널을 먼저 읽어본다. 그 코드가 보통 충분히 읽을 만 한 시작점이고, 변수명만 바꾸면 그대로 쓸 수 있다.

# square_compile.py
import torch

def square(a):
    return torch.square(a)

opt = torch.compile(square)
opt(torch.randn(10000, 10000).cuda())

# 실행 시 환경변수로 켜면 console 에 Triton 커널 dump
TORCH_LOGS=output_code python square_compile.py

# torch.compile 이 만들어낸 Triton 코드 (약식)
@triton.jit
def triton_(in_ptr0, out_ptr0, xnumel,
            XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex  = xoffset + tl.arange(0, XBLOCK)[:]
    xmask   = xindex < xnumel
    tmp0    = tl.load(in_ptr0 + xindex, xmask)
    tmp1    = tmp0 * tmp0       # ← elementwise mul
    tl.store(out_ptr0 + xindex, tmp1, xmask)

코드의 변수명은 tmp0, tmp1 같이 익명이지만, 구조는 명확하다 — load → 연산 → store. 그리고 강의의 다음 관찰이 진짜 흥미로운 지점.

PyTorch profiler 까지가 “커널 이름과 시간” 을 본다면, NCU 는 한 단계 더 — 한 커널 안에서 메모리 latency 가 발목을 잡고 있는지, register 가 부족한지, occupancy 가 모자라서 SM 이 비어 있는지를 직접 알려준다. 사용법은 한 줄.

ncu python train.py

그리고 시각 모드는 --set full -o trace 를 붙여서 실행한 뒤 Nsight Compute UI 에 trace 를 연다.

강의의 마지막 메시지 중 하나. CUDA 부터 시작하지 마라. torch → Triton → CUDA 의 사다리를 NCU 가 hint 를 주는 시점에 한 단계씩 내려간다. 그 hint 가 trigger 가 된다.

이 사다리에서 한 가지 더 짚어야 하는 점 — Numba 는 강의에서 의도적으로 추천 경로에서 빠진다. 개념적으로는 CUDA 와 동일한 모델인데 (grid, block, threadIdx), 강의의 인지적 부담을 줄이기 위해 “이 클래스에서는 C++/CUDA 또는 Triton 만 쓴다” 고 못 박는다. 산업에서도 Triton 과 직접 CUDA 가 사실상 표준이 됐다.

강의에서 6개월 뒤에 다시 돌아왔을 때 가장 빨리 복원해야 하는 사실들과 — 직접 손에 박아야 하는 코드 자료들.

손에 새기기 — 실습 시퀀스

1. async timing 패턴 베이스라인 — pytorch_square.py 를 그대로 돌려 torch.square, a*a, a**2 의 GPU 시간을 측정한다. 작은 차이가 어떤 op 디스패치에서 오는지 표로 확인.
2. load_inline 으로 첫 커널 — load_inline.py 를 돌리고 ./load_inline_cuda/ 안의 자동 생성된 main.cpp 와 build.ninja 를 직접 읽는다. pybind 가 어떻게 박혀 있는지 손에 잡힐 때까지.
3. Chrome trace — 같은 코드에 torch.profiler.profile() 을 둘러 trace.json 을 export 하고 chrome://tracing 에 드래그한다. memcpy H→D 와 커널 사이의 갭을 직접 본다.
4. Triton interpret 로 디버깅 — @triton.jit(interpret=True) 와 breakpoint() 로 row-wise square 커널을 한 step 씩 들여다본다. tl.arange, tl.load 의 결과 모양을 직접 print.
5. BLOCK_SIZE sweep — Triton square 커널을 BLOCK_SIZE ∈ {64, 128, 256, 512, 1024, 2048} 로 측정. torch.square 와 비교. 그래프가 U자나 단조 감소가 아니면 다른 변수가 섞인 것.
6. torch.compile 의 Triton 코드 읽기 — TORCH_LOGS=output_code 로 torch.square 와 torch.square(torch.square(x)) 두 케이스를 dump. fusion 이 일어났는지 코드 안에서 직접 확인.
7. NCU 한번 돌려보기 — 자기 GPU 가 있다면 ncu python pytorch_square.py. 첫 hint 가 무엇인지 적어둔다. cloud 환경에서는 NCU 가 막혀 있을 수 있으니 자기 데스크탑이나 NCU 가 허용된 인스턴스에서.

L001 의 도구들이 시리즈 안에서 어떻게 다시 호출되는지를 묶어둔다 — 같은 도구가 다른 문맥에서 재등장하는 패턴이 GPU Mode 시리즈 전체의 학습 곡선이다.

학습 노트로 정리하면서 의도적으로 비워둔 자리들 — 강의 안에서 부분적으로만 등장한 주제, 또는 후속 강의에서 본격적으로 다뤄지는 주제.

• Triton 커널과 같은 일을 하는 직접 CUDA 커널의 PTX 비교 — Mark 가 강의 중 비교를 시도하다가 시간 부족으로 못 마쳤다. 같은 row-wise square 를 Triton 과 직접 CUDA 로 짜고 두 PTX 를 직접 diff 해보면 어디가 다른지 확인 가능.
• dynamic shape 와 graph break — torch.compile 의 fusion 이 깨지는 조건들. 이 강의에서는 “복잡한 코드는 graph break 로 부분만 컴파일된다” 정도로만 언급됐다. L053 에서 본격적으로 다룬다.
• nsys 와 Nsight Systems 의 위치 — 강의에서 nsys_square.py 가 repo 에 있지만 본문에서는 NCU 만 본다. system-level 타임라인(NCCL, CPU, multi-GPU)은 nsys 의 영역이고 NCU 는 단일 커널 깊이파기. 분리 잘 해두기.
• cudaGraphs — Mark 가 마지막 Q&A 에서 “언제 cudaGraphs 를 쓸지” 가 별도 강의 주제로 적합하다고 답했다. launch overhead 가 dominant 한 학습 루프에서 큰 차이를 낸다.
• 실제 프로덕션 LLM 추론 코드에서의 적용 — 강의 자체는 toy square 커널 위에서 도구를 깐다. SAM-fast / GPT-fast 같은 사례에서 같은 도구 시퀀스가 어떻게 쓰였는지 — Mark 가 마지막에 잠깐 언급. 별도로 추적할 필요.