Skip to content

Fall back to cudaMalloc on devices without stream-ordered allocation - #2528

Open
drzraf wants to merge 1 commit into
NVIDIA:mainfrom
drzraf:fea/stream-ordered-alloc-fallback
Open

Fall back to cudaMalloc on devices without stream-ordered allocation#2528
drzraf wants to merge 1 commit into
NVIDIA:mainfrom
drzraf:fea/stream-ordered-alloc-fallback

Conversation

@drzraf

@drzraf drzraf commented Aug 29, 2026

Copy link
Copy Markdown

CAGRA search allocates its device-side dataset descriptor and its filter payload with cudaMallocAsync. That API is backed by the CUDA stream-ordered memory allocator, which is optional hardware/driver functionality: where cudaDevAttrMemoryPoolsSupported reads 0 -- every Maxwell-generation GPU, and some virtualized or older-driver configurations -- the call fails and the very first search throws:

compute_distance.hpp:244: call='cudaMallocAsync(&ptr, size, stream)',
Reason=cudaErrorNotSupported:operation not supported

Add cuvs::util::malloc_async_compat / free_async_compat, which use the stream-ordered allocator where it exists and plain cudaMalloc / cudaFree where it does not. The fallback is stronger than what the caller asks for, not weaker: cudaMalloc is synchronous with respect to the device, so the memory is trivially usable by work subsequently issued into the stream, which is the guarantee cudaMallocAsync gives; and the stream is synchronised before cudaFree, so the "free once the preceding work in this stream completes" contract of cudaFreeAsync is met explicitly rather than incidentally.

What is lost is pool reuse and the absence of implicit synchronization, so the helpers are only appropriate for allocations made a handful of times per search. That is the case for both call sites: the dataset descriptor is allocated once per descriptor and cached, and the filter payload once per filter.

Devices that support memory pools are unaffected: the wrapper forwards to cudaMallocAsync exactly as before.

On why these two sites do not simply use RMM, as most of cuVS does: both are lazily initialized caches whose allocation happens inside a mutex-guarded accessor that receives only a stream -- state::eval(rmm::cuda_stream_view) and state::dev_ptr(cudaStream_t) -- with no raft::resources or memory resource in scope. Routing them through RMM means plumbing a resource into those caches and moving the ownership of the deallocation with it, which is a larger change with its own lifetime questions. This commit deliberately restores the intended semantics without touching ownership; if the RMM route is preferred, it is a straightforward follow-up.

This change is independent of rapidsai/librtcx#17 and #2527 : it is not tied to a compute capability or to a toolkit version, and applies to any GPU whose driver reports no memory-pool support.

Verified on a Maxwell device (sm_50, CUDA 12.6): before, the first CAGRA search throws the error above; after, all CAGRA build and search paths run, and recall against exact brute force is unchanged.

CAGRA search allocates its device-side dataset descriptor and its filter payload
with cudaMallocAsync. That API is backed by the CUDA stream-ordered memory
allocator, which is optional hardware/driver functionality: where
cudaDevAttrMemoryPoolsSupported reads 0 -- every Maxwell-generation GPU, and
some virtualized or older-driver configurations -- the call fails and the very
first search throws:

  compute_distance.hpp:244: call='cudaMallocAsync(&ptr, size, stream)',
  Reason=cudaErrorNotSupported:operation not supported

Add cuvs::util::malloc_async_compat / free_async_compat, which use the
stream-ordered allocator where it exists and plain cudaMalloc / cudaFree where
it does not. The fallback is stronger than what the caller asks for, not weaker:
cudaMalloc is synchronous with respect to the device, so the memory is trivially
usable by work subsequently issued into the stream, which is the guarantee
cudaMallocAsync gives; and the stream is synchronised before cudaFree, so the
"free once the preceding work in this stream completes" contract of
cudaFreeAsync is met explicitly rather than incidentally.

What is lost is pool reuse and the absence of implicit synchronisation, so the
helpers are only appropriate for allocations made a handful of times per search.
That is the case for both call sites: the dataset descriptor is allocated once
per descriptor and cached, and the filter payload once per filter.

Devices that support memory pools are unaffected: the wrapper forwards to
cudaMallocAsync exactly as before.

On why these two sites do not simply use RMM, as most of cuVS does: both are
lazily initialised caches whose allocation happens inside a mutex-guarded
accessor that receives only a stream -- state::eval(rmm::cuda_stream_view) and
state::dev_ptr(cudaStream_t) -- with no raft::resources or memory resource in
scope. Routing them through RMM means plumbing a resource into those caches and
moving the ownership of the deallocation with it, which is a larger change with
its own lifetime questions. This commit deliberately restores the intended
semantics without touching ownership; if the RMM route is preferred, it is a
straightforward follow-up.

This change is independent of everything else in this series: it is not tied to
a compute capability or to a toolkit version, and applies to any GPU whose
driver reports no memory-pool support.

Verified on a Maxwell device (sm_50, CUDA 12.6): before, the first CAGRA search
throws the error above; after, all CAGRA build and search paths run, and recall
against exact brute force is unchanged.
@drzraf
drzraf requested a review from a team as a code owner August 29, 2026 21:18
@copy-pr-bot

copy-pr-bot Bot commented Aug 29, 2026

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant