From fe5698253f3eb34facbf142b03928c9b193a1a22 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 16:00:35 +0000 Subject: [PATCH 1/2] Derive the filtering rate in filtered multi-partition CAGRA search Single-partition CAGRA search computes the filtering rate from the filter bitset when the caller does not set one, and MULTI_CTA enlarges itopk_size for selective filters accordingly. The multi-partition search skipped this, so a selective per-partition filter searched with the unadjusted itopk_size and could return fewer than k neighbors although every partition held more than k rows passing the filter. cuvs-lucene then routes the query back to per-segment search, which made TestMultiSegmentGPUFilterConcurrency flaky. Derive the rate from the per-partition bitsets, sizing the shared plan for the most selective partition that accepts any row. Also limit topK in TestCuVSGaps.testVectorSearchWithAlternatingDocuments to the number of documents that have vectors. --- cpp/src/neighbors/cagra.cuh | 30 +++- .../ann_cagra/test_float_uint32_t.cu | 148 ++++++++++++++++++ .../com/nvidia/cuvs/lucene/TestCuVSGaps.java | 4 +- 3 files changed, 180 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 0e1454d987..c781d1e427 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -606,10 +606,38 @@ void search( cuvs::neighbors::filtering::none_sample_filter>( res, params, indices, queries, partition_ids, neighbors, distances, partition_bitsets); } else { + // As in the single-partition bitset path, derive the filtering rate from the bitsets when the + // caller did not set one, so that the search plan enlarges itopk_size for selective filters. + // Without this, a selective filter leaves too few valid candidates and the search can return + // fewer than k neighbors. One plan serves every partition, so size it for the most selective + // partition that still accepts at least one row; a partition that accepts nothing contributes + // no neighbors whatever the itopk_size, and an unfiltered partition has a rate of zero. + search_params params_copy = params; + if (params.filtering_rate < 0.0) { + float filtering_rate = 0.0f; + for (size_t i = 0; i < partition_bitsets.size() && i < indices.size(); i++) { + const auto& v = partition_bitsets[i]; + if (v.data() == nullptr || v.size() == 0) { continue; } + const int64_t n_rows = indices[i]->dataset().n_rows(); + if (n_rows == 0) { continue; } + // Count only this partition's rows; the bitset may be padded past n_rows. + const auto num_set_bits = + cuvs::core::bitset_view(const_cast(v.data()), + std::min(n_rows, v.size())) + .count(res); + if (num_set_bits == 0) { continue; } + filtering_rate = + std::max(filtering_rate, static_cast(n_rows - num_set_bits) / n_rows); + } + const float min_filtering_rate = 0.0; + const float max_filtering_rate = 0.999; + params_copy.filtering_rate = + std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate); + } using bitset_filter_t = cuvs::neighbors::filtering::bitset_filter; cagra::detail::search_multi_partition( res, - params, + params_copy, indices, queries, partition_ids, diff --git a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu index 5dccce21da..2b445df7bc 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -7,6 +7,8 @@ #include "../ann_cagra.cuh" +#include +#include #include #include @@ -135,6 +137,152 @@ TEST(AnnCagraMultiPartition, MultiKernelRejected) search_params); } +namespace { +// Runs a filtered multi-partition MULTI_CTA search where every partition keeps only ~4% of its rows +// (still more than k each) and reports the number of queries that returned fewer than k neighbors +// and the recall against an exact search over the kept rows. +void search_with_selective_filter(float filtering_rate, int& short_queries, double& recall) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int num_partitions = 4, part_size = 500, dim = 128, n_queries = 500, k = 10; + // Keep one row in `keep_every` (~4% of each partition, i.e. ~21 rows per partition). + constexpr int keep_every = 24; + constexpr int n_rows = num_partitions * part_size; + + rmm::device_uvector database(static_cast(n_rows) * dim, stream); + rmm::device_uvector queries(static_cast(n_queries) * dim, stream); + raft::random::RngState r(1234ULL); + InitDataset(handle, database.data(), n_rows, dim, cuvs::distance::DistanceType::L2Expanded, r); + InitDataset(handle, queries.data(), n_queries, dim, cuvs::distance::DistanceType::L2Expanded, r); + raft::resource::sync_stream(handle); + + cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + index_params.graph_build_params = graph_build_params::nn_descent_params( + index_params.intermediate_graph_degree, index_params.metric); + + std::vector> part_indices; + // An index only holds a view, so any padded copy has to outlive it. + std::vector> part_padded; + part_padded.reserve(num_partitions); + std::vector> part_bitsets; + part_bitsets.reserve(num_partitions); + std::vector> partition_bitsets; + std::vector removed_local; + for (int64_t i = 0; i < part_size; i++) { + if (i % keep_every != 0) { removed_local.push_back(i); } + } + auto removed = + raft::make_device_vector(handle, static_cast(removed_local.size())); + raft::update_device(removed.data_handle(), removed_local.data(), removed_local.size(), stream); + for (int p = 0; p < num_partitions; p++) { + auto view = raft::make_device_matrix_view( + database.data() + static_cast(p) * part_size * dim, part_size, dim); + part_padded.emplace_back(handle, view); + auto const& padded = part_padded.back().view; + part_indices.push_back(cagra::build(handle, index_params, padded)); + auto& part_index = part_indices.back(); + part_index = cagra::update_dataset(handle, std::move(part_index), padded); + part_bitsets.emplace_back(handle, removed.view(), part_size); + partition_bitsets.push_back(part_bitsets.back().view()); + } + std::vector*> index_ptrs; + for (auto& idx : part_indices) { + index_ptrs.push_back(&idx); + } + + const size_t out_size = static_cast(n_queries) * k; + rmm::device_uvector partition_ids_dev(out_size, stream); + rmm::device_uvector neighbors_dev(out_size, stream); + rmm::device_uvector distances_dev(out_size, stream); + + // A small itopk_size, as used for small k, on the MULTI_CTA algorithm that AUTO picks for a + // handful of queries over a handful of partitions. + cagra::search_params search_params; + search_params.algo = search_algo::MULTI_CTA; + search_params.itopk_size = k; + search_params.filtering_rate = filtering_rate; + + cagra::search( + handle, + search_params, + index_ptrs, + raft::make_device_matrix_view(queries.data(), n_queries, dim), + raft::make_device_matrix_view(partition_ids_dev.data(), n_queries, k), + raft::make_device_matrix_view(neighbors_dev.data(), n_queries, k), + raft::make_device_matrix_view(distances_dev.data(), n_queries, k), + partition_bitsets); + + std::vector partition_ids(out_size); + std::vector neighbors(out_size); + std::vector distances(out_size); + std::vector host_database(static_cast(n_rows) * dim); + std::vector host_queries(static_cast(n_queries) * dim); + raft::update_host(partition_ids.data(), partition_ids_dev.data(), out_size, stream); + raft::update_host(neighbors.data(), neighbors_dev.data(), out_size, stream); + raft::update_host(distances.data(), distances_dev.data(), out_size, stream); + raft::update_host(host_database.data(), database.data(), host_database.size(), stream); + raft::update_host(host_queries.data(), queries.data(), host_queries.size(), stream); + raft::resource::sync_stream(handle); + + // Exact top-k over the rows that pass the filter, for recall. + std::vector kept_rows; + for (int64_t g = 0; g < n_rows; g++) { + if ((g % part_size) % keep_every == 0) { kept_rows.push_back(g); } + } + ASSERT_GT(kept_rows.size(), static_cast(k)); + + short_queries = 0; + size_t matches = 0; + for (int q = 0; q < n_queries; q++) { + std::vector> exact; + for (auto g : kept_rows) { + float d = 0; + for (int j = 0; j < dim; j++) { + const float diff = host_queries[q * dim + j] - host_database[g * dim + j]; + d += diff * diff; + } + exact.emplace_back(d, g); + } + std::partial_sort(exact.begin(), exact.begin() + k, exact.end()); + std::vector expected; + for (int j = 0; j < k; j++) { + expected.push_back(exact[j].second); + } + + int valid = 0; + for (int j = 0; j < k; j++) { + const size_t i = static_cast(q) * k + j; + if (distances[i] == std::numeric_limits::max()) { continue; } + ASSERT_LT(partition_ids[i], static_cast(num_partitions)); + ASSERT_LT(neighbors[i], static_cast(part_size)); + ASSERT_EQ(neighbors[i] % keep_every, 0u) << "filtered-out row returned"; + valid++; + const int64_t g = static_cast(partition_ids[i]) * part_size + neighbors[i]; + if (std::find(expected.begin(), expected.end(), g) != expected.end()) { matches++; } + } + if (valid < k) { short_queries++; } + } + recall = static_cast(matches) / out_size; +} +} // namespace + +// A selective per-partition filter must not leave the search short of neighbors. Single-partition +// search derives the filtering rate from the bitset and enlarges itopk_size accordingly; the +// multi-partition search must do the same, or MULTI_CTA can return fewer than k valid neighbors +// even though every partition holds more than k rows that pass the filter. +TEST(AnnCagraMultiPartition, SelectiveFilterReturnsKNeighbors) +{ + int short_queries = 0; + double recall = 0; + search_with_selective_filter(-1.0f, short_queries, recall); + EXPECT_EQ(short_queries, 0) << "queries returning fewer than k neighbors"; + EXPECT_GE(recall, 0.95); +} + // The shared plan descriptor and the cross-partition select_k direction are derived from // indices[0], so all partitions must share one metric; a mismatch must be rejected. TEST(AnnCagraMultiPartition, MixedMetricRejected) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java index 31958ef0db..e9bdfea9aa 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSGaps.java @@ -114,7 +114,9 @@ public void testVectorSearchWithAlternatingDocuments() throws IOException { // Use the first vector (from document 0) as query float[] queryVector = dataset[0]; - int topK = random.nextInt(5, TOP_K_LIMIT); + // Only even-numbered documents have vectors, so never ask for more results than exist. + int vectorCount = (datasetSize + 1) / 2; + int topK = Math.min(random.nextInt(5, TOP_K_LIMIT), vectorCount); GPUKnnFloatVectorQuery query = new GPUKnnFloatVectorQuery("vector", queryVector, topK, null, topK, 1); From 0601e226bda5da2b04ce704c145c5ee03638f774 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Thu, 15 Oct 2026 17:45:16 +0000 Subject: [PATCH 2/2] Factor the CAGRA filtering-rate clamp into a helper --- cpp/src/neighbors/cagra.cuh | 28 +++++++++++++--------------- 1 file changed, 13 insertions(+), 15 deletions(-) diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index c781d1e427..0ad4f2647f 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -420,6 +420,16 @@ void search_with_filtering(raft::resources const& res, res, params, idx, queries, neighbors, distances, sample_filter); } +namespace detail { +/** Clamp a filtering rate derived from a filter to the range the search plan accepts. */ +inline float clamp_filtering_rate(float filtering_rate) +{ + const float min_filtering_rate = 0.0f; + const float max_filtering_rate = 0.999f; + return std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate); +} +} // namespace detail + template ( @@ -503,13 +510,7 @@ void search(raft::resources const& res, dynamic_cast(sample_filter_ref); search_params params_copy = params; if (params.filtering_rate < 0.0) { - const float min_filtering_rate = 0.0f; - const float max_filtering_rate = 0.999f; - params_copy.filtering_rate = - sample_filter.filtering_rate < 0.0f - ? 0.0f - : std::min(std::max(sample_filter.filtering_rate, min_filtering_rate), - max_filtering_rate); + params_copy.filtering_rate = detail::clamp_filtering_rate(sample_filter.filtering_rate); } auto sample_filter_copy = sample_filter; return search_with_filtering( @@ -629,10 +630,7 @@ void search( filtering_rate = std::max(filtering_rate, static_cast(n_rows - num_set_bits) / n_rows); } - const float min_filtering_rate = 0.0; - const float max_filtering_rate = 0.999; - params_copy.filtering_rate = - std::min(std::max(filtering_rate, min_filtering_rate), max_filtering_rate); + params_copy.filtering_rate = detail::clamp_filtering_rate(filtering_rate); } using bitset_filter_t = cuvs::neighbors::filtering::bitset_filter; cagra::detail::search_multi_partition(