Skip to content

RFC: Support CUDA Stream Protocol #2040

@leofang

Description

@leofang

cuda.core is an official CUDA Python project: https://nvidia.github.io/cuda-python/cuda-core/latest/index.html. It offers a pythonic, self-contained, lightweight, and official interface over the CUDA programming model. For new Python projects, we encourage them to just use cuda.core.<experimental>.Stream.

For existing Python projects such as PyTorch, transitioning to cuda.core may or may not be immediately feasible. As a result, we encourage projects that already expose a CUDA stream to Python to follow the CUDA Stream protocol:
https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol
and add a __cuda_stream__ method to the stream class, so as to improve interoperability without introducing extra ExternalStream-like types.

Here is a PyTorch example of how it'd be used interoperably with cuda.core:
https://github.com/NVIDIA/cuda-python/blob/c4f4ffe83d246eafb6adf1574e5a7c86bbcef944/cuda_core/examples/pytorch_example.py

Previously, this issue has been discussed in #1770 and #1775. Unfortunately, we still end up exposing a stream class as part of RMM public API. I consider it as a mistake that we should rectify. But it can be done and tracked separately from this effort.

cc @kkraus14 @aterrel @rparolin for vis

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    Status

    To-do

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions