Softmax Operation — Design & Implementation Notes

1. 개요

Softmax 는 주어진 입력 벡터를 확률 분포로 변환하는 함수

 

2. 수치 안정화 기법

실제 구현에서는 오버플로우를 막기 위해 max-trick 의 사용, 

 

3. CUDA 커널 구조

각 row(M) 마다 병렬 처리

• Warp reduction 으로 max, sum 등을 구함
• 공유 메모리에 warp 별 결과를 모아 한 번 더 reduce
• 두 패스 구조
  • 첫 패스에서 row-wise max, sum 을 계산
  • 두 번째 패스에서 normalize
• Forwad 와 Backward 모두 같은 패턴 사용

Forward 커널

1) row_max = reduce_max(x[row])
2) for each element:
       e = exp(x - row_max)
       sum += e
3) for each element:
       y = e / sum   (or log(y/sum))

Backward 커널

if log=False:
    dot = sum(dY * Y)
    dX  = (dY - dot) * Y
else (log=True):
    sum_dy = sum(dY)
    dX     = dY - sum_dy * Y

 

주요 고려 사항

• Mask 지원: 일부 구현에서는 attention mask를 더해 사용 (x + mask).
• Scale 인자: temperature scaling을 지원 (e.g. Transformer attention).
• 수치 안정성: 반드시 max-trick 적용.
• 성능 최적화:
  • Warp-level shuffle (__shfl_down_sync)을 사용해 빠른 리덕션.
  • Shared memory를 최소화하여 블록 병렬성을 유지.