cuda.core: allow updating Buffer deallocation streams by Andy-Jost · Pull Request #2602 · NVIDIA/cuda-python

Andy-Jost · 2026-08-11T18:57:41Z

Summary

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

Depends on cuda.core: capture complete Buffer deallocation recipe at creation #2526 and must be rebased onto main after that PR merges.
Focused diff relative to #2526's branch

juenglin · 2026-08-13T18:36:02Z

- 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 "


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

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

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

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.

... because the DevicePtrHandle holds an independent stream handle.
What does invoking Stream.close mean then?

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.

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.

juenglin · 2026-08-14T20:51:09Z

+ # The free operation runs after the copy on that stream.
+ buffer.close()
+ buffer = None
+ consumer_stream.sync()


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?

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.

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.

juenglin · 2026-08-14T20:52:48Z

- 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 "


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

juenglin

LGTM!

github-actions · 2026-08-17T00:16:46Z

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

Andy-Jost added this to the cuda.core 1.2.0 milestone Aug 11, 2026

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 self-assigned this Aug 11, 2026

Andy-Jost requested a review from juenglin August 11, 2026 18:58

This comment has been minimized.
Sign in to view

juenglin reviewed Aug 13, 2026
View reviewed changes

Andy-Jost force-pushed the ajost/buffer-set-deallocation-stream branch from eef14b0 to 575c18fCompare August 14, 2026 15:35

Andy-Jost mentioned this pull request Aug 14, 2026
cuda.core: reject operations on closed resource objects #2627
Closed

Andy-Jost requested review from juenglin and rparolin August 14, 2026 18:19

Andy-Jost mentioned this pull request Aug 14, 2026
cuda.core: reject operations on closed resource objects #2635
Merged
2 tasks

Andy-Jost added 3 commits August 14, 2026 13:21

cuda.core: allow updating Buffer deallocation streams
25459a4

docs(cuda.core): add deallocation stream transfer example
82a8088
Show how to move a Buffer between streams with an event so its eventual deallocation remains correctly ordered.

test(cuda.core): use shared memory resource instrumentation
7b04eac
Keep the deallocation-stream tests aligned with the centralized test helper merged in NVIDIA#2624.

Andy-Jost force-pushed the ajost/buffer-set-deallocation-stream branch from 11cb846 to 7b04eacCompare August 14, 2026 20:23

juenglin reviewed Aug 14, 2026
View reviewed changes

juenglin approved these changes Aug 14, 2026
View reviewed changes

Andy-Jost merged commit db28018 into NVIDIA:mainAug 14, 2026
109 checks passed

Andy-Jost deleted the ajost/buffer-set-deallocation-stream branch August 14, 2026 22:49

This comment has been minimized.
Sign in to view

Uh oh!

Conversation

Andy-Jost commented Aug 11, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

Related Work

Uh oh!

This comment has been minimized.

Uh oh!

Uh oh!

juenglinAug 13, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-JostAug 14, 2026

Choose a reason for hiding this comment

Uh oh!

juenglinAug 14, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-JostAug 14, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

juenglinAug 14, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Andy-JostAug 14, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

juenglinAug 14, 2026

Choose a reason for hiding this comment

Uh oh!

Andy-JostAug 14, 2026• edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

juenglinAug 14, 2026

Choose a reason for hiding this comment

Uh oh!

juenglinAug 14, 2026

Choose a reason for hiding this comment

Uh oh!

juenglin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

github-actionsBot commented Aug 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Andy-Jost commented Aug 11, 2026•
edited
Loading

Andy-JostAug 14, 2026•
edited
Loading

juenglinAug 14, 2026•
edited
Loading

Andy-JostAug 14, 2026•
edited
Loading