[onnxruntime] ONNX Runtime CUDA ArgMax/ArgMin 연산 최적화: 대규모 행렬 병렬 처리 개선

PR 링크: microsoft/onnxruntime#32092 상태: Merged | 변경: +979 / -18

들어가며

ONNX Runtime의 arg_min_max_last_axis() 함수는 기존에 행당 하나의 CUDA 스레드를 할당하여 직렬로 스캔하는 방식을 사용했습니다. 이 방식은 행의 길이가 짧은 경우에는 효율적이지만, 대규모 언어 모델의 로짓 계산과 같이 행의 길이가 매우 긴(예: [1, 202048]) 경우, 단 하나의 스레드가 전체 행을 스캔해야 하므로 GPU 자원이 극도로 낭비되는 병목 현상이 발생합니다. 본 PR은 이러한 문제를 해결하기 위해 협력적 병렬 처리(Cooperative Reduction) 경로를 도입했습니다.

코드 분석

reduction_functions.cu: 협력적 커널 도입

기존의 직렬 커널은 행당 하나의 스레드만 사용했지만, 새로운 arg_min_max_last_axis_cooperative_kernel은 행을 여러 워프(Warp)와 블록(Block)으로 나누어 처리합니다.

Before (기존 직렬 방식):

__global__ void arg_min_max_last_axis_kernel(const TIn* input, int64_t* output, int m, int n) {
  const int row = blockIdx.x * blockDim.x + threadIdx.x;
  if (row >= m) return;
  // ... 단일 스레드에서 전체 행을 순차적으로 스캔
}

After (개선된 협력적 방식):

__global__ void arg_min_max_last_axis_cooperative_kernel(...) {
  // ... 워프 수준의 셔플(Shuffle)과 블록 수준의 공유 메모리(Shared Memory)를 활용한 병렬 리덕션
  // 여러 블록이 참여하는 경우, 글로벌 버퍼를 통해 최종 결과를 병합
  const int count = atomicAdd(block_done_counts_buffer + row_id, 1);
  is_last_block_done = (count == (num_blocks_in_grid_row - 1));
}

핵심은 WARP_SHFL_DOWN을 이용한 워프 내 병합, 공유 메모리를 이용한 블록 내 병합, 그리고 atomicAdd를 이용한 그리드 수준의 최종 병합입니다. 이를 통해 긴 행을 여러 스레드가 나누어 처리함으로써 GPU 점유율을 극대화했습니다.

왜 이게 좋은가

성능 개선 수치

H200 GPU 환경에서 [1, 202048] 크기의 행렬에 대해 기존 약 5ms가 소요되던 작업이 약 6μs로 단축되어 약 800배 이상의 성능 향상을 보였습니다.

교훈

1. 데이터 의존성 제거: 기존의 id += step 방식은 int 오버플로우로 인해 무한 루프에 빠질 위험이 있었습니다. 이를 remaining과 delta를 이용한 차분 방식으로 변경하여 안전성을 확보했습니다.
2. 하이브리드 전략: 모든 경우에 병렬 처리가 좋은 것은 아닙니다. 128개 컬럼 미만의 좁은 행은 오버헤드 때문에 기존 직렬 커널이 더 빠릅니다. 이 PR은 임계값(Threshold)을 설정하여 상황에 맞는 최적의 커널을 디스패치합니다.
3. 메모리 효율성: cub::DeviceReduce와 같은 범용 라이브러리보다 훨씬 적은 스크래치 메모리(약 26배 감소)를 사용하면서도 단일 커널 실행으로 성능을 최적화했습니다.

리뷰 피드백 반영

리뷰어들은 int 오버플로우 문제를 지적했습니다. 이에 따라 인덱스 계산 로직을 id + delta 대신 delta < remaining 조건을 활용한 차분 방식으로 수정하여, INT_MAX에 달하는 대규모 행렬에서도 안전하게 동작하도록 개선했습니다.

참고 자료

• https://pytorch.org/docs/stable/generated/torch.compile.html
• https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#warp-shuffle-functions

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