How Deep Learning Computes

Custom RMSNorm Extension을 PyTorch Transformer에 연결하기

마지막 수정:

pytorchcudaextensionrmsnormbenchmark

이 카드는 기존 CUDA 트랙과 PyTorch Transformer path가 만나는 지점이다.

먼저 PyTorch reference RMSNorm이 있다.

PyTorch RMSNorm
  -> correctness reference

그다음 CUDA extension RMSNorm을 같은 interface로 감싼다.

CustomRMSNorm(nn.Module)
  -> autograd.Function
  -> C++ binding
  -> CUDA kernel

마지막으로 Transformer block 안의 RMSNorm만 교체한다.

Block(norm=PyTorchRMSNorm)
Block(norm=CustomCUDARMSNorm)

비교는 correctness와 speed를 함께 본다.

확인

  • custom kernel을 모델에 넣을 때 correctness를 먼저 확인해야 하는 이유는 무엇인가?
  • RMSNorm만 빨라져도 전체 step time이 크게 줄지 않을 수 있는 이유는 무엇인가?
  • 이 실험은 CUDA 트랙과 PyTorch 트랙을 어떻게 연결하는가?
먼저 보면 좋은 카드PyTorch RMSNorm과 Residual PathTransformer block에서 normalization과 residual connection이 activation 흐름을 어떻게 안정화하는지 코드로 구현한다.먼저 보면 좋은 카드PyTorch JIT Extension으로 RMSNorm 연결하기torch.utils.cpp_extension.load로 RMSNorm CUDA kernel을 Colab에서 즉석 빌드하고 torch.Tensor workflow에 연결한다.먼저 보면 좋은 카드Optimized RMSNorm BenchmarkPath 2의 naive RMSNorm을 block/warp reduction으로 개선하고 PyTorch extension에서 비교한다.다음에 보기 좋은 카드nano-vLLM RMSNorm Kernel SwapPyTorch RMSNorm을 custom/optimized RMSNorm kernel로 교체하고 model output과 serving metric을 비교한다.다음에 보기 좋은 카드PyTorch Transformer Benchmark Report 작성하기RMSNorm과 attention profiling, 최적화 결과를 하나의 재현 가능한 실험 리포트로 정리한다.같은 경로의 다음 카드Manual PyTorch Attention Benchmarkmatmul, causal mask, softmax, value matmul로 직접 작성한 attention을 profiling 기준선으로 삼는다.