Implement cuda.bench.State.set_stream(provider)

[oleksandr-pavlyk]

The stream provide type is expected to implement __cuda_stream__ protocol, and not be an instance of type cuda.bench.CudaStream.

The reason is that cuda.bench.CudaStream type is not user-constructible. It is returned by cuda.bench.State.get_stream method. It stores reference to the stream stored in nvbench::state referenced from cuda.bench.State.

Disallowing instances of cuda.bench.CudaStream as argument to set_stream avoids footguns.

Closes #440

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f45da714-8fd8-40a7-baf8-f2be78bc3c3c

📥 Commits

Reviewing files that changed from the base of the PR and between e57ecad and 4287b99.

📒 Files selected for processing (3)

• python/cuda/bench/__init__.pyi
• python/src/py_nvbench.cpp
• python/test/test_cuda_bench.py

---

📝 Walkthrough

Summary by CodeRabbit

• New Features
  • Added support for configuring benchmark state with external CUDA stream providers.
  • Providers can expose CUDA stream handles through the standard __cuda_stream__ protocol.
  • Stream providers remain available while configured and can be replaced as needed.
  • Added validation and clear errors for unsupported or invalid providers.
  • External stream handles are now reflected in benchmark results.
  • Benchmark execution callbacks now run with the Python GIL held.

Walkthrough

The PR adds the _SupportsCudaStream protocol and State.set_stream. The binding validates external stream providers, converts version-0 protocol results to CUDA streams, preserves provider lifetime, and adds CPU-only coverage.

Changes

External CUDA stream support

Layer / File(s)	Summary
Stream provider contract python/cuda/bench/__init__.pyi	Defines _SupportsCudaStream with __cuda_stream__ and exposes State.set_stream.
Stream extraction, binding, lifetime, and validation python/src/py_nvbench.cpp, python/test/test_cuda_bench.py	Validates providers, rejects CudaStream, registers changed handles, preserves provider lifetime, propagates handles to launches, and tests valid and invalid inputs.
Execution callback documentation python/src/py_nvbench.cpp	Documents that State.exec, run_all_benchmarks, and launcher callbacks run with the Python GIL held.
API documentation validation python/test/test_cuda_bench.py	Validates the State.set_stream docstring.

Assessment against linked issues

Objective	Addressed	Explanation
Add cuda.bench.State.set_stream(stream_provider) for non-CudaStream objects supporting __cuda_stream__. [#440]	✅
Set the stream only when the protocol handle differs from the handle held by the underlying state. [#440]	✅

---

_{Comment @coderabbitai help to get the list of available commands.}

[oleksandr-pavlyk]

@PointKernel This PR fixes the limitation mentioned in NVIDIA/cuvs#2371 (comment) that prompted filing for #440

[oleksandr-pavlyk]

Use of cuda.bench.State.set_stream removes operations for setting stream provided by cuda.bench.Launch for CuPy to use from the launchable

Example code

# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import sys

import cuda.bench as bench
import cupy


@bench.register()
@bench.axis.int64_power_of_two("Elements", [22,24,26])
def elementwise_square(state: bench.State):
    # Provide cuda.bench stream to use for launching
    state.set_stream(cupy.cuda.get_current_stream())

    size = state.get_int64("Elements")
    x = cupy.random.randint(low=-16000, high=16000, size=size)
    y = cupy.empty_like(x)

    state.add_element_count(size)
    state.add_global_memory_reads(x.nbytes)
    state.add_global_memory_writes(y.nbytes)

    def launcher(launch: bench.Launch):
        cupy.square(x, out=y)

    state.exec(launcher, batched=False)


if __name__ == "__main__":
    bench.run_all_benchmarks(sys.argv)

compare to example that sets stream provided by cuda.bench.Launch:

# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
# SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

import sys

import cuda.bench as bench
import cupy


@bench.register()
@bench.axis.int64_power_of_two("Elements", [22,24,26])
def elementwise_square(state: bench.State):

    size = state.get_int64("Elements")
    x = cupy.random.randint(low=-16000, high=16000, size=size)
    y = cupy.empty_like(x)

    state.add_element_count(size)
    state.add_global_memory_reads(x.nbytes)
    state.add_global_memory_writes(y.nbytes)

    def launcher(launch: bench.Launch):
        with cupy.cuda.Stream.from_external(launch.get_stream()):
            cupy.square(x, out=y)

    state.exec(launcher, batched=False)


if __name__ == "__main__":
    bench.run_all_benchmarks(sys.argv)