Skip to content

cuda.core: allow updating Buffer deallocation streams - #2602

Merged
Andy-Jost merged 3 commits into
NVIDIA:mainfrom
Andy-Jost:ajost/buffer-set-deallocation-stream
Aug 14, 2026
Merged

cuda.core: allow updating Buffer deallocation streams#2602
Andy-Jost merged 3 commits into
NVIDIA:mainfrom
Andy-Jost:ajost/buffer-set-deallocation-stream

Conversation

@Andy-Jost

@Andy-Jost Andy-Jost commented Aug 11, 2026

Copy link
Copy Markdown
Contributor

Summary

Closes #2600.

Add Buffer.set_deallocation_stream so callers can replace the stream that orders eventual deallocation without closing the buffer. This supports transferring an allocation's lifetime to another stream while preserving the existing Buffer object.

Changes

  • Add the public Buffer.set_deallocation_stream API and generated type stub.
  • Route initialization, explicit close overrides, and the new setter through the same validation and native DeallocationStream capture path.
  • Document the caller's responsibility to order prior allocation and access work before deallocation on the replacement stream.
  • Add focused tests covering initialization, setter, and close configuration followed by explicit close or garbage collection.

Related Work

@Andy-Jost Andy-Jost added this to the cuda.core 1.2.0 milestone Aug 11, 2026
@Andy-Jost Andy-Jost added enhancement Any code-related improvements P1 Medium priority - Should do cuda.core Everything related to the cuda.core module labels Aug 11, 2026
@Andy-Jost Andy-Jost self-assigned this Aug 11, 2026
@Andy-Jost
Andy-Jost requested a review from juenglin August 11, 2026 18:58
@github-actions

This comment has been minimized.

Comment thread cuda_core/cuda/core/_memory/_buffer.pyx Outdated
Comment thread cuda_core/docs/source/api.rst
stream = Stream._from_handle(Stream, h_stream) if h_stream else default_stream()
if not h_stream:
print(
"Warning: no deallocation stream was recorded; falling back to "

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should it say "no deallocation stream was recorded or the deallocation stream was closed"?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It's saying no stream is available for ordering the destruction. Not necessarily that a stream was provided and was closed.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, but wouldn't this warning also appear when the deallocation stream was closed?

@Andy-Jost Andy-Jost Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Two points: This warning does not occur if a valid stream is passed to set_deallocation_stream() and that stream is later closed, because the DevicePtrHandle holds an independent stream handle. It can occur if an already closed stream is passed.

The follow-up in #2635 should address the second case: closed streams will be rejected throughout the API, including by set_deallocation_stream(). If the proper checks are in place throughout the codebase, this condition should never occur, which is why I flagged it as an internal error.

@juenglin juenglin Aug 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

... because the DevicePtrHandle holds an independent stream handle.

What does invoking Stream.close mean then?

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A StreamHandle is a std::shared_ptr<CUstream>. A Python Stream class holds one of those, and Stream.close resets the shared pointer, dropping that reference. Inside a DevicePtrHandle is an independent StreamHandle referring to the deallocation stream. Whatever the user-facing Python Stream object does, the embedded reference cannot be invalidated.

Comment thread cuda_core/cuda/core/_memory/_buffer.pyx
Show how to move a Buffer between streams with an event so its eventual deallocation remains correctly ordered.
Keep the deallocation-stream tests aligned with the centralized test helper merged in NVIDIA#2624.
@Andy-Jost
Andy-Jost force-pushed the ajost/buffer-set-deallocation-stream branch from 11cb846 to 7b04eac Compare August 14, 2026 20:23
# The free operation runs after the copy on that stream.
buffer.close()
buffer = None
consumer_stream.sync()

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This looks wrong: you close a buffer then sync a stream that has pending copy for the buffer. This sync could fail with a RuntimeError, couldn't it?

@Andy-Jost Andy-Jost Aug 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The buffer in question has stream-ordered deallocation semantics, so closing it places a call to cuMemFreeAsync onto consumer_stream. The idea is that transferring the buffer to the stream "currently" using it will order its free after the work using it.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah. For the other methods for stream-ordered operations stream is a mandatory parameter. That no stream is passed to close here caused me to think it is not an async operation.

stream = Stream._from_handle(Stream, h_stream) if h_stream else default_stream()
if not h_stream:
print(
"Warning: no deallocation stream was recorded; falling back to "

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, but wouldn't this warning also appear when the deallocation stream was closed?

@juenglin juenglin left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM!

@Andy-Jost
Andy-Jost merged commit db28018 into NVIDIA:main Aug 14, 2026
109 checks passed
@Andy-Jost
Andy-Jost deleted the ajost/buffer-set-deallocation-stream branch August 14, 2026 22:49
@github-actions

This comment has been minimized.

1 similar comment
@github-actions

Copy link
Copy Markdown
Doc Preview CI
Preview removed because the pull request was closed or merged.

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

Labels

cuda.core Everything related to the cuda.core module enhancement Any code-related improvements P1 Medium priority - Should do

Projects

None yet

Development

Successfully merging this pull request may close these issues.

cuda.core: allow Buffer deallocation stream to be updated before teardown

2 participants