How Deep Learning Computes

nano-vLLM RMSNorm Kernel Swap

마지막 수정:

inferencenano-vllmrmsnormkernel

RMSNorm은 첫 kernel swap 대상으로 좋다.

이유는 작고 correctness 비교가 쉽기 때문이다.

PyTorch RMSNorm
  -> custom CUDA RMSNorm
  -> optimized CUDA RMSNorm

교체 절차는 다음이다.

1. 같은 input에서 output close 확인
2. standalone RMSNorm latency 비교
3. nano-vLLM model forward에 연결
4. TTFT/TPOT/throughput 전후 비교

RMSNorm이 단독으로 빨라져도 serving 전체 개선은 작을 수 있다. 이 차이를 기록하는 것이 이 카드의 핵심이다.

확인

  • RMSNorm은 왜 첫 kernel swap 대상으로 적당한가?
  • standalone speedup과 serving speedup이 다를 수 있는 이유는 무엇인가?
  • output close 확인은 어떤 dtype tolerance로 봐야 하는가?
먼저 보면 좋은 카드nano-vLLM Kernel Integration Pointsnano-vLLM serving loop에서 RMSNorm, attention, sampling 같은 kernel 교체 지점을 찾는다.먼저 보면 좋은 카드Custom RMSNorm Extension을 PyTorch Transformer에 연결하기기존 CUDA RMSNorm kernel을 PyTorch Transformer block의 RMSNorm 위치에 연결하고 전후 성능을 비교한다.다음에 보기 좋은 카드nano-vLLM Prefill/Decode Benchmark Reportkernel swap 전후의 TTFT, TPOT, throughput, memory를 prefill/decode workload별로 정리한다.같은 경로의 다음 카드nano-vLLM Attention Kernel Swapmanual/SDPA/optimized attention을 nano-vLLM prefill/decode loop에 연결하고 전후 성능을 비교한다.