From dcebe5db62d1b09f940aaf397da552023fa68071 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 18:59:10 +0000 Subject: [PATCH 1/2] Add IVF-RaBitQ test for queries that coincide with a list centroid --- cpp/tests/CMakeLists.txt | 2 +- .../ann_ivf_rabitq/test_exact_match.cu | 169 ++++++++++++++++++ 2 files changed, 170 insertions(+), 1 deletion(-) create mode 100644 cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index eee4150eed..66fff7c5c8 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -195,7 +195,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_IVF_RABITQ_TEST - PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu + PATH neighbors/ann_ivf_rabitq/test_float_int64_t.cu neighbors/ann_ivf_rabitq/test_exact_match.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu new file mode 100644 index 0000000000..d0fd00ef95 --- /dev/null +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_exact_match.cu @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::ivf_rabitq { + +struct ivf_rabitq_exact_match_inputs { + uint32_t dim; + uint32_t bits_per_dim; + search_mode mode; +}; + +inline auto operator<<(std::ostream& os, const ivf_rabitq_exact_match_inputs& p) -> std::ostream& +{ + const char* mode = p.mode == search_mode::LUT16 ? "LUT16" + : p.mode == search_mode::LUT32 ? "LUT32" + : p.mode == search_mode::QUANT4 ? "QUANT4" + : "QUANT8"; + return os << "ivf_rabitq_exact_match_inputs {.dim = " << p.dim + << ", .bits_per_dim = " << p.bits_per_dim << ", .mode = " << mode << "}"; +} + +/** + * Queries that coincide with a list centroid. + * + * The dataset has as many distinct rows as there are lists, and each row is repeated many times, so + * each list holds the copies of one row and its centroid is that row. Each query is one of the + * distinct rows, so the probe distance ||q - c||^2 to its own list is zero up to rounding. Search + * must return copies of the query at distance ~0, for every search mode. + */ +class ivf_rabitq_exact_match_test : public ::testing::TestWithParam { + protected: + void run() + { + constexpr int64_t kDistinct = 32; + constexpr int64_t kCopies = 128; + constexpr int64_t kRows = kDistinct * kCopies; + constexpr int64_t kK = 10; + constexpr uint32_t kLists = kDistinct; + + auto ps = GetParam(); + auto dim = static_cast(ps.dim); + auto stream = raft::resource::get_cuda_stream(handle_); + + std::mt19937 gen(42); + std::uniform_real_distribution uniform(1.0f, 2.0f); + auto h_queries = raft::make_host_matrix(kDistinct, dim); + for (int64_t i = 0; i < kDistinct; i++) { + for (int64_t j = 0; j < dim; j++) { + h_queries(i, j) = uniform(gen); + } + } + auto h_dataset = raft::make_host_matrix(kRows, dim); + for (int64_t i = 0; i < kRows; i++) { + for (int64_t j = 0; j < dim; j++) { + h_dataset(i, j) = h_queries(i % kDistinct, j); + } + } + + // The squared distance from each distinct row to the nearest other distinct row. + std::vector gap(kDistinct, std::numeric_limits::max()); + for (int64_t a = 0; a < kDistinct; a++) { + for (int64_t b = 0; b < kDistinct; b++) { + if (a == b) { continue; } + float d = 0; + for (int64_t j = 0; j < dim; j++) { + float diff = h_queries(a, j) - h_queries(b, j); + d += diff * diff; + } + gap[a] = std::min(gap[a], d); + } + } + + auto d_dataset = raft::make_device_matrix(handle_, kRows, dim); + auto d_queries = raft::make_device_matrix(handle_, kDistinct, dim); + raft::copy(d_dataset.data_handle(), h_dataset.data_handle(), h_dataset.size(), stream); + raft::copy(d_queries.data_handle(), h_queries.data_handle(), h_queries.size(), stream); + + index_params ipams; + ipams.n_lists = kLists; + ipams.bits_per_dim = ps.bits_per_dim; + auto idx = build(handle_, ipams, raft::make_const_mdspan(d_dataset.view())); + + search_params sparams; + sparams.n_probes = kLists; + sparams.mode = ps.mode; + auto neighbors = raft::make_device_matrix(handle_, kDistinct, kK); + auto distances = raft::make_device_matrix(handle_, kDistinct, kK); + search(handle_, + sparams, + idx, + raft::make_const_mdspan(d_queries.view()), + neighbors.view(), + distances.view()); + + std::vector h_neighbors(kDistinct * kK); + std::vector h_distances(kDistinct * kK); + raft::copy(h_neighbors.data(), neighbors.data_handle(), h_neighbors.size(), stream); + raft::copy(h_distances.data(), distances.data_handle(), h_distances.size(), stream); + raft::resource::sync_stream(handle_); + + // Every row has kCopies > kK copies, so all kK neighbors must be exact copies of the query. + int64_t n_bad = 0; + for (int64_t q = 0; q < kDistinct; q++) { + bool bad = false; + for (int64_t j = 0; j < kK; j++) { + auto neighbor = h_neighbors[q * kK + j]; + auto dist = h_distances[q * kK + j]; + if (neighbor < 0 || neighbor >= kRows || neighbor % kDistinct != q || + !(std::abs(dist) <= 0.25f * gap[q])) { + bad = true; + } + } + if (bad) { + n_bad++; + auto msg = ::testing::Message() << ps << ": query " << q << " (gap " << gap[q] << "):"; + for (int64_t j = 0; j < kK; j++) { + msg << " " << h_neighbors[q * kK + j] << "(row " << h_neighbors[q * kK + j] % kDistinct + << ", " << h_distances[q * kK + j] << ")"; + } + ADD_FAILURE() << msg; + } + } + EXPECT_EQ(n_bad, 0) << ps << ": " << n_bad << " of " << kDistinct + << " queries did not return copies of themselves"; + } + + raft::resources handle_; +}; + +inline auto exact_match_inputs() -> std::vector +{ + std::vector xs; + for (uint32_t dim : {64, 100, 128}) { + for (uint32_t bits_per_dim : {1, 3}) { + for (auto mode : + {search_mode::LUT16, search_mode::LUT32, search_mode::QUANT4, search_mode::QUANT8}) { + xs.push_back({dim, bits_per_dim, mode}); + } + } + } + return xs; +} + +TEST_P(ivf_rabitq_exact_match_test, build_search) { this->run(); } // NOLINT +INSTANTIATE_TEST_SUITE_P(IvfRabitq, + ivf_rabitq_exact_match_test, + ::testing::ValuesIn(exact_match_inputs())); // NOLINT + +} // namespace cuvs::neighbors::ivf_rabitq From c799bd980b9fc3b69240552bcc5ffaecc9c9afb4 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 19:02:54 +0000 Subject: [PATCH 2/2] Clamp IVF-RaBitQ probe distances to non-negative values --- cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu index e27d34557a..fe2e04f761 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -19,9 +19,11 @@ #include #include #include +#include #include #include #include +#include #include #include @@ -1016,6 +1018,14 @@ void IVFGPU::PrepareClusterSearchInputs( &beta, searcher.get_centroid_distances(), num_centroids); + // Rounding in the expanded form can make ||q - c||^2 slightly negative when the query is on (or + // very near) a centroid. The search kernels take sqrt(||q - c||^2) for the error bound, which is + // NaN for a negative value, and a NaN lower bound drops every vector of that list, so clamp to 0. + raft::linalg::unaryOp(searcher.get_centroid_distances(), + searcher.get_centroid_distances(), + batch_size * num_centroids, + raft::plug_const_op{0.f}, + searcher_stream.get()); // Step 4: select top-nprobe clusters per query auto d_raft_vals = raft::make_device_matrix(searcher_handle, batch_size, nprobe);