From 171e2cacc18590ba88501d1f23f93eb9f5dc4893 Mon Sep 17 00:00:00 2001 From: Felipe Aramburu Date: Thu, 5 Mar 2026 10:51:26 -0600 Subject: [PATCH 1/2] made stream pool for memory space non blocking main --- src/memory/memory_space.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/memory/memory_space.cpp b/src/memory/memory_space.cpp index f87e227..f82422e 100644 --- a/src/memory/memory_space.cpp +++ b/src/memory/memory_space.cpp @@ -52,7 +52,8 @@ memory_space::memory_space(const gpu_memory_space_config& config) : make_default_gpu_memory_resource(config.device_id, config.memory_capacity)), _stream_pool{[&]() -> std::unique_ptr { rmm::cuda_set_device_raii guard{rmm::cuda_device_id(config.device_id)}; - return std::make_unique(16); + return std::make_unique(16, + rmm::cuda_stream_pool::flags::non_blocking); }()} { if (!_allocator) { throw std::invalid_argument("At least one allocator must be provided"); } From e4aebae6fe99968d78bbf7aef1e8f90e87f102db Mon Sep 17 00:00:00 2001 From: Amin Aramoon Date: Fri, 6 Mar 2026 10:27:56 -0800 Subject: [PATCH 2/2] fix bug --- src/memory/memory_space.cpp | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/memory/memory_space.cpp b/src/memory/memory_space.cpp index f82422e..0a43900 100644 --- a/src/memory/memory_space.cpp +++ b/src/memory/memory_space.cpp @@ -52,8 +52,7 @@ memory_space::memory_space(const gpu_memory_space_config& config) : make_default_gpu_memory_resource(config.device_id, config.memory_capacity)), _stream_pool{[&]() -> std::unique_ptr { rmm::cuda_set_device_raii guard{rmm::cuda_device_id(config.device_id)}; - return std::make_unique(16, - rmm::cuda_stream_pool::flags::non_blocking); + return std::make_unique(16, rmm::cuda_stream::flags::non_blocking); }()} { if (!_allocator) { throw std::invalid_argument("At least one allocator must be provided"); }