Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5521)

Summary:
Pull Request resolved: https://github.com/facebookresearch/faiss/pull/5521

Previous diff: D114685755. Now aarch64 passes.

`GpuIndexCagra` had two multi-GPU build entry points totalling 16 positional
arguments, neither of which fits the `Index` API. This removes one, folds the
other into `train()`, and speeds up the build.

**Usage.** Listing more than one device in the config selects the multi-GPU
build; `train()` routes to it:

```python
devices = faiss.Int32Vector()
for i in range(8):
    devices.push_back(i)

an = faiss.AllNeighborsCagraConfig()
an.n_clusters = 16                  # 0 = auto: max(2 * n_devices, 4)
an.overlap_factor = 2               # must be >= 2
an.ivf_pq_search_batch_size = 8192  # 0 = cuVS default; caps IVF-PQ memory

config = faiss.GpuIndexCagraConfig()
config.graph_degree = 32
config.intermediate_graph_degree = 32
config.build_algo = faiss.graph_build_algo_IVF_PQ
config.devices = devices            # >1 device selects the multi-GPU build
config.all_neighbors_params = an

index = faiss.GpuIndexCagra(res, d, faiss.METRIC_L2, config)
index.train(xb)          # xb must stay alive until copyTo() completes

cpu_index = faiss.IndexHNSWCagra()
cpu_index.base_level_only = True
index.copyTo(cpu_index)  # required: the GPU index is not searchable on this path
```

Float32-only, does not copy `x`, and leaves `index_` empty so `copyTo()` is the
only valid follow-up. Single-GPU behaviour is unchanged when `devices` is empty.

**Deleted `trainMultiGpu`.** Sharded SNMG CAGRA produces per-shard graphs with
zero cross-shard edges by construction, so it needs post-hoc stitching just to
be usable, and it lost to the `all_neighbors` path on both build time and
recall. It had no callers outside the benchmark and one test.

**Folded `trainAllNeighbors` into `train()`.** Its 6 trailing bare scalars now
live in the constructor-time config struct, which is the established Faiss GPU
convention:

| Old argument | Now |
| --- | --- |
| `devices` | `GpuIndexCagraConfig::devices`, also the dispatch predicate |
| `build_algo` (0/1/2) | `GpuIndexCagraConfig::build_algo` |
| `refinement_rate` | `AllNeighborsCagraConfig::refinement_rate` |
| `n_clusters`, `overlap_factor`, `ivfpq_search_batch` | new `AllNeighborsCagraConfig` |

This also kills a live footgun: the old `int build_algo` used an encoding
(0=NN-descent, 1=brute-force, 2=IVF-PQ) that disagreed with the
`graph_build_algo` enum in the same header (0=IVF_PQ, 1=NN_DESCENT). Both
callers set `config.build_algo` and then passed an unrelated int, and the config
field was silently dead. `BRUTE_FORCE` is appended to `graph_build_algo` (at the
end, so existing values do not renumber) and the config field is now the single
source of truth.

**Graph pruning no longer copies the graph off the device.** Previously only
detour counting ran on the GPU and pruning plus reverse-graph construction ran
on the host, which was 35% of total build time.

This part is **temporary and self-removing**. cuVS has no public API to prune a
graph that is already on the device: its exported `helpers::optimize` takes host
matrices, and the dispatch behind it erases the mdspan accessor to host memory,
so cuVS's own device code path is unreachable from outside. Until that is fixed
upstream, this reaches into cuVS's internal headers when the build has them
available, and otherwise falls back to the public host API -- correct either
way, just slower on the fallback. **Open-source builds get the fallback**, since
cuVS does not install those headers; the `train()` docs say so.

An upstream patch adding a `device_matrix_view` overload to
`cuvs::neighbors::cagra::helpers::optimize` is prepared and will be submitted to
rapidsai/cuvs. When it ships, the internal include, the build flag guarding it,
and the fallback branch all get deleted and every build gets the fast path.

At 50M vectors on 8 GPUs the optimize step goes from 119.1s to 1.65s (72x) and
end-to-end build->serialize from 8.0 to 5.4 minutes, recall unchanged.

**Collapsed the benchmark to one path.** With the stitching approaches gone,
`bench_approaches.py` is a single-path tool for validating and tuning the
production build, and reports a per-phase build breakdown plus an efSearch sweep
of recall and QPS.

Deliberately *not* merged into the config:

- `ivf_pq_params` / `ivf_pq_search_params` are still not consulted on this path.
  cuVS derives them from the dataset shape and those derived values beat the
  static defaults in the faiss structs, so only the knobs cuVS cannot infer are
  overridden.
- `faiss::cagra_build_algo` only has `{IVF_PQ, NN_DESCENT}`, so a `BRUTE_FORCE`
  config would silently degrade to NN-descent on the single-GPU path;
  `train_ex()` now rejects it there. `GpuIndexBinaryCagra` shares this config
  struct and has no multi-GPU build, so it rejects `devices.size() > 1` rather
  than silently building on one device.

Behaviour changes: the default `build_algo` on the multi-GPU path is now
`IVF_PQ` (the config default) rather than NN-descent (the old argument default),
and `IndexHNSW::num_base_level_search_entrypoints` goes 32 -> 256, which
measured better on both recall and QPS at every efSearch.

Reviewed By: junjieqi

Differential Revision: D116187873

fbshipit-source-id: 265be7dc0670de262cb1a937ca4b2f530578e237

CMakeLists.txt

@@ -125,6 +125,10 @@ if(FAISS_ENABLE_CUVS AND NOT TARGET cuvs::cuvs)
    find_package(cuvs)
  endif()
 
+if(FAISS_ENABLE_CUVS AND NOT TARGET rmm::rmm)
+  find_package(rmm REQUIRED)
+endif()
+
 add_subdirectory(faiss)
 
 if(FAISS_ENABLE_GPU)

faiss/IndexHNSW.cpp

@@ -17,7 +17,6 @@
 #include <limits>
 #include <memory>
 #include <queue>
-#include <random>
 
 #include <cstdint>
 #include "faiss/Index.h"
@@ -1433,13 +1432,11 @@ void IndexHNSWCagra::search(
                 // first real candidate will always be strictly better.
                 nearest_d[i] = C::neutral();
 
-                std::random_device rd;
-                std::mt19937 gen(rd());
-                std::uniform_int_distribution<idx_t> distrib(
-                        0, this->ntotal - 1);
+                // Seeded per query so entrypoints are reproducible.
+                SplitMix64RandomGenerator gen(i);
 
                 for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
-                    auto idx = distrib(gen);
+                    idx_t idx = gen.rand_int64() % this->ntotal;
                     auto distance = (*dis)(idx);
                     if (C::cmp(nearest_d[i], distance)) {
                         nearest[i] = static_cast<storage_idx_t>(idx);
@@ -1498,12 +1495,11 @@ void IndexHNSWCagra::range_search(
             // real candidate will always be strictly better.
             float nearest_d = C::neutral();
 
-            std::random_device rd;
-            std::mt19937 gen(rd());
-            std::uniform_int_distribution<idx_t> distrib(0, ntotal - 1);
+            // For reproducible entrypoint.
+            SplitMix64RandomGenerator gen(i);
 
             for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
-                auto idx = distrib(gen);
+                idx_t idx = gen.rand_int64() % ntotal;
                 auto distance = (*dis)(idx);
                 // C::cmp(nearest_d, distance) is true iff distance is
                 // strictly better than the current nearest_d.

faiss/IndexHNSW.h

@@ -253,7 +253,7 @@ struct IndexHNSWCagra : IndexHNSW {
     /// searches only the base level knn graph of the HNSW index.
     /// This parameter selects the entry point by randomly selecting
     /// some points and using the best one.
-    int num_base_level_search_entrypoints = 32;
+    int num_base_level_search_entrypoints = 256;
 
     void add(idx_t n, const float* x) override;
 

faiss/gpu/CMakeLists.txt

@@ -363,7 +363,7 @@ else()
 
 
   find_package(CUDAToolkit REQUIRED)
-  target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs> $<$<BOOL:${FAISS_ENABLE_CUVS}>:OpenMP::OpenMP_CXX>)
+  target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs> $<$<BOOL:${FAISS_ENABLE_CUVS}>:rmm::rmm> $<$<BOOL:${FAISS_ENABLE_CUVS}>:OpenMP::OpenMP_CXX>)
   target_compile_options(faiss_gpu_objs PRIVATE
     $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all
     --expt-extended-lambda --expt-relaxed-constexpr

faiss/gpu/GpuIndexBinaryCagra.cu

@@ -78,6 +78,13 @@ std::shared_ptr<GpuResources> GpuIndexBinaryCagra::getResources() {
 }
 
 void GpuIndexBinaryCagra::train(idx_t n, const uint8_t* x) {
+    // The config is shared with the float index; there is no binary
+    // multi-GPU build, so reject rather than silently ignoring the request.
+    FAISS_THROW_IF_MSG(
+            cagraConfig_.devices.size() > 1,
+            "binary CAGRA has no multi-GPU build; "
+            "GpuIndexCagraConfig::devices must name at most one device");
+
     DeviceScope scope(cagraConfig_.device);
     if (this->is_trained) {
         FAISS_ASSERT(index_);

faiss/gpu/GpuIndexCagra.h

@@ -47,7 +47,9 @@ enum class graph_build_algo {
     /// Use NN-Descent to build all-neighbors knn graph
     NN_DESCENT,
     /// Use iterative search to build knn graph
-    ITERATIVE_SEARCH
+    ITERATIVE_SEARCH,
+    /// Exact knn graph via tiled brute force
+    BRUTE_FORCE
 };
 
 /// A type for specifying how PQ codebooks are created.
@@ -173,6 +175,34 @@ struct IVFPQSearchCagraConfig {
     uint32_t max_internal_batch_size = 4096;
 };
 
+/// Knobs for the multi-GPU build path, selected by listing more than one
+/// device in GpuIndexCagraConfig::devices. Degrees, build_algo and metric
+/// still come from GpuIndexCagraConfig.
+struct AllNeighborsCagraConfig {
+    /// Number of overlapping clusters the dataset is partitioned into.
+    /// 0 selects max(2 * number of devices, 4).
+    size_t n_clusters = 0;
+
+    /// Clusters each vector is assigned to. Must be >= 2: with 1 there are no
+    /// cross-cluster edges and recall collapses. 0 selects 2.
+    size_t overlap_factor = 2;
+
+    /// Bounds IVF-PQ search memory during the knn build. cuVS's default can
+    /// exhaust device memory around 100M vectors; 8192 avoids that and does
+    /// not change recall. 0 keeps cuVS's default.
+    uint32_t ivf_pq_search_batch_size = 0;
+
+    /// IVF-PQ refinement multiplier. The refine pass runs on the CPU, so cost
+    /// scales with this. Raising it above 1.0 measured both slower and less
+    /// accurate, so it is separate from the shared refine_rate.
+    float refinement_rate = 1.0f;
+
+    /// Size the IVF-PQ index for one cluster rather than the whole dataset.
+    /// cuVS reuses these params for every cluster without rescaling, so
+    /// sizing from the full dataset badly over-partitions each one.
+    bool ivf_pq_size_from_cluster = true;
+};
+
 struct GpuIndexCagraConfig : public GpuIndexConfig {
     /// Degree of input graph for pruning.
     size_t intermediate_graph_degree = 128;
@@ -190,6 +220,27 @@ struct GpuIndexCagraConfig : public GpuIndexConfig {
 
     /// Whether to use MST optimization to guarantee graph connectivity.
     bool guarantee_connectivity = false;
+
+    /// Devices to build on. More than one selects the multi-GPU build in
+    /// train(); see AllNeighborsCagraConfig and train() for its restrictions.
+    /// Empty or one device uses GpuIndexConfig::device as usual.
+    std::vector<int> devices;
+
+    AllNeighborsCagraConfig all_neighbors_params;
+
+    /// Build the HNSW upper levels on the GPU during copyTo() instead of by
+    /// CPU insertion. Roughly 10x faster, but measured worse recall than
+    /// base_level_only, which is cheaper still. Off by default.
+    bool gpu_hnsw_upper_levels = false;
+
+    /// intermediate_graph_degree for the per-level subgraph builds.
+    /// 0 = twice the upper-level degree.
+    size_t gpu_hnsw_intermediate_degree = 0;
+
+    /// MST connectivity pass on the per-level subgraphs. On by default:
+    /// upper levels are walked greedily with no backtracking, so a
+    /// disconnected component is a trap the descent cannot escape.
+    bool gpu_hnsw_guarantee_connectivity = true;
 };
 
 enum class search_algo {
@@ -284,46 +335,6 @@ struct GpuIndexCagra : public GpuIndex {
     /// in the index instance
     void copyTo(faiss::IndexHNSWCagra* index) const;
 
-    /// Train CAGRA using multiple GPUs by sharding the dataset.
-    /// Uses cuVS native SNMG (single-node multi-GPU) CAGRA build.
-    /// Each device builds one shard in parallel via OpenMP.
-    /// Float32 only. After training, call copyTo() to produce a CPU
-    /// IndexHNSWCagra with full HNSW upper levels.
-    /// The training data pointer must remain valid until copyTo() completes.
-    /// stitch_mode: 0=CPU HNSW (Approach C), 1=GPU brute-force (Approach B)
-    void trainMultiGpu(
-            idx_t n,
-            const float* x,
-            std::vector<GpuResourcesProvider*>& providers,
-            std::vector<int>& devices,
-            idx_t stitch_per_shard = 0,
-            int stitch_k = 2,
-            int stitch_mode = 0);
-
-    /// Build a unified CAGRA graph using cuVS all_neighbors
-    /// (multi-GPU kNN graph construction with overlapping clusters) followed
-    /// by cagra::optimize (graph pruning). Produces a single unified graph
-    /// without stitching. The training data pointer must remain valid until
-    /// copyTo() completes.
-    /// build_algo: 0=NN-descent (default), 1=brute-force, 2=IVF-PQ
-    /// refinement_rate: IVF-PQ refinement multiplier (only used when
-    /// build_algo==2). Higher values trade build time for recall; the cuVS
-    /// default is 2.0.
-    /// ivfpq_search_batch: cap the IVF-PQ search `max_internal_batch_size` used
-    /// during the all_neighbors kNN build (build_algo==2). 0 = cuVS default
-    /// (128*1024), which can OOM at 100M; a smaller value (e.g. 8192) bounds
-    /// the GPU search workspace with no effect on results (recall-neutral).
-    void trainAllNeighbors(
-            idx_t n,
-            const float* x,
-            std::vector<int>& devices,
-            int n_clusters = 0,
-            int overlap_factor = 0,
-            bool multi_gpu_optimize = false,
-            int build_algo = 0,
-            float refinement_rate = 2.0f,
-            int ivfpq_search_batch = 0);
-
     void reset() override;
 
     std::vector<idx_t> get_knngraph() const;
@@ -357,6 +368,18 @@ struct GpuIndexCagra : public GpuIndex {
             idx_t* labels,
             const SearchParameters* search_params) const override;
 
+    /// Multi-GPU build path taken by train() when cagraConfig_.devices lists
+    /// more than one device: cuVS `all_neighbors` knn graph construction over
+    /// overlapping clusters, followed by graph pruning. Leaves the result in
+    /// merged_knngraph_ for copyTo(); index_ stays empty.
+    void trainAllNeighbors_(idx_t n, const float* x);
+
+    /// Populate HNSW levels >= 1 of `index` by building a CAGRA graph over
+    /// each level's node subset on the GPU. Requires the level table to be
+    /// prepared and the storage populated. Returns the max level.
+    void buildHnswUpperLevelsGpu_(faiss::IndexHNSWCagra* index, int max_lvl)
+            const;
+
     void copyToMultiGpu_(faiss::IndexHNSWCagra* index) const;
 
     /// Our configuration options
@@ -372,7 +395,7 @@ struct GpuIndexCagra : public GpuIndex {
             std::shared_ptr<CuvsCagra<int8_t>>>
             index_;
 
-    /// Multi-GPU state: populated by trainMultiGpu(), used by copyTo()
+    /// Multi-GPU state: populated by trainAllNeighbors_(), used by copyTo()
     std::vector<idx_t> merged_knngraph_;
     idx_t merged_knngraph_degree_ = 0;
     const float* multi_gpu_dataset_ = nullptr;

faiss/gpu/test/bench_approaches.py

@@ -5,15 +5,16 @@
 # LICENSE file in the root directory of this source tree.
 
 """
-Benchmark 4 approaches for multi-GPU CAGRA -> HNSW:
-  A. IndexShards (independent HNSW per shard, no stitching)
-  B. Unified + GPU brute-force sampled stitching (exact cross-shard NNs)
-  C. Unified + CPU HNSW stitching (approximate cross-shard NNs)
-  D. all_neighbors + optimize (multi-GPU overlapping clusters, no stitching)
+Benchmark the multi-GPU CAGRA -> HNSW build: cuVS all_neighbors builds a knn
+graph over overlapping clusters across all GPUs, cagra::optimize prunes it into
+a single unified graph, and copyTo produces a CPU IndexHNSWCagra.
+
+Measures build/copyTo/serialize wall-clock, index size, and recall@10 vs
+brute-force ground truth at several efSearch values.
 
 Usage:
   buck run @//mode/opt fbcode//faiss/gpu/test:bench_approaches -- \\
-      --data /path/to/vectors.npy --approaches D --multi-gpu-optimize
+      --data /path/to/vectors.npy
 """
 
 import argparse
@@ -42,6 +43,9 @@ sys.stdout = sys.stderr
 _t0 = time.time()
 
 
+EF_VALUES = [16, 32, 64, 128, 256, 512]
+
+
 def compute_recall(I_test, I_gt, k):
     nq = I_test.shape[0]
     return np.mean([len(set(I_test[i]) & set(I_gt[i])) / k for i in range(nq)])
@@ -98,164 +102,18 @@ def load_from_hive(
     return xb[:i]
 
 
-def approach_a_indexshards(xb, d, num_gpus, graph_degree=32, save_dir=None):
-    """Build N independent CAGRA->HNSW, wrap in IndexShards."""
-    n = xb.shape[0]
-    shard_size = (n + num_gpus - 1) // num_gpus
-    timings = {}
-
-    t0 = time.time()
-    shards = []
-    for g in range(num_gpus):
-        start = g * shard_size
-        end = min(start + shard_size, n)
-        shard_data = xb[start:end]
-        if len(shard_data) == 0:
-            continue
-
-        res = faiss.StandardGpuResources()
-        config = faiss.GpuIndexCagraConfig()
-        config.graph_degree = graph_degree
-        config.intermediate_graph_degree = graph_degree * 2
-        config.build_algo = faiss.graph_build_algo_NN_DESCENT
-        idx = faiss.GpuIndexCagra(res, d, faiss.METRIC_L2, config)
-        idx.train(shard_data)
-
-        cpu_idx = faiss.IndexHNSWCagra()
-        idx.copyTo(cpu_idx)
-        shards.append(cpu_idx)
-    timings["build"] = time.time() - t0
-
-    t0 = time.time()
-    index = faiss.IndexShards(d, True, True)
-    for s in shards:
-        index.add_shard(s)
-    timings["assemble"] = time.time() - t0
-
-    out_dir = save_dir or tempfile.mkdtemp()
-    t0 = time.time()
-    total_bytes = 0
-    for i, s in enumerate(shards):
-        path = os.path.join(out_dir, f"A_shard_{i}_{n // 1_000_000}M.faiss")
-        faiss.write_index(s, path)
-        total_bytes += os.path.getsize(path)
-    timings["serialize"] = time.time() - t0
-    timings["file_size_gb"] = total_bytes / 1e9
-
-    return index, shards, timings
+BUILD_ALGOS = ("ivf_pq", "nn_descent", "brute_force")
 
 
-def approach_b_unified_gpu_stitch(
-    xb,
-    d,
-    num_gpus,
-    graph_degree=32,
-    stitch_per_shard=100000,
-    stitch_k=16,
-    save_dir=None,
-):
-    """Unified SNMG build + GPU brute-force sampled stitching (exact NNs)."""
-    n = xb.shape[0]
-    timings = {}
-
-    resources = faiss.GpuResourcesVector()
-    devices = faiss.Int32Vector()
-    res_list = []
-    for i in range(num_gpus):
-        res = faiss.StandardGpuResources()
-        res_list.append(res)
-        resources.push_back(res)
-        devices.push_back(i)
-
-    config = faiss.GpuIndexCagraConfig()
-    config.graph_degree = graph_degree
-    config.intermediate_graph_degree = graph_degree * 2
-    config.build_algo = faiss.graph_build_algo_NN_DESCENT
-    index = faiss.GpuIndexCagra(res_list[0], d, faiss.METRIC_L2, config)
-
-    t0 = time.time()
-    index.trainMultiGpu(
-        n,
-        faiss.swig_ptr(xb),
-        resources,
-        devices,
-        stitch_per_shard,
-        stitch_k,
-        1,
-    )
-    timings["build"] = time.time() - t0
-
-    t0 = time.time()
-    cpu_index = faiss.IndexHNSWCagra()
-    index.copyTo(cpu_index)
-    timings["copyTo"] = time.time() - t0
-
-    out_dir = save_dir or tempfile.mkdtemp()
-    path = os.path.join(out_dir, f"B_unified_{n // 1_000_000}M.faiss")
-    t0 = time.time()
-    faiss.write_index(cpu_index, path)
-    timings["file_size_gb"] = os.path.getsize(path) / 1e9
-    timings["serialize"] = time.time() - t0
-
-    return cpu_index, timings
+def cagra_build_algo(name):
+    return getattr(faiss, f"graph_build_algo_{name.upper()}")
 
 
-def approach_c_unified_cpu_stitch(
-    xb,
-    d,
-    num_gpus,
-    graph_degree=32,
-    stitch_per_shard=0,
-    stitch_k=16,
-    save_dir=None,
-):
-    """Unified SNMG build + CPU HNSW stitching (approximate NNs)."""
-    n = xb.shape[0]
-    timings = {}
-
-    resources = faiss.GpuResourcesVector()
-    devices = faiss.Int32Vector()
-    res_list = []
-    for i in range(num_gpus):
-        res = faiss.StandardGpuResources()
-        res_list.append(res)
-        resources.push_back(res)
-        devices.push_back(i)
-
-    config = faiss.GpuIndexCagraConfig()
-    config.graph_degree = graph_degree
-    config.intermediate_graph_degree = graph_degree * 2
-    config.build_algo = faiss.graph_build_algo_NN_DESCENT
-    index = faiss.GpuIndexCagra(res_list[0], d, faiss.METRIC_L2, config)
-
-    t0 = time.time()
-    index.trainMultiGpu(
-        n,
-        faiss.swig_ptr(xb),
-        resources,
-        devices,
-        stitch_per_shard,
-        stitch_k,
-        0,
-    )
-    timings["build"] = time.time() - t0
-
-    t0 = time.time()
-    cpu_index = faiss.IndexHNSWCagra()
-    index.copyTo(cpu_index)
-    timings["copyTo"] = time.time() - t0
-
-    out_dir = save_dir or tempfile.mkdtemp()
-    path = os.path.join(out_dir, f"C_unified_{n // 1_000_000}M.faiss")
-    t0 = time.time()
-    faiss.write_index(cpu_index, path)
-    timings["file_size_gb"] = os.path.getsize(path) / 1e9
-    timings["serialize"] = time.time() - t0
-
-    return cpu_index, timings
+def index_filename(n):
+    return f"cagra_hnsw_{n // 1_000_000}M.faiss"
 
 
-def approach_d_all_neighbors(
+def build_cagra_hnsw(
     xb,
     d,
     num_gpus,
@@ -263,52 +121,68 @@ def approach_d_all_neighbors(
     save_dir=None,
     n_clusters=0,
     overlap_factor=0,
-    multi_gpu_optimize=False,
-    build_algo=0,
+    build_algo="ivf_pq",
     base_level_only=False,
     intermediate_graph_degree=48,
-    refinement_rate=2.0,
+    refinement_rate=1.0,
     ivfpq_search_batch=0,
+    guarantee_connectivity=False,
+    ivfpq_size_from_cluster=True,
+    gpu_hnsw_upper_levels=False,
+    gpu_hnsw_igd=0,
+    gpu_hnsw_guarantee_connectivity=True,
+    ef_construction=0,
+    entrypoints=0,
 ):
-    """all_neighbors + cagra::optimize → unified CAGRA graph. No stitching."""
+    """Multi-GPU all_neighbors build -> optimize -> CPU IndexHNSWCagra."""
+    # The multi-GPU build does not copy the dataset, so xb must stay alive and
+    # contiguous until copyTo() below has run.
+    xb = np.ascontiguousarray(xb, dtype=np.float32)
     n = xb.shape[0]
     timings = {}
 
+    res_list = [faiss.StandardGpuResources() for _ in range(num_gpus)]
+
     devices = faiss.Int32Vector()
-    res_list = []
     for i in range(num_gpus):
-        res = faiss.StandardGpuResources()
-        res_list.append(res)
         devices.push_back(i)
 
+    all_neighbors = faiss.AllNeighborsCagraConfig()
+    all_neighbors.n_clusters = n_clusters
+    all_neighbors.overlap_factor = overlap_factor
+    all_neighbors.ivf_pq_search_batch_size = ivfpq_search_batch
+    all_neighbors.refinement_rate = refinement_rate
+    all_neighbors.ivf_pq_size_from_cluster = ivfpq_size_from_cluster
+
     config = faiss.GpuIndexCagraConfig()
     config.graph_degree = graph_degree
     config.intermediate_graph_degree = intermediate_graph_degree
-    config.build_algo = faiss.graph_build_algo_NN_DESCENT
+    config.build_algo = cagra_build_algo(build_algo)
+    config.guarantee_connectivity = guarantee_connectivity
+    config.gpu_hnsw_upper_levels = gpu_hnsw_upper_levels
+    config.gpu_hnsw_intermediate_degree = gpu_hnsw_igd
+    config.gpu_hnsw_guarantee_connectivity = gpu_hnsw_guarantee_connectivity
+    config.devices = devices
+    config.all_neighbors_params = all_neighbors
+
     index = faiss.GpuIndexCagra(res_list[0], d, faiss.METRIC_L2, config)
 
     t0 = time.time()
-    index.trainAllNeighbors(
-        n,
-        faiss.swig_ptr(xb),
-        devices,
-        n_clusters,
-        overlap_factor,
-        multi_gpu_optimize,
-        build_algo,
-        refinement_rate,
-        ivfpq_search_batch,
-    )
+    index.train(xb)
     timings["build"] = time.time() - t0
 
     t0 = time.time()
     cpu_index = faiss.IndexHNSWCagra()
     cpu_index.base_level_only = base_level_only
+    if ef_construction > 0:
+        cpu_index.hnsw.efConstruction = ef_construction
+    if entrypoints > 0:
+        cpu_index.num_base_level_search_entrypoints = entrypoints
     index.copyTo(cpu_index)
     timings["copyTo"] = time.time() - t0
 
     out_dir = save_dir or tempfile.mkdtemp()
-    path = os.path.join(out_dir, f"D_allneighbors_{n // 1_000_000}M.faiss")
+    path = os.path.join(out_dir, index_filename(n))
     t0 = time.time()
     faiss.write_index(cpu_index, path)
     timings["file_size_gb"] = os.path.getsize(path) / 1e9
@@ -318,18 +192,12 @@ def approach_d_all_neighbors(
 
 
 def _set_ef(index, ef):
-    if isinstance(index, faiss.IndexShards):
-        for i in range(index.count()):
-            shard = faiss.downcast_index(index.at(i))
-            if hasattr(shard, "hnsw"):
-                shard.hnsw.efSearch = ef
-    elif hasattr(index, "hnsw"):
-        index.hnsw.efSearch = ef
+    index.hnsw.efSearch = ef
 
 
 def eval_recall(index, xq, Igt, k=10, ef_values=None):
     if ef_values is None:
-        ef_values = [64, 128, 256]
+        ef_values = EF_VALUES
     results = {}
 
     for ef in ef_values:
@@ -340,10 +208,30 @@ def eval_recall(index, xq, Igt, k=10, ef_values=None):
     return results
 
 
+def eval_qps(index, xq, k=10, ef_values=None, repeat=5):
+    """Full-batch throughput at each efSearch, best of `repeat` runs."""
+    if ef_values is None:
+        ef_values = EF_VALUES
+    nq = xq.shape[0]
+    results = {}
+
+    for ef in ef_values:
+        _set_ef(index, ef)
+        index.search(xq, k)  # warmup
+        best = float("inf")
+        for _ in range(repeat):
+            t0 = time.time()
+            index.search(xq, k)
+            best = min(best, time.time() - t0)
+        results[ef] = nq / best
+
+    return results
+
+
 def eval_kcycles(index, xq, k=10, ef_values=None, warmup=3, repeat=5):
     """Measure kcycles/query using vench CycleCounter."""
     if ef_values is None:
-        ef_values = [64, 128, 256]
+        ef_values = EF_VALUES
 
     try:
         from vector_search.vench.perf_cycles import CycleCounter
@@ -361,22 +249,15 @@ def eval_kcycles(index, xq, k=10, ef_values=None, warmup=3, repeat=5):
     prev_threads = faiss.omp_get_max_threads()
     faiss.omp_set_num_threads(1)
 
-    measure_index = index
-    if isinstance(index, faiss.IndexShards) and index.count() > 0:
-        measure_index = faiss.IndexShards(index.d, False, True)
-        for i in range(index.count()):
-            measure_index.add_shard(index.at(i))
-
     for ef in ef_values:
         _set_ef(index, ef)
-        _set_ef(measure_index, ef)
         for _ in range(warmup):
-            measure_index.search(xq, k)
+            index.search(xq, k)
 
         best_kcycles = float("inf")
         for _ in range(repeat):
             c0 = cc.read()
-            measure_index.search(xq, k)
+            index.search(xq, k)
             c1 = cc.read()
             kc = (c1 - c0) / 1000.0 / nq
             best_kcycles = min(best_kcycles, kc)
@@ -388,7 +269,7 @@ def eval_kcycles(index, xq, k=10, ef_values=None, warmup=3, repeat=5):
 
 def main():
     parser = argparse.ArgumentParser(
-        description="Benchmark multi-GPU CAGRA->HNSW approaches"
+        description="Benchmark the multi-GPU CAGRA->HNSW build"
     )
     parser.add_argument(
         "--data",
@@ -405,8 +286,6 @@ def main():
     )
     parser.add_argument("--num-gpus", type=int, default=0)
     parser.add_argument("--graph-degree", type=int, default=32)
-    parser.add_argument("--stitch-k", type=int, default=16)
-    parser.add_argument("--stitch-per-shard", type=int, default=100000)
     parser.add_argument(
         "--n-clusters",
         type=int,
@@ -420,15 +299,66 @@ def main():
         help="all_neighbors overlap_factor (0=default 2)",
     )
     parser.add_argument(
-        "--multi-gpu-optimize",
-        action="store_true",
-        help="Partition detour counting across GPUs",
+        "--build-algo",
+        choices=BUILD_ALGOS,
+        default="ivf_pq",
+        help="graph_build_algo used for the kNN graph (brute_force is "
+        "multi-GPU only and O(N^2 D) per cluster)",
     )
     parser.add_argument(
-        "--build-algo",
+        "--ivfpq-size-from-cluster",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="Derive IVF-PQ params from the per-cluster subproblem. "
+        "--no-ivfpq-size-from-cluster sizes them from the full dataset, "
+        "which over-partitions every cluster",
+    )
+    parser.add_argument(
+        "--guarantee-connectivity",
+        action="store_true",
+        help="Run the MST pass in cagra::optimize so the pruned graph is "
+        "guaranteed connected (cuVS default is on; costs build time)",
+    )
+    parser.add_argument(
+        "--entrypoints-sweep",
+        type=str,
+        default="",
+        help="Comma-separated num_base_level_search_entrypoints values to "
+        "sweep at search time (base-level-only only), e.g. 32,256,1024",
+    )
+    parser.add_argument(
+        "--gpu-hnsw-upper-levels",
+        action="store_true",
+        help="Build HNSW levels >=1 as GPU CAGRA subgraphs in copyTo instead "
+        "of incremental CPU insertion (base-level-only off only)",
+    )
+    parser.add_argument(
+        "--gpu-hnsw-igd",
         type=int,
         default=0,
-        help="0=nn_descent, 1=brute_force, 2=ivf_pq",
+        help="intermediate_graph_degree for the per-level GPU subgraphs "
+        "(0 = 2x the upper-level degree)",
+    )
+    parser.add_argument(
+        "--gpu-hnsw-guarantee-connectivity",
+        action=argparse.BooleanOptionalAction,
+        default=True,
+        help="MST connectivity pass on the per-level GPU subgraphs "
+        "(on by default; greedy descent needs connectivity)",
+    )
+    parser.add_argument(
+        "--ef-construction",
+        type=int,
+        default=0,
+        help="Override hnsw.efConstruction used when building upper levels "
+        "(0 = faiss default of 40)",
+    )
+    parser.add_argument(
+        "--entrypoints",
+        type=int,
+        default=0,
+        help="Override num_base_level_search_entrypoints (0 = faiss default "
+        "of 256); only affects base-level-only search",
     )
     parser.add_argument(
         "--base-level-only",
@@ -444,9 +374,10 @@ def main():
     parser.add_argument(
         "--refinement-rate",
         type=float,
-        default=2.0,
-        help="IVF-PQ refinement multiplier (build-algo=2 only); "
-        "higher trades build time for recall (cuVS default 2.0)",
+        default=1.0,
+        help="IVF-PQ refinement multiplier (build-algo=2 only). Sets "
+        "candidate_k = k * rate; the refine pass runs on the host so cost is "
+        "linear in this. Measured best at 1.0 (default)",
     )
     parser.add_argument(
         "--ivfpq-search-batch",
@@ -456,15 +387,6 @@ def main():
         "build (build-algo=2). 0=cuVS default (131072); smaller (e.g. 8192) "
         "bounds GPU search workspace to avoid OOM at 100M (recall-neutral)",
     )
-    parser.add_argument(
-        "--approaches",
-        type=str,
-        default="A,B,C,D",
-        help=(
-            "A (IndexShards), B (GPU stitch), "
-            "C (CPU stitch), D (all_neighbors)"
-        ),
-    )
     parser.add_argument(
         "--index-dir",
         type=str,
@@ -585,274 +507,101 @@ def main():
         f"graph_degree={args.graph_degree}"
     )
     print(
-        f"        stitch_k={args.stitch_k}, "
-        f"stitch_per_shard={args.stitch_per_shard}"
-    )
-    print(
-        f"        [D] n_clusters={args.n_clusters or 'auto'}, "
+        f"        n_clusters={args.n_clusters or 'auto'}, "
         f"overlap_factor={args.overlap_factor or 'auto'}, "
         f"intermediate_graph_degree={args.intermediate_graph_degree}, "
         f"build_algo={args.build_algo}, refinement_rate={args.refinement_rate}"
     )
     print()
 
-    approaches_to_run = [a.strip().upper() for a in args.approaches.split(",")]
-    results_table = []
     save_dir = args.index_dir
     os.makedirs(save_dir, exist_ok=True)
 
     if args.kcycles_only:
-        print(f"Loading persisted indices from {save_dir} (kcycles-only mode)")
-        n_tag = f"{n // 1_000_000}M"
-        for approach in approaches_to_run:
-            if approach == "A":
-                pre = "A_shard_"
-                suf = f"_{n_tag}.faiss"
-                shard_files = sorted(
-                    f
-                    for f in os.listdir(save_dir)
-                    if f.startswith(pre) and f.endswith(suf)
-                )
-                if not shard_files:
-                    print(f"  No A files for {n_tag}")
-                    continue
-                shards = [
-                    faiss.read_index(os.path.join(save_dir, f))
-                    for f in shard_files
-                ]
-                idx = faiss.IndexShards(d, True, True)
-                for s in shards:
-                    idx.add_shard(s)
-                label = f"A: IndexShards ({len(shard_files)})"
-            else:
-                prefixes = {
-                    "B": "B_unified",
-                    "C": "C_unified",
-                    "D": "D_allneighbors",
-                }
-                prefix = prefixes.get(approach)
-                if not prefix:
-                    prefix = f"{approach}_unified"
-                path = os.path.join(save_dir, f"{prefix}_{n_tag}.faiss")
-                if not os.path.exists(path):
-                    print(f"  No {prefix} for {n_tag}")
-                    continue
-                idx = faiss.read_index(path)
-                label = f"{approach}: unified"
-
-            recall = eval_recall(idx, xq, Igt, k)
-            kcycles = eval_kcycles(idx, xq, k)
-            for ef in sorted(recall.keys()):
-                kc = kcycles.get(ef, 0)
-                print(
-                    f"  {label:25s} ef={ef:>3d}  recall={recall[ef]:.4f}  "
-                    f"kcyc/q={kc:.0f}"
-                )
-            del idx
+        print(f"Loading persisted index from {save_dir} (kcycles-only mode)")
+        path = os.path.join(save_dir, index_filename(n))
+        if not os.path.exists(path):
+            print(f"ERROR: no index at {path}", file=sys.stderr)
+            sys.exit(1)
+        idx = faiss.read_index(path)
+        recall = eval_recall(idx, xq, Igt, k)
+        qps = eval_qps(idx, xq, k)
+        kcycles = eval_kcycles(idx, xq, k)
+        for ef in sorted(recall.keys()):
+            print(
+                f"  ef={ef:>4d}  recall@{k}={recall[ef]:.4f}  "
+                f"qps={qps[ef]:,.1f}  kcyc/q={kcycles.get(ef, 0):.0f}"
+            )
         return
 
-    print(f"Indices will be saved to: {save_dir}")
+    print(f"Index will be saved to: {save_dir}")
 
-    if "A" in approaches_to_run:
-        print("=" * 60)
-        print("APPROACH A: IndexShards (independent HNSW per shard)")
-        print("=" * 60)
-        index_a, shards_a, timings_a = approach_a_indexshards(
-            xb,
-            d,
-            num_gpus,
-            args.graph_degree,
-            save_dir=save_dir,
-        )
-        recall_a = eval_recall(index_a, xq, Igt, k)
-        kcycles_a = eval_kcycles(index_a, xq, k)
-        total_build_a = timings_a["build"] + timings_a.get("assemble", 0)
-        print(f"  Build:     {timings_a['build']:.1f}s")
-        print(
-            f"  Serialize: {timings_a['serialize']:.1f}s "
-            f"({timings_a['file_size_gb']:.2f} GB, {num_gpus} files)"
-        )
-        for ef in sorted(recall_a.keys()):
-            kc = kcycles_a.get(ef, 0)
-            print(
-                f"  efSearch={ef}: recall@{k}={recall_a[ef]:.4f}  "
-                f"kcycles/q={kc:.0f}"
-            )
-        results_table.append(
-            (
-                "A: IndexShards",
-                total_build_a,
-                timings_a["serialize"],
-                timings_a["file_size_gb"],
-                recall_a,
-                kcycles_a,
-            )
-        )
-        del index_a, shards_a
-        print()
-
-    if "B" in approaches_to_run:
-        print("=" * 60)
-        print(
-            f"APPROACH B: Unified + GPU brute-force stitch "
-            f"(sps={args.stitch_per_shard}, k={args.stitch_k})"
-        )
-        print("=" * 60)
-        index_b, timings_b = approach_b_unified_gpu_stitch(
-            xb,
-            d,
-            num_gpus,
-            args.graph_degree,
-            args.stitch_per_shard,
-            args.stitch_k,
-            save_dir=save_dir,
-        )
-        recall_b = eval_recall(index_b, xq, Igt, k)
-        kcycles_b = eval_kcycles(index_b, xq, k)
-        total_build_b = timings_b["build"] + timings_b["copyTo"]
-        print(f"  Build (SNMG+GPU-stitch): {timings_b['build']:.1f}s")
-        print(f"  copyTo:                  {timings_b['copyTo']:.1f}s")
-        print(
-            f"  Serialize:               {timings_b['serialize']:.1f}s "
-            f"({timings_b['file_size_gb']:.2f} GB)"
-        )
-        for ef in sorted(recall_b.keys()):
-            kc = kcycles_b.get(ef, 0)
-            print(
-                f"  efSearch={ef}: recall@{k}={recall_b[ef]:.4f}  "
-                f"kcycles/q={kc:.0f}"
-            )
-        results_table.append(
-            (
-                "B: GPU-brute",
-                total_build_b,
-                timings_b["serialize"],
-                timings_b["file_size_gb"],
-                recall_b,
-                kcycles_b,
-            )
-        )
-        del index_b
-        print()
-
-    if "C" in approaches_to_run:
-        print("=" * 60)
-        print(f"APPROACH C: Unified + CPU stitch (sps=0, k={args.stitch_k})")
-        print("=" * 60)
-        index_c, timings_c = approach_c_unified_cpu_stitch(
-            xb,
-            d,
-            num_gpus,
-            args.graph_degree,
-            0,
-            args.stitch_k,
-            save_dir=save_dir,
-        )
-        recall_c = eval_recall(index_c, xq, Igt, k)
-        kcycles_c = eval_kcycles(index_c, xq, k)
-        total_build_c = timings_c["build"] + timings_c["copyTo"]
-        print(f"  Build (SNMG+CPU-stitch): {timings_c['build']:.1f}s")
-        print(f"  copyTo:                  {timings_c['copyTo']:.1f}s")
-        print(
-            f"  Serialize:               {timings_c['serialize']:.1f}s "
-            f"({timings_c['file_size_gb']:.2f} GB)"
-        )
-        for ef in sorted(recall_c.keys()):
-            kc = kcycles_c.get(ef, 0)
-            print(
-                f"  efSearch={ef}: recall@{k}={recall_c[ef]:.4f}  "
-                f"kcycles/q={kc:.0f}"
-            )
-        results_table.append(
-            (
-                "C: CPU-stitch",
-                total_build_c,
-                timings_c["serialize"],
-                timings_c["file_size_gb"],
-                recall_c,
-                kcycles_c,
-            )
-        )
-        del index_c
-        print()
-
-    if "D" in approaches_to_run:
-        print("=" * 60)
-        print("APPROACH D: all_neighbors + cagra::optimize (no stitching)")
-        print("=" * 60)
-        index_d, timings_d = approach_d_all_neighbors(
-            xb,
-            d,
-            num_gpus,
-            args.graph_degree,
-            save_dir=save_dir,
-            n_clusters=args.n_clusters,
-            overlap_factor=args.overlap_factor,
-            multi_gpu_optimize=args.multi_gpu_optimize,
-            build_algo=args.build_algo,
-            base_level_only=args.base_level_only,
-            intermediate_graph_degree=args.intermediate_graph_degree,
-            refinement_rate=args.refinement_rate,
-            ivfpq_search_batch=args.ivfpq_search_batch,
-        )
-        recall_d = eval_recall(index_d, xq, Igt, k)
-        kcycles_d = eval_kcycles(index_d, xq, k)
-        total_build_d = timings_d["build"] + timings_d["copyTo"]
-        print(f"  Build (allneighbors+optimize): {timings_d['build']:.1f}s")
-        print(f"  copyTo:                        {timings_d['copyTo']:.1f}s")
-        ser = timings_d["serialize"]
-        gb = timings_d["file_size_gb"]
-        print(f"  Serialize:                     {ser:.1f}s " f"({gb:.2f} GB)")
-        # Index-build wall-clock, isolated from data load (Koski/Manifold) and
-        # ground-truth: this is the headline build->serialize number.
-        index_total_d = (
-            timings_d["build"] + timings_d["copyTo"] + timings_d["serialize"]
-        )
-        print(
-            f"  >>> INDEX build->serialize total "
-            f"(excl. data load + ground truth): "
-            f"{index_total_d:.1f}s ({index_total_d / 60:.1f} min)"
-        )
-        for ef in sorted(recall_d.keys()):
-            kc = kcycles_d.get(ef, 0)
-            print(
-                f"  efSearch={ef}: recall@{k}="
-                f"{recall_d[ef]:.4f}  kcyc/q={kc:.0f}"
-            )
-        results_table.append(
-            (
-                "D: all_neighbors",
-                total_build_d,
-                timings_d["serialize"],
-                timings_d["file_size_gb"],
-                recall_d,
-                kcycles_d,
-            )
-        )
-        del index_d
-        print()
-
-    print("=" * 60)
-    print("SUMMARY")
-    print("=" * 60)
-    header = (
-        f"{'Approach':<20s} {'Build':>7s} {'Ser':>5s} {'GB':>5s}"
-        f" {'ef':>4s} {'recall':>7s} {'kcyc/q':>7s}"
+    index, timings = build_cagra_hnsw(
+        xb,
+        d,
+        num_gpus,
+        args.graph_degree,
+        save_dir=save_dir,
+        n_clusters=args.n_clusters,
+        overlap_factor=args.overlap_factor,
+        build_algo=args.build_algo,
+        base_level_only=args.base_level_only,
+        intermediate_graph_degree=args.intermediate_graph_degree,
+        refinement_rate=args.refinement_rate,
+        ivfpq_search_batch=args.ivfpq_search_batch,
+        guarantee_connectivity=args.guarantee_connectivity,
+        ivfpq_size_from_cluster=args.ivfpq_size_from_cluster,
+        gpu_hnsw_upper_levels=args.gpu_hnsw_upper_levels,
+        gpu_hnsw_igd=args.gpu_hnsw_igd,
+        gpu_hnsw_guarantee_connectivity=args.gpu_hnsw_guarantee_connectivity,
+        ef_construction=args.ef_construction,
+        entrypoints=args.entrypoints,
     )
-    print(header)
-    print("-" * len(header))
-    for name, build, ser, size, recalls, kcycles in results_table:
-        for ef in [64, 128, 256]:
-            r = recalls.get(ef, 0)
-            kc = kcycles.get(ef, 0)
-            bld = f"{build:.0f}" if ef == 64 else ""
-            sr = f"{ser:.0f}" if ef == 64 else ""
-            sz = f"{size:.1f}" if ef == 64 else ""
-            nm = name if ef == 64 else ""
-            print(
-                f"{nm:<20s} {bld:>7s} {sr:>5s} {sz:>5s}"
-                f" {ef:>4d} {r:>7.4f} {kc:>7.0f}"
-            )
+    sweep_eps = [
+        int(v) for v in args.entrypoints_sweep.split(",") if v.strip()
+    ]
+    built_ep = index.num_base_level_search_entrypoints
+    for ep in sweep_eps:
+        index.num_base_level_search_entrypoints = ep
+        r = eval_recall(index, xq, Igt, k)
+        q = eval_qps(index, xq, k)
+        print(f"\n  --- num_base_level_search_entrypoints={ep} ---")
+        print(f"  {'efSearch':>8s} {'recall@' + str(k):>10s} {'QPS':>12s}")
+        for ef in sorted(r.keys()):
+            print(f"  {ef:>8d} {r[ef]:>10.4f} {q[ef]:>12,.1f}")
+    index.num_base_level_search_entrypoints = built_ep
+
+    recall = eval_recall(index, xq, Igt, k)
+    qps = eval_qps(index, xq, k)
+    kcycles = eval_kcycles(index, xq, k)
+
+    print("=" * 60)
+    print("RESULTS")
+    print("=" * 60)
+    print(f"  Build (all_neighbors+optimize): {timings['build']:.1f}s")
+    print(f"  copyTo:                         {timings['copyTo']:.1f}s")
+    print(
+        f"  Serialize:                      {timings['serialize']:.1f}s "
+        f"({timings['file_size_gb']:.2f} GB)"
+    )
+    # Index-build wall-clock, isolated from data load (Koski/Manifold) and
+    # ground-truth: this is the headline build->serialize number.
+    index_total = (
+        timings["build"] + timings["copyTo"] + timings["serialize"]
+    )
+    print(
+        f"  >>> INDEX build->serialize total "
+        f"(excl. data load + ground truth): "
+        f"{index_total:.1f}s ({index_total / 60:.1f} min)"
+    )
+    print(f"  {'efSearch':>8s} {'recall@' + str(k):>10s} {'QPS':>12s} "
+          f"{'us/query':>10s} {'kcyc/q':>9s}")
+    for ef in sorted(recall.keys()):
+        q = qps[ef]
+        print(
+            f"  {ef:>8d} {recall[ef]:>10.4f} {q:>12,.1f} "
+            f"{1e6 / q:>10.1f} {kcycles.get(ef, 0):>9.0f}"
+        )
 
 
 if __name__ == "__main__":

faiss/gpu/test/test_cagra.py

@@ -245,6 +245,49 @@ class TestIDMapCagra(unittest.TestCase):
         self.do_IDMapCagra(faiss.METRIC_INNER_PRODUCT, faiss.Int8)
 
 
+@unittest.skipIf(
+    "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in"
+)
+class TestCagraConfig(unittest.TestCase):
+    """Config surface for the multi-GPU build, checkable without 2 GPUs."""
+
+    def test_multi_gpu_knobs_survive_swig(self):
+        devices = faiss.Int32Vector()
+        devices.push_back(0)
+        devices.push_back(1)
+
+        all_neighbors = faiss.AllNeighborsCagraConfig()
+        all_neighbors.n_clusters = 16
+        all_neighbors.overlap_factor = 3
+        all_neighbors.ivf_pq_search_batch_size = 8192
+
+        config = faiss.GpuIndexCagraConfig()
+        config.devices = devices
+        config.all_neighbors_params = all_neighbors
+
+        self.assertEqual(config.devices.size(), 2)
+        self.assertEqual(config.devices.at(1), 1)
+        stored = config.all_neighbors_params
+        self.assertEqual(
+            (
+                stored.n_clusters,
+                stored.overlap_factor,
+                stored.ivf_pq_search_batch_size,
+            ),
+            (16, 3, 8192),
+        )
+
+    def test_brute_force_rejected_without_multiple_devices(self):
+        ds = datasets.SyntheticDataset(32, 0, 1000, 1)
+        config = faiss.GpuIndexCagraConfig()
+        config.build_algo = faiss.graph_build_algo_BRUTE_FORCE
+        res = faiss.StandardGpuResources()
+        index = faiss.GpuIndexCagra(res, ds.d, faiss.METRIC_L2, config)
+
+        with self.assertRaises(RuntimeError):
+            index.train(ds.get_database())
+
+
 @unittest.skipIf(
     "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in"
 )
@@ -253,53 +296,8 @@ class TestIDMapCagra(unittest.TestCase):
 )
 class TestMultiGpuCagra(unittest.TestCase):
 
-    def test_multi_gpu_build_and_search(self):
-        ds = datasets.SyntheticDataset(128, 0, 100_000, 1000)
-        xb = ds.get_database()
-        xq = ds.get_queries()
-        k = 10
-
-        gt_index = faiss.IndexFlatL2(ds.d)
-        gt_index.add(xb)
-        Dref, Iref = gt_index.search(xq, k)
-
-        num_gpus = min(faiss.get_num_gpus(), 4)
-        resources = faiss.GpuResourcesVector()
-        devices = faiss.Int32Vector()
-        res_list = []
-        for i in range(num_gpus):
-            res = faiss.StandardGpuResources()
-            res_list.append(res)
-            resources.push_back(res)
-            devices.push_back(i)
-
-        config = faiss.GpuIndexCagraConfig()
-        config.graph_degree = 32
-        index = faiss.GpuIndexCagra(res_list[0], ds.d, faiss.METRIC_L2, config)
-        index.trainMultiGpu(ds.nb, faiss.swig_ptr(xb), resources, devices, 0, 2)
-
-        cpu_index = faiss.IndexHNSWCagra()
-        index.copyTo(cpu_index)
-        self.assertEqual(cpu_index.ntotal, ds.nb)
-
-        cpu_index.hnsw.efSearch = 128
-        Dnew, Inew = cpu_index.search(xq, k)
-
-        recall = np.mean(
-            [len(set(Inew[i]) & set(Iref[i])) / k for i in range(ds.nq)]
-        )
-        self.assertGreater(
-            recall, 0.80, f"Multi-GPU recall@{k} too low: {recall:.4f}"
-        )
-
-        # Serialization roundtrip
-        data = faiss.serialize_index(cpu_index)
-        loaded = faiss.deserialize_index(data)
-        loaded.hnsw.efSearch = 128
-        Dnew2, Inew2 = loaded.search(xq, k)
-        np.testing.assert_array_equal(Inew, Inew2)
-
     def test_all_neighbors_build(self):
+        """train() on a multi-device config builds a unified CAGRA graph."""
         ds = datasets.SyntheticDataset(32, 0, 50_000, 100)
         xb = ds.get_database()
         xq = ds.get_queries()
@@ -309,19 +307,22 @@ class TestMultiGpuCagra(unittest.TestCase):
         gt_index.add(xb)
         Dref, Iref = gt_index.search(xq, k)
 
-        num_gpus = min(faiss.get_num_gpus(), 4)
         devices = faiss.Int32Vector()
-        for i in range(num_gpus):
+        for i in range(min(faiss.get_num_gpus(), 4)):
             devices.push_back(i)
 
-        res = faiss.StandardGpuResources()
+        all_neighbors = faiss.AllNeighborsCagraConfig()
+
         config = faiss.GpuIndexCagraConfig()
         config.graph_degree = 32
         config.intermediate_graph_degree = 48
+        config.build_algo = faiss.graph_build_algo_NN_DESCENT
+        config.devices = devices
+        config.all_neighbors_params = all_neighbors
+
+        res = faiss.StandardGpuResources()
         index = faiss.GpuIndexCagra(res, ds.d, faiss.METRIC_L2, config)
-        index.trainAllNeighbors(
-            ds.nb, faiss.swig_ptr(xb), devices, 0, 0, True, 0
-        )
+        index.train(xb)
 
         cpu_index = faiss.IndexHNSWCagra()
         cpu_index.base_level_only = True
@@ -337,3 +338,9 @@ class TestMultiGpuCagra(unittest.TestCase):
         self.assertGreater(
             recall, 0.70, f"all_neighbors recall@{k} too low: {recall:.4f}"
         )
+
+        # Serialization roundtrip
+        loaded = faiss.deserialize_index(faiss.serialize_index(cpu_index))
+        loaded.hnsw.efSearch = 128
+        _, Inew2 = loaded.search(xq, k)
+        np.testing.assert_array_equal(Inew, Inew2)

pyproject-gpu-cuvs.toml

@@ -45,6 +45,7 @@ dependencies = [
   "nvidia-curand>=10.3.7,<11",
   "nvidia-nvjitlink>=13.2,<14",
   "libcuvs-cu13>=26.06,<27",
+  "librmm-cu13>=26.06,<27",
 ]
 
 [project.urls]
@@ -114,6 +115,7 @@ test-requires = [
   "nvidia-curand>=10.3.7,<11",
   "nvidia-nvjitlink>=13.2,<14",
   "libcuvs-cu13>=26.06,<27",
+  "librmm-cu13>=26.06,<27",
 ]
 test-command = "python -m pytest {project}/tests/test_wheel_smoke_gpu.py -v"