reduction은 작은 예제지만 내용은 작지 않다

sum, max, mean 같은 reduction은 겉보기에 단순하다. 많은 값을 하나로 줄이는 것이기 때문이다. 하지만 GPU에서는 이 단순한 문제가 굉장히 좋은 학습 재료가 된다. shared memory, warp primitive, synchronization, memory access, multi-stage aggregation이 모두 들어 있기 때문이다.

왜 reduction이 어렵게 느껴질까

GPU는 병렬 장치다. 그런데 reduction은 많은 값을 결국 적은 값으로 줄여야 한다. 즉, 병렬성을 살리면서도 점점 결과를 모아야 한다는 점에서 구조적으로 긴장이 있다.

이 때문에 reduction kernel을 설계할 때는 보통 두 가지를 동시에 생각해야 한다.

  • 각 thread가 충분히 많은 일을 하도록 만들기
  • 그 결과를 너무 비싸지 않게 합치기

전형적인 block-level reduction 흐름

가장 흔한 구조는 이렇다.

  1. 각 thread가 여러 원소를 읽어 partial sum을 만든다
  2. partial sum을 shared memory에 쓴다
  3. block 내부에서 단계적으로 합친다
  4. 마지막 warp 수준은 shuffle로 줄일 수 있다
  5. block 결과는 별도로 다시 합친다

이 구조를 이해하면 softmax, layernorm, attention 일부 연산도 훨씬 읽기 쉬워진다.

memory access도 여전히 중요하다

reduction이라고 해서 메모리 문제가 사라지는 것은 아니다. 오히려 input을 읽는 첫 단계에서 coalescing이 잘 되어 있는지, thread 하나가 여러 원소를 처리하는지, grid-stride loop를 쓰는지 같은 요소가 성능에 크게 영향을 준다.

즉, reduction은 aggregation 문제이면서 동시에 memory throughput 문제다.

block 내부 reduction에서 synchronization

shared memory를 쓰는 순간 synchronization이 필요하다. 이 비용은 필수지만, 구조를 잘 짜지 않으면 지나치게 커질 수 있다. 그래서 reduction 최적화에서는 어느 지점부터 warp-level primitive로 넘어갈지 판단하는 것이 중요하다.

보통 마지막 32개 수준에서는 warp synchronous하게 다루는 것이 더 가볍다.

multi-block reduction은 또 다른 문제다

입력이 크면 block 하나로는 끝나지 않는다. 그러면 각 block이 partial result를 만들고, 그 결과들을 다시 reduction해야 한다. 이때는 다음과 같은 선택지가 생긴다.

  • 두 단계 kernel launch
  • atomic 사용
  • cooperative group 같은 더 고급 구조

어떤 방법이 좋은지는 결과 크기와 병목 구조에 따라 달라진다.

reduction을 잘 이해하면 얻는 것

reduction은 단순 합계 예제 이상이다. 이 패턴을 잘 이해하면:

  • softmax의 max/sum reduction
  • layernorm의 mean/variance 계산
  • attention 내부 일부 집계
  • statistics 계산 kernel

같은 구조를 훨씬 쉽게 읽게 된다.

정리

reduction kernel은 GPU 최적화의 핵심 요소를 모아놓은 예제다.

  • input 읽기 단계의 memory pattern
  • shared memory 협업
  • warp-level primitive 사용
  • synchronization 비용
  • multi-stage aggregation 구조

이 패턴에 익숙해질수록 실제 딥러닝 kernel을 볼 때 훨씬 덜 막히게 된다.

다음 글에서는 이 reduction 패턴이 실제로 어떻게 softmax kernel에 들어가는지 구체적으로 본다.