From 2f945b616f30bee171671d5072211cb89751b2a5 Mon Sep 17 00:00:00 2001
From: Kevin Boyd <kboyd@nvidia.com>
Date: Mon, 27 Jul 2026 11:05:36 -0400
Subject: [PATCH] Reject cross-device streams in acquireExternalStream

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.
---
 src/utils/device.cpp          | 52 +++++++++++++++++++++++++++++++----
 tests/test_fmcs_primitives.cu |  6 ++--
 2 files changed, 50 insertions(+), 8 deletions(-)

diff --git a/src/utils/device.cpp b/src/utils/device.cpp
index 6e447ba6..f217b213 100644
--- a/src/utils/device.cpp
+++ b/src/utils/device.cpp
@@ -37,18 +37,60 @@ WithDevice::~WithDevice() {
   cudaCheckErrorNoThrow(cudaSetDevice(original_device_id_));
 }
 
+namespace {
+
+//! Returns true if `stream` belongs to the current device.
+//! cudaStreamQuery accepts streams from any device in the process, so it cannot answer this on its
+//! own.
+bool streamIsOnCurrentDevice(cudaStream_t stream) {
+#if CUDART_VERSION >= 12080
+  int currentDevice = -1;
+  if (cudaGetDevice(&currentDevice) != cudaSuccess) {
+    cudaGetLastError();
+    return false;
+  }
+  int streamDevice = -1;
+  if (cudaStreamGetDevice(stream, &streamDevice) != cudaSuccess) {
+    cudaGetLastError();
+    return false;
+  }
+  return streamDevice == currentDevice;
+#else
+  // cudaStreamGetDevice needs CUDA 12.8. Below that, probe with an event instead: 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 if the stream belongs to another one.
+  cudaEvent_t probe = nullptr;
+  if (cudaEventCreateWithFlags(&probe, cudaEventDisableTiming) != cudaSuccess) {
+    cudaGetLastError();
+    return false;
+  }
+  const cudaError_t recordErr = cudaEventRecord(probe, stream);
+  cudaCheckErrorNoThrow(cudaEventDestroy(probe));
+  if (recordErr != cudaSuccess) {
+    cudaGetLastError();
+    return false;
+  }
+  return true;
+#endif
+}
+
+}  // namespace
+
 std::optional<cudaStream_t> acquireExternalStream(std::uintptr_t streamPtr) {
   auto stream = reinterpret_cast<cudaStream_t>(streamPtr);
   if (streamPtr == 0) {
     return stream;
   }
   cudaError_t err = cudaStreamQuery(stream);
-  if (err == cudaSuccess || err == cudaErrorNotReady) {
-    return stream;
+  if (err != cudaSuccess && err != cudaErrorNotReady) {
+    // Clear the sticky error state
+    cudaGetLastError();
+    return std::nullopt;
+  }
+  if (!streamIsOnCurrentDevice(stream)) {
+    return std::nullopt;
   }
-  // Clear the sticky error state
-  cudaGetLastError();
-  return std::nullopt;
+  return stream;
 }
 
 size_t getDeviceFreeMemory() {
diff --git a/tests/test_fmcs_primitives.cu b/tests/test_fmcs_primitives.cu
index 324371c8..bb4338f3 100644
--- a/tests/test_fmcs_primitives.cu
+++ b/tests/test_fmcs_primitives.cu
@@ -13,9 +13,9 @@
 #include <set>
 #include <vector>
 
-#include "fmcs_cuda/fmcs_seed.cuh"
-#include "fmcs_cuda/fmcs_seed_queue.cuh"
-#include "mcs_common/mcs_cooperative_copy.cuh"
+#include "src/mcs/fmcs_cuda/fmcs_seed.cuh"
+#include "src/mcs/fmcs_cuda/fmcs_seed_queue.cuh"
+#include "src/mcs/mcs_common/mcs_cooperative_copy.cuh"
 #include "src/utils/device_vector.h"
 
 namespace {