From 456b98dabc73bc863387b5745100526249a0069e Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Sat, 10 Oct 2026 17:39:22 +0000 Subject: [PATCH 1/3] 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 2/3] 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 cd97df7c1110244564e4bf392006b70144304552 Mon Sep 17 00:00:00 2001 From: Bradley Dice Date: Mon, 12 Oct 2026 02:28:40 +0000 Subject: [PATCH 3/3] Correct the descriptor cache comment: the LRU cache does not deduplicate keys --- cpp/src/neighbors/detail/cagra/factory.cuh | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/factory.cuh b/cpp/src/neighbors/detail/cagra/factory.cuh index 3107f841ae..f35fc1ffb2 100644 --- a/cpp/src/neighbors/detail/cagra/factory.cuh +++ b/cpp/src/neighbors/detail/cagra/factory.cuh @@ -178,10 +178,12 @@ auto dataset_descriptor_init_with_cache(const raft::resources& res, 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). + // each create a descriptor and insert it under the same key (the LRU cache does not deduplicate + // keys), so concurrent searches can end up with different descriptors. When the last owner of + // one releases it, its device memory is freed with `cudaFreeAsync` on the stream that + // initialized it, which is not ordered after a kernel on another thread's stream that 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)) {