How Deep Learning Computes

PyTorch Profiling과 Kernel Optimization

이 경로는 PyTorch 학습 코드를 측정하고, 기존 CUDA 트랙을 실제 모델에 연결하는 최적화 경로입니다.

목표는 “빠를 것 같다”가 아니라, profiler와 benchmark로 병목을 찾고, PyTorch 레벨 최적화와 custom CUDA extension의 전후 차이를 기록하는 것입니다.

첫 실습은 RMSNorm입니다. RMSNorm은 작고 통제하기 쉬워 PyTorch reference, custom CUDA extension, optimized CUDA kernel을 비교하기 좋습니다.

두 번째 실습은 attention입니다. 직접 작성한 PyTorch attention, torch.nn.functional.scaled_dot_product_attention, optimized attention backend를 같은 shape와 dtype에서 비교합니다. 특히 attention은 prefill-like workload와 decode-like workload를 나눠야 vLLM과 KV cache로 자연스럽게 이어집니다.

이 path의 최종 산출물은 코드만이 아니라 재현 가능한 benchmark report입니다.

baseline profile
  -> RMSNorm benchmark
  -> custom CUDA RMSNorm integration
  -> manual attention benchmark
  -> PyTorch SDPA benchmark
  -> prefill vs decode workload split
  -> optimized attention comparison
  -> end-to-end report
  1. PyTorch Single-GPU Training Loop — forward, loss, backward, optimizer update, mixed precision을 포함한 단일 GPU 학습 루프를 만든다.
  2. CUDA Event Timing과 Benchmark — CUDA 최적화 전에 kernel 구간을 일관되게 측정하는 법을 익힌다.
  3. PyTorch Profiler로 Baseline 측정하기 — torch.profiler와 CUDA event로 Transformer forward/backward의 시간과 메모리 병목을 기록한다.
  4. PyTorch RMSNorm Benchmark Baseline — 직접 구현한 PyTorch RMSNorm을 기준선으로 삼아 latency, memory, correctness tolerance를 측정한다.
  5. PyTorch JIT Extension으로 RMSNorm 연결하기 — torch.utils.cpp_extension.load로 RMSNorm CUDA kernel을 Colab에서 즉석 빌드하고 torch.Tensor workflow에 연결한다.
  6. Optimized RMSNorm Benchmark — Path 2의 naive RMSNorm을 block/warp reduction으로 개선하고 PyTorch extension에서 비교한다.
  7. Custom RMSNorm Extension을 PyTorch Transformer에 연결하기 — 기존 CUDA RMSNorm kernel을 PyTorch Transformer block의 RMSNorm 위치에 연결하고 전후 성능을 비교한다.
  8. Manual PyTorch Attention Benchmark — matmul, causal mask, softmax, value matmul로 직접 작성한 attention을 profiling 기준선으로 삼는다.
  9. PyTorch SDPA Attention Benchmark — torch.nn.functional.scaled_dot_product_attention을 manual attention 기준선과 비교한다.
  10. Attention Workload: Prefill vs Decode — attention benchmark를 prefill-like와 decode-like workload로 나누어 측정한다.
  11. torch.compile에서 Triton과 CUDA까지 — PyTorch 코드 최적화를 시작할 때 torch.compile, Triton, CUDA를 어떤 순서로 생각하면 좋은지 이해한다.
  12. torch.compile과 Fused Op 비교 — eager 실행, torch.compile, fused operation 후보를 비교하며 PyTorch 레벨 최적화의 한계를 확인한다.
  13. Optimized Attention Benchmark — manual attention, PyTorch SDPA, optimized attention backend를 prefill/decode workload에서 비교한다.
  14. PyTorch Transformer Benchmark Report 작성하기 — RMSNorm과 attention profiling, 최적화 결과를 하나의 재현 가능한 실험 리포트로 정리한다.