From 13725b301aaab5c85b8a0e1b7f9d84e2ea2e70c3 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Fri, 6 Feb 2026 05:28:57 -0800 Subject: [PATCH 1/9] add a new max_node_id parameter to the CAGRA search API, allowing users to constrain random seed node selection to a subset of the dataset. This is useful when the graph is smaller than the dataset, such as during iterative build with compression. --- cpp/include/cuvs/neighbors/cagra.hpp | 7 +++++++ .../neighbors/detail/cagra/device_common.hpp | 6 ++++-- .../cagra/search_multi_cta_kernel-inl.cuh | 9 ++++++--- .../cagra/search_single_cta_kernel-inl.cuh | 20 +++++++++++++------ 4 files changed, 31 insertions(+), 11 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 9b9e8eb0e6..ef55507869 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -247,6 +247,13 @@ struct search_params : cuvs::neighbors::search_params { /** Bit mask used for initial random seed node selection. */ uint64_t rand_xor_mask = 0x128394; + /** + * Maximum node ID for random seed selection. + * When > 0, random seeds are constrained to [0, max_node_id) instead of [0, dataset_size). + * This is useful when the graph is smaller than the dataset (e.g., iterative build with compression). + * Default 0 means no constraint (use dataset_size). + */ + uint32_t max_node_id = 0; /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ bool persistent = false; /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ diff --git a/cpp/src/neighbors/detail/cagra/device_common.hpp b/cpp/src/neighbors/detail/cagra/device_common.hpp index 8a5bb6ba1f..df22b28081 100644 --- a/cpp/src/neighbors/detail/cagra/device_common.hpp +++ b/cpp/src/neighbors/detail/cagra/device_common.hpp @@ -102,11 +102,13 @@ RAFT_DEVICE_INLINE_FUNCTION void compute_distance_to_random_nodes( IndexT* __restrict__ traversed_hash_ptr, const uint32_t traversed_hash_bitlen, const uint32_t block_id = 0, - const uint32_t num_blocks = 1) + const uint32_t num_blocks = 1, + const IndexT max_node_id = 0) { const auto team_size_bits = dataset_desc.team_size_bitshift_from_smem(); const auto max_i = raft::round_up_safe(num_pickup, warp_size >> team_size_bits); const auto compute_distance = dataset_desc.compute_distance_impl; + const IndexT seed_index_limit = max_node_id > 0 ? max_node_id : dataset_desc.size; for (uint32_t i = threadIdx.x >> team_size_bits; i < max_i; i += (blockDim.x >> team_size_bits)) { const bool valid_i = (i < num_pickup); @@ -122,7 +124,7 @@ RAFT_DEVICE_INLINE_FUNCTION void compute_distance_to_random_nodes( if (seed_ptr && (gid < num_seeds)) { seed_index = seed_ptr[gid]; } else { - seed_index = device::xorshift64(gid ^ rand_xor_mask) % dataset_desc.size; + seed_index = device::xorshift64(gid ^ rand_xor_mask) % seed_index_limit; } } diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index c8b885dffe..916767c00b 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -193,7 +193,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const uint32_t min_iteration, const uint32_t max_iteration, uint32_t* const num_executed_iterations, /* stats */ - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + const typename DATASET_DESCRIPTOR_T::INDEX_T max_node_id = 0) { using DATA_T = typename DATASET_DESCRIPTOR_T::DATA_T; using INDEX_T = typename DATASET_DESCRIPTOR_T::INDEX_T; @@ -281,7 +282,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( local_traversed_hashmap_ptr, traversed_hash_bitlen, block_id, - num_blocks); + num_blocks, + max_node_id); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -627,7 +629,8 @@ void select_and_run(const dataset_descriptor_host& dat ps.min_iterations, ps.max_iterations, num_executed_iterations, - sample_filter); + sample_filter, + ps.max_node_id); } } // namespace multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 404817e582..956f9ac3e4 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -702,7 +702,8 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( const std::uint32_t small_hash_bitlen, const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id, - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + const typename DATASET_DESCRIPTOR_T::INDEX_T max_node_id = 0) { using LOAD_T = device::LOAD_128BIT_T; @@ -791,7 +792,10 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( local_visited_hashmap_ptr, hash_bitlen, (INDEX_T*)nullptr, - 0); + 0, + 0, + 1, + max_node_id); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -1124,7 +1128,8 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const std::uint32_t hash_bitlen, const std::uint32_t small_hash_bitlen, const std::uint32_t small_hash_reset_interval, - SAMPLE_FILTER_T sample_filter) + SAMPLE_FILTER_T sample_filter, + const typename DATASET_DESCRIPTOR_T::INDEX_T max_node_id = 0) { const auto query_id = blockIdx.y; search_core Date: Fri, 6 Feb 2026 06:51:31 -0800 Subject: [PATCH 2/9] Changed the max node id parameter name to graph_size for clarity; removed max_node_id from the search parameters structure --- cpp/include/cuvs/neighbors/cagra.hpp | 7 ------- cpp/src/neighbors/detail/cagra/device_common.hpp | 4 ++-- .../detail/cagra/search_multi_cta_kernel-inl.cuh | 6 +++--- .../detail/cagra/search_single_cta_kernel-inl.cuh | 13 ++++++------- 4 files changed, 11 insertions(+), 19 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index ef55507869..9b9e8eb0e6 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -247,13 +247,6 @@ struct search_params : cuvs::neighbors::search_params { /** Bit mask used for initial random seed node selection. */ uint64_t rand_xor_mask = 0x128394; - /** - * Maximum node ID for random seed selection. - * When > 0, random seeds are constrained to [0, max_node_id) instead of [0, dataset_size). - * This is useful when the graph is smaller than the dataset (e.g., iterative build with compression). - * Default 0 means no constraint (use dataset_size). - */ - uint32_t max_node_id = 0; /** Whether to use the persistent version of the kernel (only SINGLE_CTA is supported a.t.m.) */ bool persistent = false; /** Persistent kernel: time in seconds before the kernel stops if no requests received. */ diff --git a/cpp/src/neighbors/detail/cagra/device_common.hpp b/cpp/src/neighbors/detail/cagra/device_common.hpp index df22b28081..8cbbe1d366 100644 --- a/cpp/src/neighbors/detail/cagra/device_common.hpp +++ b/cpp/src/neighbors/detail/cagra/device_common.hpp @@ -103,12 +103,12 @@ RAFT_DEVICE_INLINE_FUNCTION void compute_distance_to_random_nodes( const uint32_t traversed_hash_bitlen, const uint32_t block_id = 0, const uint32_t num_blocks = 1, - const IndexT max_node_id = 0) + const IndexT graph_size = 0) { const auto team_size_bits = dataset_desc.team_size_bitshift_from_smem(); const auto max_i = raft::round_up_safe(num_pickup, warp_size >> team_size_bits); const auto compute_distance = dataset_desc.compute_distance_impl; - const IndexT seed_index_limit = max_node_id > 0 ? max_node_id : dataset_desc.size; + const IndexT seed_index_limit = graph_size > 0 ? graph_size : dataset_desc.size; for (uint32_t i = threadIdx.x >> team_size_bits; i < max_i; i += (blockDim.x >> team_size_bits)) { const bool valid_i = (i < num_pickup); diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index 916767c00b..a2e9a43ff0 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -194,7 +194,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const uint32_t max_iteration, uint32_t* const num_executed_iterations, /* stats */ SAMPLE_FILTER_T sample_filter, - const typename DATASET_DESCRIPTOR_T::INDEX_T max_node_id = 0) + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { using DATA_T = typename DATASET_DESCRIPTOR_T::DATA_T; using INDEX_T = typename DATASET_DESCRIPTOR_T::INDEX_T; @@ -283,7 +283,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( traversed_hash_bitlen, block_id, num_blocks, - max_node_id); + graph_size); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -630,7 +630,7 @@ void select_and_run(const dataset_descriptor_host& dat ps.max_iterations, num_executed_iterations, sample_filter, - ps.max_node_id); + static_cast(graph.extent(0))); } } // namespace multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 956f9ac3e4..5d465c25b5 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -703,7 +703,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( const std::uint32_t small_hash_reset_interval, const std::uint32_t query_id, SAMPLE_FILTER_T sample_filter, - const typename DATASET_DESCRIPTOR_T::INDEX_T max_node_id = 0) + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { using LOAD_T = device::LOAD_128BIT_T; @@ -795,7 +795,7 @@ RAFT_DEVICE_INLINE_FUNCTION void search_core( 0, 0, 1, - max_node_id); + graph_size); __syncthreads(); _CLK_REC(clk_compute_1st_distance); @@ -1129,7 +1129,7 @@ RAFT_KERNEL __launch_bounds__(1024, 1) search_kernel( const std::uint32_t small_hash_bitlen, const std::uint32_t small_hash_reset_interval, SAMPLE_FILTER_T sample_filter, - const typename DATASET_DESCRIPTOR_T::INDEX_T max_node_id = 0) + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { const auto query_id = blockIdx.y; search_core(graph.extent(0))); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } From 70a69d922ec757a7093cecde07eb5a87f3c44de2 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Fri, 6 Feb 2026 08:55:07 -0800 Subject: [PATCH 3/9] wrote test --- .../detail/cagra/search_multi_kernel.cuh | 16 +- cpp/tests/CMakeLists.txt | 1 + .../bug_graph_smaller_than_dataset.cu | 159 ++++++++++++++++++ 3 files changed, 171 insertions(+), 5 deletions(-) create mode 100644 cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh index f7d353d864..0ee7439ac6 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh @@ -104,7 +104,8 @@ RAFT_KERNEL random_pickup_kernel( typename DATASET_DESCRIPTOR_T::DISTANCE_T* const result_distances_ptr, // [num_queries, ldr] const std::uint32_t ldr, // (*) ldr >= num_pickup typename DATASET_DESCRIPTOR_T::INDEX_T* const visited_hashmap_ptr, // [num_queries, 1 << bitlen] - const std::uint32_t hash_bitlen) + const std::uint32_t hash_bitlen, + const typename DATASET_DESCRIPTOR_T::INDEX_T graph_size = 0) { using DATA_T = typename DATASET_DESCRIPTOR_T::DATA_T; using INDEX_T = typename DATASET_DESCRIPTOR_T::INDEX_T; @@ -119,6 +120,8 @@ RAFT_KERNEL random_pickup_kernel( dataset_desc = dataset_desc->setup_workspace(smem, queries_ptr, query_id); __syncthreads(); + const INDEX_T seed_index_limit = graph_size > 0 ? graph_size : dataset_desc->size; + INDEX_T best_index_team_local; DISTANCE_T best_norm2_team_local = utils::get_max_value(); for (unsigned i = 0; i < num_distilation; i++) { @@ -128,7 +131,7 @@ RAFT_KERNEL random_pickup_kernel( } else { // Chose a seed node randomly seed_index = - device::xorshift64((global_team_index ^ rand_xor_mask) * (i + 1)) % dataset_desc->size; + device::xorshift64((global_team_index ^ rand_xor_mask) * (i + 1)) % seed_index_limit; } DISTANCE_T norm2 = dataset_desc->compute_distance(seed_index, true); @@ -166,7 +169,8 @@ void random_pickup(const dataset_descriptor_host& data std::size_t ldr, // (*) ldr >= num_pickup IndexT* visited_hashmap_ptr, // [num_queries, 1 << bitlen] std::uint32_t hash_bitlen, - cudaStream_t cuda_stream) + cudaStream_t cuda_stream, + IndexT graph_size = 0) { const auto block_size = 256u; const auto num_teams_per_threadblock = block_size / dataset_desc.team_size; @@ -185,7 +189,8 @@ void random_pickup(const dataset_descriptor_host& data result_distances_ptr, ldr, visited_hashmap_ptr, - hash_bitlen); + hash_bitlen, + graph_size); } template @@ -826,7 +831,8 @@ struct search result_buffer_allocation_size, hashmap.data(), hash_bitlen, - stream); + stream, + static_cast(this->dataset_size)); unsigned iter = 0; while (1) { diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 9fc620b4cb..3643b2e12d 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -156,6 +156,7 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu neighbors/ann_cagra/bug_multi_cta_crash.cu + neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu GPUS 1 PERCENT 100 ) diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu new file mode 100644 index 0000000000..8a29779bbd --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -0,0 +1,159 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include + +#include + +#include +#include +#include +#include +#include + +#include + +namespace cuvs::neighbors::cagra { + +/** + * @brief Test verifying graph.extent(0) is used for random seed selection + * + * This test ensures that CAGRA search kernels correctly use graph.extent(0) + * (graph size) rather than dataset.size for random seed node selection. + * + * The bug: random seed selection previously used dataset_desc.size, which + * could cause OOB access if the graph size differed from dataset size + * (e.g., in CAGRA-Q iterative builds with compression). + * + * The fix: kernels now receive graph.extent(0) as graph_size parameter, + * ensuring seeds are always within valid graph node range [0, graph_size). + */ +class cagra_graph_smaller_than_dataset_test : public ::testing::Test { + public: + using data_type = float; + using index_type = uint32_t; + + protected: + void run() + { + // Create a dataset with 10000 points + constexpr int64_t n_dataset = 10000; + constexpr int64_t n_dim = 128; + constexpr int64_t n_queries = 100; + constexpr int64_t k = 10; + + // Build index normally + auto dataset = raft::make_device_matrix(res, n_dataset, n_dim); + raft::random::RngState r(1234ULL); + raft::random::uniform(res, r, dataset.data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); + + cagra::index_params index_params; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + + auto index = cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); + raft::resource::sync_stream(res); + + // Get the graph from the index + auto original_graph = index.graph(); + ASSERT_EQ(original_graph.extent(0), n_dataset); + + // Recreate the bug scenario: LARGE dataset, SMALL graph + // (like iterative_build_graph does in intermediate iterations) + constexpr int64_t n_graph = n_dataset / 2; // Only 5000 nodes in graph + + // Step 1: Build index on SMALL subset (5000 points) + auto small_dataset_view = raft::make_device_matrix_view( + dataset.data_handle(), n_graph, n_dim); + + cagra::index_params small_index_params; + small_index_params.graph_degree = 32; + auto small_index = cagra::build(res, small_index_params, small_dataset_view); + raft::resource::sync_stream(res); + + // Step 2: Update to FULL dataset (10000 points) but keep small graph (5000 nodes) + // This creates the exact bug scenario: dataset.size=10000, graph.extent(0)=5000 + small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); + + // Verify the mismatch - THIS IS THE BUG SCENARIO! + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 5000 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 10000 points + ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! + + // Create queries + auto queries = raft::make_device_matrix(res, n_queries, n_dim); + raft::random::uniform(res, r, queries.data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); + + // Allocate output + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + + // Setup search params + cagra::search_params search_params; + search_params.itopk_size = 64; + search_params.search_width = 1; + search_params.max_iterations = 10; + search_params.algo = cagra::search_algo::SINGLE_CTA; + + // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS + // Before fix: random seeds use dataset.size (10000) -> tries to access graph[7000] -> CRASH! + // After fix: random seeds use graph.extent(0) (5000) -> only accesses graph[0-4999] -> SAFE! + cagra::search(res, + search_params, + small_index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(res); + + // Verify results are valid (neighbors should be < graph size) + auto neighbors_host = raft::make_host_matrix(n_queries, k); + raft::copy(neighbors_host.data_handle(), + neighbors.data_handle(), + n_queries * k, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + // All neighbor indices should be valid (< n_graph) + for (int64_t i = 0; i < n_queries * k; i++) { + ASSERT_LT(neighbors_host.data_handle()[i], n_graph) + << "Neighbor index " << neighbors_host.data_handle()[i] + << " is >= graph size " << n_graph; + } + + // Test with MULTI_CTA algorithm as well (also had the same bug) + search_params.algo = cagra::search_algo::MULTI_CTA; + + cagra::search(res, + search_params, + small_index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view()); + + raft::resource::sync_stream(res); + + // Verify again + raft::copy(neighbors_host.data_handle(), + neighbors.data_handle(), + n_queries * k, + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + for (int64_t i = 0; i < n_queries * k; i++) { + ASSERT_LT(neighbors_host.data_handle()[i], n_graph) + << "Neighbor index " << neighbors_host.data_handle()[i] + << " is >= graph size " << n_graph << " (MULTI_CTA)"; + } + } + + private: + raft::resources res; +}; + +TEST_F(cagra_graph_smaller_than_dataset_test, search_with_smaller_graph) { this->run(); } + +} // namespace cuvs::neighbors::cagra From f428e54bf64ad1ba1f67ae01be6ce68742ac85ed Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Fri, 6 Feb 2026 08:55:34 -0800 Subject: [PATCH 4/9] minor pre-commit changes --- .../detail/cagra/search_multi_kernel.cuh | 2 +- .../bug_graph_smaller_than_dataset.cu | 45 ++++++++++--------- 2 files changed, 24 insertions(+), 23 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh index 0ee7439ac6..045c63fe59 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_kernel.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index 8a29779bbd..b06c1cba92 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -19,14 +19,14 @@ namespace cuvs::neighbors::cagra { /** * @brief Test verifying graph.extent(0) is used for random seed selection - * + * * This test ensures that CAGRA search kernels correctly use graph.extent(0) * (graph size) rather than dataset.size for random seed node selection. - * + * * The bug: random seed selection previously used dataset_desc.size, which * could cause OOB access if the graph size differed from dataset size * (e.g., in CAGRA-Q iterative builds with compression). - * + * * The fix: kernels now receive graph.extent(0) as graph_size parameter, * ensuring seeds are always within valid graph node range [0, graph_size). */ @@ -47,12 +47,13 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Build index normally auto dataset = raft::make_device_matrix(res, n_dataset, n_dim); raft::random::RngState r(1234ULL); - raft::random::uniform(res, r, dataset.data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); + raft::random::uniform( + res, r, dataset.data_handle(), n_dataset * n_dim, data_type(-1), data_type(1)); cagra::index_params index_params; index_params.graph_degree = 32; index_params.intermediate_graph_degree = 64; - + auto index = cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); raft::resource::sync_stream(res); @@ -63,28 +64,29 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Recreate the bug scenario: LARGE dataset, SMALL graph // (like iterative_build_graph does in intermediate iterations) constexpr int64_t n_graph = n_dataset / 2; // Only 5000 nodes in graph - + // Step 1: Build index on SMALL subset (5000 points) auto small_dataset_view = raft::make_device_matrix_view( dataset.data_handle(), n_graph, n_dim); - + cagra::index_params small_index_params; small_index_params.graph_degree = 32; - auto small_index = cagra::build(res, small_index_params, small_dataset_view); + auto small_index = cagra::build(res, small_index_params, small_dataset_view); raft::resource::sync_stream(res); - + // Step 2: Update to FULL dataset (10000 points) but keep small graph (5000 nodes) // This creates the exact bug scenario: dataset.size=10000, graph.extent(0)=5000 small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); - + // Verify the mismatch - THIS IS THE BUG SCENARIO! - ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 5000 nodes - ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 10000 points - ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 5000 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 10000 points + ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! // Create queries auto queries = raft::make_device_matrix(res, n_queries, n_dim); - raft::random::uniform(res, r, queries.data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); + raft::random::uniform( + res, r, queries.data_handle(), n_queries * n_dim, data_type(-1), data_type(1)); // Allocate output auto neighbors = raft::make_device_matrix(res, n_queries, k); @@ -92,10 +94,10 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Setup search params cagra::search_params search_params; - search_params.itopk_size = 64; - search_params.search_width = 1; + search_params.itopk_size = 64; + search_params.search_width = 1; search_params.max_iterations = 10; - search_params.algo = cagra::search_algo::SINGLE_CTA; + search_params.algo = cagra::search_algo::SINGLE_CTA; // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS // Before fix: random seeds use dataset.size (10000) -> tries to access graph[7000] -> CRASH! @@ -120,13 +122,12 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // All neighbor indices should be valid (< n_graph) for (int64_t i = 0; i < n_queries * k; i++) { ASSERT_LT(neighbors_host.data_handle()[i], n_graph) - << "Neighbor index " << neighbors_host.data_handle()[i] - << " is >= graph size " << n_graph; + << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph; } // Test with MULTI_CTA algorithm as well (also had the same bug) search_params.algo = cagra::search_algo::MULTI_CTA; - + cagra::search(res, search_params, small_index, @@ -145,8 +146,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { for (int64_t i = 0; i < n_queries * k; i++) { ASSERT_LT(neighbors_host.data_handle()[i], n_graph) - << "Neighbor index " << neighbors_host.data_handle()[i] - << " is >= graph size " << n_graph << " (MULTI_CTA)"; + << "Neighbor index " << neighbors_host.data_handle()[i] << " is >= graph size " << n_graph + << " (MULTI_CTA)"; } } From 13b1f77b8000053d0c8a2170b390f9707fc6fbc1 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Mon, 16 Feb 2026 05:54:19 -0800 Subject: [PATCH 5/9] addressed comments regarding type cast --- cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh | 2 +- cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index a2e9a43ff0..92331b0eb6 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -630,7 +630,7 @@ void select_and_run(const dataset_descriptor_host& dat ps.max_iterations, num_executed_iterations, sample_filter, - static_cast(graph.extent(0))); + static_cast(graph.extent(0))); } } // namespace multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index 5d465c25b5..32bdcf07ad 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -2348,7 +2348,7 @@ control is returned in this thread (in persistent_runner_t constructor), so we'r small_hash_bitlen, small_hash_reset_interval, sample_filter, - static_cast(graph.extent(0))); + static_cast(graph.extent(0))); RAFT_CUDA_TRY(cudaPeekAtLastError()); } } From cb60c6b4cfb618f0f0980d200e91d00895d84c4c Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Fri, 20 Feb 2026 01:32:09 -0800 Subject: [PATCH 6/9] Pre-commit style fix --- cpp/src/neighbors/detail/cagra/device_common.hpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/device_common.hpp b/cpp/src/neighbors/detail/cagra/device_common.hpp index 3d3a25c8eb..0b75de6bab 100644 --- a/cpp/src/neighbors/detail/cagra/device_common.hpp +++ b/cpp/src/neighbors/detail/cagra/device_common.hpp @@ -103,11 +103,11 @@ RAFT_DEVICE_INLINE_FUNCTION void compute_distance_to_random_nodes( const uint32_t traversed_hash_bitlen, const uint32_t block_id = 0, const uint32_t num_blocks = 1, - const IndexT graph_size = 0) + const IndexT graph_size = 0) { const auto team_size_bits = dataset_desc.team_size_bitshift_from_smem(); const auto max_i = raft::round_up_safe(num_pickup, warp_size >> team_size_bits); - const auto compute_distance = dataset_desc.compute_distance_impl; + const auto compute_distance = dataset_desc.compute_distance_impl; const IndexT seed_index_limit = graph_size > 0 ? graph_size : dataset_desc.size; for (uint32_t i = threadIdx.x >> team_size_bits; i < max_i; i += (blockDim.x >> team_size_bits)) { From b27ad1b0eefb7ef73f12c9eb08e77258898e6766 Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 24 Feb 2026 02:28:28 -0800 Subject: [PATCH 7/9] style fix --- cpp/tests/CMakeLists.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 83e35eb034..7ed9c08764 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -155,7 +155,8 @@ ConfigureTest( ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_TEST_BUGS - PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu neighbors/ann_cagra/bug_multi_cta_crash.cu + PATH neighbors/ann_cagra/bug_extreme_inputs_oob.cu + neighbors/ann_cagra/bug_multi_cta_crash.cu neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu neighbors/ann_cagra/bug_issue_93_reproducer.cu GPUS 1 From 3843a7b281ef44d4e7d93c62526bb3e7ad43b99e Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 3 Mar 2026 04:53:56 -0800 Subject: [PATCH 8/9] smaller test --- .../neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index b06c1cba92..b9d3d92b03 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -38,8 +38,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { protected: void run() { - // Create a dataset with 10000 points - constexpr int64_t n_dataset = 10000; + // Create a dataset with 1000 points + constexpr int64_t n_dataset = 1000; constexpr int64_t n_dim = 128; constexpr int64_t n_queries = 100; constexpr int64_t k = 10; From 86f8b8668add0bd92f14e9ea1324b0084333ba0f Mon Sep 17 00:00:00 2001 From: Irina Reshodko Date: Tue, 3 Mar 2026 04:57:16 -0800 Subject: [PATCH 9/9] comment fix --- .../ann_cagra/bug_graph_smaller_than_dataset.cu | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu index b9d3d92b03..adeb774a8b 100644 --- a/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu +++ b/cpp/tests/neighbors/ann_cagra/bug_graph_smaller_than_dataset.cu @@ -63,9 +63,9 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { // Recreate the bug scenario: LARGE dataset, SMALL graph // (like iterative_build_graph does in intermediate iterations) - constexpr int64_t n_graph = n_dataset / 2; // Only 5000 nodes in graph + constexpr int64_t n_graph = n_dataset / 2; // Only 500 nodes in graph - // Step 1: Build index on SMALL subset (5000 points) + // Step 1: Build index on SMALL subset (500 points) auto small_dataset_view = raft::make_device_matrix_view( dataset.data_handle(), n_graph, n_dim); @@ -74,13 +74,13 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { auto small_index = cagra::build(res, small_index_params, small_dataset_view); raft::resource::sync_stream(res); - // Step 2: Update to FULL dataset (10000 points) but keep small graph (5000 nodes) - // This creates the exact bug scenario: dataset.size=10000, graph.extent(0)=5000 + // Step 2: Update to FULL dataset (1000 points) but keep small graph (500 nodes) + // This creates the exact bug scenario: dataset.size=1000, graph.extent(0)=500 small_index.update_dataset(res, raft::make_const_mdspan(dataset.view())); // Verify the mismatch - THIS IS THE BUG SCENARIO! - ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 5000 nodes - ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 10000 points + ASSERT_EQ(small_index.graph().extent(0), n_graph); // Graph has 500 nodes + ASSERT_EQ(small_index.size(), n_dataset); // Dataset has 1000 points ASSERT_NE(small_index.graph().extent(0), small_index.size()); // Mismatch! // Create queries @@ -100,8 +100,8 @@ class cagra_graph_smaller_than_dataset_test : public ::testing::Test { search_params.algo = cagra::search_algo::SINGLE_CTA; // THIS SHOULD NOT CRASH OR CAUSE OOB ACCESS - // Before fix: random seeds use dataset.size (10000) -> tries to access graph[7000] -> CRASH! - // After fix: random seeds use graph.extent(0) (5000) -> only accesses graph[0-4999] -> SAFE! + // Before fix: random seeds use dataset.size (1000) -> tries to access graph[700] -> CRASH! + // After fix: random seeds use graph.extent(0) (500) -> only accesses graph[0-499] -> SAFE! cagra::search(res, search_params, small_index,