From 735180369eeb52838ce74569b3983def1e02e5d3 Mon Sep 17 00:00:00 2001 From: Malte Foerster Date: Wed, 27 May 2026 10:45:23 +0000 Subject: [PATCH 1/5] fix workspace vs large workspace usage in optimize --- cpp/src/neighbors/detail/cagra/graph_core.cuh | 20 ++++++------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 0e2ad6d769..76992cc7d5 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -1626,11 +1626,12 @@ void prune_graph_gpu( namespace bli = cuvs::spatial::knn::detail::utils; auto [copy_stream, enable_prefetch] = bli::get_prefetch_stream(res); auto workspace_mr = raft::resource::get_workspace_resource_ref(res); + auto large_workspace_mr = raft::resource::get_large_workspace_resource_ref(res); // Single-batch read-only iterator for the input graph (graph_size rows fit in one batch). bli::batch_load_iterator< raft::mdspan, raft::row_major, AccessorKnnGraph>> - d_input_graph(res, knn_graph, graph_size, copy_stream, workspace_mr); + d_input_graph(res, knn_graph, graph_size, copy_stream, large_workspace_mr); auto input_view = (*d_input_graph).view(); bli::batch_load_iterator< @@ -1770,19 +1771,10 @@ void optimize( prune_graph_gpu(res, knn_graph, new_graph); } - // reverse graph creation will always use the GPU - // using default workspace resource for random access - // otherwise will be managed memory which is slow upon first access - auto d_rev_graph = raft::make_device_mdarray(res, raft::make_extents(0, 0)); - try { - d_rev_graph = raft::make_device_mdarray( - res, default_ws_mr, raft::make_extents(graph_size, output_graph_degree)); - } catch (const std::exception& e) { - RAFT_LOG_DEBUG( - "Failed to create device matrix for reverse graph, switching to large workspace resource"); - d_rev_graph = raft::make_device_mdarray( - res, large_tmp_mr, raft::make_extents(graph_size, output_graph_degree)); - } + // reverse graph creation will always use the GPU / large workspace resource + auto d_rev_graph = raft::make_device_mdarray( + res, large_tmp_mr, raft::make_extents(graph_size, output_graph_degree)); + // This should use the default workspace resource for random access / atomics auto d_rev_graph_count = raft::make_device_mdarray( res, default_ws_mr, raft::make_extents(graph_size)); From dfb5c6f8f3e8ee925af4923c9009d0b8cb341303 Mon Sep 17 00:00:00 2001 From: Malte Foerster Date: Wed, 27 May 2026 10:48:52 +0000 Subject: [PATCH 2/5] limit kmeans workspace usage to 80% of free --- cpp/src/cluster/detail/kmeans_balanced.cuh | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index 26ce110158..ae033c5107 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -204,8 +204,11 @@ auto calc_minibatch_size(const raft::resources& handle, // If we need to convert to MathT, space required for the converted batch. if (!needs_conversion) { mem_per_row += sizeof(MathT) * dim; } - // Heuristic: calculate the minibatch size in order to use at most 1GB of memory. - IdxT minibatch_size = (1 << 30) / mem_per_row; + // Heuristic: calculate the minibatch size in order to use at most 80% or 1GB workspace memory. + auto available_ws_size = + std::min(raft::resource::get_workspace_free_bytes(handle) * 0.8, 1 << 30); + + IdxT minibatch_size = available_ws_size / mem_per_row; minibatch_size = 64 * raft::div_rounding_up_safe(minibatch_size, IdxT{64}); minibatch_size = std::min(minibatch_size, n_rows); return std::make_tuple(minibatch_size, mem_per_row); From 54ba9c8c93fe984acbcc8a28ec665ca89700a3e3 Mon Sep 17 00:00:00 2001 From: Malte Foerster Date: Wed, 27 May 2026 10:51:14 +0000 Subject: [PATCH 3/5] limit ivf-pq batchsize estimates to 80% of free workspace --- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index 08fcd1f09a..cc3266c0eb 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -780,7 +780,12 @@ inline auto get_max_fine_batch_size(raft::resources const& res, return static_cast(bs) * (other + (is_local_topk_feasible(k, n_probes, bs) ? buffers_fused : buffers_non_fused)); }; - auto max_ws_size = raft::resource::get_workspace_free_bytes(res); + + // Use only 40% of the available workspace memory as the coarse batch allocations have not been + // allocated yet. Also limit the maximum to 1GB to avoid conflicts in case the workspace is shared + // with the large workspace. + auto max_ws_size = + std::min(raft::resource::get_workspace_free_bytes(res) * 0.4, 1 << 30); if (ws_size(max_batch_size) > max_ws_size) { uint32_t smaller_batch_size = raft::bound_by_power_of_two(max_batch_size); // gradually reduce the batch size until we fit into the max size limit. @@ -837,11 +842,16 @@ inline auto get_max_coarse_batch_size(raft::resources const& res, // Transient allocations during coarse search (select_clusters): qc_distances + cluster_dists. auto transient_per_query = static_cast(n_lists + n_probes) * qc_elem_size; auto total_per_query = persistent_per_query + transient_per_query; - auto max_per_ws = raft::resource::get_workspace_free_bytes(res) / total_per_query; + + // Use only 40% of the available workspace memory as the fine batch allocations have not been + // allocated yet. Also limit the maximum to 1GB to avoid conflicts in case the workspace is shared + // with the large workspace. + auto available_ws_size = + std::min(raft::resource::get_workspace_free_bytes(res) * 0.4, 1 << 30); + auto max_per_ws = available_ws_size / total_per_query; return std::max( 1, - std::min(max_per_ws / 2, - std::min(params.max_internal_batch_size, n_queries))); + std::min(max_per_ws, std::min(params.max_internal_batch_size, n_queries))); } template From 0ed83f38d70364f2aedb91a11730326f75a8361c Mon Sep 17 00:00:00 2001 From: Malte Foerster Date: Wed, 27 May 2026 15:34:18 +0000 Subject: [PATCH 4/5] coderabbit suggestion --- cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh index cc3266c0eb..1dd3d7663d 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_search.cuh @@ -784,8 +784,9 @@ inline auto get_max_fine_batch_size(raft::resources const& res, // Use only 40% of the available workspace memory as the coarse batch allocations have not been // allocated yet. Also limit the maximum to 1GB to avoid conflicts in case the workspace is shared // with the large workspace. - auto max_ws_size = - std::min(raft::resource::get_workspace_free_bytes(res) * 0.4, 1 << 30); + const auto free_ws_size = raft::resource::get_workspace_free_bytes(res); + const auto max_ws_size = + std::min((free_ws_size * size_t{40}) / size_t{100}, size_t{1} << 30); if (ws_size(max_batch_size) > max_ws_size) { uint32_t smaller_batch_size = raft::bound_by_power_of_two(max_batch_size); // gradually reduce the batch size until we fit into the max size limit. @@ -846,8 +847,9 @@ inline auto get_max_coarse_batch_size(raft::resources const& res, // Use only 40% of the available workspace memory as the fine batch allocations have not been // allocated yet. Also limit the maximum to 1GB to avoid conflicts in case the workspace is shared // with the large workspace. - auto available_ws_size = - std::min(raft::resource::get_workspace_free_bytes(res) * 0.4, 1 << 30); + const auto free_ws_size = raft::resource::get_workspace_free_bytes(res); + const auto available_ws_size = + std::min((free_ws_size * size_t{40}) / size_t{100}, size_t{1} << 30); auto max_per_ws = available_ws_size / total_per_query; return std::max( 1, From 05edd7bc4bcdf9619f832d8cacde10e01fcf1a0d Mon Sep 17 00:00:00 2001 From: Malte Foerster Date: Wed, 27 May 2026 15:51:47 +0000 Subject: [PATCH 5/5] coderabbit suggestion and decerased limit for kmeans --- cpp/src/cluster/detail/kmeans_balanced.cuh | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index ae033c5107..7fac255810 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -204,13 +204,17 @@ auto calc_minibatch_size(const raft::resources& handle, // If we need to convert to MathT, space required for the converted batch. if (!needs_conversion) { mem_per_row += sizeof(MathT) * dim; } - // Heuristic: calculate the minibatch size in order to use at most 80% or 1GB workspace memory. - auto available_ws_size = - std::min(raft::resource::get_workspace_free_bytes(handle) * 0.8, 1 << 30); + // Heuristic: calculate the minibatch size in order to use at most 80% or 512MB workspace memory. + // We go below 1GB here as the allocation is mostly done in a single chunk which + // is problematic if e.g. a pool allocator manages its own chunks <= 1GB. + const auto free_ws_size = raft::resource::get_workspace_free_bytes(handle); + const auto available_ws_size = + std::min((free_ws_size * size_t{8}) / size_t{10}, size_t{1} << 29); - IdxT minibatch_size = available_ws_size / mem_per_row; - minibatch_size = 64 * raft::div_rounding_up_safe(minibatch_size, IdxT{64}); - minibatch_size = std::min(minibatch_size, n_rows); + IdxT minibatch_size = std::max(IdxT{1}, static_cast(available_ws_size / mem_per_row)); + + minibatch_size = raft::round_down_safe(minibatch_size, IdxT{64}); + minibatch_size = std::min(minibatch_size, n_rows); return std::make_tuple(minibatch_size, mem_per_row); }