[onnxruntime] ONNX Runtime CUDA: int64 CumSum 연산 9배 가속화 최적화 분석

PR 링크: microsoft/onnxruntime#32238 상태: Merged | 변경: +175 / -8

들어가며

ONNX Runtime의 CumSum 연산은 누적 합계를 계산하는 필수적인 연산입니다. 기존 CUDA 구현체는 각 출력 요소를 독립적으로 계산하는 방식을 사용했는데, 이는 특히 batch-1과 같이 lane이 적은 경우 O(context)의 연산을 반복적으로 수행하게 되어 심각한 성능 저하를 야기했습니다. 본 PR은 CUB(CUDA Unbound) 라이브러리의 BlockScan을 도입하여, int64 타입의 low-lane 입력에 대해 연산 복잡도를 선형적으로 개선함으로써 최대 9.16배의 성능 향상을 달성했습니다.

코드 분석

1. onnxruntime/core/providers/cuda/math/cumsum_impl.cu

핵심 변경 사항은 _CumSumInt64BlockKernel 커널의 도입입니다. 기존 방식은 각 스레드가 독립적으로 전체 누적을 다시 계산했지만, 새로운 구현은 CUDA 블록 단위로 협력하여 prefix sum을 계산합니다.

// Before: 기존 방식은 각 스레드가 독립적으로 계산 (생략)

// After: CUB BlockScan을 활용한 협력적 커널
template <int BlockSize>
__global__ void _CumSumInt64BlockKernel(...) {
  using BlockScan = cub::BlockScan<uint64_t, BlockSize>;
  __shared__ typename BlockScan::TempStorage temp_storage;
  // ... (중략)
  for (int64_t tile = 0; tile < width; tile += BlockSize) {
    // ...
    BlockScan(temp_storage).InclusiveSum(value, prefix, aggregate);
    // ...
  }
}

또한, CumSumInt64Impl에서 width와 lanes를 기반으로 기존 커널과 새 커널을 분기 처리하여 최적화된 경로를 선택하도록 했습니다.

2. onnxruntime/core/providers/cuda/math/cumsum.cc

ComputeInternal 함수에서 int64_t 타입에 대해 새로 정의된 CumSumInt64Impl을 호출하도록 디스패치 로직을 수정했습니다.

-    CumSumImpl(Stream(ctx), ...);
+    ORT_RETURN_IF_ERROR(CumSumInt64Impl(Stream(ctx), ..., GetDeviceProp().multiProcessorCount));

왜 이게 좋은가

이번 최적화의 핵심은 연산 복잡도의 개선입니다. 기존의 one-thread-per-output 방식은 width가 커질수록 연산량이 제곱으로 증가하는 구조였으나, BlockScan을 활용하면 tile 단위로 aggregate 값을 넘겨주며 선형적으로 합계를 구할 수 있습니다.

성능 지표 (RTX A1000 기준):

• Width 8192: 180.398 us -> 19.692 us (9.16x 가속)
• Width 4096: 78.223 us -> 9.656 us (8.10x 가속)

교훈:

1. 알고리즘적 개선: GPU 커널 최적화 시 단순히 스레드 수를 늘리는 것보다, CUB와 같은 라이브러리를 활용해 블록 내 스레드 간의 데이터 공유(Shared Memory)를 최적화하는 것이 훨씬 효과적입니다.
2. 엣지 케이스 처리: 리뷰 과정에서 int 오버플로우 가능성을 지적받아 int64_t로 인덱스 타입을 변경한 점은 대규모 GPU 메모리 환경에서 매우 중요한 안정성 확보 사례입니다.
3. 회귀 방지: width < 4인 경우 기존 커널을 사용하도록 게이트를 설정하여 작은 입력에서의 성능 저하를 방지했습니다.

리뷰어 피드백 반영

리뷰어 tianleiwu는 tile 인덱스 계산 시 int 오버플로우 가능성을 지적했습니다. 이에 따라 int64_t로 타입을 변경하여 대용량 데이터 처리 시의 안정성을 확보했습니다. 또한 Copilot의 제안에 따라 width=2인 경우를 테스트 케이스에 추가하여, 최적화된 경로뿐만 아니라 기존 경로의 무결성도 함께 검증했습니다.

참고 자료

• https://nvlabs.github.io/cub/structcub_1_1_block_scan.html
• https://onnxruntime.ai/docs/api/c/index.html

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.