From 9f4ea20b12ededda454be61f130c6f1a54edde33 Mon Sep 17 00:00:00 2001 From: Michael Norris Date: Fri, 14 Aug 2026 08:47:08 -0700 Subject: [PATCH] Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500) Summary: `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. Differential Revision: D114685755 --- CMakeLists.txt | 4 + faiss/IndexHNSW.cpp | 16 +- faiss/IndexHNSW.h | 2 +- faiss/gpu/CMakeLists.txt | 2 +- faiss/gpu/GpuIndexBinaryCagra.cu | 7 + faiss/gpu/GpuIndexCagra.cu | 1020 ++++++++-------------------- faiss/gpu/GpuIndexCagra.h | 107 +-- faiss/gpu/test/bench_approaches.py | 683 ++++++------------- faiss/gpu/test/test_cagra.py | 103 +-- pyproject-gpu-cuvs.toml | 2 + 10 files changed, 653 insertions(+), 1293 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 5b090f70d8..0a85cce027 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/faiss/IndexHNSW.cpp b/faiss/IndexHNSW.cpp index 42f8548a65..c67d0e941d 100644 --- a/faiss/IndexHNSW.cpp +++ b/faiss/IndexHNSW.cpp @@ -17,7 +17,6 @@ #include #include #include -#include #include #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 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(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 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. diff --git a/faiss/IndexHNSW.h b/faiss/IndexHNSW.h index 8e3d7f2906..405e14d68a 100644 --- a/faiss/IndexHNSW.h +++ b/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; diff --git a/faiss/gpu/CMakeLists.txt b/faiss/gpu/CMakeLists.txt index c896a85a29..2387eada10 100644 --- a/faiss/gpu/CMakeLists.txt +++ b/faiss/gpu/CMakeLists.txt @@ -363,7 +363,7 @@ else() find_package(CUDAToolkit REQUIRED) - target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$:cuvs::cuvs> $<$:OpenMP::OpenMP_CXX>) + target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$:cuvs::cuvs> $<$:rmm::rmm> $<$:OpenMP::OpenMP_CXX>) target_compile_options(faiss_gpu_objs PRIVATE $<$:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr diff --git a/faiss/gpu/GpuIndexBinaryCagra.cu b/faiss/gpu/GpuIndexBinaryCagra.cu index 46b6efb636..d2d560f430 100644 --- a/faiss/gpu/GpuIndexBinaryCagra.cu +++ b/faiss/gpu/GpuIndexBinaryCagra.cu @@ -78,6 +78,13 @@ std::shared_ptr 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_); diff --git a/faiss/gpu/GpuIndexCagra.cu b/faiss/gpu/GpuIndexCagra.cu index acd356c5aa..0eac9f8752 100644 --- a/faiss/gpu/GpuIndexCagra.cu +++ b/faiss/gpu/GpuIndexCagra.cu @@ -23,15 +23,12 @@ #include #include -#include #include #include #include #include #include -#include #include -#include #include #include @@ -40,56 +37,31 @@ #include #include #include -#include -#include +// clang-format off +#if __has_include() +// Some cuVS versions put helpers::optimize in its own header. +#include +#endif +// TEMPORARY. cuVS has no public API to prune a graph that is already on the +// device, so we reach into its internals when they happen to be available. +// Delete this and the branch it guards once that API ships upstream. +#if defined(FAISS_CAGRA_DEVICE_OPTIMIZE) +#include +#endif +// clang-format on namespace { -template -__global__ void kern_detour_count( - const uint32_t* __restrict__ knn_graph, - const uint32_t graph_size, - const uint32_t graph_degree, - const uint32_t output_degree, - const uint32_t batch_size, - const uint32_t offset, - uint8_t* __restrict__ detour_count) { - __shared__ uint32_t smem[MAX_DEGREE]; - - const uint64_t iA = blockIdx.x + (uint64_t)offset; - if (iA >= graph_size) - return; - - for (uint32_t k = threadIdx.x; k < graph_degree; k += blockDim.x) { - smem[k] = 0; - if (knn_graph[k + (uint64_t)graph_degree * iA] == (uint32_t)iA) - smem[k] = graph_degree; - } - __syncthreads(); - - for (uint32_t kAD = 0; kAD < graph_degree - 1; kAD++) { - const uint64_t iD = knn_graph[kAD + (uint64_t)graph_degree * iA]; - if (iD >= graph_size) - continue; - for (uint32_t kDB = threadIdx.x; kDB < graph_degree; - kDB += blockDim.x) { - const uint64_t iB_cand = - knn_graph[kDB + (uint64_t)graph_degree * iD]; - for (uint32_t kAB = kAD + 1; kAB < graph_degree; kAB++) { - const uint64_t iB = - knn_graph[kAB + (uint64_t)graph_degree * iA]; - if (iB == iB_cand) { - atomicAdd(smem + kAB, 1); - break; - } - } - } - __syncthreads(); - } - - for (uint32_t k = threadIdx.x; k < graph_degree; k += blockDim.x) { - detour_count[k + (uint64_t)graph_degree * iA] = - static_cast(min(smem[k], 255u)); +// all_neighbors emits int64 indices, optimize consumes uint32. Narrowing on +// device lets the graph stay there. +__global__ void kern_narrow_indices( + const int64_t* __restrict__ src, + uint32_t* __restrict__ dst, + size_t count) { + size_t stride = (size_t)gridDim.x * blockDim.x; + for (size_t i = (size_t)blockIdx.x * blockDim.x + threadIdx.x; i < count; + i += stride) { + dst[i] = static_cast(src[i]); } } @@ -109,6 +81,19 @@ GpuIndexCagra::GpuIndexCagra( } void GpuIndexCagra::train_ex(idx_t n, const void* x, NumericType numeric_type) { + if (cagraConfig_.devices.size() > 1) { + FAISS_THROW_IF_NOT_MSG( + numeric_type == NumericType::Float32, + "multi-GPU CAGRA build only supports Float32"); + trainAllNeighbors_(n, static_cast(x)); + return; + } + + FAISS_THROW_IF_MSG( + cagraConfig_.build_algo == graph_build_algo::BRUTE_FORCE, + "graph_build_algo::BRUTE_FORCE is only available on the multi-GPU " + "build path (set GpuIndexCagraConfig::devices)"); + numeric_type_ = numeric_type; bool index_is_initialized = !std::holds_alternative(index_); @@ -370,461 +355,99 @@ void GpuIndexCagra::searchImpl_( search_params); } -void GpuIndexCagra::trainMultiGpu( - idx_t n, - const float* x, - std::vector& providers, - std::vector& devices, - idx_t stitch_per_shard, - int stitch_k, - int stitch_mode) { +void GpuIndexCagra::trainAllNeighbors_(idx_t n, const float* x) { FAISS_THROW_IF_MSG(is_trained, "index is already trained"); - FAISS_THROW_IF_MSG(devices.empty(), "must provide at least one GPU"); - FAISS_THROW_IF_NOT_MSG(stitch_k >= 1, "stitch_k must be >= 1"); - - numeric_type_ = NumericType::Float32; - int num_shards = static_cast(devices.size()); - idx_t shard_size = (n + num_shards - 1) / num_shards; - - int total_cross = (num_shards > 1) ? stitch_k * (num_shards - 1) : 0; - - auto t_phase = std::chrono::high_resolution_clock::now(); - auto t_start = t_phase; - - // Use cuVS native multi-GPU CAGRA build via SNMG - raft::device_resources_snmg clique(devices); - - cuvs::neighbors::mg_index_params - mg_params; - mg_params.mode = cuvs::neighbors::SHARDED; - mg_params.intermediate_graph_degree = - cagraConfig_.intermediate_graph_degree; - mg_params.graph_degree = cagraConfig_.graph_degree; - mg_params.guarantee_connectivity = cagraConfig_.guarantee_connectivity; - mg_params.metric = metricFaissToCuvs(this->metric_type, false); - - if (cagraConfig_.build_algo == graph_build_algo::IVF_PQ) { - cuvs::neighbors::cagra::graph_build_params::ivf_pq_params - graph_build_params; - if (cagraConfig_.ivf_pq_params) { - graph_build_params.build_params.n_lists = - cagraConfig_.ivf_pq_params->n_lists; - graph_build_params.build_params.kmeans_n_iters = - cagraConfig_.ivf_pq_params->kmeans_n_iters; - graph_build_params.build_params.kmeans_trainset_fraction = - cagraConfig_.ivf_pq_params->kmeans_trainset_fraction; - graph_build_params.build_params.pq_bits = - cagraConfig_.ivf_pq_params->pq_bits; - graph_build_params.build_params.pq_dim = - cagraConfig_.ivf_pq_params->pq_dim; - graph_build_params.build_params.codebook_kind = - static_cast( - cagraConfig_.ivf_pq_params->codebook_kind); - graph_build_params.build_params.force_random_rotation = - cagraConfig_.ivf_pq_params->force_random_rotation; - graph_build_params.build_params.conservative_memory_allocation = - cagraConfig_.ivf_pq_params->conservative_memory_allocation; - } - if (cagraConfig_.ivf_pq_search_params) { - graph_build_params.search_params.n_probes = - cagraConfig_.ivf_pq_search_params->n_probes; - graph_build_params.search_params.lut_dtype = - cagraConfig_.ivf_pq_search_params->lut_dtype; - graph_build_params.search_params.preferred_shmem_carveout = - cagraConfig_.ivf_pq_search_params->preferred_shmem_carveout; - } - graph_build_params.build_params.metric = - metricFaissToCuvs(this->metric_type, false); - graph_build_params.refinement_rate = cagraConfig_.refine_rate; - mg_params.graph_build_params = graph_build_params; - if (mg_params.graph_degree == mg_params.intermediate_graph_degree) { - mg_params.intermediate_graph_degree = 1.5 * mg_params.graph_degree; - } - } else { - cuvs::neighbors::cagra::graph_build_params::nn_descent_params - graph_build_params(mg_params.intermediate_graph_degree); - graph_build_params.max_iterations = cagraConfig_.nn_descent_niter; - graph_build_params.metric = metricFaissToCuvs(this->metric_type, false); - mg_params.graph_build_params = graph_build_params; - } - - auto dataset = raft::make_host_matrix_view( - x, n, static_cast(this->d)); - - auto mg_idx = cuvs::neighbors::cagra::build(clique, mg_params, dataset); - - auto t_now = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainMultiGpu] SNMG CAGRA build: %.2f seconds\n", - std::chrono::duration(t_now - t_phase).count()); - t_phase = t_now; - - // Extract per-shard graphs into expanded layout: [CAGRA neighbors | - // cross-shard slots] - idx_t graph_degree = static_cast(cagraConfig_.graph_degree); - idx_t expanded_degree = graph_degree + total_cross; - merged_knngraph_.assign(n * expanded_degree, -1); - - for (int s = 0; s < num_shards; s++) { - idx_t offset = static_cast(s) * shard_size; - idx_t actual_size = std::min(shard_size, n - offset); - - auto& shard_idx = mg_idx.ann_interfaces_[s].index_.value(); - auto device_graph = shard_idx.graph(); - - const auto& dev_res = - raft::resource::set_current_device_to_rank(clique, s); - FAISS_THROW_IF_NOT_FMT( - device_graph.extent(1) == graph_degree, - "Shard %d has graph_degree %ld, expected %ld", - s, - (long)device_graph.extent(1), - (long)graph_degree); - - std::vector host_graph(actual_size * graph_degree); - raft::resource::sync_stream(dev_res); - thrust::copy( - thrust::device_ptr(device_graph.data_handle()), - thrust::device_ptr( - device_graph.data_handle() + - actual_size * graph_degree), - host_graph.data()); - -#pragma omp parallel for - for (idx_t i = 0; i < actual_size; i++) { - for (idx_t j = 0; j < graph_degree; j++) { - merged_knngraph_[(offset + i) * expanded_degree + j] = - static_cast(host_graph[i * graph_degree + j]) + - offset; - } - } - } - - merged_knngraph_degree_ = expanded_degree; - - t_now = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainMultiGpu] Graph extract + merge: %.2f seconds\n", - std::chrono::duration(t_now - t_phase).count()); - t_phase = t_now; - - // Cross-shard stitching: appends cross-shard edges after CAGRA neighbors. - // Mode 0: CPU HNSW search (approximate). Mode 1: GPU brute-force (exact). - if (num_shards > 1) { - if (stitch_mode == 1 && stitch_per_shard == 0) { - stitch_per_shard = 100000; - fprintf(stderr, - " [trainMultiGpu] GPU brute-force mode: defaulting " - "stitch_per_shard to %ld (all-vectors is O(N^2))\n", - (long)stitch_per_shard); - } - - fprintf(stderr, - " [trainMultiGpu] Stitching: mode=%s, stitch_per_shard=%ld, " - "stitch_k=%d, total_cross=%d, expanded_degree=%ld\n", - stitch_mode == 1 ? "GPU-brute-force" : "CPU-HNSW", - (long)stitch_per_shard, - stitch_k, - total_cross, - (long)expanded_degree); - - if (stitch_mode == 1) { - // GPU brute-force stitching: exact cross-shard neighbors - for (int j = 0; j < num_shards; j++) { - idx_t j_offset = static_cast(j) * shard_size; - idx_t j_size = std::min(shard_size, n - j_offset); - if (j_size == 0) - continue; - - GpuIndexFlatConfig flat_config; - flat_config.device = devices[j]; - GpuIndexFlatL2 flat_idx(providers[j], this->d, flat_config); - flat_idx.add(j_size, x + j_offset * this->d); - - for (int i = 0; i < num_shards; i++) { - if (i == j) - continue; - idx_t i_offset = static_cast(i) * shard_size; - idx_t i_size = std::min(shard_size, n - i_offset); - if (i_size == 0) - continue; - - int target_pos = (j < i) ? j : j - 1; - idx_t slot_base = graph_degree + - static_cast(target_pos) * stitch_k; - - idx_t num_queries = i_size; - const float* query_data = x + i_offset * this->d; - std::vector sampled_indices; - std::vector sampled_vectors; - - if (stitch_per_shard > 0 && stitch_per_shard < i_size) { - num_queries = stitch_per_shard; - sampled_indices.resize(i_size); - std::iota( - sampled_indices.begin(), - sampled_indices.end(), - idx_t(0)); - std::mt19937 rng(42 + i * num_shards + j); - for (idx_t s = 0; s < stitch_per_shard; s++) { - std::uniform_int_distribution dist( - s, i_size - 1); - std::swap( - sampled_indices[s], - sampled_indices[dist(rng)]); - } - sampled_indices.resize(stitch_per_shard); - - sampled_vectors.resize(stitch_per_shard * this->d); -#pragma omp parallel for - for (idx_t s = 0; s < stitch_per_shard; s++) { - memcpy(sampled_vectors.data() + s * this->d, - x + - (i_offset + sampled_indices[s]) * - this->d, - this->d * sizeof(float)); - } - query_data = sampled_vectors.data(); - } - - std::vector distances(num_queries * stitch_k); - std::vector labels(num_queries * stitch_k); - flat_idx.search( - num_queries, - query_data, - stitch_k, - distances.data(), - labels.data()); - -#pragma omp parallel for - for (idx_t s = 0; s < num_queries; s++) { - idx_t orig_idx = (!sampled_indices.empty()) - ? sampled_indices[s] - : s; - idx_t global_id = i_offset + orig_idx; - for (int k = 0; k < stitch_k; k++) { - idx_t local_nb = labels[s * stitch_k + k]; - merged_knngraph_ - [global_id * expanded_degree + slot_base + - k] = (local_nb >= 0) - ? (local_nb + j_offset) - : local_nb; - } - } - } - fprintf(stderr, - " Shard %d/%d GPU-stitched\n", - j + 1, - num_shards); - } - } else { - // CPU HNSW stitching (Approach C) - auto M_temp = graph_degree / 2; - for (int j = 0; j < num_shards; j++) { - idx_t j_offset = static_cast(j) * shard_size; - idx_t j_size = std::min(shard_size, n - j_offset); - if (j_size == 0) - continue; - - IndexHNSWCagra temp_idx; - temp_idx.d = this->d; - temp_idx.metric_type = this->metric_type; - temp_idx.base_level_only = true; - temp_idx.num_base_level_search_entrypoints = 32; - if (this->metric_type == METRIC_L2) { - temp_idx.storage = new IndexFlatL2(this->d); - } else { - temp_idx.storage = new IndexFlatIP(this->d); - } - temp_idx.own_fields = true; - temp_idx.keep_max_size_level0 = true; - temp_idx.hnsw.reset(); - temp_idx.hnsw.assign_probas.clear(); - temp_idx.hnsw.cum_nneighbor_per_level.clear(); - temp_idx.hnsw.set_default_probas(M_temp, 1.0 / log(M_temp)); - temp_idx.init_level0 = false; - temp_idx.hnsw.prepare_level_tab(j_size, false); - temp_idx.storage->add(j_size, x + j_offset * this->d); - temp_idx.ntotal = j_size; - -#pragma omp parallel for - for (idx_t i = 0; i < j_size; i++) { - size_t begin, end; - temp_idx.hnsw.neighbor_range(i, 0, &begin, &end); - for (size_t k = begin; k < end; k++) { - idx_t global_nb = merged_knngraph_ - [(j_offset + i) * expanded_degree + - (k - begin)]; - temp_idx.hnsw.neighbors[k] = global_nb - j_offset; - } - } - - temp_idx.hnsw.efSearch = 64; - - for (int i = 0; i < num_shards; i++) { - if (i == j) - continue; - idx_t i_offset = static_cast(i) * shard_size; - idx_t i_size = std::min(shard_size, n - i_offset); - if (i_size == 0) - continue; - - int target_pos = (j < i) ? j : j - 1; - idx_t slot_base = graph_degree + - static_cast(target_pos) * stitch_k; - - idx_t num_queries = i_size; - const float* query_data = x + i_offset * this->d; - std::vector sampled_indices; - std::vector sampled_vectors; - - if (stitch_per_shard > 0 && stitch_per_shard < i_size) { - num_queries = stitch_per_shard; - sampled_indices.resize(i_size); - std::iota( - sampled_indices.begin(), - sampled_indices.end(), - idx_t(0)); - std::mt19937 rng(42 + i * num_shards + j); - for (idx_t s = 0; s < stitch_per_shard; s++) { - std::uniform_int_distribution dist( - s, i_size - 1); - std::swap( - sampled_indices[s], - sampled_indices[dist(rng)]); - } - sampled_indices.resize(stitch_per_shard); - - sampled_vectors.resize(stitch_per_shard * this->d); -#pragma omp parallel for - for (idx_t s = 0; s < stitch_per_shard; s++) { - memcpy(sampled_vectors.data() + s * this->d, - x + - (i_offset + sampled_indices[s]) * - this->d, - this->d * sizeof(float)); - } - query_data = sampled_vectors.data(); - } - - std::vector distances(num_queries * stitch_k); - std::vector labels(num_queries * stitch_k); - temp_idx.search( - num_queries, - query_data, - stitch_k, - distances.data(), - labels.data()); - -#pragma omp parallel for - for (idx_t s = 0; s < num_queries; s++) { - idx_t orig_idx = (!sampled_indices.empty()) - ? sampled_indices[s] - : s; - idx_t global_id = i_offset + orig_idx; - for (int k = 0; k < stitch_k; k++) { - idx_t local_nb = labels[s * stitch_k + k]; - merged_knngraph_ - [global_id * expanded_degree + slot_base + - k] = (local_nb >= 0) - ? (local_nb + j_offset) - : local_nb; - } - } - } - fprintf(stderr, - " Shard %d/%d stitched\n", - j + 1, - num_shards); - } - } // end else (CPU HNSW) - } - - t_now = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainMultiGpu] Stitching: %.2f seconds\n", - std::chrono::duration(t_now - t_phase).count()); - fprintf(stderr, - " [trainMultiGpu] Total: %.2f seconds\n", - std::chrono::duration(t_now - t_start).count()); - - multi_gpu_dataset_ = x; - this->is_trained = true; - this->ntotal = n; -} -void GpuIndexCagra::trainAllNeighbors( - idx_t n, - const float* x, - std::vector& devices, - int n_clusters_override, - int overlap_factor_override, - bool multi_gpu_optimize, - int build_algo, - float refinement_rate, - int ivfpq_search_batch) { - FAISS_THROW_IF_MSG(is_trained, "index is already trained"); - FAISS_THROW_IF_MSG(devices.empty(), "must provide at least one GPU"); + const auto& devices = cagraConfig_.devices; + const auto& an_config = cagraConfig_.all_neighbors_params; numeric_type_ = NumericType::Float32; idx_t graph_degree = static_cast(cagraConfig_.graph_degree); idx_t intermediate_degree = static_cast(cagraConfig_.intermediate_graph_degree); + const size_t knn_count = (size_t)n * intermediate_degree; auto t_start = std::chrono::high_resolution_clock::now(); auto t_phase = t_start; - // Phase 1+2: Build kNN graph (multi-GPU) then D2H + cast. - // clique and d_indices are scoped so SNMG resources are fully destroyed - // before Phase 3, avoiding CUDA context interference with single-GPU - // optimize. - auto h_knn = - raft::make_host_matrix(n, intermediate_degree); + // Allocated off the clique so it survives the clique's destruction below. + CUDA_VERIFY(cudaSetDevice(devices[0])); + raft::device_resources single_gpu_res; + auto d_knn = raft::make_device_matrix( + single_gpu_res, n, intermediate_degree); + + // Build the knn graph across all devices. The clique is scoped: optimize + // hangs intermittently if multi-GPU state is still alive when it runs. { raft::device_resources_snmg clique(devices); cuvs::neighbors::all_neighbors::all_neighbors_params an_params; int num_gpus = static_cast(devices.size()); - an_params.n_clusters = n_clusters_override > 0 - ? static_cast(n_clusters_override) - : std::max(num_gpus * 2, 4); - an_params.overlap_factor = overlap_factor_override > 0 - ? static_cast(overlap_factor_override) - : 2; + an_params.n_clusters = an_config.n_clusters > 0 + ? an_config.n_clusters + : static_cast(std::max(num_gpus * 2, 4)); + an_params.overlap_factor = an_config.overlap_factor; an_params.metric = metricFaissToCuvs(this->metric_type, false); const char* algo_name = "nn_descent"; - if (build_algo == 1) { - // Brute-force: exact kNN via tiled GEMM, O(N²D) per cluster - cuvs::neighbors::all_neighbors::graph_build_params:: - brute_force_params bf_params; - an_params.graph_build_params = bf_params; - algo_name = "brute_force"; - } else if (build_algo == 2) { - // IVF-PQ: approximate kNN with refinement - auto dataset_ext = raft::make_extents( - n, static_cast(this->d)); - cuvs::neighbors::all_neighbors::graph_build_params::ivf_pq_params - ivfpq_params(dataset_ext); - ivfpq_params.refinement_rate = refinement_rate; - // Bound the IVF-PQ search workspace to avoid GPU OOM at scale. The - // cuVS default max_internal_batch_size is 128*1024, whose search - // buffers can exceed H100 memory for dense 100M clusters. Capping - // it batches the search into smaller query chunks (results - // unchanged). - if (ivfpq_search_batch > 0) { - ivfpq_params.search_params.max_internal_batch_size = - static_cast(ivfpq_search_batch); + switch (cagraConfig_.build_algo) { + case graph_build_algo::BRUTE_FORCE: { + // Exact kNN via tiled GEMM, O(N^2 D) per cluster + cuvs::neighbors::all_neighbors::graph_build_params:: + brute_force_params bf_params; + an_params.graph_build_params = bf_params; + algo_name = "brute_force"; + break; } - an_params.graph_build_params = ivfpq_params; - algo_name = "ivf_pq"; - } else { - // Default: NN-descent - cuvs::neighbors::all_neighbors::graph_build_params:: - nn_descent_params nn_params; - nn_params.graph_degree = intermediate_degree; - nn_params.max_iterations = cagraConfig_.nn_descent_niter; - an_params.graph_build_params = nn_params; + case graph_build_algo::IVF_PQ: { + // cuVS derives good IVF-PQ params from the dataset shape, + // so only override what it cannot infer. See + // AllNeighborsCagraConfig::ivf_pq_size_from_cluster. + idx_t sizing_rows = n; + if (an_config.ivf_pq_size_from_cluster) { + sizing_rows = std::max( + 1, + (idx_t)((double)n * an_params.overlap_factor / + (double)an_params.n_clusters)); + } + auto dataset_ext = raft::make_extents( + sizing_rows, static_cast(this->d)); + cuvs::neighbors::all_neighbors::graph_build_params:: + ivf_pq_params ivfpq_params(dataset_ext); + ivfpq_params.refinement_rate = an_config.refinement_rate; + // Bounds search memory; batches the search rather than + // changing its result. + if (an_config.ivf_pq_search_batch_size > 0) { + ivfpq_params.search_params.max_internal_batch_size = + an_config.ivf_pq_search_batch_size; + } + fprintf(stderr, + " [trainAllNeighbors] IVF-PQ sized from %s " + "(%ld rows): n_lists=%u n_probes=%u " + "kmeans_trainset_fraction=%.4f refine=%.2f\n", + an_config.ivf_pq_size_from_cluster ? "cluster" + : "full dataset", + (long)sizing_rows, + ivfpq_params.build_params.n_lists, + ivfpq_params.search_params.n_probes, + ivfpq_params.build_params.kmeans_trainset_fraction, + ivfpq_params.refinement_rate); + an_params.graph_build_params = ivfpq_params; + algo_name = "ivf_pq"; + break; + } + case graph_build_algo::NN_DESCENT: { + cuvs::neighbors::all_neighbors::graph_build_params:: + nn_descent_params nn_params; + nn_params.graph_degree = intermediate_degree; + nn_params.max_iterations = cagraConfig_.nn_descent_niter; + an_params.graph_build_params = nn_params; + break; + } + default: + FAISS_THROW_MSG( + "multi-GPU CAGRA build supports build_algo IVF_PQ, " + "NN_DESCENT or BRUTE_FORCE"); } auto dataset = raft::make_host_matrix_view( @@ -842,7 +465,9 @@ void GpuIndexCagra::trainAllNeighbors( an_params.n_clusters, an_params.overlap_factor, algo_name, - build_algo == 2 ? refinement_rate : 0.0f); + cagraConfig_.build_algo == graph_build_algo::IVF_PQ + ? an_config.refinement_rate + : 0.0f); cuvs::neighbors::all_neighbors::build( clique, an_params, dataset, d_indices.view()); @@ -853,258 +478,76 @@ void GpuIndexCagra::trainAllNeighbors( std::chrono::duration(t_now - t_phase).count()); t_phase = t_now; - // D2H + cast int64 -> uint32 - std::vector h_indices_i64(n * intermediate_degree); - raft::copy( - h_indices_i64.data(), - d_indices.data_handle(), - n * intermediate_degree, - raft::resource::get_cuda_stream(clique)); - raft::resource::sync_stream(clique); - -#pragma omp parallel for - for (idx_t i = 0; i < n * intermediate_degree; i++) { - h_knn.data_handle()[i] = static_cast(h_indices_i64[i]); - } + CUDA_VERIFY(cudaSetDevice(devices[0])); + constexpr int kThreads = 256; + const int blocks = (int)std::min( + (knn_count + kThreads - 1) / kThreads, 65535); + kern_narrow_indices<<>>( + d_indices.data_handle(), d_knn.data_handle(), knn_count); + CUDA_VERIFY(cudaGetLastError()); + CUDA_VERIFY(cudaDeviceSynchronize()); } // clique + d_indices destroyed here, freeing all GPU resources auto t_now = std::chrono::high_resolution_clock::now(); fprintf(stderr, - " [trainAllNeighbors] D2H + cast: %.2f seconds\n", + " [trainAllNeighbors] narrow to uint32 (device): %.2f seconds\n", std::chrono::duration(t_now - t_phase).count()); t_phase = t_now; - // Phase 3: Optimize kNN graph into CAGRA graph - auto h_cagra = raft::make_host_matrix(n, graph_degree); - memset(h_cagra.data_handle(), 0xff, n * graph_degree * sizeof(uint32_t)); - - if (multi_gpu_optimize && devices.size() > 1) { - // S2: Multi-GPU detour counting + CPU pruning/reverse/merge - int num_gpus = static_cast(devices.size()); - fprintf(stderr, - " [trainAllNeighbors] Multi-GPU optimize: %d GPUs, " - "n=%ld, %ld->%ld\n", - num_gpus, - (long)n, - (long)intermediate_degree, - (long)graph_degree); - - // Phase 3a: Multi-GPU detour counting - FAISS_THROW_IF_NOT_MSG( - intermediate_degree <= 1024, - "intermediate_graph_degree must be <= 1024 for " - "multi-GPU optimize"); - - auto h_detour = raft::make_host_matrix( - n, intermediate_degree); - memset(h_detour.data_handle(), - 0xff, - n * intermediate_degree * sizeof(uint8_t)); - - idx_t chunk = (n + num_gpus - 1) / num_gpus; - -#pragma omp parallel for num_threads(num_gpus) - for (int g = 0; g < num_gpus; g++) { - idx_t my_start = g * chunk; - idx_t my_end = std::min(my_start + chunk, n); - idx_t my_n = my_end - my_start; - if (my_n <= 0) - continue; - - CUDA_VERIFY(cudaSetDevice(devices[g])); - cudaStream_t stream; - CUDA_VERIFY(cudaStreamCreate(&stream)); - - uint32_t* d_knn_graph; - CUDA_VERIFY(cudaMalloc( - &d_knn_graph, - (size_t)n * intermediate_degree * sizeof(uint32_t))); - CUDA_VERIFY(cudaMemcpyAsync( - d_knn_graph, - h_knn.data_handle(), - (size_t)n * intermediate_degree * sizeof(uint32_t), - cudaMemcpyHostToDevice, - stream)); - - uint8_t* d_detour; - CUDA_VERIFY(cudaMalloc( - &d_detour, - (size_t)n * intermediate_degree * sizeof(uint8_t))); - CUDA_VERIFY(cudaMemsetAsync( - d_detour, - 0xff, - (size_t)n * intermediate_degree * sizeof(uint8_t), - stream)); - - // Launch detour counting for this GPU's node range - constexpr uint32_t BATCH = 262144; - dim3 threads(32, 1, 1); - uint32_t total_batches = - (static_cast(my_n) + BATCH - 1) / BATCH; - - for (uint32_t ib = 0; ib < total_batches; ib++) { - uint32_t batch_start = - static_cast(my_start) + ib * BATCH; - uint32_t batch_n = std::min( - BATCH, static_cast(my_end) - batch_start); - dim3 blocks(batch_n, 1, 1); - kern_detour_count<1024><<>>( - d_knn_graph, - static_cast(n), - static_cast(intermediate_degree), - static_cast(graph_degree), - BATCH, - batch_start, - d_detour); - } - - CUDA_VERIFY(cudaGetLastError()); - CUDA_VERIFY(cudaMemcpyAsync( - h_detour.data_handle() + my_start * intermediate_degree, - d_detour + my_start * intermediate_degree, - (size_t)my_n * intermediate_degree * sizeof(uint8_t), - cudaMemcpyDeviceToHost, - stream)); - - CUDA_VERIFY(cudaStreamSynchronize(stream)); - CUDA_VERIFY(cudaFree(d_knn_graph)); - CUDA_VERIFY(cudaFree(d_detour)); - CUDA_VERIFY(cudaStreamDestroy(stream)); - } - - auto t_opt1 = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainAllNeighbors] multi-GPU detour counting: " - "%.2f seconds\n", - std::chrono::duration(t_opt1 - t_phase).count()); - - // Phase 3b: CPU pruning (select top-graph_degree by lowest detour - // count) - auto* output_ptr = h_cagra.data_handle(); -#pragma omp parallel for - for (idx_t i = 0; i < n; i++) { - idx_t pk = 0; - uint32_t num_detour = 0; - for (uint32_t l = 0; - l < (uint32_t)intermediate_degree && pk < graph_degree; - l++) { - uint32_t next_num_detour = std::numeric_limits::max(); - for (idx_t k = 0; k < intermediate_degree; k++) { - uint32_t dc = - h_detour.data_handle()[i * intermediate_degree + k]; - if (dc > num_detour) - next_num_detour = std::min(dc, next_num_detour); - if (dc != num_detour) - continue; - - uint32_t candidate = - h_knn.data_handle()[i * intermediate_degree + k]; - bool dup = false; - for (idx_t dk = 0; dk < pk; dk++) { - if (candidate == output_ptr[i * graph_degree + dk]) { - dup = true; - break; - } - } - if (!dup && candidate < (uint32_t)n) { - output_ptr[i * graph_degree + pk] = candidate; - pk++; - } - if (pk >= graph_degree) - break; - } - if (pk >= graph_degree) - break; - if (next_num_detour == std::numeric_limits::max()) - break; - num_detour = next_num_detour; - } - } - - auto t_opt2 = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainAllNeighbors] CPU pruning: %.2f seconds\n", - std::chrono::duration(t_opt2 - t_opt1).count()); - - // Phase 3c: Build reverse graph and merge into output - // (matches graph_core.cuh Phase 4+5 logic) - std::vector rev_graph(n * graph_degree, UINT32_MAX); - std::vector rev_count(n, 0); - - // Build reverse adjacency: for each edge A→B, record B←A - for (idx_t i = 0; i < n; i++) { - for (idx_t k = 0; k < graph_degree; k++) { - uint32_t dest = output_ptr[i * graph_degree + k]; - if (dest >= (uint32_t)n) - continue; - uint32_t slot = rev_count[dest]; - if (slot < (uint32_t)graph_degree) { - rev_graph[dest * graph_degree + slot] = - static_cast(i); - rev_count[dest]++; - } - } - } - - // Merge reverse edges into output graph (replace tail edges) -#pragma omp parallel for - for (idx_t i = 0; i < n; i++) { - uint32_t num_protected = std::max(graph_degree / 2, 1); - uint32_t kr = - std::min(rev_count[i], static_cast(graph_degree)); - while (kr > 0) { - kr--; - uint32_t rev_node = rev_graph[i * graph_degree + kr]; - if (rev_node >= (uint32_t)n) - continue; - - // Check if already in neighbor list - bool found = false; - for (idx_t j = 0; j < graph_degree; j++) { - if (output_ptr[i * graph_degree + j] == rev_node) { - found = true; - break; - } - } - if (found) - continue; - - // Shift tail edges and insert at protected boundary - for (idx_t j = graph_degree - 1; j > (idx_t)num_protected; - j--) { - output_ptr[i * graph_degree + j] = - output_ptr[i * graph_degree + j - 1]; - } - output_ptr[i * graph_degree + num_protected] = rev_node; - } - } + // Phase 3: prune the knn graph into a CAGRA graph. + const size_t out_count = (size_t)n * graph_degree; + std::vector h_cagra(out_count); - auto t_opt3 = std::chrono::high_resolution_clock::now(); - fprintf(stderr, - " [trainAllNeighbors] reverse graph: %.2f seconds\n", - std::chrono::duration(t_opt3 - t_opt2).count()); - } else { - // Single-GPU optimize (original path) - cudaSetDevice(devices[0]); - cudaDeviceSynchronize(); - raft::device_resources single_gpu_res; - cuvs::neighbors::cagra::helpers::optimize( - single_gpu_res, h_knn.view(), h_cagra.view()); - } +#if !defined(FAISS_CAGRA_DEVICE_OPTIMIZE) + // Public API: host matrices only, so the graph makes a round trip it does + // not need. Some cuVS versions also read the input as host memory whatever + // accessor you hand them, so passing device pointers here is not an option. + auto h_knn = + raft::make_host_matrix(n, intermediate_degree); + raft::copy( + h_knn.data_handle(), + d_knn.data_handle(), + knn_count, + raft::resource::get_cuda_stream(single_gpu_res)); + raft::resource::sync_stream(single_gpu_res); + + auto h_cagra_view = raft::make_host_matrix_view( + h_cagra.data(), n, graph_degree); + cuvs::neighbors::cagra::helpers::optimize( + single_gpu_res, h_knn.view(), h_cagra_view); + const char* optimize_where = "host (public API)"; +#else + auto d_cagra = raft::make_device_matrix( + single_gpu_res, n, graph_degree); + + cuvs::neighbors::cagra::detail::graph::optimize( + single_gpu_res, + d_knn.view(), + d_cagra.view(), + cagraConfig_.guarantee_connectivity); + raft::resource::sync_stream(single_gpu_res); + raft::copy( + h_cagra.data(), + d_cagra.data_handle(), + out_count, + raft::resource::get_cuda_stream(single_gpu_res)); + raft::resource::sync_stream(single_gpu_res); + const char* optimize_where = "device-resident"; +#endif t_now = std::chrono::high_resolution_clock::now(); fprintf(stderr, - " [trainAllNeighbors] optimize: %.2f seconds\n", + " [trainAllNeighbors] optimize (%s): %.2f seconds\n", + optimize_where, std::chrono::duration(t_now - t_phase).count()); t_phase = t_now; - // Phase 4: Store as merged_knngraph_ for copyTo() - merged_knngraph_.resize(n * graph_degree); + merged_knngraph_.resize(out_count); merged_knngraph_degree_ = graph_degree; #pragma omp parallel for - for (idx_t i = 0; i < n * graph_degree; i++) { - merged_knngraph_[i] = static_cast(h_cagra.data_handle()[i]); + for (idx_t i = 0; i < (idx_t)out_count; i++) { + merged_knngraph_[i] = static_cast(h_cagra[i]); } multi_gpu_dataset_ = x; @@ -1112,6 +555,9 @@ void GpuIndexCagra::trainAllNeighbors( this->ntotal = n; t_now = std::chrono::high_resolution_clock::now(); + fprintf(stderr, + " [trainAllNeighbors] D2H + widen: %.2f seconds\n", + std::chrono::duration(t_now - t_phase).count()); fprintf(stderr, " [trainAllNeighbors] Total: %.2f seconds\n", std::chrono::duration(t_now - t_start).count()); @@ -1411,6 +857,121 @@ void GpuIndexCagra::copyTo(faiss::IndexHNSWCagra* index) const { index->init_level0 = true; } +void GpuIndexCagra::buildHnswUpperLevelsGpu_( + faiss::IndexHNSWCagra* index, + int max_lvl) const { + auto& hnsw = index->hnsw; + const idx_t n_train = this->ntotal; + // Levels >= 1 all carry M neighbours; level 0 carries 2*M. + const int upper_degree = (int)hnsw.nb_neighbors(1); + + // BRUTE_FORCE only exists on the multi-GPU path. + auto sub_algo = cagraConfig_.build_algo == graph_build_algo::IVF_PQ + ? faiss::cagra_build_algo::IVF_PQ + : faiss::cagra_build_algo::NN_DESCENT; + + const int sub_intermediate = cagraConfig_.gpu_hnsw_intermediate_degree > 0 + ? (int)cagraConfig_.gpu_hnsw_intermediate_degree + : 2 * upper_degree; + + // Search starts at entry_point and immediately reads its neighbours at + // max_level, so the two must agree. Track the highest populated level + // rather than assuming max_lvl has any nodes. + int top_lvl = 0; + HNSW::storage_idx_t top_entry = 0; + + std::vector ids; + std::vector sub_x; + for (int lvl = 1; lvl <= max_lvl; lvl++) { + ids.clear(); + for (idx_t i = 0; i < n_train; i++) { + if (hnsw.levels[i] - 1 >= lvl) { + ids.push_back(i); + } + } + const idx_t n_lvl = (idx_t)ids.size(); + if (n_lvl < 1) { + continue; + } + top_lvl = lvl; + top_entry = static_cast(ids[0]); + if (n_lvl == 1) { + // Single node: no edges to add, but it is still a valid entry. + continue; + } + + if (n_lvl <= upper_degree + 1) { + // Too small for a meaningful proximity graph: connect all-to-all. + for (idx_t a = 0; a < n_lvl; a++) { + size_t begin, end; + hnsw.neighbor_range(ids[a], lvl, &begin, &end); + size_t slot = begin; + for (idx_t b = 0; b < n_lvl && slot < end; b++) { + if (b != a) { + hnsw.neighbors[slot++] = ids[b]; + } + } + } + continue; + } + + sub_x.resize((size_t)n_lvl * this->d); +#pragma omp parallel for + for (idx_t a = 0; a < n_lvl; a++) { + memcpy(sub_x.data() + (size_t)a * this->d, + multi_gpu_dataset_ + (size_t)ids[a] * this->d, + this->d * sizeof(float)); + } + + auto t0 = std::chrono::high_resolution_clock::now(); + std::vector sub_graph; + { + CuvsCagra sub( + this->resources_.get(), + this->d, + sub_intermediate, + upper_degree, + sub_algo, + cagraConfig_.nn_descent_niter, + /*store_dataset=*/false, + this->metric_type, + this->metric_arg, + INDICES_64_BIT, + std::nullopt, + std::nullopt, + cagraConfig_.refine_rate, + cagraConfig_.gpu_hnsw_guarantee_connectivity); + sub.train(n_lvl, sub_x.data()); + sub_graph = sub.get_knngraph(); + } + + // Map subgraph-local ids back to global ids. +#pragma omp parallel for + for (idx_t a = 0; a < n_lvl; a++) { + size_t begin, end; + hnsw.neighbor_range(ids[a], lvl, &begin, &end); + for (size_t j = begin; j < end; j++) { + idx_t local = sub_graph[(size_t)a * upper_degree + (j - begin)]; + hnsw.neighbors[j] = + (local >= 0 && local < n_lvl) ? ids[local] : -1; + } + } + + fprintf(stderr, + " [copyTo] level %d: %ld nodes, igd=%d gc=%d, " + "CAGRA subgraph %.2f s\n", + lvl, + (long)n_lvl, + sub_intermediate, + (int)cagraConfig_.gpu_hnsw_guarantee_connectivity, + std::chrono::duration( + std::chrono::high_resolution_clock::now() - t0) + .count()); + } + hnsw.entry_point = top_entry; + hnsw.max_level = top_lvl; +} + void GpuIndexCagra::copyToMultiGpu_(faiss::IndexHNSWCagra* index) const { FAISS_ASSERT(this->is_trained && index && !merged_knngraph_.empty()); FAISS_THROW_IF_NOT_MSG( @@ -1446,12 +1007,23 @@ void GpuIndexCagra::copyToMultiGpu_(faiss::IndexHNSWCagra* index) const { auto n_train = this->ntotal; index->init_level0 = false; - if (!index->base_level_only) { - index->add(n_train, multi_gpu_dataset_); - } else { + if (index->base_level_only) { index->hnsw.prepare_level_tab(n_train, false); index->storage->add(n_train, multi_gpu_dataset_); index->ntotal = n_train; + } else if (cagraConfig_.gpu_hnsw_upper_levels) { + auto t0 = std::chrono::high_resolution_clock::now(); + int max_lvl = index->hnsw.prepare_level_tab(n_train, false); + index->storage->add(n_train, multi_gpu_dataset_); + index->ntotal = n_train; + buildHnswUpperLevelsGpu_(index, max_lvl); + fprintf(stderr, + " [copyTo] GPU upper levels: %.2f seconds\n", + std::chrono::duration( + std::chrono::high_resolution_clock::now() - t0) + .count()); + } else { + index->add(n_train, multi_gpu_dataset_); } #pragma omp parallel for diff --git a/faiss/gpu/GpuIndexCagra.h b/faiss/gpu/GpuIndexCagra.h index 6f0ae93be2..0ec56b4e14 100644 --- a/faiss/gpu/GpuIndexCagra.h +++ b/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 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& providers, - std::vector& 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& 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 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>> index_; - /// Multi-GPU state: populated by trainMultiGpu(), used by copyTo() + /// Multi-GPU state: populated by trainAllNeighbors_(), used by copyTo() std::vector merged_knngraph_; idx_t merged_knngraph_degree_ = 0; const float* multi_gpu_dataset_ = nullptr; diff --git a/faiss/gpu/test/bench_approaches.py b/faiss/gpu/test/bench_approaches.py index b62ebd7252..5070f9f698 100644 --- a/faiss/gpu/test/bench_approaches.py +++ b/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 @@ _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 - - -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 +BUILD_ALGOS = ("ivf_pq", "nn_descent", "brute_force") -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 = {} +def cagra_build_algo(name): + return getattr(faiss, f"graph_build_algo_{name.upper()}") - 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) +def index_filename(n): + return f"cagra_hnsw_{n // 1_000_000}M.faiss" - 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 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", + "--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( + "--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="Partition detour counting across GPUs", + 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( - "--build-algo", + "--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 - return - - print(f"Indices 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"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" 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, + f" ef={ef:>4d} recall@{k}={recall[ef]:.4f} " + f"qps={qps[ef]:,.1f} kcyc/q={kcycles.get(ef, 0):.0f}" ) - ) - del index_a, shards_a - print() + return - 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(f"Index will be saved to: {save_dir}") + + 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, + ) + 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("SUMMARY") + print("RESULTS") print("=" * 60) - header = ( - f"{'Approach':<20s} {'Build':>7s} {'Ser':>5s} {'GB':>5s}" - f" {'ef':>4s} {'recall':>7s} {'kcyc/q':>7s}" - ) - 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}" - ) + 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__": diff --git a/faiss/gpu/test/test_cagra.py b/faiss/gpu/test/test_cagra.py index 7cbd07447f..e18a5b6a65 100644 --- a/faiss/gpu/test/test_cagra.py +++ b/faiss/gpu/test/test_cagra.py @@ -248,58 +248,56 @@ def test_IDMapCagra_IP_Int8(self): @unittest.skipIf( "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in" ) -@unittest.skipIf( - faiss.get_num_gpus() < 2, "need at least 2 GPUs for multi-GPU test" -) -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) +class TestCagraConfig(unittest.TestCase): + """Config surface for the multi-GPU build, checkable without 2 GPUs.""" - num_gpus = min(faiss.get_num_gpus(), 4) - resources = faiss.GpuResourcesVector() + def test_multi_gpu_knobs_survive_swig(self): 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) + 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.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) + 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), + ) - cpu_index = faiss.IndexHNSWCagra() - index.copyTo(cpu_index) - self.assertEqual(cpu_index.ntotal, ds.nb) + 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) - cpu_index.hnsw.efSearch = 128 - Dnew, Inew = cpu_index.search(xq, k) + with self.assertRaises(RuntimeError): + index.train(ds.get_database()) - 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) +@unittest.skipIf( + "CUVS" not in faiss.get_compile_options(), "only if cuVS is compiled in" +) +@unittest.skipIf( + faiss.get_num_gpus() < 2, "need at least 2 GPUs for multi-GPU test" +) +class TestMultiGpuCagra(unittest.TestCase): 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 @@ def test_all_neighbors_build(self): 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 @@ def test_all_neighbors_build(self): 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) diff --git a/pyproject-gpu-cuvs.toml b/pyproject-gpu-cuvs.toml index 9e1b1d5d60..21899ede08 100644 --- a/pyproject-gpu-cuvs.toml +++ b/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"