How Deep Learning Computes

Transformer를 위한 CUDA Naive Kernels

이 경로는 CUDA 좌표계를 익힌 뒤, Transformer에 필요한 기본 연산을 직접 CUDA kernel로 옮기는 단계다.

목표는 최적화가 아니다. 먼저 CPU reference와 비교해서 kernel이 맞는지 확인한다. transpose, GEMM, GEMV, softmax, RMSNorm은 모두 나중에 최적화와 vLLM 이해로 이어지는 기본 재료다.

마지막에는 naive RMSNorm kernel을 PyTorch JIT extension으로 연결한다. 이 카드가 장난감 .cu 실행과 실제 PyTorch tensor workflow 사이의 다리 역할을 한다.

  1. CPU Reference와 검증 루프 — CUDA kernel을 빠르게 만들기 전에 CPU reference와 비교해 정확성을 먼저 확인하는 습관을 만든다.
  2. Matrix Transpose Kernel — 입력 matrix의 row, col을 바꿔 out[col, row]에 쓰는 첫 non-elementwise memory movement kernel.
  3. Naive GEMM Kernel — thread 하나가 C[row, col] 하나를 맡고 A의 row와 B의 column을 dot product로 곱한다.
  4. Naive GEMV Kernel — matrix-vector multiplication을 CUDA로 구현하고 LLM decode에서 GEMV가 왜 중요한지 연결한다.
  5. Naive Softmax Kernel — row-wise max, sum(exp), normalize를 직접 구현하며 softmax가 단순 elementwise kernel이 아님을 확인한다.
  6. Reduction Pattern — 여러 input 값을 하나의 통계량으로 줄이는 sum, max, mean 패턴이 softmax와 RMSNorm의 공통 기반임을 이해한다.
  7. RMSNorm Forward Kernel — 최신 LLM에서 흔히 쓰이는 RMSNorm forward를 row-wise sum of squares와 normalize-scale로 구현한다.
  8. RMSNorm Backward Kernel — backward를 하나만 다룬다면 RMSNorm으로 grad_x와 grad_weight 계산에 필요한 reduction 패턴을 확인한다.
  9. Transformer Kernel Map — transpose, GEMM, GEMV, softmax, reduction, RMSNorm이 Transformer 안에서 어디에 쓰이는지 연결한다.
  10. PyTorch JIT Extension으로 RMSNorm 연결하기 — torch.utils.cpp_extension.load로 RMSNorm CUDA kernel을 Colab에서 즉석 빌드하고 torch.Tensor workflow에 연결한다.