normalization kernel은 왜 중요할까

transformer 계열 모델을 보면 normalization은 거의 빠지지 않는다. layernorm이나 RMSNorm은 수식만 보면 크지 않아 보여도, 실제 학습과 추론에서 반복적으로 등장하기 때문에 누적 비용이 상당하다. 그래서 normalization kernel은 GPU 최적화에서 자주 다루는 주제다.

layernorm은 어떤 계산을 하나

layernorm은 보통 한 row 혹은 feature vector에 대해 평균과 분산을 구하고, 이를 이용해 normalize한 뒤 affine transform을 적용한다. 즉, 내부적으로는:

  • mean reduction
  • variance 계산
  • normalization
  • scale/bias 적용

같은 흐름이 들어 있다.

이 자체만 봐도 reduction과 elementwise 연산이 함께 섞여 있다는 점에서 softmax와 닮아 있다.

왜 memory-bound가 되기 쉬운가

layernorm은 matmul처럼 거대한 arithmetic intensity를 갖는 연산이 아니다. 많은 경우 memory traffic에 비해 계산이 그리 크지 않기 때문에 bandwidth-bound가 되기 쉽다.

즉, layernorm 최적화의 핵심은 엄청난 수학 최적화보다:

  • row를 몇 번 읽는가
  • intermediate를 얼마나 덜 남기는가
  • reduction을 얼마나 효율적으로 하느냐

에 더 가깝다.

RMSNorm은 무엇이 다를까

RMSNorm은 mean을 빼지 않고 root mean square 기반으로 normalize한다. 수식적으로는 조금 단순해지고, 경우에 따라 구현도 더 가벼울 수 있다. 하지만 GPU 관점에서 중요한 점은 여전히 비슷하다.

  • reduction 구조가 있다
  • memory traffic가 중요하다
  • row 단위 병렬화가 중요하다

즉, layernorm과 RMSNorm은 디테일은 다르지만 같은 class의 kernel로 보는 편이 좋다.

normalization kernel을 최적화할 때 보는 것

  • 한 row를 몇 번 스캔하는가
  • reduction을 warp 수준에서 끝낼 수 있는가
  • block 수준 협업이 필요한가
  • weight/bias 적용까지 fuse할 수 있는가

이 질문들이 normalization kernel 설계의 핵심이다.

정리

layernorm과 RMSNorm은 작아 보이는 연산이지만, GPU kernel 관점에서는 매우 좋은 공부 재료다. reduction, memory-bound 특성, fusion 가능성, row-wise 병렬화가 모두 들어 있기 때문이다.

다음 글에서는 vectorized load/store와 alignment를 보면서, 메모리 접근 폭을 어떻게 더 효율적으로 맞출 수 있는지 본다.