Follow up fixes for pinning hardware device

libs/qec/lib/decoder.cpp

@@ -154,6 +154,47 @@ static int read_cuda_device_id(const cudaqx::heterogeneous_map &params) {
   return value;
 }
 
+/// Selects the construction device and restores the previous device unless
+/// commit() is called. Makes failed plugin construction transactional: if the
+/// plugin ctor throws, the calling thread is left on its original device
+/// instead of leaking the pin to whichever device was selected for the attempt.
+class ConstructionDevicePin {
+public:
+  explicit ConstructionDevicePin(int target) : target_(target) {
+    cudaError_t err = cudaGetDevice(&previous_);
+    if (err != cudaSuccess)
+      throw std::runtime_error(
+          "cuda_device_id " + std::to_string(target_) +
+          " could not be selected because cudaGetDevice() failed: " +
+          cudaGetErrorString(err));
+    err = cudaSetDevice(target_);
+    if (err != cudaSuccess)
+      throw std::runtime_error("cudaSetDevice(" + std::to_string(target_) +
+                               ") failed: " + cudaGetErrorString(err));
+    selected_ = true;
+  }
+
+  ~ConstructionDevicePin() {
+    if (selected_ && !committed_ && previous_ != target_)
+      (void)cudaSetDevice(previous_);
+  }
+
+  ConstructionDevicePin(const ConstructionDevicePin &) = delete;
+  ConstructionDevicePin &operator=(const ConstructionDevicePin &) = delete;
+
+  // Keep the pin: one thread owns one decoder, so leaving the constructing
+  // thread on the target device lets later allocations and kernel launches --
+  // including lazy ones inside decode() -- land there with no per-call
+  // machinery.
+  void commit() { committed_ = true; }
+
+private:
+  int target_ = -1;
+  int previous_ = -1;
+  bool selected_ = false;
+  bool committed_ = false;
+};
+
 std::unique_ptr<decoder>
 decoder::get(const std::string &name, const decoder_init &init,
              const cudaqx::heterogeneous_map &param_map) {
@@ -168,14 +209,7 @@ decoder::get(const std::string &name, const decoder_init &init,
   const int cuda_device_id = read_cuda_device_id(param_map);
   if (cuda_device_id < 0)
     return iter->second(init, param_map);
-  // Pin the constructing thread persistently (no restore): one thread owns
-  // one decoder, so every later allocation and kernel launch on this thread
-  // -- including lazy allocations inside a plugin's decode() -- lands on the
-  // requested device with no per-call machinery.
-  cudaError_t err = cudaSetDevice(cuda_device_id);
-  if (err != cudaSuccess)
-    throw std::runtime_error("cudaSetDevice(" + std::to_string(cuda_device_id) +
-                             ") failed: " + cudaGetErrorString(err));
+  ConstructionDevicePin device_pin(cuda_device_id);
   // The key is consumed here; strip it so plugins that strictly validate
   // their parameter keys do not reject it.
   cudaqx::heterogeneous_map plugin_params;
@@ -184,6 +218,7 @@ decoder::get(const std::string &name, const decoder_init &init,
       plugin_params.insert(kv.first, kv.second);
   auto d = iter->second(init, plugin_params);
   d->cuda_device_id_ = cuda_device_id;
+  device_pin.commit();
   return d;
 }
 

libs/qec/lib/hardware_guards.h

@@ -14,6 +14,22 @@
 
 namespace cudaq::qec::detail_affinity {
 
+/// Point the calling thread at \p target before work that allocates or
+/// launches on it (no restore; set-if-different). No-op for target < 0.
+/// Throws on failure: never silently decode on the wrong GPU.
+inline void set_cuda_device_for_decode(int target) {
+  if (target < 0)
+    return;
+  int current = -1;
+  if (cudaGetDevice(&current) == cudaSuccess && current == target)
+    return;
+  cudaError_t err = cudaSetDevice(target);
+  if (err != cudaSuccess)
+    throw std::runtime_error("set_cuda_device_for_decode: cudaSetDevice(" +
+                             std::to_string(target) +
+                             ") failed: " + cudaGetErrorString(err));
+}
+
 /// RAII: set the calling thread's CUDA device, restore the previous device on
 /// scope exit. No-op for target < 0. Lib-private and header-only so decoder
 /// plugins built as separate .so files can reuse it (PR2 extends this header

libs/qec/lib/realtime/decoding-server-cqr/CMakeLists.txt

@@ -149,6 +149,10 @@ target_link_libraries(cudaq-qec-decoding-server
     # violates the realtime-server dependency-closure contract.
     cudaq-qec-decoders
     cudaq-qec-realtime-decoding
+  PRIVATE
+    # DecodingSession worker threads pin themselves to their decoder's
+    # cuda_device_id (cudaSetDevice).
+    CUDA::cudart
 )
 
 if(CUDAQ_GPU_ROCE_AVAILABLE)

libs/qec/lib/realtime/decoding-server-cqr/DecodingServer.cpp

@@ -28,12 +28,38 @@ using cudaq::qec::decoding::config::DecoderTransport;
 // Constructors
 // ---------------------------------------------------------------------------
 
+/// gpu_roce runs the whole pipeline -- rings, dispatch scheduler, device-side
+/// graph fire -- on ONE GPU: the one the FPGA/NIC is affine to
+/// (HOLOLINK_GPU_ID). A decoder pinned elsewhere would split graph capture
+/// and graph launch across devices, which CUDA graphs cannot do. Both knobs
+/// name the same topology fact, so they must agree.
+int reconcile_gpu_roce_device(std::optional<int> env_gpu_id, int decoder_pin) {
+  if (env_gpu_id && *env_gpu_id < 0)
+    throw std::runtime_error("HOLOLINK_GPU_ID must be >= 0 (got " +
+                             std::to_string(*env_gpu_id) + ")");
+  if (env_gpu_id && decoder_pin >= 0 && *env_gpu_id != decoder_pin)
+    throw std::runtime_error(
+        "gpu_roce device conflict: HOLOLINK_GPU_ID=" +
+        std::to_string(*env_gpu_id) + " but the decoder is pinned to " +
+        std::to_string(decoder_pin) +
+        " (cuda_device_id). The FPGA-affine GPU and the decoder pin must be "
+        "the same device.");
+  if (env_gpu_id)
+    return *env_gpu_id;
+  return decoder_pin >= 0 ? decoder_pin : 0;
+}
+
 std::unique_ptr<ITransceiver>
-DecodingServer::make_transport(DecoderTransport transport_type) {
+DecodingServer::make_transport(DecoderTransport transport_type,
+                               int pinned_cuda_device) {
   switch (transport_type) {
   case DecoderTransport::gpu_roce:
 #ifdef CUDAQ_GPU_ROCE_AVAILABLE
-    return std::make_unique<GpuRoceTransceiver>(GpuRoceConfig::from_env());
+  {
+    auto cfg = GpuRoceConfig::from_env();
+    cfg.gpu_id = reconcile_gpu_roce_device(cfg.gpu_id_env, pinned_cuda_device);
+    return std::make_unique<GpuRoceTransceiver>(cfg);
+  }
 #else
     throw std::runtime_error(
         "gpu_roce transport requested but CUDAQ_GPU_ROCE_AVAILABLE is not set. "
@@ -65,7 +91,12 @@ DecodingServer::DecodingServer(const std::string &config_yaml) {
   registry_.load_from_config(config, config_yaml);
   register_handlers();
 
-  auto t = make_transport(registry_.required_transport());
+  const auto &boot_sessions = registry_.sessions();
+  const int pinned_cuda_device =
+      boot_sessions.size() == 1
+          ? boot_sessions.begin()->second->dec->get_cuda_device_id()
+          : -1;
+  auto t = make_transport(registry_.required_transport(), pinned_cuda_device);
   ITransceiver *raw = t.get();
   owned_transports_.push_back(std::move(t));
   function_transport_[kEnqueueSyndromesFunctionId] = raw;
@@ -102,7 +133,12 @@ DecodingServer::DecodingServer(std::unique_ptr<ITransceiver> transport,
   function_transport_[kEnqueueSyndromesFunctionId] = raw;
   function_transport_[kGetCorrectionsFunctionId] = raw;
   function_transport_[kResetDecoderFunctionId] = raw;
-  init(config_yaml);
+  try {
+    init(config_yaml);
+  } catch (...) {
+    registry_.stop_workers();
+    throw;
+  }
 }
 
 DecodingServer::DecodingServer(
@@ -113,7 +149,14 @@ DecodingServer::DecodingServer(
   function_transport_[kEnqueueSyndromesFunctionId] = raw;
   function_transport_[kGetCorrectionsFunctionId] = raw;
   function_transport_[kResetDecoderFunctionId] = raw;
-  registry_.load_from_config(config, "configure_decoders()");
+  try {
+    registry_.load_from_config(config, "configure_decoders()");
+  } catch (...) {
+    // Members destroy in reverse order (transports before registry); join any
+    // already-started workers while the transports still exist.
+    registry_.stop_workers();
+    throw;
+  }
   register_handlers();
 }
 
@@ -122,7 +165,12 @@ DecodingServer::DecodingServer(std::vector<std::unique_ptr<ITransceiver>> owned,
                                const std::string &config_yaml)
     : owned_transports_(std::move(owned)),
       function_transport_(std::move(function_transport)) {
-  init(config_yaml);
+  try {
+    init(config_yaml);
+  } catch (...) {
+    registry_.stop_workers();
+    throw;
+  }
 }
 
 DecodingServer::~DecodingServer() {

libs/qec/lib/realtime/decoding-server-cqr/DecodingServer.h

@@ -15,12 +15,19 @@
 
 #include <atomic>
 #include <memory>
+#include <optional>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
 namespace cudaq::qec::decoding_server {
 
+/// Resolve the single GPU a gpu_roce pipeline runs on from the two knobs that
+/// can name it: HOLOLINK_GPU_ID (FPGA/NIC affinity; nullopt when unset) and
+/// the decoder's cuda_device_id (-1 when unpinned). Throws when both are set
+/// and disagree; unset env defers to the pin; neither set -> 0.
+int reconcile_gpu_roce_device(std::optional<int> env_gpu_id, int decoder_pin);
+
 /// Maps function_id → non-owning ITransceiver pointer.
 /// Ownership lives in DecodingServer::owned_transports_.
 using TransportMap = std::unordered_map<uint32_t, ITransceiver *>;
@@ -74,7 +81,8 @@ class DecodingServer {
   /// until CpuRoceTransceiverAdapter / GpuRoceTransceiverAdapter are
   /// available via CUDAQ_REALTIME.
   static std::unique_ptr<ITransceiver>
-  make_transport(cudaq::qec::decoding::config::DecoderTransport transport_type);
+  make_transport(cudaq::qec::decoding::config::DecoderTransport transport_type,
+                 int pinned_cuda_device);
 
   // Destruction order matters: the GPU RoCE scheduler (inside
   // owned_transports_) holds a cudaGraphExec_t captured from a session's

libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.cpp

@@ -8,10 +8,12 @@
 
 #include "DecodingSession.h"
 #include "RpcWireFormat.h"
+#include "../../hardware_guards.h"
 #include "cudaq/qec/logger.h"
 
 #include <chrono>
 #include <cstring>
+#include <future>
 #include <stdexcept>
 #include <vector>
 
@@ -61,7 +63,30 @@ DecodingSession::create(std::unique_ptr<cudaq::qec::decoder> decoder,
 }
 
 void DecodingSession::start_worker() {
-  worker = std::thread([this] { worker_loop(); });
+  // The pin must happen ON the worker thread (CUDA device selection is
+  // thread-local), but a failure is a startup error that belongs to the
+  // caller: hand it back through a promise so load_from_config aborts the
+  // server instead of a worker silently decoding on the wrong device.
+  std::promise<void> pinned;
+  auto pin_result = pinned.get_future();
+  worker = std::thread([this, &pinned] {
+    try {
+      cudaq::qec::detail_affinity::set_cuda_device_for_decode(
+          dec->get_cuda_device_id());
+      pinned.set_value();
+    } catch (...) {
+      pinned.set_exception(std::current_exception());
+      return; // never serve work from a mispinned thread
+    }
+    worker_loop();
+  });
+  try {
+    pin_result.get();
+  } catch (...) {
+    if (worker.joinable())
+      worker.join();
+    throw;
+  }
 }
 
 bool DecodingSession::try_enqueue(WorkItem item) {

libs/qec/lib/realtime/decoding-server-cqr/DecodingSession.h

@@ -106,7 +106,10 @@ struct DecodingSession {
   create(std::unique_ptr<cudaq::qec::decoder> decoder,
          SyndromeMappingTable mapping_table);
 
-  /// Start the FIFO worker thread.  Must be called after create().
+  /// Start the FIFO worker thread.  Must be called after create().  The
+  /// worker pins itself to the decoder's cuda_device_id before serving work;
+  /// a pin failure throws HERE (one worker owns one decoder -- a worker on
+  /// the wrong device must never serve).
   void start_worker();
 
   /// Signal shutdown and join the worker (drains any queued items first).

libs/qec/lib/realtime/decoding-server-cqr/GpuRoceTransceiver.cpp

@@ -117,7 +117,9 @@ GpuRoceConfig GpuRoceConfig::from_env() {
   c.device_name = env_str("HOLOLINK_DEVICE");
   c.peer_ip = env_str("HOLOLINK_PEER_IP");
   c.remote_qp = env_u32("HOLOLINK_REMOTE_QP", 0);
-  c.gpu_id = env_int("HOLOLINK_GPU_ID", 0);
+  if (std::getenv("HOLOLINK_GPU_ID"))
+    c.gpu_id_env = env_int("HOLOLINK_GPU_ID", 0);
+  c.gpu_id = c.gpu_id_env.value_or(0);
   c.frame_size = env_size("HOLOLINK_FRAME_SIZE", 384);
   c.page_size = env_size("HOLOLINK_PAGE_SIZE", 0); // 0 → derived below
   c.num_pages = env_size("HOLOLINK_NUM_PAGES", 64);

libs/qec/lib/realtime/decoding-server-cqr/GpuRoceTransceiver.h

@@ -15,6 +15,7 @@
 #include <atomic>
 #include <cstddef>
 #include <cstdint>
+#include <optional>
 #include <string>
 #include <thread>
 
@@ -38,6 +39,9 @@ struct GpuRoceConfig {
   std::string device_name; ///< HOLOLINK_DEVICE   (IB netdev, e.g. "mlx5_0")
   uint32_t remote_qp{0};   ///< HOLOLINK_REMOTE_QP (FPGA/emulator QP number)
   int gpu_id{0};           ///< HOLOLINK_GPU_ID
+  /// Set iff HOLOLINK_GPU_ID was present in the environment (the FPGA/NIC
+  /// affinity is a topology fact; absence defers to the decoder's pin).
+  std::optional<int> gpu_id_env;
   size_t frame_size{384};  ///< HOLOLINK_FRAME_SIZE (max RPC frame bytes)
   size_t page_size{0};     ///< HOLOLINK_PAGE_SIZE (0 → derived from frame_size)
   size_t num_pages{64};    ///< HOLOLINK_NUM_PAGES (ring depth)

libs/qec/lib/realtime/qec_realtime_session.cpp

@@ -10,6 +10,8 @@
 
 #include "qec_realtime_session.h"
 
+#include "../hardware_guards.h"
+
 #include "cudaq/qec/logger.h"
 #include "cudaq/qec/realtime/decoder_rpc_ids.h"
 #include "cudaq/qec/realtime/graph_resources.h"
@@ -135,16 +137,11 @@ cudaq::qec::decoder *get_decoder_or_throw(std::int64_t decoder_id) {
 static void apply_decoder_cuda_device(cudaq::qec::decoder *dec) {
   if (!dec)
     return;
-  const int id = dec->get_cuda_device_id();
-  if (id < 0)
-    return;
-  int cur = -1;
-  if (cudaGetDevice(&cur) == cudaSuccess && cur == id)
-    return;
-  cudaError_t err = cudaSetDevice(id);
-  if (err != cudaSuccess)
-    CUDA_QEC_WARN("apply_decoder_cuda_device: cudaSetDevice({}) failed: {}", id,
-                  cudaGetErrorString(err));
+  // Throws on failure (fail fast): host dispatch surfaces it as an error
+  // response and graph initialization aborts, rather than continuing on
+  // whichever device happened to be current.
+  cudaq::qec::detail_affinity::set_cuda_device_for_decode(
+      dec->get_cuda_device_id());
 }
 
 // Two-ring response writer: the request stays in `rx_slot` (read-only); the
@@ -282,7 +279,9 @@ void reset_decoder_host(const void *rx_slot, void *tx_slot, std::size_t) {
     const auto *body = reinterpret_cast<const rpc::ResetRequestPayload *>(
         static_cast<const std::uint8_t *>(rx_slot) +
         sizeof(cudaq::realtime::RPCHeader));
-    get_decoder_or_throw(body->decoder_id)->reset_decoder();
+    auto *decoder = get_decoder_or_throw(body->decoder_id);
+    apply_decoder_cuda_device(decoder);
+    decoder->reset_decoder();
     write_response(tx_slot, rx_slot, 0);
   } catch (...) {
     write_response(tx_slot, rx_slot, -2);