Skip to content

Reject cross-device streams in acquireExternalStream - #235

Merged
scal444 merged 1 commit into
NVIDIA-BioNeMo:mainfrom
scal444:fix-cross-device-stream-check
Jul 27, 2026
Merged

Reject cross-device streams in acquireExternalStream#235
scal444 merged 1 commit into
NVIDIA-BioNeMo:mainfrom
scal444:fix-cross-device-stream-check

Conversation

@scal444

@scal444 scal444 commented Jul 27, 2026

Copy link
Copy Markdown
Collaborator

acquireExternalStream validated a caller-supplied stream handle with cudaStreamQuery, which reports whether the stream is live, not which device owns it. CUDA permits querying a stream that belongs to another device in the same process, so a foreign stream returned cudaSuccess and was accepted. The cross-device rejection the function advertises was therefore never actually happening.

Resolve the stream's owning device and compare it against the current device. cudaStreamGetDevice does this directly but requires CUDA 12.8, above our supported floor, so below that fall back to probing with an event: an event and the stream it is recorded into must share a CUDA context, so recording an event created on the current device fails with cudaErrorInvalidResourceHandle when the stream belongs to another device.

The driver API (cuStreamGetCtx) was avoided deliberately: it changes signature in CUDA 13 and would add a libcuda link dependency.

DeviceTest.AcquireExternalStreamWrongDevice covers this but only runs with 2+ GPUs, which is why it went unnoticed until it failed in multi-GPU CI.

acquireExternalStream validated a caller-supplied stream handle with
cudaStreamQuery, which reports whether the stream is live, not which
device owns it. CUDA permits querying a stream that belongs to another
device in the same process, so a foreign stream returned cudaSuccess and
was accepted. The cross-device rejection the function advertises was
therefore never actually happening.

Resolve the stream's owning device and compare it against the current
device. cudaStreamGetDevice does this directly but requires CUDA 12.8,
above our supported floor, so below that fall back to probing with an
event: an event and the stream it is recorded into must share a CUDA
context, so recording an event created on the current device fails with
cudaErrorInvalidResourceHandle when the stream belongs to another device.

The driver API (cuStreamGetCtx) was avoided deliberately: it changes
signature in CUDA 13 and would add a libcuda link dependency.

DeviceTest.AcquireExternalStreamWrongDevice covers this but only runs
with 2+ GPUs, which is why it went unnoticed until it failed in
multi-GPU CI.
@scal444
scal444 requested a review from evasnow1992 July 27, 2026 16:43
@greptile-apps

greptile-apps Bot commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR strengthens external CUDA stream validation and aligns test includes with the repository’s include convention.

  • Resolves a stream’s owning device directly with CUDA 12.8 or newer.
  • Uses an event-based device-context probe on older CUDA versions.
  • Rejects valid foreign-device streams while preserving default and same-device streams.
  • Converts FMCS test includes to repository-root-relative paths.

Confidence Score: 5/5

The PR appears safe to merge with no actionable defects identified.

The new ownership checks preserve valid default and same-device stream handling while rejecting foreign-device streams, and the changed test includes resolve through the established project-root include path.

Important Files Changed

Filename Overview
src/utils/device.cpp Adds compatible cross-device stream ownership validation without exposing a concrete correctness regression.
tests/test_fmcs_primitives.cu Updates local includes to the repository-root-relative convention supported by the target’s include configuration.

Reviews (1): Last reviewed commit: "Reject cross-device streams in acquireEx..." | Re-trigger Greptile

@evasnow1992 evasnow1992 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

Thank you for fixing the test failure.

@scal444
scal444 merged commit ac7714d into NVIDIA-BioNeMo:main Jul 27, 2026
10 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants