softmax는 단순해 보이지만 실제로는 꽤 복잡하다

softmax는 수식만 보면 특별히 어려워 보이지 않는다. 각 값을 exp로 바꾸고 합으로 나누면 된다. 하지만 GPU kernel 관점에서는 여러 까다로운 요소가 한 번에 들어오는 연산이다.

  • max reduction이 필요하다
  • sum reduction이 또 필요하다
  • 메모리를 여러 번 읽고 쓸 위험이 크다
  • numerical stability를 고려해야 한다

그래서 softmax는 GPU 시스템을 공부할 때 매우 좋은 연습 문제다.

왜 max reduction이 먼저 필요한가

softmax를 그대로 계산하면 큰 값에서 exp가 터질 수 있다. 그래서 보통 row의 최대값을 먼저 구하고, 각 원소에서 그 값을 빼준 뒤 exp를 계산한다. 이 과정은 numerical stability를 위한 기본 패턴이다.

즉, softmax는 단순 elementwise kernel이 아니라, row-wise reduction과 elementwise transformation이 결합된 구조다.

메모리 관점에서 왜 부담이 큰가

naive하게 구현하면 row를 여러 번 읽게 된다.

  1. max를 구하기 위해 한 번 읽고
  2. exp를 계산하며 다시 읽고
  3. sum을 구하며 또 읽고
  4. 최종 출력에 쓰기 위해 한 번 더 처리한다

이런 구조는 GPU에서 곧바로 bandwidth 문제로 이어진다. 즉, softmax의 핵심은 exp 계산이 아니라 메모리 흐름을 얼마나 잘 정리하느냐인 경우가 많다.

row 단위 병렬화가 왜 중요할까

softmax는 보통 row 단위로 처리하는 경우가 많다. 예를 들어 attention score matrix의 각 row에 대해 softmax를 적용한다. 이때 중요한 것은 row 길이에 따라 block 구성과 reduction 전략이 달라질 수 있다는 점이다.

row가 짧으면 warp 수준에서 끝낼 수 있고, 길면 block 전체 reduction이 필요할 수 있다. 결국 softmax도 데이터 shape에 맞는 kernel 설계가 중요하다.

reduction과 elementwise를 어떻게 묶을까

좋은 softmax kernel은 reduction 단계와 exp/normalize 단계를 완전히 따로 떼기보다, memory traffic를 줄이는 방식으로 최대한 묶으려 한다. 여기서 shared memory나 warp-level primitive가 중요한 역할을 한다.

즉, softmax 최적화는 단순히 더 빠른 exp를 찾는 작업이 아니라:

  • row max를 어떻게 줄일지
  • row sum을 어떻게 줄일지
  • intermediate를 메모리에 어떻게 덜 남길지

를 설계하는 문제다.

왜 딥러닝 kernel의 축소판처럼 느껴질까

softmax를 잘 보면 현대 딥러닝 kernel의 특징이 꽤 많이 들어 있다.

  • reduction이 있다
  • memory-bound가 되기 쉽다
  • numerical stability가 중요하다
  • shape에 따라 최적 전략이 달라진다

그래서 softmax를 이해하면 layernorm, attention, fused normalization 계열도 더 쉽게 읽을 수 있다.

정리

softmax kernel은 GPU 시스템 학습에서 좋은 이유가 분명하다.

  • reduction을 실제 연산에 연결해준다
  • memory traffic 문제를 체감하게 해준다
  • numerical stability까지 같이 보게 만든다

다음 글에서는 layernorm과 RMSNorm을 보면서 비슷하지만 다른 normalization kernel의 구조를 비교해본다.