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

@@ -158,6 +158,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
@@ -239,6 +256,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

@@ -735,6 +735,49 @@ 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]  filters        array of `num_partitions` filters, one per partition (or NULL for a
+ *                            fully unfiltered search). `filters[i]` applies to partition `i`: use
+ *                            {.type=NO_FILTER, .addr=0} for no filter on that partition, or
+ *                            {.type=BITSET, .addr=ptr} where ptr is a uintptr_t-cast
+ *                            DLManagedTensor* holding that partition's own bitset (one bit per
+ *                            vector in that partition; standard 32-bit packing).
+ */
+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* filters);
+
 /**
  * @}
  */

c/include/cuvs/neighbors/common.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
  */
 
@@ -34,6 +34,12 @@ enum cuvsFilterType {
 /**
  * @brief Struct to hold address of cuvs::neighbors::prefilter and its type
  *
+ * `addr` points to a filter object owned by the caller; the library performs no caching of the
+ * underlying bitset across search calls. Allocating and populating the device bitset may be more
+ * expensive than a single filtered search, so callers that issue repeated searches against the same
+ * filter (e.g. many queries over one index) should build the bitset once and reuse the same
+ * cuvsFilter across those calls rather than rebuild it per search. Reusing the bitset is essential
+ * for realizing the full throughput of filtered search.
  */
 typedef struct {
   uintptr_t addr;

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

@@ -10,11 +10,13 @@
 #include <raft/core/memory_tracking_resources.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>
@@ -39,6 +41,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 cuvsResourcesCreateWithMemoryTracking(cuvsResources_t* res,
                                                              const char* csv_path,
                                                              int64_t sample_interval_ms)
@@ -153,17 +168,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);
   });
 }
 
@@ -185,9 +200,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

@@ -263,6 +263,96 @@ 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* filters)
+{
+  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);
+
+  // One bitset view per partition; an empty view (null ptr) means "no filter for that partition".
+  using bitset_view_t   = cuvs::core::bitset_view<std::uint32_t, int64_t>;
+  using bitset_mdspan_t = raft::device_vector_view<std::uint32_t, int64_t, raft::row_major>;
+  std::vector<bitset_view_t> partition_bitsets(
+    num_partitions, bitset_view_t(static_cast<std::uint32_t*>(nullptr), static_cast<int64_t>(0)));
+  if (filters != nullptr) {
+    for (uint32_t i = 0; i < num_partitions; i++) {
+      if (filters[i].type == NO_FILTER) {
+        continue;  // leave the empty (accept-all) view for this partition
+      } else if (filters[i].type == BITSET) {
+        auto* bitset_tensor = reinterpret_cast<DLManagedTensor*>(filters[i].addr);
+        RAFT_EXPECTS(
+          bitset_tensor != nullptr, "BITSET filter addr must be non-null (partition %u)", i);
+        auto bitset_mds = cuvs::core::from_dlpack<bitset_mdspan_t>(bitset_tensor);
+        partition_bitsets[i] =
+          bitset_view_t(bitset_mds, static_cast<int64_t>(bitset_mds.size()) * 32);
+      } else {
+        RAFT_FAIL("Unsupported filter type for multi-partition search (partition %u): %d",
+                  i,
+                  (int)filters[i].type);
+      }
+    }
+  }
+
+  cuvs::neighbors::cagra::search(*res_ptr,
+                                 search_params,
+                                 idx_vec,
+                                 queries_view,
+                                 partition_ids_view,
+                                 neighbors_view,
+                                 distances_view,
+                                 partition_bitsets);
+}
+
+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* filters)
+{
+  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, filters);
+  } 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, filters);
+  } 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 +780,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* filters)
+{
+  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, filters);
+    } else if (index_dtype.code == kDLFloat && index_dtype.bits == 16) {
+      _search_multi_partition<half>(
+        res, *params, num_partitions, indices, queries, partition_ids, neighbors, distances, filters);
+    } 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, filters);
+    } 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, filters);
+    } 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,