From dd45b5dab614441753e8bed926bfcb7ae357eb29 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 15:41:08 +0000 Subject: [PATCH 01/27] Make k-means centroid updates deterministic Sum the rows and weights of each cluster in a fixed order instead of with atomics, and use cub::DeviceReduce for the centroid shift and the weight sum. Identical fits now return bitwise-identical centroids, so the convergence check no longer stops at a different iteration from run to run. --- .../cluster/detail/kmeans_centroid_sums.cuh | 348 ++++++++++++++++++ cpp/src/cluster/detail/kmeans_common.cuh | 86 +++-- 2 files changed, 398 insertions(+), 36 deletions(-) create mode 100644 cpp/src/cluster/detail/kmeans_centroid_sums.cuh diff --git a/cpp/src/cluster/detail/kmeans_centroid_sums.cuh b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh new file mode 100644 index 0000000000..9e0fdb378b --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh @@ -0,0 +1,348 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +#include +#include +#include + +/** + * Deterministic per-cluster weighted row sums for the k-means centroid update. + * + * Accumulating the rows of each cluster with atomics (as raft::linalg::reduce_rows_by_key does) + * makes the floating-point summation order, and therefore the centroids, vary from run to run. + * Through the convergence check this can change the number of Lloyd iterations, so two identical + * fits can return noticeably different centroids. Here the summation order is a function of the + * input only: + * + * 1. Row indices are stably sorted by label, so each cluster becomes a contiguous segment. + * 2. Each segment is cut into pieces of at most `kRowsPerPiece` rows. One thread block reduces a + * piece with a fixed assignment of rows to threads and a fixed-order combine. + * 3. One thread block per cluster adds up that cluster's piece sums in the same fixed way, and + * adds the result to the output. + */ +namespace cuvs::cluster::kmeans::detail::sums_by_label { + +constexpr int kBlockThreads = 256; +constexpr int kRowsPerPiece = 256; + +/** Number of threads along the columns: smallest power of two >= n_cols (at most kBlockThreads). */ +inline int column_threads(int64_t n_cols) +{ + int tc = 1; + while (tc < n_cols && tc < kBlockThreads) { + tc <<= 1; + } + return tc; +} + +/** + * Sums `values(i, j)` over i in [begin, end) for each column j in [0, n_cols), using a fixed + * assignment of items to threads and a fixed-order combine, so the result does not depend on + * scheduling. The result for column j is passed to `write(j, sum)` by a single thread. + */ +template +__device__ void fixed_order_column_sums( + IndexT begin, IndexT end, IndexT n_cols, int col_threads, ValueOp values, WriteOp write) +{ + __shared__ DataT smem[kBlockThreads]; + const int lanes = kBlockThreads / col_threads; + const int tc = threadIdx.x % col_threads; + const int lane = threadIdx.x / col_threads; + for (IndexT col0 = 0; col0 < n_cols; col0 += col_threads) { + const IndexT j = col0 + tc; + DataT acc = DataT{0}; + if (j < n_cols) { +#pragma unroll 4 + for (IndexT i = begin + lane; i < end; i += lanes) { + acc += values(i, j); + } + } + smem[threadIdx.x] = acc; + __syncthreads(); + if (lane == 0 && j < n_cols) { + DataT sum = smem[tc]; + for (int l = 1; l < lanes; ++l) { + sum += smem[l * col_threads + tc]; + } + write(j, sum); + } + __syncthreads(); + } +} + +template +RAFT_KERNEL init_sort_kernel(LabelsIterator labels, IndexT n_rows, uint32_t* keys, IndexT* rows) +{ + for (IndexT i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n_rows; + i += static_cast(blockDim.x) * gridDim.x) { + keys[i] = static_cast(labels[i]); + rows[i] = i; + } +} + +/** Finds each cluster's segment of the sorted labels and the number of pieces it is cut into. */ +template +RAFT_KERNEL segment_kernel(const uint32_t* sorted_keys, + IndexT n_rows, + IndexT n_clusters, + IndexT* cluster_offsets, + IndexT* piece_counts) +{ + auto lower_bound = [=](IndexT c) { + IndexT lo = 0, hi = n_rows; + while (lo < hi) { + IndexT mid = lo + (hi - lo) / 2; + if (static_cast(sorted_keys[mid]) < c) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + }; + for (IndexT c = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; c <= n_clusters; + c += static_cast(blockDim.x) * gridDim.x) { + IndexT begin = lower_bound(c); + cluster_offsets[c] = begin; + if (c < n_clusters) { + IndexT end = lower_bound(c + 1); + piece_counts[c] = raft::ceildiv(end - begin, kRowsPerPiece); + } else { + piece_counts[c] = 0; + } + } +} + +template +RAFT_KERNEL piece_sums_kernel(const DataT* X, + IndexT n_cols, + const DataT* weights, + const IndexT* sorted_rows, + const IndexT* cluster_offsets, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + DataT* piece_sums) +{ + const IndexT n_pieces = piece_offsets[n_clusters]; + for (IndexT piece = blockIdx.x; piece < n_pieces; piece += gridDim.x) { + // The cluster of this piece is the last c with piece_offsets[c] <= piece. + IndexT lo = 0, hi = n_clusters; + while (hi - lo > 1) { + IndexT mid = lo + (hi - lo) / 2; + if (piece_offsets[mid] <= piece) { + lo = mid; + } else { + hi = mid; + } + } + const IndexT c = lo; + const IndexT begin = cluster_offsets[c] + (piece - piece_offsets[c]) * kRowsPerPiece; + const IndexT end = std::min(begin + kRowsPerPiece, cluster_offsets[c + 1]); + DataT* out = piece_sums + static_cast(piece) * (n_cols + 1); + fixed_order_column_sums( + begin, + end, + n_cols, + col_threads, + [=](IndexT i, IndexT j) { + const IndexT row = sorted_rows[i]; + const DataT x = X[static_cast(row) * n_cols + j]; + return weights != nullptr ? x * weights[row] : x; + }, + [=](IndexT j, DataT sum) { out[j] = sum; }); + // The sum of weights is stored after the column sums. + fixed_order_column_sums( + begin, + end, + IndexT{1}, + 1, + [=](IndexT i, IndexT) { return weights != nullptr ? weights[sorted_rows[i]] : DataT{1}; }, + [=](IndexT, DataT sum) { out[n_cols] = sum; }); + } +} + +template +RAFT_KERNEL cluster_sums_kernel(const DataT* piece_sums, + IndexT n_cols, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + bool reset_sums, + DataT* centroid_sums, + DataT* weight_per_cluster) +{ + for (IndexT c = blockIdx.x; c < n_clusters; c += gridDim.x) { + const IndexT p_begin = piece_offsets[c]; + const IndexT p_end = piece_offsets[c + 1]; + const IndexT stride = n_cols + 1; + DataT* out = centroid_sums + static_cast(c) * n_cols; + fixed_order_column_sums( + p_begin, + p_end, + n_cols, + col_threads, + [=](IndexT p, IndexT j) { return piece_sums[static_cast(p) * stride + j]; }, + [=](IndexT j, DataT sum) { out[j] = reset_sums ? sum : out[j] + sum; }); + fixed_order_column_sums( + p_begin, + p_end, + IndexT{1}, + 1, + [=](IndexT p, IndexT) { return piece_sums[static_cast(p) * stride + n_cols]; }, + [=](IndexT, DataT sum) { + weight_per_cluster[c] = reset_sums ? sum : weight_per_cluster[c] + sum; + }); + } +} + +inline size_t align_bytes(size_t bytes) { return raft::alignTo(bytes, 256); } + +/** + * @brief Deterministic weighted sum of the rows of X per label, and of the weights per label. + * + * centroid_sums[c, :] (+)= sum_{i : labels[i] == c} weights[i] * X[i, :] + * weight_per_cluster[c] (+)= sum_{i : labels[i] == c} weights[i] + * + * The result is bitwise reproducible for identical inputs on the same device. + * + * @param[in] handle RAFT resources handle + * @param[in] X Row-major input [n_rows x n_cols] + * @param[in] weights Per-row weights [n_rows] + * @param[in] labels Label of each row, in [0, n_clusters) + * @param[in] n_rows Number of rows of X + * @param[in] n_cols Number of columns of X + * @param[in] n_clusters Number of labels + * @param[inout] centroid_sums Row-major output [n_clusters x n_cols] + * @param[inout] weight_per_cluster Output [n_clusters] + * @param[inout] workspace Scratch buffer, resized as needed + * @param[in] reset_sums Overwrite the outputs if true, add to them otherwise + */ +template +void weighted_sums_by_label(raft::resources const& handle, + const DataT* X, + const DataT* weights, + LabelsIterator labels, + IndexT n_rows, + IndexT n_cols, + IndexT n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); + if (n_clusters <= 0) { return; } + + int end_bit = 1; + while (end_bit < 32 && (uint64_t{1} << end_bit) < static_cast(n_clusters)) { + ++end_bit; + } + const IndexT max_pieces = raft::ceildiv(n_rows, kRowsPerPiece) + n_clusters; + + // Workspace layout + size_t sort_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(nullptr, + sort_bytes, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + n_rows, + 0, + end_bit, + stream)); + size_t scan_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + scan_bytes, + static_cast(nullptr), + static_cast(nullptr), + n_clusters + 1, + stream)); + const size_t keys_bytes = align_bytes(sizeof(uint32_t) * n_rows); + const size_t rows_bytes = align_bytes(sizeof(IndexT) * n_rows); + const size_t offsets_bytes = align_bytes(sizeof(IndexT) * (n_clusters + 1)); + const size_t pieces_bytes = + align_bytes(sizeof(DataT) * static_cast(max_pieces) * (n_cols + 1)); + const size_t temp_bytes = align_bytes(std::max(sort_bytes, scan_bytes)); + workspace.resize(2 * keys_bytes + 2 * rows_bytes + 3 * offsets_bytes + pieces_bytes + temp_bytes, + stream); + + char* ptr = workspace.data(); + auto take = [&ptr](size_t bytes) { + char* p = ptr; + ptr += bytes; + return p; + }; + auto* keys_in = reinterpret_cast(take(keys_bytes)); + auto* keys_out = reinterpret_cast(take(keys_bytes)); + auto* rows_in = reinterpret_cast(take(rows_bytes)); + auto* rows_out = reinterpret_cast(take(rows_bytes)); + auto* cluster_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_counts = reinterpret_cast(take(offsets_bytes)); + auto* piece_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_sums = reinterpret_cast(take(pieces_bytes)); + void* temp = take(temp_bytes); + + if (n_rows > 0) { + const int grid = static_cast( + std::min(raft::ceildiv(n_rows, kBlockThreads), int64_t{65536})); + init_sort_kernel<<>>(labels, n_rows, keys_in, rows_in); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + temp, sort_bytes, keys_in, keys_out, rows_in, rows_out, n_rows, 0, end_bit, stream)); + } + { + const int grid = static_cast( + std::min(raft::ceildiv(n_clusters + 1, kBlockThreads), int64_t{65536})); + segment_kernel<<>>( + keys_out, n_rows, n_clusters, cluster_offsets, piece_counts); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + temp, scan_bytes, piece_counts, piece_offsets, n_clusters + 1, stream)); + + const int col_threads = column_threads(n_cols); + if (max_pieces > 0) { + const int grid = static_cast(std::min(max_pieces, int64_t{65536})); + piece_sums_kernel<<>>(X, + n_cols, + weights, + rows_out, + cluster_offsets, + piece_offsets, + n_clusters, + col_threads, + piece_sums); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + { + const int grid = static_cast(std::min(n_clusters, int64_t{65536})); + cluster_sums_kernel<<>>(piece_sums, + n_cols, + piece_offsets, + n_clusters, + col_threads, + reset_sums, + centroid_sums, + weight_per_cluster); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + +} // namespace cuvs::cluster::kmeans::detail::sums_by_label diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 78edf8759e..ac00770d58 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -6,6 +6,7 @@ #include "../../distance/distance.cuh" #include "../../distance/top_1_nn.cuh" +#include "kmeans_centroid_sums.cuh" #include #include #include @@ -121,6 +122,18 @@ struct WeightedDistanceOp { } }; +template +struct SquaredDiffOp { + const DataT* a; + const DataT* b; + + __device__ DataT operator()(IndexT offset) const + { + DataT diff = a[offset] - b[offset]; + return diff * diff; + } +}; + template void copyClusterLabels(raft::resources const& handle, const MinClusterAndDistanceResult& result, @@ -194,8 +207,17 @@ void weightSum( DataT wt_sum_h = DataT{0}; if constexpr (raft::is_device_mdspan_v) { - raft::linalg::mapThenSumReduce( - d_wt_sum.data_handle(), n_samples, raft::identity_op{}, stream.get(), weight.data_handle()); + // cub::DeviceReduce is run-to-run deterministic, unlike the atomics in mapThenSumReduce. + size_t temp_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceReduce::Sum( + nullptr, temp_bytes, weight.data_handle(), d_wt_sum.data_handle(), n_samples, stream.get())); + rmm::device_uvector temp(temp_bytes, stream); + RAFT_CUDA_TRY(cub::DeviceReduce::Sum(temp.data(), + temp_bytes, + weight.data_handle(), + d_wt_sum.data_handle(), + n_samples, + stream.get())); if (check_positive) { raft::copy(&wt_sum_h, d_wt_sum.data_handle(), 1, stream); raft::resource::sync_stream(handle); @@ -600,8 +622,11 @@ void countSamplesInCluster(raft::resources const& handle, * @brief Compute centroid adjustments (weighted sums and counts per cluster) * * This helper function computes: - * 1. Weighted sum of samples per cluster using reduce_rows_by_key - * 2. Sum of weights per cluster using reduce_cols_by_key + * 1. Weighted sum of samples per cluster + * 2. Sum of weights per cluster + * + * Both use a fixed summation order (see sums_by_label::weighted_sums_by_label), so that + * identical fits give bitwise-identical centroids. * * @tparam DataT Data type for samples and weights * @tparam IndexT Index type @@ -632,31 +657,17 @@ void compute_centroid_adjustments( rmm::device_uvector& workspace, bool reset_sums = true) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - auto n_samples = X.extent(0); - - workspace.resize(n_samples, stream); - - raft::linalg::reduce_rows_by_key(X.data_handle(), - X.extent(1), - cluster_labels, - sample_weights.data_handle(), - workspace.data(), - X.extent(0), - X.extent(1), - n_clusters, - centroid_sums.data_handle(), - stream, - reset_sums); - - raft::linalg::reduce_cols_by_key(sample_weights.data_handle(), - cluster_labels, - weight_per_cluster.data_handle(), - static_cast(1), - static_cast(n_samples), - n_clusters, - stream, - reset_sums); + sums_by_label::weighted_sums_by_label(handle, + X.data_handle(), + sample_weights.data_handle(), + cluster_labels, + X.extent(0), + X.extent(1), + n_clusters, + centroid_sums.data_handle(), + weight_per_cluster.data_handle(), + workspace, + reset_sums); } /** * @brief Finalize centroids by dividing accumulated sums by counts. @@ -715,13 +726,16 @@ void compute_centroid_shift(raft::resources const& handle, raft::device_matrix_view new_centroids, raft::device_scalar_view sqrd_norm_out) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - raft::linalg::mapThenSumReduce(sqrd_norm_out.data_handle(), - old_centroids.size(), - raft::sqdiff_op{}, - stream, - old_centroids.data_handle(), - new_centroids.data_handle()); + // cub::DeviceReduce is run-to-run deterministic, unlike the atomics in mapThenSumReduce. + rmm::device_uvector workspace(0, raft::resource::get_cuda_stream(handle)); + computeClusterCost( + handle, + cuda::counting_iterator(0), + static_cast(old_centroids.size()), + workspace, + sqrd_norm_out, + SquaredDiffOp{old_centroids.data_handle(), new_centroids.data_handle()}, + raft::add_op{}); } /** From 11b7b47eda963491e713598aac29c7dccc055c39 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 15:51:26 +0000 Subject: [PATCH 02/27] Test that repeated k-means fits are bitwise identical --- cpp/tests/cluster/kmeans.cu | 89 +++++++++++++++++++++++++++++++++++++ 1 file changed, 89 insertions(+) diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 23fc248917..a6e5b4c0f6 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -23,6 +23,7 @@ #include +#include #include #include @@ -495,6 +496,82 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(handle, n_clusters, n_features); + raft::random::RngState rng(p.rng_state.seed); + raft::random::uniform(handle, rng, d_init.data_handle(), n_clusters * n_features, T(-1), T(1)); + + auto d_sw = d_sw_view(); + std::optional> h_sw = std::nullopt; + auto h_sample_weight = raft::make_host_vector(testparams.weighted ? n_samples : 0); + if (testparams.weighted) { + std::fill_n(h_sample_weight.data_handle(), n_samples, T(1)); + h_sw = std::make_optional(raft::make_const_mdspan(h_sample_weight.view())); + } + + constexpr int n_runs = 3; + for (int host : {0, 1}) { + std::vector> centroids(n_runs, std::vector(n_clusters * n_features)); + std::vector inertias(n_runs); + std::vector n_iters(n_runs); + for (int r = 0; r < n_runs; ++r) { + auto d_c = raft::make_device_matrix(handle, n_clusters, n_features); + raft::copy(d_c.data_handle(), d_init.data_handle(), n_clusters * n_features, stream); + T inertia = 0; + if (host) { + auto pb = p; + pb.device_buffer_samples = testparams.device_buffer_samples; + int64_t n_iter = 0; + cuvs::cluster::kmeans::fit( + handle, + pb, + raft::make_const_mdspan(h_X->view()), + h_sw, + raft::make_device_matrix_view(d_c.data_handle(), n_clusters, n_features), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } else { + int n_iter = 0; + cuvs::cluster::kmeans::fit(handle, + p, + raft::make_const_mdspan(d_X->view()), + d_sw, + d_c.view(), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } + raft::copy(centroids[r].data(), d_c.data_handle(), n_clusters * n_features, stream); + raft::resource::sync_stream(handle, stream); + inertias[r] = inertia; + } + for (int r = 1; r < n_runs; ++r) { + EXPECT_EQ(n_iters[0], n_iters[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ(inertias[0], inertias[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ( + 0, std::memcmp(centroids[0].data(), centroids[r].data(), centroids[0].size() * sizeof(T))) + << (host ? "host" : "device") << " run " << r << ": centroids differ"; + } + } + } + T fitHostWithInitSize(int64_t init_size_value) { int n_features = testparams.n_col; @@ -709,6 +786,18 @@ TEST_P(KmeansFitBatchedTestD, Result) runZeroCost(); } +TEST_P(KmeansFitBatchedTestF, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + +TEST_P(KmeansFitBatchedTestD, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + INSTANTIATE_TEST_CASE_P(KmeansFitBatchedTests, KmeansFitBatchedTestF, ::testing::ValuesIn(batched_inputsf2)); From d7de50bed08b9540bf6ca19a535b6cdfce3f49f1 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 15:16:57 +0000 Subject: [PATCH 03/27] Test searching an NRT reader whose segments are being merged --- .../TestSearchSegmentsOpenedForMerge.java | 169 ++++++++++++++++++ 1 file changed, 169 insertions(+) create mode 100644 java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java new file mode 100644 index 0000000000..679eb7c56e --- /dev/null +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestSearchSegmentsOpenedForMerge.java @@ -0,0 +1,169 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs.lucene; + +import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; +import static org.apache.lucene.tests.util.TestUtil.alwaysKnnVectorsFormat; + +import java.io.IOException; +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import org.apache.lucene.document.Document; +import org.apache.lucene.document.Field; +import org.apache.lucene.document.KnnFloatVectorField; +import org.apache.lucene.document.StringField; +import org.apache.lucene.index.CodecReader; +import org.apache.lucene.index.ConcurrentMergeScheduler; +import org.apache.lucene.index.DirectoryReader; +import org.apache.lucene.index.FilterMergePolicy; +import org.apache.lucene.index.IndexWriter; +import org.apache.lucene.index.IndexWriterConfig; +import org.apache.lucene.index.MergeTrigger; +import org.apache.lucene.index.NoMergePolicy; +import org.apache.lucene.index.SegmentCommitInfo; +import org.apache.lucene.index.SegmentInfos; +import org.apache.lucene.index.VectorSimilarityFunction; +import org.apache.lucene.search.IndexSearcher; +import org.apache.lucene.search.KnnFloatVectorQuery; +import org.apache.lucene.search.Query; +import org.apache.lucene.search.TopDocs; +import org.apache.lucene.store.Directory; +import org.apache.lucene.tests.util.LuceneTestCase; +import org.apache.lucene.tests.util.LuceneTestCase.SuppressSysoutChecks; +import org.junit.BeforeClass; +import org.junit.Test; + +/** + * IndexWriter pools the {@link org.apache.lucene.index.SegmentReader} it opens to merge a segment + * and hands that same reader (or a clone sharing its core) to near-real-time readers opened while + * the merge is still running. Searching such a reader must work even though its vectors reader was + * opened with a merge {@link org.apache.lucene.store.IOContext}. + */ +@SuppressSysoutChecks(bugUrl = "") +public class TestSearchSegmentsOpenedForMerge extends LuceneTestCase { + + private static final int NUM_SEGMENTS = 3; + private static final int DOCS_PER_SEGMENT = 64; + private static final int DIMENSION = 16; + + @BeforeClass + public static void beforeClass() { + assumeTrue("cuVS not supported", isSupported()); + } + + /** Forced merge of all segments that blocks once its input readers are open. */ + private static final class BlockingMergePolicy extends FilterMergePolicy { + final CountDownLatch mergeReadersOpened = new CountDownLatch(1); + final CountDownLatch releaseMerge = new CountDownLatch(1); + + BlockingMergePolicy() { + super(NoMergePolicy.INSTANCE); + } + + @Override + public MergeSpecification findForcedMerges( + SegmentInfos segmentInfos, + int maxSegmentCount, + Map segmentsToMerge, + MergeContext mergeContext) { + List segments = new ArrayList<>(); + for (SegmentCommitInfo info : segmentInfos) { + if (segmentsToMerge.containsKey(info) + && !mergeContext.getMergingSegments().contains(info)) { + segments.add(info); + } + } + if (segments.size() < 2) { + return null; + } + MergeSpecification spec = new MergeSpecification(); + spec.add( + new OneMerge(segments) { + @Override + public CodecReader wrapForMerge(CodecReader reader) throws IOException { + // Called by IndexWriter#mergeMiddle once the input readers have been opened (with a + // merge IOContext) and pooled. + mergeReadersOpened.countDown(); + try { + if (!releaseMerge.await(5, TimeUnit.MINUTES)) { + throw new IOException("merge was never released"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IOException(e); + } + return reader; + } + }); + return spec; + } + + @Override + public MergeSpecification findMerges( + MergeTrigger mergeTrigger, SegmentInfos segmentInfos, MergeContext mergeContext) { + return null; + } + } + + @Test + public void testSearchNRTReaderDuringMerge() throws Exception { + BlockingMergePolicy mergePolicy = new BlockingMergePolicy(); + IndexWriterConfig config = + new IndexWriterConfig() + .setCodec(alwaysKnnVectorsFormat(new CuVS2510GPUVectorsFormat())) + .setMergePolicy(mergePolicy) + .setMergeScheduler(new ConcurrentMergeScheduler()) + .setMaxBufferedDocs(IndexWriterConfig.DISABLE_AUTO_FLUSH) + .setRAMBufferSizeMB(64); + + float[][] vectors = new float[NUM_SEGMENTS * DOCS_PER_SEGMENT][]; + try (Directory dir = newDirectory(); + IndexWriter writer = new IndexWriter(dir, config)) { + for (int seg = 0; seg < NUM_SEGMENTS; seg++) { + for (int i = 0; i < DOCS_PER_SEGMENT; i++) { + int id = seg * DOCS_PER_SEGMENT + i; + vectors[id] = new float[DIMENSION]; + for (int d = 0; d < DIMENSION; d++) { + vectors[id][d] = random().nextFloat(); + } + Document doc = new Document(); + doc.add(new StringField("id", Integer.toString(id), Field.Store.YES)); + doc.add( + new KnnFloatVectorField("vector", vectors[id], VectorSimilarityFunction.EUCLIDEAN)); + writer.addDocument(doc); + } + writer.commit(); + } + + try { + writer.forceMerge(1, false); + assertTrue( + "merge did not start", mergePolicy.mergeReadersOpened.await(5, TimeUnit.MINUTES)); + + // The segments being merged are still live, so this NRT reader shares the cores the merge + // opened. + try (DirectoryReader reader = DirectoryReader.open(writer)) { + assertEquals(NUM_SEGMENTS, reader.leaves().size()); + IndexSearcher searcher = new IndexSearcher(reader); + int topK = 10; + float[] target = vectors[random().nextInt(vectors.length)]; + List queries = + List.of( + new GPUKnnFloatVectorQuery("vector", target, topK, null, topK, 1), + new KnnFloatVectorQuery("vector", target, topK)); + for (Query query : queries) { + TopDocs hits = searcher.search(query, topK); + assertEquals(query.toString(), topK, hits.scoreDocs.length); + } + } + } finally { + mergePolicy.releaseMerge.countDown(); + } + } + } +} From 04ab2a81c05333e40b514641750398ebc7981443 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 15:17:39 +0000 Subject: [PATCH 04/27] Load GPU indexes on first search for readers opened for a merge --- ...ia-cuvs-lucene-cuvs2510gpuvectorsreader.md | 52 +++++++----- .../cuvs/lucene/CuVS2510GPUVectorsReader.java | 84 ++++++++++++++----- 2 files changed, 92 insertions(+), 44 deletions(-) diff --git a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md index 5f3ff89fe9..a19381a5a5 100644 --- a/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md +++ b/fern/pages/lucene_api/lucene-api-com-nvidia-cuvs-lucene-cuvs2510gpuvectorsreader.md @@ -35,7 +35,7 @@ Initializes the `CuVS2510GPUVectorsReader`, checks and loads the index. | --- | --- | | `IOException` | I/O exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:87`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:90`_ ### CuVS2510GPUVectorsReader @@ -59,7 +59,7 @@ Initializes the reader with the cache owned by its vectors format. | --- | --- | | `IOException` | I/O exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:100`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:103`_ ### readSimilarityFunction @@ -85,7 +85,7 @@ an instance of VectorSimilarityFunction | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:228`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:231`_ ### readVectorEncoding @@ -111,7 +111,7 @@ the vector encoding | --- | --- | | `IOException` | | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:243`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:246`_ ### getFieldEntry @@ -132,7 +132,7 @@ cuVS index for it. the meta information for the field, or `null` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:306`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:309`_ ### openCagraIndexForMerge @@ -165,7 +165,7 @@ a freshly loaded CAGRA index, or `null` if this segment has none for the field | --- | --- | | `IOException` | I/O exception | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:325`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:328`_ ### close @@ -175,7 +175,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Closes the resources. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:394`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:430`_ ### checkIntegrity @@ -185,7 +185,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Checks consistency of this reader. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:412`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:453`_ ### getFloatVectorValues @@ -195,7 +195,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Returns the FloatVectorValues for the given field. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:421`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:462`_ ### getByteVectorValues @@ -207,7 +207,7 @@ Returns the ByteVectorValues for the given field. This is not supported. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:431`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:472`_ ### search @@ -217,7 +217,7 @@ _Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVecto Returns the k nearest neighbor documents using cuVS's CAGRA or brute force algorithm for this field, to the given vector. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:455`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:496`_ ### search @@ -229,7 +229,7 @@ Return the k nearest neighbor documents as determined by comparison of their vec This is not supported. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:615`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:656`_ ### readEntry @@ -257,16 +257,16 @@ an instance of FieldEntry | --- | --- | | `IOException` | I/O Exceptions | -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:643`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:684`_ ### getCagraIndexForField ```java -public CagraIndex getCagraIndexForField(String field) +public CagraIndex getCagraIndexForField(String field) throws IOException ``` -Returns the `CagraIndex` for the given field, or `null` if unavailable -(e.g., during a merge or when the field is missing). +Returns the `CagraIndex` for the given field, or `null` if this segment has none +(e.g., when the field is missing or the segment only holds a brute force index). **Parameters** @@ -278,7 +278,13 @@ Returns the `CagraIndex` for the given field, or `null` if unavailable the CAGRA index, or `null` -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:695`_ +**Throws** + +| Type | Description | +| --- | --- | +| `IOException` | I/O exception | + +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:737`_ ### getFilterBitsetCache @@ -288,7 +294,7 @@ FilterBitsetCache getFilterBitsetCache() Returns the filter cache owned by the vectors format that created this reader. -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:705`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:746`_ ### getFieldInfos @@ -302,7 +308,7 @@ Gets the instance of FieldInfos. the instance of FieldInfos -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:714`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:755`_ ### getCuvsIndexes @@ -314,9 +320,9 @@ Gets the map of `GPUIndex` objects. **Returns** -the map of GPU index objects +the map of GPU index objects, or `null` if this reader was opened for a merge and has not been searched yet -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:723`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:765`_ ### getFieldEntries @@ -330,6 +336,6 @@ Gets the map of FieldEntry objects that hold the meta information for the field. the map of FieldEntry objects -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:732`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:774`_ -_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:59`_ +_Source: `java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java:60`_ diff --git a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java index 1efa2d1174..949e09ed22 100644 --- a/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java +++ b/java/cuvs-lucene/src/main/java/com/nvidia/cuvs/lucene/CuVS2510GPUVectorsReader.java @@ -41,6 +41,7 @@ import org.apache.lucene.index.VectorSimilarityFunction; import org.apache.lucene.internal.hppc.IntObjectHashMap; import org.apache.lucene.search.KnnCollector; +import org.apache.lucene.store.AlreadyClosedException; import org.apache.lucene.store.ChecksumIndexInput; import org.apache.lucene.store.DataInput; import org.apache.lucene.store.IOContext; @@ -64,7 +65,9 @@ public class CuVS2510GPUVectorsReader extends KnnVectorsReader { private final FlatVectorsReader flatVectorsReader; private final FieldInfos fieldInfos; private final IntObjectHashMap fields; - private final IntObjectHashMap cuvsIndices; + // Loaded eagerly unless this reader was opened for a merge; see getCuVSIndices(). + private volatile IntObjectHashMap cuvsIndices; + private boolean closed; private final IndexInput cuvsIndexInput; private final FilterBitsetCache filterBitsetCache; @@ -130,13 +133,13 @@ public CuVS2510GPUVectorsReader(SegmentReadState state, FlatVectorsReader flatRe var ioContext = state.context.withReadAdvice(ReadAdvice.SEQUENTIAL); cuvsIndexInput = openCuVSInput(state, versionMeta, ioContext); /* - * Only load indexes on the GPU when this reader is opening for searches. - * Do not load indexes on the GPU when this reader is opening during merge calls. - * With this approach we reduce device memory usage by approximately 50% during merges. + * Do not load indexes on the GPU when this reader is opened for a merge, which reads them + * through openCagraIndexForMerge instead. This reduces device memory usage by approximately + * 50% during merges. IndexWriter pools the reader it opens for a merge, though, and hands it + * (or a clone sharing this core) to near-real-time readers opened while the merge runs, so + * such a reader can still be searched; it then loads its indexes on first use. */ - if (state.context.context().equals(Context.MERGE)) { - cuvsIndices = null; - } else { + if (!state.context.context().equals(Context.MERGE)) { cuvsIndices = loadCuVSIndices(); } success = true; @@ -347,11 +350,44 @@ CagraIndex openCagraIndexForMerge(String field) throws IOException { */ private IntObjectHashMap loadCuVSIndices() throws IOException { var indices = new IntObjectHashMap(); - for (var field : fields) { - var fieldEntry = field.value; - int fieldNumber = field.key; - var cuvsIndex = loadCuVSIndex(fieldEntry); - indices.put(fieldNumber, cuvsIndex); + boolean success = false; + try { + for (var field : fields) { + var fieldEntry = field.value; + int fieldNumber = field.key; + var cuvsIndex = loadCuVSIndex(fieldEntry); + indices.put(fieldNumber, cuvsIndex); + } + success = true; + return indices; + } finally { + if (success == false) { + var loaded = stream(indices.values().iterator()).map(cursor -> cursor.value); + IOUtils.closeWhileHandlingException(loaded::iterator); + } + } + } + + /** + * Returns the GPU indexes of this segment, loading them first if this reader was opened for a + * merge and has not been searched yet. + * + * @return the map of {@link GPUIndex} objects, keyed by field number + * @throws IOException I/O exception + */ + private IntObjectHashMap getCuVSIndices() throws IOException { + var indices = cuvsIndices; + if (indices == null) { + synchronized (this) { + if (closed) { + throw new AlreadyClosedException("this reader is closed"); + } + indices = cuvsIndices; + if (indices == null) { + indices = loadCuVSIndices(); + cuvsIndices = indices; + } + } } return indices; } @@ -394,10 +430,15 @@ private GPUIndex loadCuVSIndex(FieldEntry fieldEntry) throws IOException { */ @Override public void close() throws IOException { + final IntObjectHashMap indices; + synchronized (this) { + closed = true; + indices = cuvsIndices; + } var closeableStream = Stream.of(flatVectorsReader, cuvsIndexInput); IOUtils.close(closeableStream::iterator); - if (cuvsIndices != null) { - var indexClosableStream = stream(cuvsIndices.values().iterator()).map(cursor -> cursor.value); + if (indices != null) { + var indexClosableStream = stream(indices.values().iterator()).map(cursor -> cursor.value); IOUtils.close(indexClosableStream::iterator); } closeCuVSResourcesInstance(); @@ -462,7 +503,7 @@ public void search(String field, float[] target, KnnCollector knnCollector, Bits } var fieldNumber = fieldInfos.fieldInfo(field).number; - GPUIndex cuvsIndex = cuvsIndices != null ? cuvsIndices.get(fieldNumber) : null; + GPUIndex cuvsIndex = getCuVSIndices().get(fieldNumber); if (cuvsIndex == null) { throw new IllegalStateException("Index not found for field:" + field); } @@ -687,17 +728,17 @@ private static void checkVersion(int versionMeta, int versionVectorData, IndexIn } /** - * Returns the {@link CagraIndex} for the given field, or {@code null} if unavailable - * (e.g., during a merge or when the field is missing). + * Returns the {@link CagraIndex} for the given field, or {@code null} if this segment has none + * (e.g., when the field is missing or the segment only holds a brute force index). * * @param field the vector field name * @return the CAGRA index, or {@code null} + * @throws IOException I/O exception */ - public CagraIndex getCagraIndexForField(String field) { - if (cuvsIndices == null) return null; + public CagraIndex getCagraIndexForField(String field) throws IOException { FieldInfo info = fieldInfos.fieldInfo(field); if (info == null) return null; - GPUIndex gpuIndex = cuvsIndices.get(info.number); + GPUIndex gpuIndex = getCuVSIndices().get(info.number); if (gpuIndex == null) return null; return gpuIndex.getCagraIndex(); } @@ -719,7 +760,8 @@ public FieldInfos getFieldInfos() { /** * Gets the map of {@link GPUIndex} objects. * - * @return the map of GPU index objects + * @return the map of GPU index objects, or {@code null} if this reader was opened for a merge and + * has not been searched yet */ public IntObjectHashMap getCuvsIndexes() { return cuvsIndices; From 390f97b0ead222dc2546252633cf5d2e4d871586 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 16:14:45 +0000 Subject: [PATCH 05/27] Make the L2 expanded self-distance clamp relative to the norms The expanded L2 distance epilogues clamp round-off of self-distances to zero when the two norms are equal and the squared distance is below an absolute tolerance. For data with a small magnitude, this also zeroes the distances between distinct points with equal norms. Spectral embeddings of graphs with several connected components are such data, which made k-means (and spectral clustering) assign points to the wrong clusters on GPUs that use the CUTLASS fused 1-NN kernel. Scale the tolerance with the norms when they are smaller than one. --- .../distance/detail/distance_ops/l2_exp.cuh | 39 +++++----- cpp/tests/distance/dist_l2_exp.cu | 49 ++++++++++++- cpp/tests/neighbors/distance_nn.cu | 73 +++++++++++++++++++ 3 files changed, 143 insertions(+), 18 deletions(-) diff --git a/cpp/src/distance/detail/distance_ops/l2_exp.cuh b/cpp/src/distance/detail/distance_ops/l2_exp.cuh index c0433c676a..d059836ed1 100644 --- a/cpp/src/distance/detail/distance_ops/l2_exp.cuh +++ b/cpp/src/distance/detail/distance_ops/l2_exp.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,6 +28,23 @@ __device__ constexpr AccT get_clamp_precision() } } +/** + * Whether an expanded L2 distance is round-off error of a self-distance and should be clamped to + * zero. Self-neighboring points have (aNorm == bNorm) == accVal, but the dot product (accVal) can + * have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. + * + * The round-off error scales with the norms, so the tolerance is relative to them for squared + * norms below one. Otherwise, distinct points with equal norms and a small magnitude (for example, + * a spectral embedding) would have their distances incorrectly clamped to zero. + */ +template +DI bool is_self_distance_round_off(AccT dist, AccT aNorm, AccT bNorm) +{ + const AccT norm_product = aNorm * bNorm; + const AccT norm_scale = norm_product < AccT(1) ? norm_product : AccT(1); + return (aNorm == bNorm) && (dist * dist < get_clamp_precision() * norm_scale); +} + // Epilogue operator for CUTLASS based kernel template struct l2_exp_cutlass_op { @@ -38,13 +55,7 @@ struct l2_exp_cutlass_op { inline __device__ AccT operator()(AccT aNorm, AccT bNorm, AccT accVal) const noexcept { AccT outVal = aNorm + bNorm - AccT(2.0) * accVal; - - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product (accVal) - * can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. - */ - outVal = - outVal * AccT(!((outVal * outVal < get_clamp_precision()) * (aNorm == bNorm))); + outVal = outVal * AccT(!is_self_distance_round_off(outVal, aNorm, bNorm)); return sqrt ? raft::sqrt(outVal * static_cast(outVal > AccT(0))) : outVal; } @@ -112,15 +123,9 @@ struct l2_exp_distance_op { for (int j = 0; j < Policy::AccColsPerTh; ++j) { AccT accVal = acc[i][j]; AccT val = regxn[i] + regyn[j] - (AccT)2.0 * accVal; - - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product - * (accVal) can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal - * instead. - */ - acc[i][j] = val * static_cast((val > AccT(0))) * - static_cast( - !((val * val < get_clamp_precision()) * (regxn[i] == regyn[j]))); + acc[i][j] = + val * static_cast((val > AccT(0))) * + static_cast(!is_self_distance_round_off(val, regxn[i], regyn[j])); } } if (sqrt) { diff --git a/cpp/tests/distance/dist_l2_exp.cu b/cpp/tests/distance/dist_l2_exp.cu index a91daff1b3..1c9dbe4e28 100644 --- a/cpp/tests/distance/dist_l2_exp.cu +++ b/cpp/tests/distance/dist_l2_exp.cu @@ -1,11 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../test_utils.cuh" #include "distance_base.cuh" +#include +#include + +#include +#include + namespace cuvs { namespace distance { @@ -162,5 +168,46 @@ INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucExpTestD, ::testing::ValuesIn( class BigMatrixEucExp : public BigMatrixDistanceTest {}; TEST_F(BigMatrixEucExp, Result) {} + +// Distinct points with equal norms must keep their (small) distances. The clamp of self-distance +// round-off to zero used to apply an absolute tolerance, which collapsed such points when the data +// has a small magnitude (e.g. a spectral embedding). +template +void test_small_scale_equal_norms(T scale) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n = 8; + std::vector h_x(n * n, T{0}); + for (int64_t i = 0; i < n; i++) { + h_x[i * n + i] = scale; + } + auto x = raft::make_device_matrix(handle, n, n); + auto d = raft::make_device_matrix(handle, n, n); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream); + + for (auto metric : {DistanceType::L2Expanded, DistanceType::L2SqrtExpanded}) { + cuvs::distance::pairwise_distance(handle, + raft::make_const_mdspan(x.view()), + raft::make_const_mdspan(x.view()), + d.view(), + metric); + std::vector h_d(n * n); + raft::update_host(h_d.data(), d.data_handle(), h_d.size(), stream); + raft::resource::sync_stream(handle, stream); + + const T expected = + metric == DistanceType::L2Expanded ? T(2) * scale * scale : std::sqrt(T(2)) * scale; + for (int64_t i = 0; i < n; i++) { + for (int64_t j = 0; j < n; j++) { + EXPECT_NEAR(h_d[i * n + j], i == j ? T(0) : expected, T(1e-3) * expected) + << "metric=" << static_cast(metric) << " i=" << i << " j=" << j; + } + } + } +} + +TEST(DistanceEucExpSmallScale, EqualNormsF) { test_small_scale_equal_norms(1e-2f); } +TEST(DistanceEucExpSmallScale, EqualNormsD) { test_small_scale_equal_norms(1e-5); } } // end namespace distance } // namespace cuvs diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 5ac80fda73..9ab376927c 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -16,8 +16,10 @@ #include #include #include +#include #include +#include namespace cuvs::neighbors { @@ -446,6 +448,77 @@ TEST(Top1nnPlan, RejectsMismatchedLaunch) plan))); } +// Distinct points with equal norms and a small magnitude (as in a spectral embedding) must not be +// treated as self-neighbors when finding the nearest neighbor. +TEST(Top1nnEqualNorms, SmallScale) +{ + using Backend = cuvs::distance::detail::Top1nnBackend; + using Kvp = raft::KeyValuePair; + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int m = 1000; + constexpr int n = 5; + constexpr int k = 5; + constexpr float s = 1e-2f; + std::vector h_x(m * k, 0.0f); + std::vector h_y(n * k, 0.0f); + for (int i = 0; i < m; i++) { + h_x[i * k + i % n] = s; + } + for (int j = 0; j < n; j++) { + h_y[j * k + j] = s; + } + auto x = raft::make_device_matrix(handle, m, k); + auto y = raft::make_device_matrix(handle, n, k); + auto x_norm = raft::make_device_vector(handle, m); + auto y_norm = raft::make_device_vector(handle, n); + auto output = raft::make_device_vector(handle, m); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream.get()); + raft::update_device(y.data_handle(), h_y.data(), h_y.size(), stream.get()); + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, m, stream.get()); + raft::linalg::rowNorm( + y_norm.data_handle(), y.data_handle(), k, n, stream.get()); + + cuvs::distance::detail::Top1nnTuning tuning{}; + for (auto backend : {Backend::Cutlass, Backend::Unfused}) { + if (!cuvs::distance::detail::is_top_1_nn_backend_available( + backend, x.data_handle(), y.data_handle(), m, n, k, DistanceType::L2Expanded)) { + continue; + } + const auto plan = cuvs::distance::probe_top_1_nn( + handle, x.data_handle(), y.data_handle(), m, n, k, tuning, DistanceType::L2Expanded, backend); + ASSERT_TRUE(plan.available); + auto workspace = raft::make_device_vector(handle, plan.workspace_bytes); + cuvs::distance::top_1_nn(handle, + output.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + tuning, + workspace.data_handle(), + workspace.size(), + false, + true, + true, + DistanceType::L2Expanded, + 0.0f, + plan); + std::vector h_out(m); + raft::update_host(h_out.data(), output.data_handle(), m, stream.get()); + raft::resource::sync_stream(handle, stream); + int n_wrong = 0; + for (int i = 0; i < m; i++) { + n_wrong += h_out[i].key != i % n; + } + EXPECT_EQ(n_wrong, 0) << "backend=" << static_cast(backend); + } +} + TEST(Top1nnPlan, ExplicitUnavailableBackendIsStrict) { raft::resources handle; From f7c813d3295ade161d222b4e2829c926fe4db9d4 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:12:45 +0000 Subject: [PATCH 06/27] Handle empty lists when building IVF-RaBitQ indexes k-means can leave lists empty, for example when the dataset has fewer distinct rows than n_lists. Building from device data, or from host data that fits on the device, then launched the quantization kernels with zero blocks and failed with cudaErrorInvalidValue. The streaming build skipped empty lists, which left their rotated centroids uninitialized. Compute only the rotated centroid for an empty list and skip the kernels that have no rows to process. Add a regression test with fewer distinct rows than lists. --- .../neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu | 54 ++++---- .../ivf_rabitq/gpu_index/quantizer_gpu.cu | 12 ++ cpp/tests/neighbors/ann_ivf_rabitq.cuh | 119 ++++++++++++++++++ .../ann_ivf_rabitq/test_float_int64_t.cu | 5 + 4 files changed, 166 insertions(+), 24 deletions(-) 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..1290175f4c 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/ivf_gpu.cu @@ -719,38 +719,44 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, cluster_idx++; } - if (batch_vectors == 0) break; // No more clusters to process + // Every cluster fits into a batch (checked above), so each batch makes progress. + RAFT_EXPECTS(cluster_idx > batch_start_cluster, "IVF-RaBitQ streaming batch is empty"); // ------------------------- // 11. Gather and transfer batch data to GPU using reordered PIDs // ------------------------- - // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) - raft::copy(batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); - raft::resource::sync_stream(handle_); + // A batch that only contains empty clusters has no data to transfer, but the clusters still + // need their rotated centroids below. + if (batch_vectors > 0) { + // First, copy the reordered PIDs for this batch to host (reuse preallocated buffer) + raft::copy( + batch_pids.data_handle(), d_flat_pids + batch_start_offset, batch_vectors, stream_); + raft::resource::sync_stream(handle_); - // Gather data on host using the reordered PIDs (reuse preallocated buffer) - // Use OpenMP parallel region to leverage persistent thread pool across batches + // Gather data on host using the reordered PIDs (reuse preallocated buffer) + // Use OpenMP parallel region to leverage persistent thread pool across batches #pragma omp parallel num_threads(num_threads) - { - size_t tid = omp_get_thread_num(); - size_t chunk_size = batch_vectors / num_threads; - size_t start = tid * chunk_size; - size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; - - for (size_t i = start; i < end; ++i) { - std::memcpy(&batch_data.data_handle()[i * num_dimensions], - &host_data[batch_pids.data_handle()[i] * num_dimensions], - num_dimensions * sizeof(float)); + { + size_t tid = omp_get_thread_num(); + size_t chunk_size = batch_vectors / num_threads; + size_t start = tid * chunk_size; + size_t end = (tid == num_threads - 1) ? batch_vectors : start + chunk_size; + + for (size_t i = start; i < end; ++i) { + std::memcpy(&batch_data.data_handle()[i * num_dimensions], + &host_data[batch_pids.data_handle()[i] * num_dimensions], + num_dimensions * sizeof(float)); + } } - } - // OpenMP implicit barrier ensures all threads complete before continuing + // OpenMP implicit barrier ensures all threads complete before continuing - // Transfer batch to GPU (reuse preallocated buffer) - raft::copy(d_batch_data.data_handle(), - batch_data.data_handle(), - batch_vectors * num_dimensions, - stream_); + // Transfer batch to GPU (reuse preallocated buffer) + raft::copy(d_batch_data.data_handle(), + batch_data.data_handle(), + batch_vectors * num_dimensions, + stream_); + } // ------------------------- // 12. Quantize each cluster in the batch @@ -758,8 +764,8 @@ void IVFGPU::construct_on_gpu_streaming(const float* host_data, size_t batch_offset = 0; for (size_t c = batch_start_cluster; c < cluster_idx; ++c) { + // Empty clusters are quantized too, so that their rotated centroids are initialized. size_t cluster_size = h_cluster_meta[c].num; - if (cluster_size == 0) continue; const float* cur_centroid = device_centroids + c * num_dimensions; float* cur_rotated_c = d_rotated_centroids.data_handle() + c * num_padded_dim; diff --git a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu index b35ac9fab9..5eef2179d4 100644 --- a/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu +++ b/cpp/src/neighbors/ivf_rabitq/gpu_index/quantizer_gpu.cu @@ -520,6 +520,9 @@ void data_transformation_batch_opt(const float* d_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); + // An empty list only needs its rotated centroid; a launch with zero blocks is invalid. + if (num_points == 0) { return; } + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -637,6 +640,9 @@ void DataQuantizerGPU::quantize_batch_opt(const float* d_data, D, handle_); + // An empty list has no codes or factors to compute. + if (num_points == 0) { return; } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), @@ -714,6 +720,9 @@ void data_transformation_batch_opt_contiguous(const float* d_contiguous_data, // 5. Save the rotated centroid: copy CP into d_rotated_c. raft::copy(d_rotated_c, d_CP, D, stream); + // An empty list only needs its rotated centroid; a launch with zero blocks is invalid. + if (num_points == 0) { return; } + // 6. Launch the single FUSED kernel for subtract, normalize, and binarize. const unsigned int FusedBlockSize = 256; // A good default, can be tuned. dim3 gridDim(num_points); @@ -758,6 +767,9 @@ void DataQuantizerGPU::quantize_batch_opt_contiguous(const float* d_contiguous_d D, handle_); + // An empty list has no codes or factors to compute. + if (num_points == 0) { return; } + rabitq_codes_and_factors_fused(d_rotated_c, d_bin_XP.data_handle(), d_XP.data_handle(), diff --git a/cpp/tests/neighbors/ann_ivf_rabitq.cuh b/cpp/tests/neighbors/ann_ivf_rabitq.cuh index 2755258600..7296e67ddf 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq.cuh +++ b/cpp/tests/neighbors/ann_ivf_rabitq.cuh @@ -396,4 +396,123 @@ inline auto var_search_mode_1_bit() -> test_cases_t #define INSTANTIATE(type, vals) \ INSTANTIATE_TEST_SUITE_P(IvfRabitq, type, ::testing::ValuesIn(vals)); /* NOLINT */ +/* Empty lists */ + +enum class dataset_input { device, host, host_streaming }; + +struct ivf_rabitq_empty_lists_inputs { + uint32_t dim; + dataset_input input; +}; + +inline auto operator<<(std::ostream& os, const ivf_rabitq_empty_lists_inputs& p) -> std::ostream& +{ + const char* input = p.input == dataset_input::device ? "device" + : p.input == dataset_input::host ? "host" + : "host_streaming"; + return os << "ivf_rabitq_empty_lists_inputs {.dim = " << p.dim << ", .input = " << input << "}"; +} + +/** + * The dataset has fewer distinct rows than there are lists. Identical rows are always assigned to + * the same list, so k-means is guaranteed to leave some lists empty. The index must still build + * from device and host input, and search must still find the nearest rows. + */ +class ivf_rabitq_empty_lists_test : public ::testing::TestWithParam { + protected: + void run() + { + constexpr int64_t kRows = 4096; + constexpr int64_t kDistinct = 8; + constexpr int64_t kQueries = 2 * kDistinct; + constexpr int64_t kK = 10; + constexpr uint32_t kLists = 32; + constexpr float kSpacing = 0.25f; + constexpr float kOffset = 0.1f; + + auto ps = GetParam(); + auto dim = static_cast(ps.dim); + auto stream = raft::resource::get_cuda_stream(handle_); + + // Row i is the point kSpacing * (i % kDistinct) on the diagonal. Query q is offset from the + // distinct row q % kDistinct by +-kOffset in every dimension, so that row is its unique nearest + // distinct row at squared L2 distance kOffset^2 * dim. + 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) = kSpacing * static_cast(i % kDistinct); + } + } + auto h_queries = raft::make_host_matrix(kQueries, dim); + for (int64_t i = 0; i < kQueries; i++) { + for (int64_t j = 0; j < dim; j++) { + h_queries(i, j) = + kSpacing * static_cast(i % kDistinct) + (i < kDistinct ? kOffset : -kOffset); + } + } + auto d_dataset = raft::make_device_matrix(handle_, kRows, dim); + auto d_queries = raft::make_device_matrix(handle_, kQueries, 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; + if (ps.input == dataset_input::host_streaming) { + ipams.force_streaming = true; + ipams.streaming_batch_size = kRows; + } + auto idx = ps.input == dataset_input::device + ? build(handle_, ipams, raft::make_const_mdspan(d_dataset.view())) + : build(handle_, ipams, raft::make_const_mdspan(h_dataset.view())); + + // Probe every list, because an empty list's centroid can be closer to a query than the + // centroid of the list that holds its neighbors. + search_params sparams; + sparams.n_probes = kLists; + auto neighbors = raft::make_device_matrix(handle_, kQueries, kK); + auto distances = raft::make_device_matrix(handle_, kQueries, kK); + search(handle_, + sparams, + idx, + raft::make_const_mdspan(d_queries.view()), + neighbors.view(), + distances.view()); + + std::vector h_neighbors(kQueries * kK); + std::vector h_distances(kQueries * 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 kRows / kDistinct > kK copies, so all kK neighbors are copies of the nearest + // distinct row. The next nearest distinct row is (kSpacing - kOffset)^2 * dim away. + const float expected_dist = kOffset * kOffset * static_cast(dim); + for (int64_t q = 0; q < kQueries; q++) { + for (int64_t j = 0; j < kK; j++) { + auto neighbor = h_neighbors[q * kK + j]; + auto dist = h_distances[q * kK + j]; + ASSERT_TRUE(neighbor >= 0 && neighbor < kRows) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + ASSERT_EQ(neighbor % kDistinct, q % kDistinct) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + ASSERT_NEAR(dist, expected_dist, 0.1f * expected_dist) + << ps << ": query " << q << ", neighbor " << j << " = " << neighbor; + } + } + } + + raft::resources handle_; +}; + +inline auto empty_lists_inputs() -> std::vector +{ + std::vector xs; + for (uint32_t dim : {1, 2, 7, 64, 100}) { + for (auto input : {dataset_input::device, dataset_input::host, dataset_input::host_streaming}) { + xs.push_back({dim, input}); + } + } + return xs; +} + } // namespace cuvs::neighbors::ivf_rabitq diff --git a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu index 42bf84e6a3..28771bd785 100644 --- a/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_rabitq/test_float_int64_t.cu @@ -18,4 +18,9 @@ INSTANTIATE(f32_f32_i64, defaults() + small_dims() + big_dims() + var_n_probes() + var_k() + var_bits_per_dim() + var_search_mode() + var_search_mode_1_bit()); +TEST_P(ivf_rabitq_empty_lists_test, build_search) { this->run(); } // NOLINT +INSTANTIATE_TEST_SUITE_P(IvfRabitq, + ivf_rabitq_empty_lists_test, + ::testing::ValuesIn(empty_lists_inputs())); // NOLINT + } // namespace cuvs::neighbors::ivf_rabitq From dc92f0aee9407ba3e93683bb019b382077d0a283 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:13:11 +0000 Subject: [PATCH 07/27] Reject single-row datasets in CAGRA graph build With one row, check_graph_degree reduced both the intermediate and the final graph degree to 0, and the kNN graph build (NN-descent or IVF-PQ) then hit an illegal memory access, which aborted the process from a stream synchronization in a destructor. Throw a clear error up front instead, and keep HnswRandomizedIT from generating single-row datasets. --- .../neighbors/detail/cagra/cagra_build.cuh | 7 +- .../com/nvidia/cuvs/CagraRandomizedIT.java | 3 +- .../com/nvidia/cuvs/CagraSingleRowIT.java | 64 +++++++++++++++++++ .../com/nvidia/cuvs/HnswRandomizedIT.java | 3 +- 4 files changed, 73 insertions(+), 4 deletions(-) create mode 100644 java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 9a22676334..7d2cdb7517 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -164,10 +164,15 @@ class ace_disk_workspace { template void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_t dataset_size) { + // Every row needs at least one neighbor other than itself. With a single row both degrees would + // drop to zero, and the kNN graph build kernels would then access memory out of bounds. + RAFT_EXPECTS(dataset_size >= 2, + "CAGRA graph build requires a dataset with at least 2 rows, got %zu", + dataset_size); if (intermediate_degree >= static_cast(dataset_size)) { RAFT_LOG_WARN( "Intermediate graph degree cannot be larger than dataset size, reducing it to %lu", - dataset_size); + dataset_size - 1); intermediate_degree = dataset_size - 1; } if (intermediate_degree < graph_degree) { diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index 9110f6ddc2..a38a12ea77 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -56,8 +56,7 @@ private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryK int TOP_K_LIMIT = 64; // nocommit This fails beyond 64 int datasetSize = - random.nextInt(DATASET_SIZE_LIMIT) - + 2; // datasetSize of 1 fails/crashed due to a bug, hence adding 2 here. + random.nextInt(DATASET_SIZE_LIMIT) + 2; // A CAGRA graph needs at least 2 rows. int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, datasetSize); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java new file mode 100644 index 0000000000..6410dafc43 --- /dev/null +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraSingleRowIT.java @@ -0,0 +1,64 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +package com.nvidia.cuvs; + +import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.carrotsearch.randomizedtesting.RandomizedRunner; +import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import com.nvidia.cuvs.CagraIndexParams.CuvsDistanceType; +import org.junit.Before; +import org.junit.Test; +import org.junit.runner.RunWith; + +/** + * A CAGRA graph needs at least two rows: with a single row there is no neighbor to link to, so + * both graph degrees collapse to zero. Building one used to launch the kNN graph kernels anyway, + * which hit an illegal memory access and aborted the JVM (exit code 134). It has to fail with an + * exception instead. + */ +@RunWith(RandomizedRunner.class) +public class CagraSingleRowIT extends CuVSTestCase { + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + initializeRandom(); + } + + private void assertSingleRowBuildFails(CagraGraphBuildAlgo algo) throws Throwable { + float[][] vectors = generateData(random, 1, 16); + try (CuVSResources resources = CheckedCuVSResources.create()) { + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(algo) + .withGraphDegree(64) + .withIntermediateGraphDegree(128) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + RuntimeException e = + assertThrows( + RuntimeException.class, + () -> + CagraIndex.newBuilder(resources) + .withDataset(vectors) + .withIndexParams(indexParams) + .build()); + assertTrue(e.getMessage(), e.getMessage().contains("at least 2 rows")); + } + } + + @Test + public void testSingleRowNnDescentBuildThrows() throws Throwable { + assertSingleRowBuildFails(CagraGraphBuildAlgo.NN_DESCENT); + } + + @Test + public void testSingleRowIvfPqBuildThrows() throws Throwable { + assertSingleRowBuildFails(CagraGraphBuildAlgo.IVF_PQ); + } +} diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java index 4ace8fc92e..e35f3ce4a3 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/HnswRandomizedIT.java @@ -49,7 +49,8 @@ private void tmpResultsTopKWithRandomValues(boolean useNativeMemoryDataset) thro int NUM_QUERIES_LIMIT = 10; int TOP_K_LIMIT = 64; // nocommit This fails beyond 64 - int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 1; + // A CAGRA graph needs at least 2 rows, so the HNSW index converted from it does too. + int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 2; int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, datasetSize); From fe5698253f3eb34facbf142b03928c9b193a1a22 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 16:00:35 +0000 Subject: [PATCH 08/27] 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 456b98dabc73bc863387b5745100526249a0069e Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:39:22 +0000 Subject: [PATCH 09/27] Fix CAGRA dataset descriptor race across threads sharing a handle dataset_descriptor_host::state::get() skipped waiting for the descriptor initialization event when the caller's stream handle equaled the stream that ran the initialization. The default raft stream is cudaStreamPerThread, whose handle is identical in every host thread but refers to a different stream in each. When several threads search with the same raft::resources, a thread could pick up the cached descriptor and launch a search kernel before the initialization kernel queued on another thread's stream had run, reading an uninitialized descriptor (cudaErrorMisalignedAddress / cudaErrorIllegalAddress in NEIGHBORS_MG_TEST). Always wait on the ready event. --- cpp/src/neighbors/detail/cagra/compute_distance.hpp | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/compute_distance.hpp b/cpp/src/neighbors/detail/cagra/compute_distance.hpp index c1d0033a01..99bd6fc78d 100644 --- a/cpp/src/neighbors/detail/cagra/compute_distance.hpp +++ b/cpp/src/neighbors/detail/cagra/compute_distance.hpp @@ -253,10 +253,13 @@ struct dataset_descriptor_host { { if (!ready.load(std::memory_order_acquire)) { eval(stream); } // value is immutable at this point. - auto& [ptr, ready_stream] = std::get(value); - if (ready_stream != stream) { - RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); - } + auto* ptr = std::get<0>(std::get(value)); + // Always order `stream` after the initialization kernel. Comparing the stream handles is not + // enough: the per-thread default stream (`cudaStreamPerThread`, which is also raft's default + // stream) has the same handle in every host thread but refers to a different stream in each. + // When several threads share a raft::resources (and thus this cached descriptor), skipping + // the wait lets a search kernel read the descriptor before it is initialized. + RAFT_CUDA_TRY(cudaStreamWaitEvent(stream.get(), ready_event, 0)); return ptr; } }; From 733caab2017f0eb147049cda46dd4933881ffe7e Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:51:57 +0000 Subject: [PATCH 10/27] Make the CAGRA descriptor cache lookup-or-insert atomic dataset_descriptor_init_with_cache() did the LRU get and set under separate locks. Threads that missed concurrently each created a descriptor, and the later insertion replaced the earlier one while other threads could already be using it. The last owner of the replaced descriptor then freed its device memory with cudaFreeAsync on its own stream, while a search kernel on another thread's stream could still be reading it. Serialize the lookup and insertion so that concurrent searches on one raft::resources share a single descriptor. --- cpp/src/neighbors/detail/cagra/factory.cuh | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index ace28b7842..3107f841ae 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -16,6 +16,8 @@ #include +#include + namespace cuvs::neighbors::cagra::detail { template , dataset_descriptor_host> value{kDefaultSize}; + /** Serializes lookup-or-insert, so that concurrent searches share a single descriptor. */ + std::mutex mutex; }; } // namespace descriptor_cache @@ -170,9 +174,15 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res, -> dataset_descriptor_host { auto key = descriptor_cache::make_key(params, dataset, metric, dataset_norms); - auto& cache = - raft::resource::get_custom_resource>(res) - ->value; + auto* store = + raft::resource::get_custom_resource>(res); + auto& cache = store->value; + // Hold the lock across the lookup and the insertion. Otherwise threads that miss concurrently + // each create a descriptor and the later insertion replaces the earlier one, which may already + // be in use by another thread; the last owner then frees it on its own stream while a kernel on + // a different stream can still be reading it. Only host-side work happens under the lock (the + // device-side initialization is deferred to the first `dev_ptr()` call). + std::lock_guard lock(store->mutex); dataset_descriptor_host desc; if (!cache.get(key, &desc)) { desc = From ccb70f13286f699f51130781fac4df4be7f7e803 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:52:35 +0000 Subject: [PATCH 11/27] Fix flaky IVF-PQ Python tests test_build_precomputed compared the neighbor indices of two separately built indexes exactly. Both indexes hold the same PQ codes, and their distances are bitwise equal, but extend fills each list in a nondeterministic order and the search breaks distance ties by that order. Compare the distances exactly and the indices of each group of equal distances as a set, searching one extra neighbor to detect a tie across the k-th position. test_ivf_pq used the default pq_dim (8 for 10 columns). With it the inner-product recall is about 0.74 and varies by about 0.01 between builds of the same data, so recall > 0.7 failed occasionally. Use one dimension per PQ subspace, which gives about 0.94. --- python/cuvs/cuvs/tests/test_ivf_pq.py | 54 ++++++++++++++++++--------- 1 file changed, 37 insertions(+), 17 deletions(-) diff --git a/python/cuvs/cuvs/tests/test_ivf_pq.py b/python/cuvs/cuvs/tests/test_ivf_pq.py index da43d22ea2..d0371c5525 100644 --- a/python/cuvs/cuvs/tests/test_ivf_pq.py +++ b/python/cuvs/cuvs/tests/test_ivf_pq.py @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # @@ -149,7 +149,13 @@ def run_ivf_pq_build_search_test( "metric", ["sqeuclidean", "inner_product", "euclidean"] ) def test_ivf_pq(inplace, dtype, metric): + # One dimension per PQ subspace. With the default pq_dim (8 for 10 + # columns), the inner-product recall on this data is only about 0.74 and + # varies by about 0.01 from build to build, so `recall > 0.7` fails in + # roughly 1 of 1000 runs. run_ivf_pq_build_search_test( + n_cols=10, + pq_dim=10, dtype=dtype, inplace=inplace, metric=metric, @@ -304,15 +310,14 @@ def test_build_precomputed(codebook_kind, metric): queries = generate_data((n_queries, n_cols), dtype) queries_device = device_ndarray(queries) - # Search regular index + # Search both indexes for one extra neighbor, so that a distance tie + # across the k-th position can be detected. search_params = ivf_pq.SearchParams(n_probes=n_probes) regular_dist, regular_idx = ivf_pq.search( - search_params, regular_index, queries_device, k + search_params, regular_index, queries_device, k + 1 ) - - # Search precomputed index precomputed_dist, precomputed_idx = ivf_pq.search( - search_params, precomputed_index, queries_device, k + search_params, precomputed_index, queries_device, k + 1 ) # Copy results to host for comparison @@ -321,20 +326,35 @@ def test_build_precomputed(codebook_kind, metric): precomputed_idx_host = precomputed_idx.copy_to_host() precomputed_dist_host = precomputed_dist.copy_to_host() - # Compare results for exact match + # Both indexes hold the same centers and PQ codes, so the distances must + # match exactly. np.testing.assert_array_equal( - regular_idx_host, - precomputed_idx_host, - err_msg="Neighbor indices should match exactly", - ) - np.testing.assert_allclose( - regular_dist_host, - precomputed_dist_host, - rtol=1e-5, - atol=1e-5, - err_msg="Distances should match closely", + regular_dist_host[:, :k], + precomputed_dist_host[:, :k], + err_msg="Distances should match exactly", ) + # The order of the vectors within each list is not deterministic (extend + # assigns the slots atomically), and the search breaks distance ties by + # that order. Neighbors with bitwise-equal distances may therefore come + # back in a different order, and a tie across the k-th position may + # select different neighbors. Compare the indices of each group of equal + # distances as a set, except for a group that crosses the k-th position. + for row in range(n_queries): + dist = regular_dist_host[row] + start = 0 + while start < k: + end = start + 1 + while end <= k and dist[end] == dist[start]: + end += 1 + if end <= k: + np.testing.assert_array_equal( + np.sort(regular_idx_host[row, start:end]), + np.sort(precomputed_idx_host[row, start:end]), + err_msg=f"Neighbor indices should match (query {row})", + ) + start = end + @pytest.mark.parametrize("codebook_kind", ["subspace", "cluster"]) @pytest.mark.parametrize("dtype", [np.float32, np.float16, np.int8, np.uint8]) From 87ca43989c348f209879c50c3d7266e1dd76921e Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 18:02:37 +0000 Subject: [PATCH 12/27] Keep the two segments apart in testMergeTwoSegsWithASingleDocPerSeg --- .../java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java | 5 ++++- .../lucene/TestLucene99AcceleratedHNSWVectorsFormat.java | 5 ++++- .../com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java | 5 ++++- 3 files changed, 12 insertions(+), 3 deletions(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java index 6db380a3ce..a09b306152 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSVectorsFormat.java @@ -18,6 +18,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.LogDocMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; @@ -42,7 +43,9 @@ protected Codec getCodec() { public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + // A random merge policy could merge the two segments before they are checked. + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(new LogDocMergePolicy()))) { Document doc1 = new Document(); doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java index 707ee35709..8f56c7f45e 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestLucene99AcceleratedHNSWVectorsFormat.java @@ -18,6 +18,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.LogDocMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.Directory; import org.apache.lucene.tests.index.BaseKnnVectorsFormatTestCase; @@ -42,7 +43,9 @@ protected Codec getCodec() { public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { float[][] f = new float[][] {randomVector(384), randomVector(384)}; try (Directory dir = newDirectory(); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + // A random merge policy could merge the two segments before they are checked. + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(new LogDocMergePolicy()))) { Document doc1 = new Document(); doc1.add(new StringField("id", "0", Field.Store.YES)); doc1.add(new KnnFloatVectorField("f", f[0], EUCLIDEAN)); diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java index 3b05beaacf..67cae2b4e9 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestQuantizedVectorsFormats.java @@ -26,6 +26,7 @@ import org.apache.lucene.index.IndexWriter; import org.apache.lucene.index.LeafReader; import org.apache.lucene.index.LeafReaderContext; +import org.apache.lucene.index.LogDocMergePolicy; import org.apache.lucene.index.VectorEncoding; import org.apache.lucene.store.ByteBuffersDirectory; import org.apache.lucene.store.Directory; @@ -75,7 +76,9 @@ public void testMergeTwoSegsWithASingleDocPerSeg() throws Exception { } try (Directory dir = newDirectory(new ByteBuffersDirectory()); - IndexWriter w = new IndexWriter(dir, newIndexWriterConfig())) { + // A random merge policy could merge the two segments before they are checked. + IndexWriter w = + new IndexWriter(dir, newIndexWriterConfig().setMergePolicy(new LogDocMergePolicy()))) { for (int i = 0; i < R; i++) { Document doc = new Document(); doc.add(new StringField("id", String.valueOf(i), Field.Store.YES)); From dcebe5db62d1b09f940aaf397da552023fa68071 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 18:59:10 +0000 Subject: [PATCH 13/27] 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 14/27] 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); From 7dcc920a109cbbfd28b46e2ec5087044eaedf9e7 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 20:12:26 +0000 Subject: [PATCH 15/27] Require run-to-run determinism explicitly in k-means reductions Use the environment-based cub::DeviceReduce with an explicit cuda::execution::determinism::run_to_run requirement for the cluster cost, the weight sum and the centroid shift. Temporary storage comes from the caller's workspace buffer through a small memory resource, so the reductions do not allocate on every call. Also let the per-cluster row sums take any data iterator and optional weights, and process the rows in passes of at most 2^22 to bound the workspace. --- cpp/src/cluster/detail/kmeans.cuh | 3 +- .../cluster/detail/kmeans_centroid_sums.cuh | 206 ++++++++++++------ cpp/src/cluster/detail/kmeans_common.cuh | 139 ++++++++---- cpp/src/cluster/detail/kmeans_mg.cuh | 3 +- 4 files changed, 237 insertions(+), 114 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 332fed16b5..4478cec2a6 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -1082,7 +1082,8 @@ void kmeans_fit( compute_centroid_shift(handle, raft::make_const_mdspan(centroids_const), raft::make_const_mdspan(new_centroids_view), - sqrdNormError.view()); + sqrdNormError.view(), + ws); std::swap(cur_centroids_ptr, new_centroids_ptr); diff --git a/cpp/src/cluster/detail/kmeans_centroid_sums.cuh b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh index 9e0fdb378b..6248e00ae1 100644 --- a/cpp/src/cluster/detail/kmeans_centroid_sums.cuh +++ b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh @@ -21,7 +21,7 @@ #include /** - * Deterministic per-cluster weighted row sums for the k-means centroid update. + * Deterministic per-cluster (weighted) row sums for the k-means centroid updates. * * Accumulating the rows of each cluster with atomics (as raft::linalg::reduce_rows_by_key does) * makes the floating-point summation order, and therefore the centroids, vary from run to run. @@ -29,11 +29,22 @@ * fits can return noticeably different centroids. Here the summation order is a function of the * input only: * - * 1. Row indices are stably sorted by label, so each cluster becomes a contiguous segment. + * 1. Row indices are stably sorted by label (cub::DeviceRadixSort), so each cluster becomes a + * contiguous segment. * 2. Each segment is cut into pieces of at most `kRowsPerPiece` rows. One thread block reduces a * piece with a fixed assignment of rows to threads and a fixed-order combine. * 3. One thread block per cluster adds up that cluster's piece sums in the same fixed way, and * adds the result to the output. + * + * The order depends only on the labels and the sizes, not on scheduling or on the launch + * configuration. + * + * Why not a CCCL reduction: cub::DeviceSegmentedReduce (run_to_run) and + * cub::DeviceReduce::ReduceByKey (run_to_run) reduce one scalar per segment or run. A + * [n_rows x n_cols] matrix would need n_clusters * n_cols segments that read the gathered rows + * column by column (strided, uncoalesced), and DeviceSegmentedReduce gives each segment to one + * thread block, which serializes large clusters when n_clusters is small. The per-cluster weight + * sum alone could use them, but it comes for free here in the same pass over the rows. */ namespace cuvs::cluster::kmeans::detail::sums_by_label { @@ -128,10 +139,11 @@ RAFT_KERNEL segment_kernel(const uint32_t* sorted_keys, } } -template -RAFT_KERNEL piece_sums_kernel(const DataT* X, +template +RAFT_KERNEL piece_sums_kernel(DataIterator X, IndexT n_cols, const DataT* weights, + bool sum_weights, const IndexT* sorted_rows, const IndexT* cluster_offsets, const IndexT* piece_offsets, @@ -162,18 +174,20 @@ RAFT_KERNEL piece_sums_kernel(const DataT* X, col_threads, [=](IndexT i, IndexT j) { const IndexT row = sorted_rows[i]; - const DataT x = X[static_cast(row) * n_cols + j]; + const auto x = static_cast(X[static_cast(row) * n_cols + j]); return weights != nullptr ? x * weights[row] : x; }, [=](IndexT j, DataT sum) { out[j] = sum; }); // The sum of weights is stored after the column sums. - fixed_order_column_sums( - begin, - end, - IndexT{1}, - 1, - [=](IndexT i, IndexT) { return weights != nullptr ? weights[sorted_rows[i]] : DataT{1}; }, - [=](IndexT, DataT sum) { out[n_cols] = sum; }); + if (sum_weights) { + fixed_order_column_sums( + begin, + end, + IndexT{1}, + 1, + [=](IndexT i, IndexT) { return weights != nullptr ? weights[sorted_rows[i]] : DataT{1}; }, + [=](IndexT, DataT sum) { out[n_cols] = sum; }); + } } } @@ -199,61 +213,47 @@ RAFT_KERNEL cluster_sums_kernel(const DataT* piece_sums, col_threads, [=](IndexT p, IndexT j) { return piece_sums[static_cast(p) * stride + j]; }, [=](IndexT j, DataT sum) { out[j] = reset_sums ? sum : out[j] + sum; }); - fixed_order_column_sums( - p_begin, - p_end, - IndexT{1}, - 1, - [=](IndexT p, IndexT) { return piece_sums[static_cast(p) * stride + n_cols]; }, - [=](IndexT, DataT sum) { - weight_per_cluster[c] = reset_sums ? sum : weight_per_cluster[c] + sum; - }); + if (weight_per_cluster != nullptr) { + fixed_order_column_sums( + p_begin, + p_end, + IndexT{1}, + 1, + [=](IndexT p, IndexT) { return piece_sums[static_cast(p) * stride + n_cols]; }, + [=](IndexT, DataT sum) { + weight_per_cluster[c] = reset_sums ? sum : weight_per_cluster[c] + sum; + }); + } } } inline size_t align_bytes(size_t bytes) { return raft::alignTo(bytes, 256); } -/** - * @brief Deterministic weighted sum of the rows of X per label, and of the weights per label. - * - * centroid_sums[c, :] (+)= sum_{i : labels[i] == c} weights[i] * X[i, :] - * weight_per_cluster[c] (+)= sum_{i : labels[i] == c} weights[i] - * - * The result is bitwise reproducible for identical inputs on the same device. - * - * @param[in] handle RAFT resources handle - * @param[in] X Row-major input [n_rows x n_cols] - * @param[in] weights Per-row weights [n_rows] - * @param[in] labels Label of each row, in [0, n_clusters) - * @param[in] n_rows Number of rows of X - * @param[in] n_cols Number of columns of X - * @param[in] n_clusters Number of labels - * @param[inout] centroid_sums Row-major output [n_clusters x n_cols] - * @param[inout] weight_per_cluster Output [n_clusters] - * @param[inout] workspace Scratch buffer, resized as needed - * @param[in] reset_sums Overwrite the outputs if true, add to them otherwise - */ -template -void weighted_sums_by_label(raft::resources const& handle, - const DataT* X, - const DataT* weights, - LabelsIterator labels, - IndexT n_rows, - IndexT n_cols, - IndexT n_clusters, - DataT* centroid_sums, - DataT* weight_per_cluster, - rmm::device_uvector& workspace, - bool reset_sums) +/** Rows per pass of `weighted_sums_by_label`; bounds the workspace and fits 32-bit indices. */ +constexpr int64_t kMaxRowsPerPass = int64_t{1} << 22; + +/** One pass of `weighted_sums_by_label` over at most kMaxRowsPerPass rows. */ +template +void weighted_sums_by_label_pass(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + uint32_t n_rows, + uint32_t n_cols, + uint32_t n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) { + using IdxT = uint32_t; cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - if (n_clusters <= 0) { return; } int end_bit = 1; while (end_bit < 32 && (uint64_t{1} << end_bit) < static_cast(n_clusters)) { ++end_bit; } - const IndexT max_pieces = raft::ceildiv(n_rows, kRowsPerPiece) + n_clusters; + const IdxT max_pieces = raft::ceildiv(n_rows, kRowsPerPiece) + n_clusters; // Workspace layout size_t sort_bytes = 0; @@ -261,8 +261,8 @@ void weighted_sums_by_label(raft::resources const& handle, sort_bytes, static_cast(nullptr), static_cast(nullptr), - static_cast(nullptr), - static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), n_rows, 0, end_bit, @@ -270,18 +270,19 @@ void weighted_sums_by_label(raft::resources const& handle, size_t scan_bytes = 0; RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, scan_bytes, - static_cast(nullptr), - static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), n_clusters + 1, stream)); const size_t keys_bytes = align_bytes(sizeof(uint32_t) * n_rows); - const size_t rows_bytes = align_bytes(sizeof(IndexT) * n_rows); - const size_t offsets_bytes = align_bytes(sizeof(IndexT) * (n_clusters + 1)); + const size_t rows_bytes = align_bytes(sizeof(IdxT) * n_rows); + const size_t offsets_bytes = align_bytes(sizeof(IdxT) * (n_clusters + 1)); const size_t pieces_bytes = align_bytes(sizeof(DataT) * static_cast(max_pieces) * (n_cols + 1)); const size_t temp_bytes = align_bytes(std::max(sort_bytes, scan_bytes)); - workspace.resize(2 * keys_bytes + 2 * rows_bytes + 3 * offsets_bytes + pieces_bytes + temp_bytes, - stream); + const size_t total_bytes = + 2 * keys_bytes + 2 * rows_bytes + 3 * offsets_bytes + pieces_bytes + temp_bytes; + if (workspace.size() < total_bytes) { workspace.resize(total_bytes, stream); } char* ptr = workspace.data(); auto take = [&ptr](size_t bytes) { @@ -291,11 +292,11 @@ void weighted_sums_by_label(raft::resources const& handle, }; auto* keys_in = reinterpret_cast(take(keys_bytes)); auto* keys_out = reinterpret_cast(take(keys_bytes)); - auto* rows_in = reinterpret_cast(take(rows_bytes)); - auto* rows_out = reinterpret_cast(take(rows_bytes)); - auto* cluster_offsets = reinterpret_cast(take(offsets_bytes)); - auto* piece_counts = reinterpret_cast(take(offsets_bytes)); - auto* piece_offsets = reinterpret_cast(take(offsets_bytes)); + auto* rows_in = reinterpret_cast(take(rows_bytes)); + auto* rows_out = reinterpret_cast(take(rows_bytes)); + auto* cluster_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_counts = reinterpret_cast(take(offsets_bytes)); + auto* piece_offsets = reinterpret_cast(take(offsets_bytes)); auto* piece_sums = reinterpret_cast(take(pieces_bytes)); void* temp = take(temp_bytes); @@ -318,11 +319,12 @@ void weighted_sums_by_label(raft::resources const& handle, temp, scan_bytes, piece_counts, piece_offsets, n_clusters + 1, stream)); const int col_threads = column_threads(n_cols); - if (max_pieces > 0) { + { const int grid = static_cast(std::min(max_pieces, int64_t{65536})); piece_sums_kernel<<>>(X, n_cols, weights, + weight_per_cluster != nullptr, rows_out, cluster_offsets, piece_offsets, @@ -345,4 +347,70 @@ void weighted_sums_by_label(raft::resources const& handle, } } +/** + * @brief Deterministic (weighted) sum of the rows of X per label, and of the weights per label. + * + * centroid_sums[c, :] (+)= sum_{i : labels[i] == c} weights[i] * X[i, :] + * weight_per_cluster[c] (+)= sum_{i : labels[i] == c} weights[i] + * + * The summation order depends only on the labels and the sizes, not on scheduling, + * so the result is bitwise reproducible. Rows are processed in passes of at most kMaxRowsPerPass + * rows, which bounds the workspace (16 bytes per row of a pass, plus the piece sums). + * + * @tparam DataT Accumulation and output type + * @tparam IndexT Index type + * @tparam DataIterator Random-access iterator over the row-major input; its values are converted + * to DataT (e.g. a transform iterator that maps int8 data to float) + * @tparam LabelsIterator Random-access iterator over the labels + * + * @param[in] handle RAFT resources handle + * @param[in] X Row-major input [n_rows x n_cols] + * @param[in] weights Per-row weights [n_rows], or nullptr for unit weights + * @param[in] labels Label of each row, in [0, n_clusters) + * @param[in] n_rows Number of rows of X + * @param[in] n_cols Number of columns of X + * @param[in] n_clusters Number of labels + * @param[inout] centroid_sums Row-major output [n_clusters x n_cols] + * @param[inout] weight_per_cluster Output [n_clusters], or nullptr to skip the weight sums + * @param[inout] workspace Scratch buffer, resized as needed + * @param[in] reset_sums Overwrite the outputs if true, add to them otherwise + */ +template +void weighted_sums_by_label(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + IndexT n_rows, + IndexT n_cols, + IndexT n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + if (n_clusters <= 0) { return; } + RAFT_EXPECTS(static_cast(n_clusters) < (uint64_t{1} << 31) && + static_cast(n_cols) < (uint64_t{1} << 31), + "weighted_sums_by_label: too many clusters or columns"); + // Passes are accumulated in a fixed order, so the result stays deterministic. There is always at + // least one pass, so that the outputs are reset even when there are no rows. + const auto total_rows = static_cast(n_rows); + int64_t start = 0; + do { + const int64_t pass_rows = std::min(kMaxRowsPerPass, total_rows - start); + weighted_sums_by_label_pass(handle, + X + start * static_cast(n_cols), + weights != nullptr ? weights + start : nullptr, + labels + start, + static_cast(pass_rows), + static_cast(n_cols), + static_cast(n_clusters), + centroid_sums, + weight_per_cluster, + workspace, + reset_sums && start == 0); + start += pass_rows; + } while (start < total_rows); +} + } // namespace cuvs::cluster::kmeans::detail::sums_by_label diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ac00770d58..efedd05b1b 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -21,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -43,7 +44,13 @@ #include #include #include +#include +#include #include +#include +#include +#include +#include #include #include @@ -122,6 +129,73 @@ struct WeightedDistanceOp { } }; +/** + * @brief A `cuda::mr::resource` that hands out temporary storage from one caller-owned buffer. + * + * The buffer only grows, and deallocation is a no-op, so repeated calls reuse the same memory + * instead of allocating (and possibly synchronizing) every time. Only one allocation may be live + * at a time. + */ +class workspace_buffer_resource { + public: + explicit workspace_buffer_resource(rmm::device_uvector& buffer) : buffer_{&buffer} {} + + void* allocate(cuda::stream_ref stream, size_t bytes, size_t alignment) + { + const size_t needed = bytes + alignment; + if (buffer_->size() < needed) { buffer_->resize(needed, stream); } + const auto ptr = reinterpret_cast(buffer_->data()); + return reinterpret_cast(raft::alignTo(ptr, alignment)); + } + void deallocate(cuda::stream_ref, void*, size_t, size_t) noexcept {} + void* allocate_sync(size_t bytes, size_t alignment) + { + return allocate(buffer_->stream(), bytes, alignment); + } + void deallocate_sync(void*, size_t, size_t) noexcept {} + bool operator==(const workspace_buffer_resource& other) const noexcept + { + return buffer_ == other.buffer_; + } + bool operator!=(const workspace_buffer_resource& other) const noexcept + { + return !(*this == other); + } + + private: + rmm::device_uvector* buffer_; +}; + +/** + * @brief Device-wide reduction that is bitwise reproducible from run to run. + * + * Requests `cuda::execution::determinism::run_to_run` explicitly (also cub's current default): + * cub combines the partial results in a fixed order on a given GPU, with no atomics. This is + * stronger than raft's mapThenSumReduce, which combines block results with atomics. + * + * `gpu_to_gpu` (an RFA-based sum that is identical across architectures) is not requested: the + * inputs of these reductions (distances from GEMM / cuTile TF32 kernels and the centroids derived + * from them) already differ between architectures, so a fit cannot be reproducible across GPUs. + * + * Temporary storage comes from `workspace`, which is resized as needed. + */ +template +void run_to_run_reduce(raft::resources const& handle, + InputIterator input, + IndexT size, + OutputT* output, + ReductionOpT reduction_op, + OutputT init, + rmm::device_uvector& workspace) +{ + auto env = + cuda::std::execution::env{cuda::stream_ref{raft::resource::get_cuda_stream(handle)}, + cuda::execution::require(cuda::execution::determinism::run_to_run), + cuda::std::execution::prop{cuda::mr::get_memory_resource, + workspace_buffer_resource{workspace}}}; + RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(input, output, size, reduction_op, init, env)); +} + template struct SquaredDiffOp { const DataT* a; @@ -207,17 +281,14 @@ void weightSum( DataT wt_sum_h = DataT{0}; if constexpr (raft::is_device_mdspan_v) { - // cub::DeviceReduce is run-to-run deterministic, unlike the atomics in mapThenSumReduce. - size_t temp_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceReduce::Sum( - nullptr, temp_bytes, weight.data_handle(), d_wt_sum.data_handle(), n_samples, stream.get())); - rmm::device_uvector temp(temp_bytes, stream); - RAFT_CUDA_TRY(cub::DeviceReduce::Sum(temp.data(), - temp_bytes, - weight.data_handle(), - d_wt_sum.data_handle(), - n_samples, - stream.get())); + rmm::device_uvector workspace(0, stream); + run_to_run_reduce(handle, + weight.data_handle(), + n_samples, + d_wt_sum.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); if (check_positive) { raft::copy(&wt_sum_h, d_wt_sum.data_handle(), 1, stream); raft::resource::sync_stream(handle); @@ -292,29 +363,9 @@ void computeClusterCost(raft::resources const& handle, MainOpT main_op, ReductionOpT reduction_op) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cuda::transform_iterator itr(input, main_op); - - size_t temp_storage_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr, - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); - - workspace.resize(temp_storage_bytes, stream); - - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(workspace.data(), - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); + run_to_run_reduce( + handle, itr, size, clusterCost.data_handle(), reduction_op, OutputT(), workspace); } template void compute_centroid_shift(raft::resources const& handle, raft::device_matrix_view old_centroids, raft::device_matrix_view new_centroids, - raft::device_scalar_view sqrd_norm_out) + raft::device_scalar_view sqrd_norm_out, + rmm::device_uvector& workspace) { - // cub::DeviceReduce is run-to-run deterministic, unlike the atomics in mapThenSumReduce. - rmm::device_uvector workspace(0, raft::resource::get_cuda_stream(handle)); - computeClusterCost( + run_to_run_reduce( handle, - cuda::counting_iterator(0), + cuda::transform_iterator( + cuda::counting_iterator(0), + SquaredDiffOp{old_centroids.data_handle(), new_centroids.data_handle()}), static_cast(old_centroids.size()), - workspace, - sqrd_norm_out, - SquaredDiffOp{old_centroids.data_handle(), new_centroids.data_handle()}, - raft::add_op{}); + sqrd_norm_out.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); } /** diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index a86c30a087..4b25221f63 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -605,7 +605,8 @@ void mnmg_fit( dev_res, rank_centroids_const, raft::make_const_mdspan(new_centroids.view()), - sqrd_norm_error_dev.view()); + sqrd_norm_error_dev.view(), + workspace); raft::copy( rank_centroids.data_handle(), new_centroids.data_handle(), n_clusters * n_features, stream); From 18855627573bd2d22a73648b09acb258e21c026f Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 22:22:22 +0000 Subject: [PATCH 16/27] Explain unavailable top_1_nn paths in NNTest skip message --- cpp/tests/neighbors/distance_nn.cu | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 5ac80fda73..589550aa45 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -10,6 +10,7 @@ #include "../../src/distance/unfused_distance_nn.cuh" #include +#include #include #include #include @@ -116,7 +117,17 @@ class NNTest : public ::testing::TestWithParam> { workspace_size = m * n * sizeof(AccT); } - if (backend_unavailable) { GTEST_SKIP() << "Requested top_1_nn path is unavailable"; } + if (backend_unavailable) { + // cuTile cubins are embedded only for specific architectures and are not forward compatible + // across minor versions, so cuTile (and FP16 Auto, which has no non-cuTile fallback) is + // expected to be unavailable on devices such as GB300 (sm_103) and GB10 (sm_121). + const auto& prop = raft::resource::get_device_properties(handle); + GTEST_SKIP() << "Requested top_1_nn path is unavailable on compute capability " << prop.major + << "." << prop.minor + << (cuvs::detail::jit_lto::cutile_launch_available_on_current_device() + ? "" + : " (this build has no cuTile cubin for this architecture)"); + } // Reset buffer if constexpr (std::is_same_v>) { From 37f4477e8b8c961a435ad2ae55616228a7aae48f Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 00:47:38 +0000 Subject: [PATCH 17/27] Use a recall-based check in TestCuVSRandomizedHNSWVectorSearch The test required every hit of a single query to be among the closest 3 * topK vectors. The search is approximate, and on this test's data (uniform random vectors with up to 2048 dimensions) a top-1 to top-3 search occasionally returns a valid neighbor just outside that window. Check all queries instead, allow 2 such misses in total, and keep hard checks that every result is a distinct document with a vector, scored by its own vector, in score order, with min(topK, vectors) results. Also raise TOP_K_LIMIT to 128: the search does not fail beyond 64. --- .../TestCuVSRandomizedHNSWVectorSearch.java | 77 ++++--------------- .../com/nvidia/cuvs/lucene/TestUtils.java | 62 +++++++++++++-- 2 files changed, 74 insertions(+), 65 deletions(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java index 2f77f19b8d..9f5f01d9c1 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestCuVSRandomizedHNSWVectorSearch.java @@ -4,17 +4,13 @@ */ package com.nvidia.cuvs.lucene; +import static com.nvidia.cuvs.lucene.TestUtils.assertSearchFindsNearestNeighbors; import static com.nvidia.cuvs.lucene.TestUtils.generateDataset; import static com.nvidia.cuvs.lucene.TestUtils.generateQueries; import static com.nvidia.cuvs.lucene.ThreadLocalCuVSResourcesProvider.isSupported; import java.io.IOException; -import java.util.ArrayList; -import java.util.Arrays; -import java.util.List; -import java.util.Map; import java.util.Random; -import java.util.TreeMap; import java.util.logging.Level; import java.util.logging.Logger; import org.apache.lucene.codecs.Codec; @@ -57,7 +53,9 @@ public class TestCuVSRandomizedHNSWVectorSearch extends LuceneTestCase { static int DATASET_SIZE_LIMIT = 1000; static int DIMENSIONS_LIMIT = 2048; static int NUM_QUERIES_LIMIT = 10; - static int TOP_K_LIMIT = 64; // TODO This fails beyond 64 + // A wider search is more accurate, not less: in the measurements described in testVectorSearch, + // no result was outside the closest 3 * topK for any topK from 8 to 200. + static int TOP_K_LIMIT = 128; static float[][] dataset; // Whether each document of the dataset was indexed with a vector; only those can be hits. static boolean[] hasVector; @@ -121,66 +119,25 @@ public void testVectorSearch() throws IOException { int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, dataset.length); - if (dataset.length < topK) topK = dataset.length; - float[][] queries = generateQueries(random, dataset[0].length, numQueries); - List> expected = generateExpectedResults(topK, dataset, queries); log.log(Level.FINE, "Dataset size: " + dataset.length + "x" + dataset[0].length); log.log(Level.FINE, "Query size: " + numQueries + "x" + queries[0].length); log.log(Level.FINE, "TopK: " + topK); - Query query = new KnnFloatVectorQuery("vector", queries[0], topK); - int correct[] = new int[topK]; - for (int i = 0; i < topK; i++) correct[i] = expected.get(0).get(i); - - ScoreDoc[] hits = searcher.search(query, topK).scoreDocs; - log.log(Level.FINE, "RESULTS: " + Arrays.toString(hits)); - log.log(Level.FINE, "EXPECTD: " + expected.get(0)); - - for (ScoreDoc hit : hits) { - log.log( - Level.FINE, "\t" + reader.storedFields().document(hit.doc).get("id") + ": " + hit.score); - } - - for (ScoreDoc hit : hits) { - int doc = Integer.parseInt(reader.storedFields().document(hit.doc).get("id")); - assertTrue("Result returned was not in topK*3: " + doc, expected.get(0).contains(doc)); - } - } - - private static List> generateExpectedResults( - int topK, float[][] dataset, float[][] queries) { - List> neighborsResult = new ArrayList<>(); - int dimensions = dataset[0].length; - - for (float[] query : queries) { - Map distances = new TreeMap<>(); - for (int j = 0; j < dataset.length; j++) { - if (hasVector[j] == false) { - continue; - } - double distance = 0; - for (int k = 0; k < dimensions; k++) { - distance += (query[k] - dataset[j][k]) * (query[k] - dataset[j][k]); - } - distances.put(j, (distance)); - } - - Map sorted = new TreeMap(distances); - log.log(Level.FINE, "EXPECTED: " + sorted); - - // Sort by distance and select the topK nearest neighbors - List neighbors = - distances.entrySet().stream() - .sorted(Map.Entry.comparingByValue()) - .map(Map.Entry::getKey) - .toList(); - neighborsResult.add(neighbors.subList(0, Math.min(topK * 3, neighbors.size()))); - } - - log.log(Level.FINE, "Expected results generated successfully."); - return neighborsResult; + // The search is approximate. Lucene searches each segment's HNSW graph with a beam as wide as + // topK, and in this data (uniform random vectors with up to 2048 dimensions, so all distances + // are within about 10% of each other) a narrow beam sometimes stops just short of the true + // neighbors. Over 27,003 random indexes from this test's distribution, with 10 queries each, a + // result fell outside the closest 3 * topK in 1.7% of queries at topK = 1, 0.21% at topK = 2, + // 0.03% at topK = 3, 3 times in 270,000 queries at topK = 5, and never at topK >= 8. Those + // results were at most rank 17, and typically only 1.3% farther away than the nearest + // neighbor. On the same segments, Lucene's own HNSW format missed in 22% of queries at + // topK = 1. Allowing 2 misses across all the queries keeps the false failure rate below 1e-5 + // per run. + int allowedMisses = 2; + assertSearchFindsNearestNeighbors( + searcher, "vector", "id", dataset, hasVector, queries, topK, allowedMisses); } @Test diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java index ce58fa908c..42a26583e0 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestUtils.java @@ -90,6 +90,38 @@ public static void assertSearchFindsNearestNeighbors( int topK, int allowedMisses) throws IOException { + assertSearchFindsNearestNeighbors( + searcher, vectorField, idField, dataset, null, queries, topK, allowedMisses); + } + + /** + * Like {@link #assertSearchFindsNearestNeighbors(IndexSearcher, String, String, float[][], + * float[][], int, int)}, for an index where only some rows of {@code dataset} were indexed with a + * vector: {@code indexed[i]} says whether row {@code i} was, and {@code null} means all of them. + * The closest vectors are ranked among the indexed rows only, since the others can never be + * returned. + * + *

Besides the approximate check, every query must return {@code min(topK, indexed rows)} + * distinct indexed documents in order of decreasing score, each scored by its own vector. + */ + public static void assertSearchFindsNearestNeighbors( + IndexSearcher searcher, + String vectorField, + String idField, + float[][] dataset, + boolean[] indexed, + float[][] queries, + int topK, + int allowedMisses) + throws IOException { + int indexedCount = dataset.length; + if (indexed != null) { + indexedCount = 0; + for (boolean isIndexed : indexed) { + if (isIndexed) indexedCount++; + } + } + int expectedResults = Math.min(topK, indexedCount); int acceptableRank = topK * 3; int misses = 0; StringBuilder found = new StringBuilder(); @@ -97,13 +129,23 @@ public static void assertSearchFindsNearestNeighbors( for (int q = 0; q < queries.length; q++) { KnnFloatVectorQuery query = new KnnFloatVectorQuery(vectorField, queries[q], topK); TopDocs results = searcher.search(query, topK); - assertEquals("results for query " + q, topK, results.scoreDocs.length); + assertEquals("results for query " + q, expectedResults, results.scoreDocs.length); - int[] ranks = distanceRanks(dataset, queries[q]); - int[] ids = new int[topK]; - int[] resultRanks = new int[topK]; - for (int i = 0; i < topK; i++) { + int[] ranks = distanceRanks(dataset, indexed, queries[q]); + int[] ids = new int[expectedResults]; + int[] resultRanks = new int[expectedResults]; + Set seen = new HashSet<>(); + for (int i = 0; i < expectedResults; i++) { ids[i] = Integer.parseInt(storedFields.document(results.scoreDocs[i].doc).get(idField)); + assertTrue( + "id " + ids[i] + " for query " + q + " has no vector", + indexed == null || indexed[ids[i]]); + assertTrue("id " + ids[i] + " returned twice for query " + q, seen.add(ids[i])); + if (i > 0) { + assertTrue( + "results for query " + q + " are not sorted by score at position " + i, + results.scoreDocs[i].score <= results.scoreDocs[i - 1].score); + } float expectedScore = VectorSimilarityFunction.EUCLIDEAN.compare(queries[q], dataset[ids[i]]); assertEquals( @@ -142,6 +184,14 @@ public static void assertSearchFindsNearestNeighbors( * @return the 1-based rank of each vector, indexed like {@code dataset} */ public static int[] distanceRanks(float[][] dataset, float[] query) { + return distanceRanks(dataset, null, query); + } + + /** + * Like {@link #distanceRanks(float[][], float[])}, ranking only the rows where {@code indexed} is + * true ({@code null} means all rows). The other rows get rank {@link Integer#MAX_VALUE}. + */ + public static int[] distanceRanks(float[][] dataset, boolean[] indexed, float[] query) { double[] distances = new double[dataset.length]; for (int i = 0; i < dataset.length; i++) { for (int j = 0; j < query.length; j++) { @@ -151,11 +201,13 @@ public static int[] distanceRanks(float[][] dataset, float[] query) { } int[] byDistance = IntStream.range(0, dataset.length) + .filter(i -> indexed == null || indexed[i]) .boxed() .sorted(Comparator.comparingDouble(i -> distances[i])) .mapToInt(Integer::intValue) .toArray(); int[] ranks = new int[dataset.length]; + Arrays.fill(ranks, Integer.MAX_VALUE); for (int rank = 0; rank < byDistance.length; rank++) { ranks[byDistance[rank]] = rank + 1; } From 5af73320da734d57203bbb28d688b9bb0b47cbd2 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 22:48:42 +0000 Subject: [PATCH 18/27] Find bin2c next to the real nvcc so conda builds enable cuTile --- cpp/cmake/modules/generate_cutile_kernels.cmake | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/cpp/cmake/modules/generate_cutile_kernels.cmake b/cpp/cmake/modules/generate_cutile_kernels.cmake index 0578dd0589..df82d51193 100644 --- a/cpp/cmake/modules/generate_cutile_kernels.cmake +++ b/cpp/cmake/modules/generate_cutile_kernels.cmake @@ -34,14 +34,23 @@ function(_cutile_detect) return() endif() + # bin2c is installed next to the real nvcc. In conda environments, CUDAToolkit_BIN_DIR is + # targets//bin, which only contains an nvcc symlink, so also search the symlink target's + # directory. + set(bin2c_dirs ${CUDAToolkit_BIN_DIR}) + if(CUDAToolkit_NVCC_EXECUTABLE) + get_filename_component(nvcc_real_path "${CUDAToolkit_NVCC_EXECUTABLE}" REALPATH) + get_filename_component(nvcc_real_dir "${nvcc_real_path}" DIRECTORY) + list(APPEND bin2c_dirs "${nvcc_real_dir}") + endif() find_program( CUVS_CUTILE_BIN2C NAMES bin2c - PATHS ${CUDAToolkit_BIN_DIR} + PATHS ${bin2c_dirs} NO_DEFAULT_PATH ) if(NOT CUVS_CUTILE_BIN2C) - message(STATUS "cuTile disabled: bin2c not found in ${CUDAToolkit_BIN_DIR}.") + message(STATUS "cuTile disabled: bin2c not found in ${bin2c_dirs}.") return() endif() From c6f44ea8ffe864f83cf88295f2e15f05332d027d Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 02:13:40 +0000 Subject: [PATCH 19/27] Do not treat pageable host memory as device memory on HMM systems Several checks decided that memory is device-accessible when cudaPointerGetAttributes() returns a non-null devicePointer. On systems with HMM, devicePointer is also set for pageable (unregistered) host memory once the process has accessed pageable memory from the GPU, so ordinary host buffers were classified as device memory: - check_pointer_residency returned host_and_device, so utils::memzero issued cudaMemsetAsync on host memory (for example the accum_sorted_sizes host_vector at the start of ivf_pq::build), and the PQ/VPQ encoders launched kernels directly on pageable host input. - PQ train_from_rows, expect_device_accessible_data_handle and the CAGRA benchmark wrapper used the same rule. GPU accesses to pageable memory under HMM are serviced through GPU page faults, and the memset at the start of ivf_pq::build was observed to never complete on H100 CI runners. Add cuvs::util::detail::is_device_accessible, which classifies memory by cudaPointerAttributes::type: device and managed memory are device accessible, registered host memory is when it is mapped, and pageable memory is only when the device reads it through the host page tables (cudaDevAttrPageableMemoryAccessUsesHostPageTables, i.e. ATS systems such as Grace), which keeps the previous behavior there. Use it at all four sites. --- cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h | 8 ++- cpp/include/cuvs/neighbors/common.hpp | 7 ++- cpp/include/cuvs/util/pointer_attributes.hpp | 58 ++++++++++++++++++++ cpp/src/neighbors/detail/ann_utils.cuh | 4 +- cpp/src/preprocessing/quantize/pq.cu | 5 +- 5 files changed, 76 insertions(+), 6 deletions(-) create mode 100644 cpp/include/cuvs/util/pointer_attributes.hpp diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index da5e5baa9d..23c9dd56bd 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -66,13 +66,15 @@ inline void maybe_log_cagra_persistent_concurrency_hint(bool persistent_search) threads_rec); } -/** Whether the GPU can dereference this pointer: device, managed or mapped host memory. */ +/** + * Whether the GPU can dereference this pointer: device, managed or mapped host memory, or pageable + * memory on systems where the device reads it through the host page tables (ATS). + */ inline auto is_device_accessible(const void* ptr) -> bool { cudaPointerAttributes attrs{}; RAFT_CUDA_TRY(cudaPointerGetAttributes(&attrs, ptr)); - return attrs.type == cudaMemoryTypeDevice || attrs.type == cudaMemoryTypeManaged || - attrs.devicePointer != nullptr; + return cuvs::util::detail::is_device_accessible(attrs); } /** diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index baa5b44da2..abb661128f 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include @@ -1115,7 +1116,11 @@ template if (ptr_attrs.type == cudaMemoryTypeDevice || ptr_attrs.type == cudaMemoryTypeManaged) { return const_cast(src.data_handle()); } - auto* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + // Host memory needs a device mapping; pageable memory only counts when the device reads it + // through the host page tables (ATS), not through HMM page faults. + auto* device_ptr = cuvs::util::detail::is_device_accessible(ptr_attrs) + ? reinterpret_cast(ptr_attrs.devicePointer) + : nullptr; RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); return device_ptr; } diff --git a/cpp/include/cuvs/util/pointer_attributes.hpp b/cpp/include/cuvs/util/pointer_attributes.hpp new file mode 100644 index 0000000000..7a2fa689e7 --- /dev/null +++ b/cpp/include/cuvs/util/pointer_attributes.hpp @@ -0,0 +1,58 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include + +#include + +#include +#include + +namespace cuvs::util::detail { + +/** + * Whether `device` reads pageable (unregistered) host memory through the host page tables (ATS, + * e.g. Grace Hopper / Grace Blackwell). The result is cached per device. + */ +inline auto pageable_memory_uses_host_page_tables(int device) -> bool +{ + static std::mutex mutex; + static std::unordered_map cache; + std::lock_guard lock(mutex); + if (auto it = cache.find(device); it != cache.end()) { return it->second; } + int uses_host_page_tables = 0; + RAFT_CUDA_TRY(cudaDeviceGetAttribute( + &uses_host_page_tables, cudaDevAttrPageableMemoryAccessUsesHostPageTables, device)); + cache.emplace(device, uses_host_page_tables != 0); + return uses_host_page_tables != 0; +} + +/** + * Whether memory described by `attr` may be dereferenced in device code on the current device. + * + * Device and managed memory are device-accessible, and registered (pinned) host memory is when it + * is mapped (`devicePointer` is set). Pageable (unregistered) host memory is device-accessible only + * when the device reads it through the host page tables (ATS). On systems with HMM, + * `cudaPointerGetAttributes` can also report a non-null `devicePointer` for pageable memory (for + * example after the first GPU access to pageable memory in the process), but device accesses to it + * are then serviced through GPU page faults, so it is treated as host-only there. + */ +inline auto is_device_accessible(const cudaPointerAttributes& attr) -> bool +{ + switch (attr.type) { + case cudaMemoryTypeDevice: + case cudaMemoryTypeManaged: return true; + case cudaMemoryTypeHost: return attr.devicePointer != nullptr; + default: { + if (attr.devicePointer == nullptr) { return false; } + int device = 0; + RAFT_CUDA_TRY(cudaGetDevice(&device)); + return pageable_memory_uses_host_page_tables(device); + } + } +} + +} // namespace cuvs::util::detail diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh index deddc47dcb..e7fdc17c3a 100644 --- a/cpp/src/neighbors/detail/ann_utils.cuh +++ b/cpp/src/neighbors/detail/ann_utils.cuh @@ -6,6 +6,7 @@ #pragma once #include +#include #include #include #include @@ -60,7 +61,8 @@ struct pointer_residency_count { auto [on_device, on_host] = pointer_residency_count::run(ptrs...); cudaPointerAttributes attr; RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); - if (attr.devicePointer || attr.type == cudaMemoryTypeDevice) { ++on_device; } + // Not `attr.devicePointer != nullptr`: with HMM it can be set for pageable host memory. + if (cuvs::util::detail::is_device_accessible(attr)) { ++on_device; } if (attr.hostPointer || attr.type == cudaMemoryTypeUnregistered) { ++on_host; } return std::make_tuple(on_device, on_host); } diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 20b8f21d36..6488dfa996 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -6,6 +6,7 @@ #include "./detail/pq.cuh" #include +#include #include #include @@ -88,7 +89,9 @@ auto train_from_rows(raft::resources const& res, { cudaPointerAttributes ptr_attrs; RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src_ptr)); - auto const* device_ptr = reinterpret_cast(ptr_attrs.devicePointer); + auto const* device_ptr = cuvs::util::detail::is_device_accessible(ptr_attrs) + ? reinterpret_cast(ptr_attrs.devicePointer) + : nullptr; if (device_ptr == nullptr) { // A host mdspan makes training subsample the rows and encoding stream them in bounded batches, // so the dense dataset is never staged on the device. From 0a27e5fa83582aec739f4dafb34576d95f170b0a Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 03:10:33 +0000 Subject: [PATCH 20/27] Query the entry node's own vector in TestAcceleratedHNSWMultiLayerRoundTrip The test looked up the query vector with vectors[entryNode], treating the graph ordinal as the insertion id. A randomized merge policy can reorder documents during forceMerge(1), so the query was then an arbitrary document's vector and the approximate binary-quantized search could legitimately miss it. Look up the entry node's stored id instead and match the hit by docID. --- .../lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java index b9f39e6c05..5ec3814cd9 100644 --- a/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java +++ b/java/cuvs-lucene/src/test/java/com/nvidia/cuvs/lucene/TestAcceleratedHNSWMultiLayerRoundTrip.java @@ -111,13 +111,17 @@ public void testThreeLevelGraphCanBeReopenedAndSearched() throws Exception { int queryNode = graph.entryNode(); assertTrue(nodesByLevel.get(2).contains(queryNode)); + // Graph nodes are vector ordinals, which equal docIDs in this dense single-segment index. + // The randomized merge policy can reorder documents during forceMerge, so the entry node's + // vector must be looked up through its stored id rather than by indexing the input array + // with the ordinal. + int queryId = Integer.parseInt(leaf.storedFields().document(queryNode).get("id")); IndexSearcher searcher = new IndexSearcher(reader); - var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[queryNode], 10), 10); + var hits = searcher.search(new KnnFloatVectorQuery(FIELD, vectors[queryId], 10), 10); assertEquals(10, hits.scoreDocs.length); - String queryNodeId = Integer.toString(queryNode); boolean foundQueryNode = false; for (var hit : hits.scoreDocs) { - foundQueryNode |= queryNodeId.equals(searcher.storedFields().document(hit.doc).get("id")); + foundQueryNode |= hit.doc == queryNode; } assertTrue("the entry-node vector must be returned for its own query", foundQueryNode); } From ae0f9c2f46da232a679508e1d6d68d964685ad06 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 03:14:13 +0000 Subject: [PATCH 21/27] Guard IVF-Flat cosine coarse search against zero-norm centers k-means gives an empty cluster a zero center, and the cosine coarse distance to it was NaN. select_k never picks NaN columns, so when fewer than n_probes lists have a finite distance the remaining probes are padded with list 0, which is then scanned several times and returns duplicated neighbors. Define the cosine distance to a zero-norm center as 1, as IVF-SQ already does, and add a regression test. --- .../neighbors/ivf_flat/ivf_flat_search.cuh | 6 +- .../ann_ivf_flat/test_float_int64_t.cu | 80 +++++++++++++++++++ 2 files changed, 85 insertions(+), 1 deletion(-) diff --git a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh index 6117935712..6bb3b2ffe8 100644 --- a/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh +++ b/cpp/src/neighbors/ivf_flat/ivf_flat_search.cuh @@ -171,7 +171,11 @@ void search_impl(raft::resources const& handle, [=] __device__(const uint32_t idx, const float dist) { const auto query = idx / n_lists; const auto cluster = idx % n_lists; - return dist / (q_norm_ptr[query] * index_center_norm_ptr[cluster]); + // A list whose k-means cluster ended up empty has a zero center. Define the cosine + // distance to a zero-norm vector as 1 (similarity 0), as IVF-SQ does: a NaN here is never + // selected, and with fewer finite lists than n_probes the probes get padded with list 0. + const float denom = q_norm_ptr[query] * index_center_norm_ptr[cluster]; + return denom > 0.0f ? dist / denom : 0.0f; }, raft::make_const_mdspan(distance_buffer_dev_view)); } diff --git a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu index fee22db3c5..a70d41a09f 100644 --- a/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu +++ b/cpp/tests/neighbors/ann_ivf_flat/test_float_int64_t.cu @@ -10,10 +10,13 @@ #include #include #include +#include +#include #include #include #include +#include namespace cuvs::neighbors::ivf_flat { @@ -48,6 +51,83 @@ TEST(IVFFlatSerialization, StreamRoundTrip) EXPECT_EQ(loaded.n_lists(), index.n_lists()); } +// k-means gives an empty cluster a zero center. With the cosine metric, the coarse search must +// still rank such lists instead of producing NaN distances: when fewer than n_probes lists have a +// finite distance, the NaN columns are never selected and the remaining probe slots are padded +// with list 0, which then is scanned several times and returns duplicated neighbors. +TEST(AnnIVFFlatTest, CosineZeroNormCenters) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + + constexpr int64_t n_rows = 10000; + constexpr int64_t n_queries = 500; + constexpr int64_t dim = 8; + constexpr uint32_t n_lists = 1024; + constexpr uint32_t n_zero_lists = 1000; + constexpr uint32_t n_probes = 40; + constexpr uint32_t k = 16; + + auto data = raft::make_device_matrix(handle, n_rows, dim); + auto queries = raft::make_device_matrix(handle, n_queries, dim); + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, data.data_handle(), data.size(), 0.1f, 2.0f); + raft::random::uniform(handle, rng, queries.data_handle(), queries.size(), 0.1f, 2.0f); + + index_params params; + params.n_lists = n_lists; + params.metric = cuvs::distance::DistanceType::CosineExpanded; + params.add_data_on_build = false; + auto idx = build(handle, params, raft::make_const_mdspan(data.view())); + + // Zero the centers of lists 1..n_zero_lists, as if k-means had left these clusters empty. Fewer + // lists than n_probes keep a non-zero center, and list 0 is one of them. + std::vector centers_h(idx.centers().size()); + raft::update_host(centers_h.data(), idx.centers().data_handle(), centers_h.size(), stream); + raft::resource::sync_stream(handle); + std::fill(centers_h.begin() + dim, centers_h.begin() + (n_zero_lists + 1) * dim, 0.0f); + raft::update_device(idx.centers().data_handle(), centers_h.data(), centers_h.size(), stream); + extend(handle, raft::make_const_mdspan(data.view()), std::nullopt, &idx); + + // All data rows are in the lists with a non-zero center, and all of those lists are probed. + search_params sparams; + sparams.n_probes = n_probes; + auto neighbors = raft::make_device_matrix(handle, n_queries, k); + auto distances = raft::make_device_matrix(handle, n_queries, k); + search(handle, + sparams, + idx, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + auto naive_neighbors = raft::make_device_matrix(handle, n_queries, k); + auto naive_distances = raft::make_device_matrix(handle, n_queries, k); + cuvs::neighbors::naive_knn(handle, + naive_distances.data_handle(), + naive_neighbors.data_handle(), + queries.data_handle(), + data.data_handle(), + n_queries, + n_rows, + dim, + k, + params.metric); + + std::vector neighbors_h(n_queries * k), naive_neighbors_h(n_queries * k); + std::vector distances_h(n_queries * k), naive_distances_h(n_queries * k); + raft::update_host(neighbors_h.data(), neighbors.data_handle(), neighbors_h.size(), stream); + raft::update_host(distances_h.data(), distances.data_handle(), distances_h.size(), stream); + raft::update_host( + naive_neighbors_h.data(), naive_neighbors.data_handle(), naive_neighbors_h.size(), stream); + raft::update_host( + naive_distances_h.data(), naive_distances.data_handle(), naive_distances_h.size(), stream); + raft::resource::sync_stream(handle); + + EXPECT_TRUE(eval_neighbours( + naive_neighbors_h, neighbors_h, naive_distances_h, distances_h, n_queries, k, 0.001, 0.99)); +} + TEST(AnnIVFFlatTest, RepeatedExtendCopyPreservesSharedListWithinCapacity) { raft::resources handle; From d410e9f530418e630067c451f5a6ce5a1bd1ef65 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 09:52:47 +0000 Subject: [PATCH 22/27] Re-enable CagraRandomizedIT with a recall-based check The test required the first 5 results of every query to be among the exact 2k+10 nearest neighbors. On uniformly random data in hundreds of dimensions the nearest neighbors are nearly equidistant, so CAGRA can legitimately miss that window now and then (#631). The other reported failures (#1048, #1468) were filtered-out rows leaking into the results through the prefilter buffer that BitSet.toLongArray() truncated, fixed in #2375. Check what must always hold instead: every query returns min(k, passing rows) distinct, valid rows that pass the prefilter, sorted, with their exact distances, and a copy of a row finds that row first. Check recall only as a mean over all queries of each dataset memory kind, with bounds set from the distribution measured in CI. Also attach the padded device dataset that CAGRA search now requires. --- .../com/nvidia/cuvs/CagraRandomizedIT.java | 275 +++++++++++++----- 1 file changed, 207 insertions(+), 68 deletions(-) diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java index a38a12ea77..4cb1bdc89e 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraRandomizedIT.java @@ -5,29 +5,55 @@ package com.nvidia.cuvs; import static com.carrotsearch.randomizedtesting.RandomizedTest.assumeTrue; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; import com.carrotsearch.randomizedtesting.RandomizedRunner; import com.nvidia.cuvs.CagraIndexParams.CagraGraphBuildAlgo; +import java.util.Arrays; import java.util.BitSet; import java.util.List; +import java.util.Map; import org.junit.Before; -import org.junit.Ignore; import org.junit.Test; import org.junit.runner.RunWith; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +/** + * Builds CAGRA indexes over random datasets and compares the search results with an exact search + * on the CPU. + * + *

CAGRA is an approximate index, and on uniformly random data in hundreds or thousands of + * dimensions the nearest neighbors are barely closer than the rest, so a single query can + * legitimately miss some or all of its exact neighbors. The test therefore checks what must hold + * for every query (the right number of distinct, valid rows that pass the prefilter, returned in + * order with their correct distances), and checks recall only as a mean over many queries. Each + * case also searches for a copy of one of its rows, which the search should return first. + * + *

This test used to require the first 5 results of every query to be among the exact {@code + * 2k + 10} nearest neighbors. That failed now and then for legitimate approximate results (#631), + * and it also caught the Java bindings leaking filtered-out rows (#1048, #1468, fixed in #2375), + * which the prefilter check here reports directly. + */ @RunWith(RandomizedRunner.class) public class CagraRandomizedIT extends CuVSTestCase { private static final Logger log = LoggerFactory.getLogger(CagraRandomizedIT.class); - @Before - public void setup() { - assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); - initializeRandom(); - log.trace("Random context initialized for test."); - } + // Lower bounds for the mean recall@k over the queries of each dataset memory kind, without and + // with a prefilter. In CI (RTX PRO 6000, CUDA 13), about 300 runs of this test measured means of + // 0.83 (standard deviation 0.022, lowest 0.76) without a prefilter and 0.94 (standard deviation + // 0.011, lowest 0.90) with one. The bounds sit 8 and 12 standard deviations below the means, so + // they only fail when results are wrong rather than approximate. + private static final double MIN_MEAN_RECALL_UNFILTERED = 0.65; + private static final double MIN_MEAN_RECALL_FILTERED = 0.8; + // Minimum fraction of the searches for a copy of a row that return that row first. All of more + // than 100,000 such searches did in CI. + private static final double MIN_SELF_MATCH_RATE = 0.99; + // Returned distances are computed in float on the GPU; allow for the difference in rounding + // from the double-precision distances computed here (at most 2.4e-7 in CI). + private static final double DISTANCE_RELATIVE_TOLERANCE = 1e-4; enum TestDatasetMemoryKind { HEAP, @@ -35,17 +61,70 @@ enum TestDatasetMemoryKind { DEVICE } - @Ignore // https://github.com/nvidia/cuvs/issues/1468 + /** Recall accumulated per dataset memory kind, with and without a prefilter. */ + private final double[][] recallSum = new double[TestDatasetMemoryKind.values().length][2]; + + private final int[][] queryCount = new int[TestDatasetMemoryKind.values().length][2]; + + private int selfQueries; + private int selfMatches; + + @Before + public void setup() { + assumeTrue("not supported on " + System.getProperty("os.name"), isLinuxSupportedArch()); + initializeRandom(); + } + @Test public void testResultsTopKWithRandomValues() throws Throwable { - TestDatasetMemoryKind[] testDatasetMemoryKinds = { - TestDatasetMemoryKind.HEAP, TestDatasetMemoryKind.NATIVE, TestDatasetMemoryKind.DEVICE - }; for (int i = 0; i < 100; i++) { - for (var datasetMemoryKind : testDatasetMemoryKinds) { + for (var datasetMemoryKind : TestDatasetMemoryKind.values()) { tmpResultsTopKWithRandomValues(datasetMemoryKind); } } + + assertTrue( + String.format( + "Only %d of %d searches for a row returned it first", selfMatches, selfQueries), + selfMatches >= MIN_SELF_MATCH_RATE * selfQueries); + for (var kind : TestDatasetMemoryKind.values()) { + for (int filtered = 0; filtered < 2; filtered++) { + int n = queryCount[kind.ordinal()][filtered]; + if (n == 0) { + continue; + } + double meanRecall = recallSum[kind.ordinal()][filtered] / n; + double minMeanRecall = + filtered == 1 ? MIN_MEAN_RECALL_FILTERED : MIN_MEAN_RECALL_UNFILTERED; + log.info( + "{} {} prefilter: mean recall@k {} over {} queries", + kind, + filtered == 1 ? "with" : "without", + meanRecall, + n); + assertTrue( + String.format( + "Mean recall@k %f over %d queries (%s, prefilter=%b) is below %f", + meanRecall, n, kind, filtered == 1, minMeanRecall), + meanRecall >= minMeanRecall); + } + } + } + + /** + * Gives the index the device-padded vectors CAGRA search needs. A device matrix whose rows are + * already padded can only be viewed, not copied. + */ + private static AutoCloseable attachPaddedDataset(CagraIndex index, CuVSMatrix deviceVectors) + throws Throwable { + if (CagraIndex.isPaddedDataset(deviceVectors)) { + var view = index.makePaddedDatasetView(deviceVectors); + index.updateDataset(view); + return view; + } + var dataset = index.makePaddedDataset(deviceVectors); + index.updateDataset(dataset); + return dataset; } private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryKind) @@ -53,65 +132,43 @@ private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryK int DATASET_SIZE_LIMIT = 10_000; int DIMENSIONS_LIMIT = 2048; int NUM_QUERIES_LIMIT = 10; - int TOP_K_LIMIT = 64; // nocommit This fails beyond 64 + // The default CAGRA search parameters (itopk_size = 64) bound the number of results. + int TOP_K_LIMIT = 64; - int datasetSize = - random.nextInt(DATASET_SIZE_LIMIT) + 2; // A CAGRA graph needs at least 2 rows. + // A CAGRA graph needs at least 2 rows. + int datasetSize = random.nextInt(DATASET_SIZE_LIMIT) + 2; int dimensions = random.nextInt(DIMENSIONS_LIMIT) + 1; int numQueries = random.nextInt(NUM_QUERIES_LIMIT) + 1; int topK = Math.min(random.nextInt(TOP_K_LIMIT) + 1, datasetSize); boolean usePrefilter = random.nextBoolean(); - if (datasetSize < topK) datasetSize = topK; - - BitSet sharedPrefilter = null; + BitSet prefilter = null; if (usePrefilter) { - sharedPrefilter = new BitSet(datasetSize); + prefilter = new BitSet(datasetSize); for (int j = 0; j < datasetSize; j++) { - sharedPrefilter.set(j, random.nextBoolean()); + prefilter.set(j, random.nextBoolean()); } } - BitSet[] prefilters = null; - if (sharedPrefilter != null) { - prefilters = new BitSet[numQueries]; - for (int i = 0; i < numQueries; i++) { - prefilters[i] = sharedPrefilter; - } - } - - // Generate a random dataset float[][] vectors = generateData(random, datasetSize, dimensions); - - // Generate random query vectors float[][] queries = generateData(random, numQueries, dimensions); - - log.debug("Dataset size: {}x{}", datasetSize, dimensions); - log.debug("Query size: {}x{}", numQueries, dimensions); - log.debug("TopK: {}", topK); - log.debug("Use memory dataset: " + datasetMemoryKind.name()); - - // Debugging: Log dataset and queries - if (log.isDebugEnabled()) { - log.debug("Dataset:"); - for (float[] row : vectors) { - log.debug(java.util.Arrays.toString(row)); - } - log.debug("Queries:"); - for (float[] query : queries) { - log.debug(java.util.Arrays.toString(query)); - } + // One more query is a copy of a row that passes the prefilter, if any does. + int selfRow = random.nextInt(datasetSize); + if (prefilter != null && !prefilter.get(selfRow)) { + int next = prefilter.nextSetBit(selfRow); + selfRow = next >= 0 ? next : prefilter.previousSetBit(selfRow); + } + if (selfRow >= 0) { + queries = Arrays.copyOf(queries, numQueries + 1); + queries[numQueries] = vectors[selfRow].clone(); } - // Sanity checks - assert vectors.length > 0 : "Dataset is empty."; - assert queries.length > 0 : "Queries are empty."; - assert dimensions > 0 : "Invalid dimensions."; - assert topK > 0 && topK <= datasetSize : "Invalid topK value."; - // Generate expected results using brute force - List> expected = generateExpectedResults(topK, vectors, queries, prefilters, log); + String caseDescription = + String.format( + "dataset %dx%d, %d queries, k=%d, prefilter=%b, %s", + datasetSize, dimensions, numQueries, topK, usePrefilter, datasetMemoryKind); + log.debug("Case: {}", caseDescription); - // Create CuVS index and query try (CuVSResources resources = CheckedCuVSResources.create()) { CagraIndexParams indexParams = new CagraIndexParams.Builder() @@ -143,35 +200,117 @@ private void tmpResultsTopKWithRandomValues(TestDatasetMemoryKind datasetMemoryK .withIndexParams(indexParams) .build(); } else { - assert datasetMemoryKind == TestDatasetMemoryKind.HEAP; index = CagraIndex.newBuilder(resources) .withDataset(vectors) .withIndexParams(indexParams) .build(); } - log.trace("Index built successfully."); - try (var queryVectors = CuVSMatrix.ofArray(queries)) { - // Execute search and retrieve results + // CAGRA search needs the device-padded vectors, whatever memory the index was built from. + try (var hostVectors = CuVSMatrix.ofArray(vectors); + var deviceVectors = hostVectors.toDevice(resources); + var indexDataset = attachPaddedDataset(index, deviceVectors); + var queryVectors = CuVSMatrix.ofArray(queries)) { CagraQuery.Builder queryBuilder = new CagraQuery.Builder(resources) .withQueryVectors(queryVectors) .withTopK(topK) .withSearchParams(new CagraSearchParams.Builder().build()); - - if (sharedPrefilter != null) { - queryBuilder.withPrefilter(sharedPrefilter, datasetSize); + if (prefilter != null) { + queryBuilder.withPrefilter(prefilter, datasetSize); + } + SearchResults results = index.search(queryBuilder.build()); + checkResults( + caseDescription, + datasetMemoryKind, + results, + vectors, + queries, + numQueries, + prefilter, + topK); + if (selfRow >= 0) { + var selfResult = results.getResults().get(numQueries); + selfQueries++; + if (!selfResult.isEmpty() && selfResult.keySet().iterator().next() == selfRow) { + selfMatches++; + } } - - CagraQuery query = queryBuilder.build(); - log.trace("Query built successfully. Executing search..."); - SearchResults results = index.search(query); - - compareResults(results, expected, topK, datasetSize, numQueries); } finally { index.close(); } } } + + /** Squared L2 distances from the query to every row, +inf for the rows the filter removes. */ + private static double[] exactDistances(float[] query, float[][] vectors, BitSet prefilter) { + double[] distances = new double[vectors.length]; + for (int j = 0; j < vectors.length; j++) { + if (prefilter != null && !prefilter.get(j)) { + distances[j] = Double.POSITIVE_INFINITY; + continue; + } + double sum = 0; + for (int d = 0; d < query.length; d++) { + double diff = (double) query[d] - (double) vectors[j][d]; + sum += diff * diff; + } + distances[j] = sum; + } + return distances; + } + + private void checkResults( + String caseDescription, + TestDatasetMemoryKind datasetMemoryKind, + SearchResults results, + float[][] vectors, + float[][] queries, + int numRandomQueries, + BitSet prefilter, + int topK) { + List> resultMaps = results.getResults(); + assertEquals(caseDescription, queries.length, resultMaps.size()); + int numPassing = prefilter == null ? vectors.length : prefilter.cardinality(); + int expectedCount = Math.min(topK, numPassing); + int filtered = prefilter == null ? 0 : 1; + + for (int q = 0; q < queries.length; q++) { + String where = caseDescription + ", query " + q; + double[] exact = exactDistances(queries[q], vectors, prefilter); + double[] sorted = exact.clone(); + Arrays.sort(sorted); + + // The results map holds one entry per distinct row, so the count also rules out duplicates. + Map result = resultMaps.get(q); + assertEquals(where + ": number of distinct results", expectedCount, result.size()); + + float previous = Float.NEGATIVE_INFINITY; + int found = 0; + for (var entry : result.entrySet()) { + int id = entry.getKey(); + float distance = entry.getValue(); + assertTrue(where + ": invalid row id " + id, id >= 0 && id < vectors.length); + assertTrue( + where + ": row " + id + " is excluded by the prefilter", + prefilter == null || prefilter.get(id)); + assertTrue(where + ": results are not sorted by distance", distance >= previous); + previous = distance; + double tolerance = DISTANCE_RELATIVE_TOLERANCE * Math.max(exact[id], 1.0); + assertTrue( + where + ": distance " + distance + " of row " + id + " should be " + exact[id], + Math.abs(distance - exact[id]) <= tolerance); + // Rows exactly as close as the k-th nearest neighbor count as found. + if (exact[id] <= sorted[expectedCount - 1]) { + found++; + } + } + // Recall is measured on the random queries only, not on the copy of a row. + if (expectedCount > 0 && q < numRandomQueries) { + recallSum[datasetMemoryKind.ordinal()][filtered] += (double) found / expectedCount; + queryCount[datasetMemoryKind.ordinal()][filtered]++; + } + } + } } From 9af2f6ba7137a12356c489b8137bf4fdbdc25d7e Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 12:13:27 +0000 Subject: [PATCH 23/27] Do not ignore a CAGRA prefilter that excludes every row CagraIndexImpl.search built the filter bitmap through Util.concatenate, which treats a BitSet with no bits set like a missing filter and sets every bit. A prefilter that excludes all rows therefore returned the excluded rows. CAGRA takes a single filter for all queries, so use its bits directly. The re-enabled CagraRandomizedIT hit this on tiny datasets whose random prefilter happened to exclude every row. --- .../nvidia/cuvs/internal/CagraIndexImpl.java | 12 ++--- .../nvidia/cuvs/CagraBuildAndSearchIT.java | 47 +++++++++++++++++++ 2 files changed, 51 insertions(+), 8 deletions(-) diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java index 491c5213fc..55ab018e18 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CagraIndexImpl.java @@ -14,7 +14,6 @@ import static com.nvidia.cuvs.internal.common.Util.CudaMemcpyKind.HOST_TO_DEVICE; import static com.nvidia.cuvs.internal.common.Util.buildMemorySegment; import static com.nvidia.cuvs.internal.common.Util.checkCuVSError; -import static com.nvidia.cuvs.internal.common.Util.concatenate; import static com.nvidia.cuvs.internal.common.Util.prepareTensor; import static com.nvidia.cuvs.internal.panama.headers_h.*; @@ -468,8 +467,6 @@ private static boolean isCagraPaddedLayout(CuVSMatrixInternal dataset) { == cagraRequiredRowWidth(logicalColumns, elementSizeBytes(dataset.dataType())); } - private static final BitSet[] EMPTY_PREFILTER_BITSET = new BitSet[0]; - /** * Invokes the native search_cagra_index via the Panama API for searching a * CAGRA index. @@ -495,9 +492,7 @@ public SearchResults search(CagraQuery query) throws Throwable { final long neighborsBytes = C_INT_BYTE_SIZE * numQueries * topK; final long distancesBytes = C_FLOAT_BYTE_SIZE * numQueries * topK; final boolean hasPreFilter = query.getPrefilter() != null; - final BitSet[] prefilters = - hasPreFilter ? new BitSet[] {query.getPrefilter()} : EMPTY_PREFILTER_BITSET; - final long prefilterDataLength = hasPreFilter ? query.getNumDocs() * prefilters.length : 0; + final long prefilterDataLength = hasPreFilter ? query.getNumDocs() : 0; final long prefilterLen = hasPreFilter ? (prefilterDataLength + 31) / 32 : 0; final long prefilterBytes = C_INT_BYTE_SIZE * prefilterLen; @@ -531,8 +526,9 @@ public SearchResults search(CagraQuery query) throws Throwable { cuvsFilter.type(prefilter, 0); // NO_FILTER cuvsFilter.addr(prefilter, 0); } else { - BitSet concatenatedFilters = concatenate(prefilters, query.getNumDocs()); - long[] filters = concatenatedFilters.toLongArray(); + // CAGRA takes a single filter for all queries. Use its bits as they are: concatenate() + // would read a BitSet with no bits set as "no filter" and return excluded rows. + long[] filters = query.getPrefilter().get(0, query.getNumDocs()).toLongArray(); var prefilterDataMemorySegment = buildMemorySegment(localArena, filters, (prefilterDataLength + 63) / 64); diff --git a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java index a4f08c47d4..4e276aeb11 100644 --- a/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java +++ b/java/cuvs-java/src/test/java/com/nvidia/cuvs/CagraBuildAndSearchIT.java @@ -692,6 +692,53 @@ public void testPrefilteringReducesResults() throws Throwable { } } + /** + * A prefilter that excludes every row leaves nothing to return. The search used to read a + * {@link BitSet} with no bits set as "no filter" and returned the excluded rows. + */ + @Test + public void testPrefilterExcludingEveryRowReturnsNothing() throws Throwable { + float[][] dataset = createSampleData(); + float[][] queries = {{0.48216683f, 0.0428398f}, {0.74021935f, 0.9209938f}}; + + CagraIndexParams indexParams = + new CagraIndexParams.Builder() + .withCagraGraphBuildAlgo(CagraGraphBuildAlgo.NN_DESCENT) + .withGraphDegree(2) + .withIntermediateGraphDegree(3) + .withMetric(CuvsDistanceType.L2Expanded) + .build(); + + try (CuVSResources resources = CheckedCuVSResources.create(); + CagraIndex index = + CagraIndex.newBuilder(resources) + .withDataset(dataset) + .withIndexParams(indexParams) + .build(); + var hostVectors = CuVSMatrix.ofArray(dataset); + var deviceVectors = hostVectors.toDevice(resources); + var indexDataset = index.makePaddedDataset(deviceVectors); + var queryVectors = CuVSMatrix.ofArray(queries)) { + index.updateDataset(indexDataset); + + for (var algo : + new CagraSearchParams.SearchAlgo[] { + CagraSearchParams.SearchAlgo.SINGLE_CTA, CagraSearchParams.SearchAlgo.MULTI_CTA + }) { + CagraQuery query = + new CagraQuery.Builder(resources) + .withTopK(2) + .withSearchParams(new CagraSearchParams.Builder().withAlgo(algo).build()) + .withQueryVectors(queryVectors) + .withPrefilter(new BitSet(dataset.length), dataset.length) + .build(); + + List> results = index.search(query).getResults(); + assertEquals(algo.toString(), List.of(Map.of(), Map.of()), results); + } + } + } + /** * Regression test for a bug in {@code CagraIndexImpl.search}: the {@code distances} output tensor * was described (dtype and buffer size) using the query vectors' data type instead of the From dc3230f697a27fe87ca089fc200d7a12cece7a3c Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 14:11:23 +0000 Subject: [PATCH 24/27] Make balanced k-means donor selection independent of scheduling The random donor selection took candidate points from a shared atomic counter, so which receiver got which donor depended on the order in which the warps ran. Derive each receiver's candidates from its index instead, with a stride that is coprime with the number of rows, so that every receiver still finds a donor whenever one exists. The pseudo-random donor offset rotated with a function-local static counter, so a fit depended on how many fits had run before it in the same process. Rotate it with the balancing step instead. --- cpp/src/cluster/detail/kmeans_balanced.cuh | 36 ++++++--- cpp/tests/cluster/kmeans_balanced.cu | 91 ++++++++++++++++++++++ 2 files changed, 116 insertions(+), 11 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_balanced.cuh b/cpp/src/cluster/detail/kmeans_balanced.cuh index bb96c9115b..881c0bd92a 100644 --- a/cpp/src/cluster/detail/kmeans_balanced.cuh +++ b/cpp/src/cluster/detail/kmeans_balanced.cuh @@ -43,6 +43,7 @@ #include #include +#include #include #include #include @@ -640,7 +641,7 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL IdxT average, MathT centroid_offset, IdxT seed, - IdxT* search_count, + IdxT stride, IdxT* update_count, MappingOpT mapping_op) { @@ -652,12 +653,17 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL IdxT i = n_rows; IdxT j = raft::laneId(); if (j == 0) { + // The candidate sequence of each receiver is derived from its index (not from a shared atomic + // counter), so the donor points do not depend on the order in which the warps run. + // `stride >= n_clusters`, so different receivers use different counters (and thus, before + // wrapping around, different candidates); `stride` is coprime with `n_rows`, so each receiver + // visits every row within `n_rows` attempts and finds a donor whenever one exists. IdxT attempt = 0; do { - auto old = atomicAdd(search_count, IdxT{1}); - auto candidate = - static_cast((static_cast(seed) * static_cast(old + 1)) % - static_cast(n_rows)); + auto counter = static_cast(attempt) * static_cast(stride) + + static_cast(receiver_cluster); + auto candidate = static_cast((static_cast(seed) * (counter + 1)) % + static_cast(n_rows)); if (static_cast(cluster_sizes[labels[candidate]]) >= average) { i = candidate; } ++attempt; } while (i >= n_rows && attempt < n_rows); @@ -768,6 +774,8 @@ __launch_bounds__((raft::WarpSize * BlockDimY)) RAFT_KERNEL * max_cluster_size > average_size * balance_upper_tolerance * balance_upper_tolerance > 1 * @param[in] centroid_offset offset from the donor cluster centroid towards a donor point + * @param[in] donor_selection how the donor points are chosen (see above) + * @param[in] balancing_step index of the current EM iteration; selects the pseudo-random offset * @param[in] mapping_op Mapping operation from T to MathT * @param[inout] device_memory memory resource to use for temporary allocations * @@ -791,6 +799,7 @@ auto adjust_centers(const raft::resources& handle, MathT balance_upper_tolerance, MathT centroid_offset, cuvs::cluster::kmeans::balanced_donor_selection donor_selection, + uint32_t balancing_step, MappingOpT mapping_op, rmm::device_async_resource_ref device_memory) -> bool { @@ -802,7 +811,6 @@ auto adjust_centers(const raft::resources& handle, 601, 659, 733, 809, 863, 941, 1013, 1069, 1151, 1223, 1291, 1373, 1451, 1511, 1583, 1657, 1733, 1811, 1889, 1987, 2053, 2129, 2213, 2287, 2357, 2423, 2531, 2617, 2687, 2741}; - static IdxT i_primes = 0; auto average = static_cast(n_rows) / static_cast(n_clusters); auto lower_threshold = average * balance_lower_tolerance; @@ -837,10 +845,12 @@ auto adjust_centers(const raft::resources& handle, auto n_pairs = static_cast(host_receiver_clusters.size()); if (n_pairs == 0) { return false; } + // Rotate the pseudo-random offset with the balancing step rather than with a process-wide + // counter, so that the donor points do not depend on how many fits ran before in the process. IdxT ofst; + size_t i_primes = balancing_step; do { - i_primes = (i_primes + 1) % kPrimes.size(); - ofst = kPrimes[i_primes]; + ofst = kPrimes[i_primes++ % kPrimes.size()]; } while (n_rows % ofst == 0); rmm::device_uvector receiver_clusters(n_pairs, stream, device_memory); @@ -851,8 +861,11 @@ auto adjust_centers(const raft::resources& handle, update_count.set_value_to_zero_async(stream); if (donor_selection == cuvs::cluster::kmeans::balanced_donor_selection::Random) { - rmm::device_scalar search_count(stream, device_memory); - search_count.set_value_to_zero_async(stream); + // The smallest stride >= n_clusters that is coprime with n_rows (see the kernel). + auto stride = std::max(n_clusters, 1); + while (std::gcd(static_cast(stride), static_cast(n_rows)) != 1) { + ++stride; + } const dim3 grid_dim(raft::ceildiv(n_clusters, static_cast(kBlockDimY)), 1, 1); adjust_centers_random_donor_kernel <<>>(centers, @@ -866,7 +879,7 @@ auto adjust_centers(const raft::resources& handle, static_cast(n_rows / n_clusters), centroid_offset, ofst, - search_count.data(), + stride, update_count.data(), mapping_op); return update_count.value(stream) > 0; // NB: rmm scalar performs the sync @@ -979,6 +992,7 @@ void balancing_em_iters(const raft::resources& handle, balance_upper_tolerance, static_cast(params.centroid_offset), params.donor_selection, + iter, mapping_op, device_memory)) { if (balancing_counter++ >= balancing_pullback) { diff --git a/cpp/tests/cluster/kmeans_balanced.cu b/cpp/tests/cluster/kmeans_balanced.cu index 3b7988e76b..185f400dff 100644 --- a/cpp/tests/cluster/kmeans_balanced.cu +++ b/cpp/tests/cluster/kmeans_balanced.cu @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../src/cluster/detail/kmeans_balanced.cuh" #include "../test_utils.h" #include @@ -12,16 +13,20 @@ #include #include #include +#include #include #include #include #include +#include #include #include +#include +#include #include #include @@ -301,4 +306,90 @@ KB_TEST((KmeansBalancedTest, KmeansBalancedTestFI8I32I32_SEP, inputsf_i32); +/* + * The balancing step must pick the same donor points for identical inputs: they may depend neither + * on the order in which the GPU runs the warps nor on how many fits ran before in the process. + * Most rows are packed in a few clusters and many clusters are empty, so that many receivers need a + * donor. + */ +void check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection donor) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + const int64_t n_rows = 10007; + const int64_t dim = 8; + const int64_t n_lists = 200; + auto dataset = raft::make_device_matrix(handle, n_rows, dim); + raft::random::RngState rng(1234ULL); + raft::random::uniform(handle, rng, dataset.data_handle(), n_rows * dim, -1.0f, 1.0f); + std::vector h_centers(n_lists * dim); + for (size_t i = 0; i < h_centers.size(); i++) { + h_centers[i] = static_cast(i % 97) / 97.0f; + } + + // 95% of the rows in clusters [0, 10), the rest spread over [10, 60); clusters [60, 200) empty. + std::vector h_labels(n_rows); + std::vector h_sizes(n_lists, 0); + for (int64_t i = 0; i < n_rows; i++) { + h_labels[i] = i < n_rows * 95 / 100 ? i % 10 : 10 + i % 50; + h_sizes[h_labels[i]]++; + } + auto labels = raft::make_device_vector(handle, n_rows); + auto sizes = raft::make_device_vector(handle, n_lists); + raft::update_device(labels.data_handle(), h_labels.data(), n_rows, stream); + raft::update_device(sizes.data_handle(), h_sizes.data(), n_lists, stream); + + const bool random = donor == cuvs::cluster::kmeans::balanced_donor_selection::Random; + const float lower_tol = random ? 0.25f : 0.7f; + const float upper_tol = random ? 4.0f : 1.5f; + const float offset = random ? 0.05f : 0.2f; + + std::vector> results; + for (int run = 0; run < 3; run++) { + auto centers = raft::make_device_matrix(handle, n_lists, dim); + raft::update_device(centers.data_handle(), h_centers.data(), h_centers.size(), stream); + bool adjusted = + cuvs::cluster::kmeans::detail::adjust_centers(handle, + centers.data_handle(), + n_lists, + dim, + dataset.data_handle(), + n_rows, + labels.data_handle(), + sizes.data_handle(), + lower_tol, + upper_tol, + offset, + donor, + /*balancing_step=*/3u, + raft::identity_op{}, + rmm::mr::get_current_device_resource_ref()); + ASSERT_TRUE(adjusted); + std::vector h(centers.size()); + raft::update_host(h.data(), centers.data_handle(), h.size(), stream); + raft::resource::sync_stream(handle); + results.push_back(std::move(h)); + } + // Both strategies move (at least) the first 60 empty clusters. + for (int64_t c = 60; c < 120; c++) { + ASSERT_NE(0, std::memcmp(&results[0][c * dim], &h_centers[c * dim], dim * sizeof(float))) + << "cluster " << c << " was not moved"; + } + for (int run = 1; run < 3; run++) { + ASSERT_EQ( + 0, std::memcmp(results[0].data(), results[run].data(), results[0].size() * sizeof(float))) + << "adjust_centers moved the centers differently for identical inputs (call " << run << ")"; + } +} + +TEST(KmeansBalancedAdjustCenters, ReproducibleSizeSorted) +{ + check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::SizeSorted); +} + +TEST(KmeansBalancedAdjustCenters, ReproducibleRandomDonor) +{ + check_adjust_centers_reproducible(cuvs::cluster::kmeans::balanced_donor_selection::Random); +} + } // namespace cuvs From 696e4612df88f9760fbfc8550f71404785ac6649 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 14:23:11 +0000 Subject: [PATCH 25/27] Move is_device_accessible out of the public include tree Define cuvs::neighbors::detail::is_device_accessible in src/neighbors/detail/pointer_attributes.cpp instead of the header-only cuvs/util/pointer_attributes.hpp. common.hpp declares it next to expect_device_accessible_data_handle, its only user in the public headers, and the per-device ATS attribute cache is now internal to libcuvs. --- cpp/CMakeLists.txt | 1 + cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h | 2 +- cpp/include/cuvs/neighbors/common.hpp | 18 ++++++++++--- cpp/src/neighbors/detail/ann_utils.cuh | 4 +-- .../neighbors/detail/pointer_attributes.cpp} | 25 ++++++++----------- cpp/src/preprocessing/quantize/pq.cu | 4 +-- 6 files changed, 30 insertions(+), 24 deletions(-) rename cpp/{include/cuvs/util/pointer_attributes.hpp => src/neighbors/detail/pointer_attributes.cpp} (58%) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 08009b731d..d6e67dbb4d 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1472,6 +1472,7 @@ if(NOT BUILD_CPU_ONLY) src/neighbors/detail/cagra/topk_for_cagra/topk.cu ${cuvs_cagra_search_cuda_inst_files} src/neighbors/detail/cagra/cagra_helpers.cpp + src/neighbors/detail/pointer_attributes.cpp src/neighbors/dynamic_batching.cu src/neighbors/composite/index.cu $<$:src/neighbors/cagra.cpp> diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 23c9dd56bd..2e202b5c89 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -74,7 +74,7 @@ inline auto is_device_accessible(const void* ptr) -> bool { cudaPointerAttributes attrs{}; RAFT_CUDA_TRY(cudaPointerGetAttributes(&attrs, ptr)); - return cuvs::util::detail::is_device_accessible(attrs); + return cuvs::neighbors::detail::is_device_accessible(attrs); } /** diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index abb661128f..2b6b2aad14 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -25,7 +25,6 @@ #include #include #include -#include #include #include @@ -1095,6 +1094,18 @@ template namespace detail { +/** + * Whether memory described by `attr` may be dereferenced in device code on the current device. + * + * Device and managed memory are device-accessible, and registered (pinned) host memory is when it + * is mapped (`devicePointer` is set). Pageable (unregistered) host memory is device-accessible only + * when the device reads it through the host page tables (ATS). On systems with HMM, + * `cudaPointerGetAttributes` can also report a non-null `devicePointer` for pageable memory (for + * example after the first GPU access to pageable memory in the process), but device accesses to it + * are then serviced through GPU page faults, so it is treated as host-only there. + */ +[[nodiscard]] auto is_device_accessible(const cudaPointerAttributes& attr) -> bool; + template [[nodiscard]] inline uint32_t mdspan_row_stride_elements(SrcT const& src) { @@ -1118,9 +1129,8 @@ template } // Host memory needs a device mapping; pageable memory only counts when the device reads it // through the host page tables (ATS), not through HMM page faults. - auto* device_ptr = cuvs::util::detail::is_device_accessible(ptr_attrs) - ? reinterpret_cast(ptr_attrs.devicePointer) - : nullptr; + auto* device_ptr = + is_device_accessible(ptr_attrs) ? reinterpret_cast(ptr_attrs.devicePointer) : nullptr; RAFT_EXPECTS(device_ptr != nullptr, "%s", error_msg); return device_ptr; } diff --git a/cpp/src/neighbors/detail/ann_utils.cuh b/cpp/src/neighbors/detail/ann_utils.cuh index e7fdc17c3a..47ea5278e8 100644 --- a/cpp/src/neighbors/detail/ann_utils.cuh +++ b/cpp/src/neighbors/detail/ann_utils.cuh @@ -6,7 +6,7 @@ #pragma once #include -#include +#include #include #include #include @@ -62,7 +62,7 @@ struct pointer_residency_count { cudaPointerAttributes attr; RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, ptr)); // Not `attr.devicePointer != nullptr`: with HMM it can be set for pageable host memory. - if (cuvs::util::detail::is_device_accessible(attr)) { ++on_device; } + if (cuvs::neighbors::detail::is_device_accessible(attr)) { ++on_device; } if (attr.hostPointer || attr.type == cudaMemoryTypeUnregistered) { ++on_host; } return std::make_tuple(on_device, on_host); } diff --git a/cpp/include/cuvs/util/pointer_attributes.hpp b/cpp/src/neighbors/detail/pointer_attributes.cpp similarity index 58% rename from cpp/include/cuvs/util/pointer_attributes.hpp rename to cpp/src/neighbors/detail/pointer_attributes.cpp index 7a2fa689e7..29e213b6c8 100644 --- a/cpp/include/cuvs/util/pointer_attributes.hpp +++ b/cpp/src/neighbors/detail/pointer_attributes.cpp @@ -2,7 +2,8 @@ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ -#pragma once + +#include #include @@ -11,13 +12,15 @@ #include #include -namespace cuvs::util::detail { +namespace cuvs::neighbors::detail { + +namespace { /** * Whether `device` reads pageable (unregistered) host memory through the host page tables (ATS, * e.g. Grace Hopper / Grace Blackwell). The result is cached per device. */ -inline auto pageable_memory_uses_host_page_tables(int device) -> bool +auto pageable_memory_uses_host_page_tables(int device) -> bool { static std::mutex mutex; static std::unordered_map cache; @@ -30,17 +33,9 @@ inline auto pageable_memory_uses_host_page_tables(int device) -> bool return uses_host_page_tables != 0; } -/** - * Whether memory described by `attr` may be dereferenced in device code on the current device. - * - * Device and managed memory are device-accessible, and registered (pinned) host memory is when it - * is mapped (`devicePointer` is set). Pageable (unregistered) host memory is device-accessible only - * when the device reads it through the host page tables (ATS). On systems with HMM, - * `cudaPointerGetAttributes` can also report a non-null `devicePointer` for pageable memory (for - * example after the first GPU access to pageable memory in the process), but device accesses to it - * are then serviced through GPU page faults, so it is treated as host-only there. - */ -inline auto is_device_accessible(const cudaPointerAttributes& attr) -> bool +} // namespace + +auto is_device_accessible(const cudaPointerAttributes& attr) -> bool { switch (attr.type) { case cudaMemoryTypeDevice: @@ -55,4 +50,4 @@ inline auto is_device_accessible(const cudaPointerAttributes& attr) -> bool } } -} // namespace cuvs::util::detail +} // namespace cuvs::neighbors::detail diff --git a/cpp/src/preprocessing/quantize/pq.cu b/cpp/src/preprocessing/quantize/pq.cu index 6488dfa996..b99acb50fd 100644 --- a/cpp/src/preprocessing/quantize/pq.cu +++ b/cpp/src/preprocessing/quantize/pq.cu @@ -5,8 +5,8 @@ #include "./detail/pq.cuh" +#include #include -#include #include #include @@ -89,7 +89,7 @@ auto train_from_rows(raft::resources const& res, { cudaPointerAttributes ptr_attrs; RAFT_CUDA_TRY(cudaPointerGetAttributes(&ptr_attrs, src_ptr)); - auto const* device_ptr = cuvs::util::detail::is_device_accessible(ptr_attrs) + auto const* device_ptr = cuvs::neighbors::detail::is_device_accessible(ptr_attrs) ? reinterpret_cast(ptr_attrs.devicePointer) : nullptr; if (device_ptr == nullptr) { From fe6517c6e9bb4951b003e05f7cda8d3f9ee72a20 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sun, 11 Oct 2026 14:11:31 +0000 Subject: [PATCH 26/27] Group IVF-PQ training rows by cluster with a stable sort The per-cluster codebook training listed each cluster's rows in the order given by an atomic counter, so the per-cluster training subsets depended on scheduling. Group the row indices with a stable sort by label instead, so each cluster's rows are listed in increasing order. --- cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh | 26 +++++++---------------- 1 file changed, 8 insertions(+), 18 deletions(-) diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh index 85024d2de2..8f6cdccd61 100644 --- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh +++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh @@ -58,6 +58,8 @@ #include #include #include +#include +#include #include #include @@ -205,17 +207,6 @@ void flat_compute_residuals( stream.get()); } -template -__launch_bounds__(BlockDim) static __global__ void fill_indices_kernel(IdxT n_rows, - IdxT* data_indices, - IdxT* data_offsets, - const uint32_t* labels) -{ - const auto i = IdxT(BlockDim) * IdxT(blockIdx.x) + IdxT(threadIdx.x); - if (i >= n_rows) { return; } - data_indices[atomicAdd(data_offsets + labels[i], 1)] = i; -} - /** * @brief Calculate cluster offsets and arrange data indices into clusters. * @@ -249,13 +240,12 @@ auto calculate_offsets_and_indices(IdxT n_rows, stream.sync(); RAFT_EXPECTS(cumsum == n_rows, "cluster sizes do not add up."); RAFT_LOG_DEBUG("Max cluster size %d", max_cluster_size); - rmm::device_uvector data_offsets_buf(n_lists, stream); - auto data_offsets = data_offsets_buf.data(); - raft::copy(data_offsets, cluster_offsets, n_lists, stream); - constexpr uint32_t n_threads = 128; // NOLINT - const IdxT n_blocks = raft::div_rounding_up_unsafe(n_rows, n_threads); - fill_indices_kernel - <<>>(n_rows, data_indices, data_offsets, labels); + // Group the row indices by cluster with a stable sort, so that the order of the rows within a + // cluster (which decides the per-cluster training subsets) is the same in every build. + rmm::device_uvector sorted_labels(n_rows, stream); + raft::copy(sorted_labels.data(), labels, n_rows, stream); + thrust::sequence(exec_policy, data_indices, data_indices + n_rows); + thrust::stable_sort_by_key(exec_policy, sorted_labels.begin(), sorted_labels.end(), data_indices); return max_cluster_size; } From 65c558f02df8f9ef0451a16ac23d13670674b2e6 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:16:28 +0000 Subject: [PATCH 27/27] CI stress testing: run all test jobs on the PR and nightly matrices Do not merge. Every test job runs regardless of changed files, the conda C++, conda Python, and wheel test jobs use the union of the pull-request and nightly test matrices, and every test job needs a trivial anchor job so that re-running the anchor re-runs all test jobs without rebuilding. --- .github/workflows/pr.yaml | 29 ++++++++++++++++++----------- 1 file changed, 18 insertions(+), 11 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index a6438c1a8d..9e49dce763 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -36,6 +36,7 @@ jobs: - wheel-tests-cuvs - devcontainer - telemetry-setup + - stress-rerun-anchor uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@main permissions: contents: read @@ -53,6 +54,14 @@ jobs: - name: Telemetry setup if: ${{ vars.TELEMETRY_ENABLED == 'true' }} uses: rapidsai/shared-actions/telemetry-dispatch-stash-base-env-vars@main + stress-rerun-anchor: + # Stress testing only: every test job needs this job, so re-running it re-runs all + # test jobs against the existing build artifacts without re-running any builds. + runs-on: ubuntu-latest + permissions: {} + steps: + - name: Anchor for test job re-runs + run: echo "Re-running the test jobs" build-details: uses: rapidsai/shared-workflows/.github/workflows/compute-build-details.yaml@main check-nightly-ci: @@ -469,7 +478,7 @@ jobs: node_type: cpu16 script: ci/build_cpp.sh conda-cpp-tests: - needs: [conda-cpp-build, changed-files] + needs: [conda-cpp-build, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -478,7 +487,6 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/conda-cpp-tests.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp # Split the C++ tests across several jobs for each CUDA/GPU configuration strategy: fail-fast: false @@ -488,6 +496,7 @@ jobs: with: build_type: pull-request script: "ci/test_cpp.sh ${{ matrix.shard }} ${{ matrix.num_shards }}" + matrix_type: "pull-request,nightly" cache-paths: .cache/cuda-jit cache-key-prefix: libcuvs-cudajit-v1 # Caches are immutable. Re-key when the JIT-relevant sources change so a completed run can @@ -548,7 +557,7 @@ jobs: # Build a conda package for each CUDA x ARCH x minimum supported Python version matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) conda-python-tests: - needs: [conda-python-build, changed-files] + needs: [conda-python-build, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -557,10 +566,10 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda with: build_type: pull-request script: ci/test_python.sh + matrix_type: "pull-request,nightly" rocky8-clib-standalone-build: needs: checks permissions: @@ -573,7 +582,7 @@ jobs: with: build_type: pull-request rocky8-clib-tests: - needs: [rocky8-clib-standalone-build, changed-files] + needs: [rocky8-clib-standalone-build, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -581,7 +590,6 @@ jobs: packages: read pull-requests: read uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_standalone_c with: build_type: pull-request container-image-repository: "rapidsai/citestwheel" @@ -595,7 +603,6 @@ jobs: packages: read pull-requests: read uses: ./.github/workflows/matrix-job.yaml - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_java || fromJSON(needs.changed-files.outputs.changed_file_groups).test_cpp # Artifacts are not published from these jobs, so it's safe to run for multiple CUDA versions. # If these jobs start producing artifacts, the names will have to differentiate between CUDA versions. with: @@ -606,7 +613,7 @@ jobs: artifact-name-prefix: "cuvs-java" file_to_upload: "java/cuvs-java/target/" conda-java-tests: - needs: conda-java-build + needs: [conda-java-build, stress-rerun-anchor] permissions: actions: read contents: read @@ -637,7 +644,7 @@ jobs: artifact-name-prefix: "cuvs-lucene" file_to_upload: "java/cuvs-lucene/target/" conda-lucene-tests: - needs: conda-lucene-build + needs: [conda-lucene-build, stress-rerun-anchor] permissions: actions: read contents: read @@ -746,7 +753,7 @@ jobs: # Build a wheel for each CUDA x ARCH x minimum supported Python version matrix_filter: group_by({CUDA_VER, ARCH}) | map(min_by(.PY_VER | split(".") | map(tonumber))) wheel-tests-cuvs: - needs: [wheel-build-cuvs, changed-files] + needs: [wheel-build-cuvs, changed-files, stress-rerun-anchor] permissions: actions: read contents: read @@ -755,10 +762,10 @@ jobs: pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main - if: fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels with: build_type: pull-request script: ci/test_wheel_cuvs.sh + matrix_type: "pull-request,nightly" devcontainer: permissions: actions: read