Use the PTDS for (most) cupy operations

[jcrist]

This configures cupy to use the per-thread default stream (PTDS) for most operations. This avoids usage of the default legacy stream in more of the codebase, allowing for improved parallelism when running across multiple threads.

This is a breaking change.

Previously any cupy operations in cuml ran in cupy's default stream (the legacy stream). We didn't synchronize the stream before returning, but that didn't matter due to the synchronization behavior of the legacy stream.

With this PR we've moved to running (most) cupy operations in the PTDS. Depending on the operation, we may not synchronize the PTDS before returning.

Most users shouldn't notice a difference and should have no issues.

Users not using threads, custom streams, or only working with host memory (e.g. numpy in/numpy out) should see no difference. Likewise any users that only use cupy's default stream (the legacy stream) in their code should see no issues.

Users doing tricky things with custom streams or threads may run into issues and require a manual sync of the PTDS (can be done with cupy.cuda.Stream.ptds.synchronize(). For example, the following workflow may run into issues:

• Run a cuml operation based on cupy in thread A, returning a cupy array
• Consume that output in thread B as a cupy array using a stream other than the legacy stream (e.g. a different PTDS or a custom stream)

For safety, you probably want to add a call to cupy.cuda.Stream.ptds.synchronize() in thread A before returning to ensure the output array is fully populated before consuming it in thread B.

Fixes #7909.

[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: 8f4c8a7a-406b-44c5-86f3-d3f728ccd5ea

📥 Commits

Reviewing files that changed from the base of the PR and between 4a86000 and f546679.

📒 Files selected for processing (4)

• python/cuml/cuml/internals/outputs.py
• python/cuml/cuml/preprocessing/label.py
• python/cuml/tests/test_prims.py
• python/cuml/tests/test_reflection.py

💤 Files with no reviewable changes (2)

• python/cuml/cuml/preprocessing/label.py
• python/cuml/tests/test_prims.py

🚧 Files skipped from review as they are similar to previous changes (2)

• python/cuml/tests/test_reflection.py
• python/cuml/cuml/internals/outputs.py

---

📝 Walkthrough

Summary by CodeRabbit

• Bug Fixes

  • Improved internal stream execution context management to ensure proper scope handling during operations.
  • Removed unnecessary synchronization calls that could impact concurrent execution efficiency.

• Tests

  • Added test to verify stream execution behavior in decorated methods.

Walkthrough

Scope cuML internal output contexts to CuPy's per-thread default stream (PTDS), apply that scope when reflect returns internal outputs, remove several explicit null-stream synchronizations, and add a test asserting PTDS is active inside reflected/internal contexts.

Changes

PTDS stream scoping and reflect integration

Layer / File(s)	Summary
PTDS import and internal-context scoping python/cuml/cuml/internals/outputs.py	Add from cupy.cuda import Stream and wrap the enter_internal_context yield with with Stream.ptds: so code run while inside the yielded internal context executes under CuPy's per-thread default stream.
Use internal-context for reflect internal-return python/cuml/cuml/internals/outputs.py	In reflect.inner, when returning internal (output_type == "cuml"), perform coercion/return from within with enter_internal_context(): instead of returning immediately, ensuring PTDS scoping is applied to the coercion/return path.
Remove explicit null-stream synchronizations python/cuml/cuml/preprocessing/label.py, python/cuml/tests/test_prims.py	Delete several cp.cuda.Stream.null.synchronize() calls: two in label_binarize/LabelBinarizer.fit and four in tests/test_prims.py around make_monotonic assertions. No other behavior or public APIs changed.
Add PTDS verification test python/cuml/tests/test_reflection.py	Add test_decorators_set_cupy_ptds which constructs a small estimator using @reflect/@run_in_internal_context and asserts that direct reflected calls, nested reflected calls, and non-reflected internal-context calls observe cp.cuda.Stream.ptds as the current stream.

🎯 3 (Moderate) | ⏱️ ~20 minutes

Suggested reviewers:

• csadorf

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'Use the PTDS for (most) cupy operations' accurately reflects the main objective of configuring CuPy to use the per-thread default stream, which is the primary change across all modified files.
Description check	✅ Passed	The description explains the PTDS configuration change, acknowledges it as a breaking change, discusses synchronization implications, and provides guidance for users—all directly related to the code changes.
Linked Issues check	✅ Passed	The PR addresses issue #7909 by configuring CuPy operations to use PTDS across common estimator methods via the internal context mechanism and decorator changes.
Out of Scope Changes check	✅ Passed	All changes are directly related to implementing PTDS for CuPy operations: internal context wrapping, synchronization call removals, and verification tests.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[jcrist]

A quick benchmark of cross-thread performance:

from time import perf_counter
from multiprocessing.pool import ThreadPool

import rmm
from cuml.linear_model import Ridge
from cuml.datasets import make_regression

rmm.mr.set_current_device_resource(
    rmm.mr.PoolMemoryResource(rmm.mr.CudaAsyncMemoryResource())
)

N = 4
X, y = make_regression(100_000, 1000)


def fit(i):
    Ridge(solver="svd").fit(X, y)


# Warmup
fit(1)

print(f"Benchmarking on {N} threads...")

# Measure sequential time
start = perf_counter()
for i in range(N):
    fit(i)
seq_time = perf_counter() - start
print(f"- Sequential: {seq_time:.2f} s")

# Measure parallel time
with ThreadPool(N) as pool:
    start = perf_counter()
    pool.map(fit, range(N))
    par_time = perf_counter() - start
print(f"- Parallel: {par_time:.2f} s")

print(f"\nParallel took {par_time / seq_time:.2f}x the time of sequential")

This PR

Benchmarking on 4 threads...
- Sequential: 1.27 s
- Parallel: 0.71 s

Parallel took 0.56x the time of sequential

On Main

$ python bench_ridge.py 
Benchmarking on 4 threads...
- Sequential: 1.28 s
- Parallel: 1.31 s

Parallel took 1.02x the time of sequential

[jcrist]

One open question is what should our synchronization behavior be?

Previously we would call handle.sync() for any libcuml call (synchronizing the PTDS that it wraps). For cupy code we didn't do any explicit synchronizing, which was fine since cupy used the legacy stream.

Currently this PR doesn't do any synchronizing for the outputs when running with cupy, but still does when running with libcuml. I don't love that we're not consistent here, but am not sure which way we should go.

• For non-device outputs the user won't notice anything, device->host copies always synchronize
• For device outputs to cupy, a user using cupy's default config (the legacy stream) won't notice anything, the legacy stream takes care of that.
• For device outputs to cudf the user won't notice anything, cudf handles any synchronization itself.

The only risk here is users running a cuml thing that uses cupy for the last step and outputs device memory in thread 1, then does post processing in thread 2 using something other than the legacy stream without synchronizing the ptds in thread 1. This case seems pretty unlikely, but if someone was doing that before this PR they might run into concurrency issues after this PR unless we add a synchronization when returning device memory.

I see two sane paths forward:

• cuml operations are always synchronous to the user, we synchronize the PTDS before returning to user space when outputting device memory. We do this uniformly, so anything running with libcuml or cupy will have the same behavior. I think we can remove most handle.sync() calls in the python codebase and rely on Stream.ptds.synchronize().
• cuml operations are asynchronous when possible (providing no guarantee of a sync) and run in the PTDS. Users using cupy or cudf won't notice, but users doing weird fancy odd things across threads may need to add a sync point when moving memory across threads. We can remove most handle.sync() calls in the python codebase.

If we opt for the 2nd route as our synchronization behavior, then I think any additional work should be done in a follow-up PR (though we might want to find a place to document this in this PR). We're already correct for that model, we just overly sync.

If we opt for the 1st route, then I'll need to add a defensive sync call here. My preference is the 2nd though - it's easier to implement, and better matches other cuda libraries (including cupy, which we're using as our main container and would be good to emulate the behavior of IMO).

[csadorf]

Blocking this while I'm investigating the posed question.

[jcrist]

After sleeping on it, I'm pretty convinced that option 2 is the best path forward. It's the simplest to implement (we provide a guarantee of the stream, but not of sync), theoretically allows for improved performance due to better GPU utilization, and matches with how cupy does things. Users doing typical things won't notice a change, but it allows better control for expert users on when a sync occurs. If we do the sync internally we don't leave that option open.

[divyegala]

I would also vote for option 2.

[csadorf]

With this change, users can no longer control the stream used by cuML's Python-level CuPy work via an outer cupy.cuda.Stream() context. That's a behavior change. I think we have three options:

1. Detect whether a non-default CuPy stream was previously selected and respect it.
2. Provide a cuML-owned API for explicitly setting the stream, e.g. with cuml.use_stream(...):.
3. Do not allow users to override this stream selection.

I think option 2 is the cleanest long-term API, but we may want a transition period where 26.06 implements option 1 to preserve existing behavior with a FutureWarning, then require the cuML-owned API in 26.08.

This PR should also add tests for the stream policy:

• the default/null-stream case should use PTDS inside enter_internal_context()
• an explicitly selected CuPy stream should have documented behavior, whether respected or intentionally ignored

One other minor thing: there are still two explicit cp.cuda.Stream.null.synchronize() calls in preprocessing/label.py that should be updated or removed. There are also four occurrences in python/cuml/tests/test_prims.py.

---

Regarding the synchronization question: I agree that option 2 is preferable. I tried to construct a demo using only cuML and CuPy calls that would pass on main and fail on this branch without an explicit sync, and I was not able to. So for the synchronization policy specifically, I think the practical risk is very low.

[jcrist]

We just synced offline and came to the following agreement:

• cuml internals use many different technologies to perform compute on device memory (libcuml, cupy, cudf, ...). Only cupy allows selecting the stream, the others run on the PTDS. For uniformity and simplicity, we think we should (for now) always use the PTDS and not provide any way to configure this.
• We should make this as a breaking change this release. It is unlikely to actually break anything, but we should at least call it out as part of the release.
• We should document our stream policy. Since our docs currently lack a place to note this, I plan to do this in a followup PR adding an "advanced usage" or "advanced topics" doc with a few other usage bits not related to streams.

[jcrist]

I've updated this PR and the description. I believe it's ready for review/merge. Once in, I'll followup with a new docs page on advanced topics.

[dantegd]

Love the PR and change, excited for the improvement!

[dantegd]

Forgot to change comment to approve in my review above

[jcrist]

/merge