C++ 템플릿 라이브러리인 CUB는 고성능 GPU 프리미티브 알고리즘을 위한 필수 도구입니다. 하지만 메모리 예측과 할당을 분리하는 기존의 2단계 API는 사용법이 번거로울 수 있습니다. 이러한 프로그래밍 모델은 유연성을 제공하지만, 종종 반복적인 상용구(Boilerplate) 코드를 양산하는 결과를 초래합니다.
본 포스팅에서는 CUDA 13.1에서 도입된 새로운 CUB 단일 호출 API로의 변화에 대해 설명합니다. 이 API는 성능 저하 없이 내부적으로 메모리를 직접 관리함으로써 개발 프로세스를 단순화합니다.
CUB란 무엇인가요?
GPU에서 스캔, 히스토그램, 정렬과 같은 표준 알고리즘을 실행해야 한다면, CUB가 가장 빠른 방법입니다. NVIDIA CCCL(CUDA Core Compute Libraries)의 핵심 구성 요소인 CUB는 성능을 저하하지 않으면서도 수동적인 CUDA 스레드 관리의 복잡성을 추상화하도록 설계되었습니다.
Thrust와 같은 라이브러리가 빠른 프로토타이핑을 위해 C++ 표준 템플릿 라이브러리(STL)와 유사한 고수준의 ‘호스트 측’ 인터페이스를 제공한다면, CUB는 일련의 ‘디바이스 측’ 프리미티브를 제공합니다. 이를 통해 개발자는 고도로 최적화된 알고리즘을 자신의 커스텀 커널에 직접 통합할 수 있습니다. CUB 사용법에 대해 더 자세히 알아보려면 NVIDIA DLI 과정인 ‘Modern CUDA C++을 활용한 가속 컴퓨팅의 기초‘를 확인해 보세요.
기존 CUB 2단계 API의 구조
CUB는 NVIDIA GPU의 연산 능력을 최대한 활용하기 위해 널리 권장되는 라이브러리입니다. 그럼에도 불구하고, 실제 사용 시에는 다소 까다롭게 느껴질 수 있는 복잡성을 수반합니다. 본 섹션에서는 이러한 기저의 메커니즘을 관점에 따라 다시 확인해볼 예정입니다.
일반적으로는 함수 프리미티브를 한 번 호출하는 것만으로 알고리즘이 실행되고 즉시 결과를 얻을 수 있는 단순한 단일 패스 실행 흐름을 가정합니다. 변수 수정이나 결과 반환과 같은 함수의 부수 효과가 다음 문장에서 바로 나타나기를 기대하는 것입니다.
그러나 CUB 실행 모델은 이러한 익숙한 패턴에서 벗어납니다. CUB 프리미티브를 호출하는 과정은 먼저 필요한 디바이스 메모리 크기를 계산하고(첫 번째 호출), 이후 메모리를 명시적으로 할당한 뒤 커널을 실행하는(두 번째 호출) 두 단계의 프로세스를 거칩니다.
다음은 전형적인 CUB 호출 방식의 예시입니다.
// FIRST CALL: determine temporary storage size
cub::DeviceScan::ExclusiveSum(nullptr, temp_storage_bytes, d_input, d_output, num_items);
// Allocate the required temporary storage
cudaMalloc(&d_temp_storage, temp_storage_bytes);
// SECOND CALL: run the actual scan
cub::DeviceScan::ExclusiveSum(d_temp_storage, temp_storage_bytes, d_input, d_output, num_items);
CUB 인터페이스는 실무적인 과제를 안겨줍니다. 프리미티브를 두 번 호출해야 하기 때문입니다. 첫 번째는 필요한 임시 메모리 양을 결정하기 위함이고, 두 번째는 할당된 저장 공간을 바탕으로 실제 알고리즘을 실행하기 위함입니다.
기존 2단계 API의 주요 단점은 메모리 추정 단계와 실행 단계 사이에서 어떤 인자가 일관되게 유지되어야 하는지 명확하지 않다는 점입니다. 위의 코드 스니펫을 참고하면, 두 단계의 함수 시그니처가 동일하기 때문에 어떤 파라미터가 내부 상태에 영향을 미치고 호출 간에 변경될 수 있는지 프로그램적으로 파악하기 어렵습니다. 예를 들어, d_input과 d_output 인자는 실제로는 두 번째 호출에서만 사용됩니다.
이러한 복잡성에도 불구하고 기존 설계는 다음과 같은 근본적인 목적을 수행합니다. 할당과 실행을 분리함으로써 사용자는 메모리 청크를 한 번 할당한 뒤 이를 여러 번 재사용하거나 서로 다른 알고리즘 간에 공유할 수 있습니다.
이러한 설계는 무시할 수 없는 일부 사용자층에게는 중요하지만, 이 기능을 활용하는 전체 사용자 기반은 다소 제한적입니다. 이것이 바로 많은 사용자가 매 사용 시 요구되는 2단계 호출을 추상화하기 위해 CUB 호출을 별도의 코드로 감싸서 사용하는 이유입니다. PyTorch가 대표적인 사례로, 매크로를 사용하여 CUB 호출을 단일 호출로 래핑하고 자동 메모리 관리를 제공합니다.
다음 소스 코드는 pytorch/pytorch GitHub 리포지토리에서 발췌한 것입니다.
// handle the temporary storage and 'twice' calls for cub API
#define CUB_WRAPPER(func, ...) do { \
size_t temp_storage_bytes = 0; \
AT_CUDA_CHECK(func(nullptr, temp_storage_bytes, __VA_ARGS__)); \
auto& caching_allocator = *::c10::cuda::CUDACachingAllocator::get(); \
auto temp_storage = caching_allocator.allocate(temp_storage_bytes); \
AT_CUDA_CHECK(func(temp_storage.get(), temp_storage_bytes, __VA_ARGS__));\
} while (false)
매크로 사용은 제어 흐름과 파라미터 전달을 모호하게 만들어 코드 이해를 어렵게 하고 디버깅을 상당히 저해한다는 단점이 있습니다.
새로운 단일 호출 CUB API
많은 프로덕션 코드베이스에서 래퍼가 널리 사용됨에 따라, 새로운 단일 호출 API를 도입하여 CUB를 확장해야 할 필요성이 인정되었습니다.
// SINGLE CALL: allocation and execution on a single step
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items);
이 예제에서 볼 수 있듯이 명시적인 메모리 할당이 필요하지 않습니다. 다만, 메모리 할당 프로세스는 여전히 내부적으로 수행된다는 점에 유의하세요. 그림 1은 임시 저장 공간 추정, 메모리 할당 및 알고리즘 호출을 포함하는 단일 호출 인터페이스가 기존 2단계 API와 비교하여 오버헤드를 전혀 발생시키지 않음을 보여줍니다.
그림 1은 기존의 2단계 ExclusiveSum 호출과 새로 도입된 단일 단계(Single-phase) 호출의 GPU 런타임을 비교한 결과입니다. x축은 다양한 입력 크기를, y축은 각 호출 방식에 따른 정규화된 실행 시간을 나타냅니다. 이 성능 데이터로부터 도출된 두 가지 주요 결론은 다음과 같습니다.
- 새로운 API는 오버헤드를 전혀 발생시키지 않습니다.
- 새로운 API에서도 메모리 할당은 여전히 수행되나, 단지 내부적으로 처리될 뿐입니다.
두 번째 결론은 새로운 API의 내부 구현을 살펴봄으로써 확인할 수 있습니다. 디바이스 프리미티브 내에 비동기 할당 방식이 내재되어 있습니다.
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items, env = {}) {
. . .
d_temp_storage = mr.allocate(stream, bytes);
mr.deallocate(stream, d_temp_storage, bytes);
. . .
}
기존의 2단계 API가 제거된 것은 아니며, 여전히 유효한 호출 방식입니다. 단일 단계 호출은 기존 API 위에 추가된 형태이며, 향후 대다수의 사용자가 이 방식을 채택할 것으로 예상됩니다.
환경 및 메모리 리소스
새로운 단일 호출 CUB API는 앞서 언급한 문제들을 해결할 뿐만 아니라, 호출되는 프리미티브의 실행 구성 역량을 확장합니다. 이 API는 ‘환경(Environment)’ 인자를 도입하여 메모리 리소스를 사용한 메모리 할당 커스텀화나, 기존 2단계 API처럼 실행을 위한 스트림 제공 기능을 지원합니다.
메모리 리소스는 메모리 할당 및 해제를 위한 새로운 유틸리티입니다. 단일 호출 API의 환경 인자에는 선택적으로 메모리 리소스를 포함할 수 있습니다. 환경 인자를 통해 메모리 리소스를 명시하지 않을 경우, API는 CCCL에서 제공하는 기본 메모리 리소스를 사용합니다. 반대로, 코드베이스에 포함된 비기본 CCCL 메모리 리소스를 전달하거나 사용자 정의 메모리 리소스를 직접 전달할 수도 있습니다.
// Use CCCL-provided memory resource type
cuda::device_memory_pool mr{cuda::devices[0]};
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items, mr);
// Create and use your custom MR
my_memory_resource my_mr{cuda::experimental::devices[0]};
// Use it with CUB
cub::DeviceScan::ExclusiveSum(d_input, d_output, num_items, my_mr);
새로운 API에서 CUDA 실행 스트림 처리는 간소화된 것이 아니라 새로운 env 변수 내에 캡슐화되었습니다. 물론 임시 할당 처리가 자동화되었더라도 이전처럼 스트림을 명시적으로 전달할 수 있습니다. CUB는 이제 타입 안정성이 보장되는 cuda::stream_ref를 제공하며, 이를 사용하는 것이 권장됩니다. 또한 하부 실행 스트림을 소유하는 cuda::stream을 전달할 수도 있습니다.
실행 옵션의 결합
단일 호출 API는 마지막 인자로 메모리 리소스나 스트림을 전달하는 것 이상의 기능을 수행합니다. 앞으로 환경 인자는 결정론적 요구 사항, 보장 사항, 사용자 정의 튜닝 등 모든 실행 관련 설정을 관리하는 지점이 될 것입니다.
단일 패스 API의 도입으로 CUB는 방대한 실행 구성 기능을 확보하게 되었습니다. 이러한 수많은 기능을 결합하는 가장 효율적인 방법은 무엇일까요?
그 해답은 새로운 env 인자에 있습니다. CUB는 cuda::std::execution을 활용하여 알고리즘을 위한 유연한 ‘제어 패널’ 역할을 하는 중앙 엔드포인트를 제공합니다. 환경 인자를 사용하면 엄격하게 정의된 함수 인자 대신, 필요한 기능을 조합하여 구성할 수 있습니다. 커스텀 스트림과 특정 메모리 풀을 쌍으로 묶거나, 엄격한 결정론적 요구 사항과 커스텀 튜닝 정책을 결합하는 등 이 모든 것이 타입 안정성이 보장된 단일 객체 내에서 처리됩니다.
cuda::stream custom_stream{cuda::device_ref{0}};
auto memory_prop = cuda::std::execution::prop{cuda::mr::get_memory_resource,
cuda::device_default_memory_pool(cuda::device_ref{0})};
auto env = cuda::std::execution::env{custom_stream.get(), memory_prop};
DeviceScan::ExclusiveSum(d_input, d_output, num_items, env);
현재 CUB에서 환경 인터페이스를 지원하는 알고리즘은 다음과 같으며, 향후 더욱 확대될 예정입니다.
- cub::DeviceReduce::Reduce
- cub::DeviceReduce::Sum
- cub::DeviceReduce::Min / Max / ArgMin / ArgMax
- cub::DeviceScan::ExclusiveSum
- cub::DeviceScan::ExclusiveScan
새로운 환경 기반 오버로드의 최신 진행 상황은 NVIDIA/cccl GitHub 리포지토리의 CUB 디바이스 프리미티브 트래킹 이슈에서 확인할 수 있습니다.
CUB 시작하기
CUB는 번거로운 2단계 패턴을 간소화된 단일 호출 인터페이스로 대체함으로써, 오버헤드 없이 상용구 코드를 제거한 현대적인 API를 제공합니다. 확장 가능한 env 인자를 활용하면 메모리 리소스, 스트림 및 기타 편의 기능을 원활하게 조합하는 통합 제어 패널을 사용할 수 있습니다. 코드베이스를 단순화하고 GPU의 연산 성능을 온전히 활용하기 위해 이 새로운 표준을 도입해 보시기 바랍니다. CUDA 13.1 이상의 버전을 다운로드하여 이 단일 호출 API를 바로 사용해 보세요.
