Multi partition cagra search

c/CMakeLists.txt

@@ -106,6 +106,7 @@ add_library(
   src/preprocessing/quantize/pq.cpp
   src/preprocessing/quantize/scalar.cpp
   src/distance/pairwise_distance.cpp
+  src/selection/select_k.cpp
 )
 add_library(cuvs::c_api ALIAS cuvs_c)
 set_target_properties(

c/include/cuvs/core/all.h

@@ -45,3 +45,5 @@
 #include <cuvs/preprocessing/quantize/binary.h>
 #include <cuvs/preprocessing/quantize/pq.h>
 #include <cuvs/preprocessing/quantize/scalar.h>
+
+#include <cuvs/selection/select_k.h>

c/include/cuvs/core/c_api.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -131,6 +131,23 @@ CUVS_EXPORT cuvsError_t cuvsStreamSync(cuvsResources_t res);
  */
 CUVS_EXPORT cuvsError_t cuvsDeviceIdGet(cuvsResources_t res, int* device_id);
 
+/**
+ * @brief Configure the temporary workspace on this resources object as an uncapped pool, backed
+ *        by the current device memory resource. After the initial reservation is allocated on
+ *        first use, subsequent calls to cuvsRMMAlloc / cuvsRMMFree on the same resources handle
+ *        hit the pool cache rather than calling cudaMallocAsync / cudaFreeAsync, reducing CUDA
+ *        context lock contention under concurrent query threads. The pool grows without shrinking:
+ *        freed allocations are returned to the pool rather than to the device, so the pool's
+ *        high-water mark only increases until the resources object is destroyed.
+ *
+ * @param[in] res                cuvsResources_t opaque C handle
+ * @param[in] initial_size_bytes initial pool reservation in bytes; size to cover the
+ *                               steady-state working set to avoid growth after warmup
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsResourcesSetWorkspacePool(cuvsResources_t res,
+                                                      size_t initial_size_bytes);
+
 /**
  * @brief Create an Initialized opaque C handle for C++ type `raft::device_resources_snmg`
  *        for multi-GPU operations
@@ -212,6 +229,19 @@ CUVS_EXPORT cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes
 CUVS_EXPORT cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
                                             int max_pool_size_percent,
                                             bool managed);
+/**
+ * @brief Switches the working memory resource to use stream-ordered asynchronous allocation
+ * (cudaMallocAsync / cudaFreeAsync). Unlike the pool resource, this resource returns memory to
+ * the stream immediately without blocking the CPU, eliminating device-wide synchronization on
+ * deallocation. This is especially beneficial when multiple CAGRA searches run concurrently on
+ * separate CUDA streams, because the internal workspace allocations no longer serialize kernel
+ * launches. Be aware that this function will change the memory resource for the whole process
+ * and the new memory resource will be used until explicitly changed.
+ *
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsRMMAsyncMemoryResourceEnable();
+
 /**
  * @brief Resets the memory resource to use the default memory resource (cuda_memory_resource)
  * @return cuvsError_t

c/include/cuvs/neighbors/cagra.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -714,6 +714,48 @@ CUVS_EXPORT cuvsError_t cuvsCagraSearch(cuvsResources_t res,
                             DLManagedTensor* distances,
                             cuvsFilter filter);
 
+/**
+ * @brief Search multiple CAGRA index partitions concurrently and return the global top-k per
+ * query.
+ *
+ * For each query row, the function searches all partitions in parallel into an internal
+ * intermediate buffer, applies per-partition distance post-processing, runs a batched top-k
+ * merge across partitions, and writes the final outputs to the caller-supplied device tensors.
+ * All work is submitted to the CUDA stream associated with @p res; use @c cuvsStreamSync to
+ * wait for completion.
+ *
+ * The index element type may be float32, float16, int8, or uint8. All partitions must share the
+ * same element type, and the queries must use that same type.
+ *
+ * @param[in]  res            cuvsResources_t opaque C handle
+ * @param[in]  params         search parameters (shared across partitions)
+ * @param[in]  num_partitions number of index partitions
+ * @param[in]  indices        array of num_partitions cuvsCagraIndex_t pointers, all of the same
+ *                            element type
+ * @param[in]  queries        DLManagedTensor* (device, same dtype as the indices, [n_queries,
+ *                            dim]); the queries matrix is searched against every partition
+ * @param[out] partition_ids  DLManagedTensor* (device, uint32, [n_queries, k]); which partition
+ *                            each returned neighbor came from
+ * @param[out] neighbors      DLManagedTensor* (device, uint32 or int64, [n_queries, k]); ordinal
+ *                            in the corresponding partition's dataset
+ * @param[out] distances      DLManagedTensor* (device, float32, [n_queries, k]); post-processed
+ *                            distance for each (query, neighbor)
+ * @param[in]  filter         filter to apply during search; use {.type=NO_FILTER, .addr=0} for
+ *                            unfiltered search, or {.type=BITSET, .addr=ptr} where ptr is a
+ *                            uintptr_t-cast DLManagedTensor* holding the concatenation of the
+ *                            per-partition bitsets (each word-aligned). Per-partition bit offsets
+ *                            are recomputed internally from the index sizes.
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraSearchMultiPartition(cuvsResources_t res,
+                                                      cuvsCagraSearchParams_t params,
+                                                      uint32_t num_partitions,
+                                                      cuvsCagraIndex_t* indices,
+                                                      DLManagedTensor* queries,
+                                                      DLManagedTensor* partition_ids,
+                                                      DLManagedTensor* neighbors,
+                                                      DLManagedTensor* distances,
+                                                      cuvsFilter filter);
+
 /**
  * @}
  */

c/include/cuvs/selection/select_k.h

@@ -0,0 +1,37 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include <cuvs/core/c_api.h>
+#include <dlpack/dlpack.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @brief Select the k smallest values from a flat device array of n candidates.
+ *
+ * Treats `in_val` as a matrix of shape [1, n] and selects the `k` smallest
+ * float values.  `out_idx` receives the int64 column positions of the selected
+ * values in [0, n), so the caller can recover per-segment identity as:
+ *
+ *   segment_index        = out_idx[j] / segment_k
+ *   position_in_segment  = out_idx[j] % segment_k
+ *
+ * @param[in]  res      cuvsResources_t handle
+ * @param[in]  in_val   DLManagedTensor* shape [1, n], float32, device memory
+ * @param[out] out_val  DLManagedTensor* shape [1, k], float32, device memory
+ * @param[out] out_idx  DLManagedTensor* shape [1, k], int64,   device memory
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsSelectK(cuvsResources_t res,
+                                    DLManagedTensor* in_val,
+                                    DLManagedTensor* out_val,
+                                    DLManagedTensor* out_idx);
+
+#ifdef __cplusplus
+}
+#endif

c/src/core/c_api.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -9,11 +9,13 @@
 #include <raft/core/device_resources_snmg.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resource/device_id.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/resource_types.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
 #include <rapids_logger/logger.hpp>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/cuda_async_memory_resource.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/managed_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
@@ -35,6 +37,19 @@ extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res)
   });
 }
 
+extern "C" cuvsError_t cuvsResourcesSetWorkspacePool(cuvsResources_t res, size_t initial_size_bytes)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto res_ptr = reinterpret_cast<raft::resources*>(res);
+    // Create an uncapped pool: pre-warms with initial_size_bytes to avoid cudaMalloc on every
+    // query, but can grow beyond that if an allocation exceeds the initial reservation.
+    raft::resource::set_workspace_resource(
+      *res_ptr,
+      rmm::mr::pool_memory_resource{rmm::mr::get_current_device_resource_ref(),
+                                    initial_size_bytes});
+  });
+}
+
 extern "C" cuvsError_t cuvsResourcesDestroy(cuvsResources_t res)
 {
   return cuvs::core::translate_exceptions([=] {
@@ -132,17 +147,17 @@ extern "C" cuvsError_t cuvsRMMAlloc(cuvsResources_t res, void** ptr, size_t byte
 {
   return cuvs::core::translate_exceptions([=] {
     auto res_ptr = reinterpret_cast<raft::resources*>(res);
-    auto mr      = rmm::mr::get_current_device_resource_ref();
-    *ptr         = mr.allocate(raft::resource::get_cuda_stream(*res_ptr), bytes);
+    auto stream  = raft::resource::get_cuda_stream(*res_ptr);
+    *ptr         = raft::resource::get_workspace_resource_ref(*res_ptr).allocate(stream, bytes);
   });
 }
 
 extern "C" cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes)
 {
   return cuvs::core::translate_exceptions([=] {
     auto res_ptr = reinterpret_cast<raft::resources*>(res);
-    auto mr      = rmm::mr::get_current_device_resource_ref();
-    mr.deallocate(raft::resource::get_cuda_stream(*res_ptr), ptr, bytes);
+    auto stream  = raft::resource::get_cuda_stream(*res_ptr);
+    raft::resource::get_workspace_resource_ref(*res_ptr).deallocate(stream, ptr, bytes);
   });
 }
 
@@ -164,9 +179,18 @@ extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_per
   });
 }
 
+extern "C" cuvsError_t cuvsRMMAsyncMemoryResourceEnable()
+{
+  return cuvs::core::translate_exceptions([=] {
+    rmm::mr::set_current_device_resource(rmm::mr::cuda_async_memory_resource{});
+  });
+}
+
 extern "C" cuvsError_t cuvsRMMMemoryResourceReset()
 {
-  return cuvs::core::translate_exceptions([=] { rmm::mr::reset_current_device_resource(); });
+  return cuvs::core::translate_exceptions([=] {
+    rmm::mr::reset_current_device_resource();
+  });
 }
 
 thread_local std::unique_ptr<rmm::mr::pinned_host_memory_resource> pinned_mr;

c/src/neighbors/cagra.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -263,6 +263,93 @@ void _search(cuvsResources_t res,
   }
 }
 
+template <typename T, typename OutIdxT>
+void _search_multi_partition(cuvsResources_t res,
+                             cuvsCagraSearchParams params,
+                             uint32_t num_partitions,
+                             cuvsCagraIndex_t* indices,
+                             DLManagedTensor* queries,
+                             DLManagedTensor* partition_ids,
+                             DLManagedTensor* neighbors,
+                             DLManagedTensor* distances,
+                             cuvsFilter filter)
+{
+  using IdxT      = uint32_t;
+  using DistanceT = float;
+  using IndexT    = cuvs::neighbors::cagra::index<T, IdxT>;
+
+  auto res_ptr       = reinterpret_cast<raft::resources*>(res);
+  auto search_params = cuvs::neighbors::cagra::search_params();
+  convert_c_search_params(params, &search_params);
+
+  std::vector<const IndexT*> idx_vec(num_partitions);
+  for (uint32_t i = 0; i < num_partitions; i++) {
+    idx_vec[i] = reinterpret_cast<const IndexT*>(indices[i]->addr);
+  }
+
+  using queries_view_t = raft::device_matrix_view<const T, int64_t, raft::row_major>;
+  using pid_view_t     = raft::device_matrix_view<uint32_t, int64_t, raft::row_major>;
+  using nbrs_view_t    = raft::device_matrix_view<OutIdxT, int64_t, raft::row_major>;
+  using dist_view_t    = raft::device_matrix_view<DistanceT, int64_t, raft::row_major>;
+
+  auto queries_view       = cuvs::core::from_dlpack<queries_view_t>(queries);
+  auto partition_ids_view = cuvs::core::from_dlpack<pid_view_t>(partition_ids);
+  auto neighbors_view     = cuvs::core::from_dlpack<nbrs_view_t>(neighbors);
+  auto distances_view     = cuvs::core::from_dlpack<dist_view_t>(distances);
+
+  if (filter.type == NO_FILTER) {
+    cuvs::neighbors::cagra::search(
+      *res_ptr, search_params, idx_vec, queries_view, partition_ids_view, neighbors_view, distances_view);
+  } else if (filter.type == BITSET) {
+    auto* bitset_tensor = reinterpret_cast<DLManagedTensor*>(filter.addr);
+    RAFT_EXPECTS(bitset_tensor != nullptr, "BITSET filter addr must be non-null");
+
+    using bitset_mdspan_t = raft::device_vector_view<std::uint32_t, int64_t, raft::row_major>;
+    auto bitset_mds       = cuvs::core::from_dlpack<bitset_mdspan_t>(bitset_tensor);
+
+    // The bitset is the concatenation of the per-partition bitsets; use its full word capacity as
+    // the logical bit length. Per-partition bit offsets and the size-consistency check against the
+    // partition sizes are handled in search_multi_partition, which already loops the partitions.
+    cuvs::core::bitset_view<std::uint32_t, int64_t> combined_bitset_view(
+      bitset_mds, static_cast<int64_t>(bitset_mds.size()) * 32);
+    cuvs::neighbors::filtering::bitset_filter<std::uint32_t, int64_t> bitset_filter_obj(
+      combined_bitset_view);
+
+    cuvs::neighbors::cagra::search(*res_ptr,
+                                   search_params,
+                                   idx_vec,
+                                   queries_view,
+                                   partition_ids_view,
+                                   neighbors_view,
+                                   distances_view,
+                                   bitset_filter_obj);
+  } else {
+    RAFT_FAIL("Unsupported filter type for multi-partition search: %d", (int)filter.type);
+  }
+}
+
+template <typename T>
+void _search_multi_partition(cuvsResources_t res,
+                             cuvsCagraSearchParams params,
+                             uint32_t num_partitions,
+                             cuvsCagraIndex_t* indices,
+                             DLManagedTensor* queries,
+                             DLManagedTensor* partition_ids,
+                             DLManagedTensor* neighbors,
+                             DLManagedTensor* distances,
+                             cuvsFilter filter)
+{
+  if (neighbors->dl_tensor.dtype.code == kDLUInt && neighbors->dl_tensor.dtype.bits == 32) {
+    _search_multi_partition<T, uint32_t>(
+      res, params, num_partitions, indices, queries, partition_ids, neighbors, distances, filter);
+  } else if (neighbors->dl_tensor.dtype.code == kDLInt && neighbors->dl_tensor.dtype.bits == 64) {
+    _search_multi_partition<T, int64_t>(
+      res, params, num_partitions, indices, queries, partition_ids, neighbors, distances, filter);
+  } else {
+    RAFT_FAIL("neighbors should be of type uint32_t or int64_t");
+  }
+}
+
 template <typename T>
 void _serialize(cuvsResources_t res,
                 const char* filename,
@@ -690,6 +777,54 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res,
   });
 }
 
+extern "C" cuvsError_t cuvsCagraSearchMultiPartition(cuvsResources_t res,
+                                                     cuvsCagraSearchParams_t params,
+                                                     uint32_t num_partitions,
+                                                     cuvsCagraIndex_t* indices,
+                                                     DLManagedTensor* queries,
+                                                     DLManagedTensor* partition_ids,
+                                                     DLManagedTensor* neighbors,
+                                                     DLManagedTensor* distances,
+                                                     cuvsFilter filter)
+{
+  return cuvs::core::translate_exceptions([=] {
+    RAFT_EXPECTS(num_partitions > 0, "num_partitions must be > 0");
+    RAFT_EXPECTS(indices != nullptr && queries != nullptr && partition_ids != nullptr &&
+                   neighbors != nullptr && distances != nullptr,
+                 "All pointer arguments must be non-null");
+
+    // Every partition index must be present, built, and share one dtype; the search dispatches on
+    // that common dtype. The queries dtype is validated against T inside from_dlpack.
+    RAFT_EXPECTS(indices[0] != nullptr, "Index at position 0 is null");
+    auto index_dtype = indices[0]->dtype;
+    for (uint32_t i = 0; i < num_partitions; i++) {
+      RAFT_EXPECTS(indices[i] != nullptr && indices[i]->addr != 0,
+                   "Index at position %u is null or not built", i);
+      RAFT_EXPECTS(
+        indices[i]->dtype.code == index_dtype.code && indices[i]->dtype.bits == index_dtype.bits,
+        "All partition indices must share the same dtype; position %u differs from position 0", i);
+    }
+
+    if (index_dtype.code == kDLFloat && index_dtype.bits == 32) {
+      _search_multi_partition<float>(
+        res, *params, num_partitions, indices, queries, partition_ids, neighbors, distances, filter);
+    } else if (index_dtype.code == kDLFloat && index_dtype.bits == 16) {
+      _search_multi_partition<half>(
+        res, *params, num_partitions, indices, queries, partition_ids, neighbors, distances, filter);
+    } else if (index_dtype.code == kDLInt && index_dtype.bits == 8) {
+      _search_multi_partition<int8_t>(
+        res, *params, num_partitions, indices, queries, partition_ids, neighbors, distances, filter);
+    } else if (index_dtype.code == kDLUInt && index_dtype.bits == 8) {
+      _search_multi_partition<uint8_t>(
+        res, *params, num_partitions, indices, queries, partition_ids, neighbors, distances, filter);
+    } else {
+      RAFT_FAIL("Unsupported multi-partition index DLtensor dtype: %d and bits: %d",
+                index_dtype.code,
+                index_dtype.bits);
+    }
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res,
                                       cuvsCagraIndexParams_t params,
                                       cuvsCagraIndex_t* indices,