Follow up fixes for pinning hardware device

docs/sphinx/examples_rst/qec/realtime_decoding.rst

@@ -155,6 +155,7 @@ arguments:
    decoders:
      - id: 0
        type: pymatching
+       cuda_device_id: 0   # optional: pin this decoder to a CUDA device
        block_size: 3
        syndrome_size: 3
        H_sparse: [ 0, -1, 1, -1, 2, -1 ]
@@ -164,6 +165,14 @@ arguments:
          error_rate_vec: [ 0.1, 0.1, 0.1 ]
          merge_strategy: smallest_weight
 
+``cuda_device_id`` pins a GPU-accelerated decoder (e.g. ``nv-qldpc-decoder``
+or ``trt_decoder``) to a specific CUDA device. The same knob is available as
+a construction parameter in C++ and Python
+(``qec.get_decoder("trt_decoder", H, cuda_device_id=1)``). The thread that
+creates a decoder is pinned to that device and is expected to drive its
+decode calls; create each pinned decoder on its own thread to place several
+decoders on different GPUs.
+
 Here is how to create and save a decoder configuration:
 
 .. tab:: Python

libs/qec/include/cudaq/qec/decoder.h

@@ -260,6 +260,12 @@ class decoder
   /// depends on D_sparse, so you must have called set_D_sparse() first.
   uint32_t get_num_msyn_per_decode() const;
 
+  /// @brief The CUDA device this decoder was pinned to at construction via
+  /// the "cuda_device_id" parameter, or -1 when no pin was requested.
+  /// Construction pins the constructing thread persistently (the thread that
+  /// creates a decoder is the thread expected to drive its decode calls).
+  int get_cuda_device_id() const { return cuda_device_id_; }
+
   /// @brief Set the observable matrix.
   void set_O_sparse(const std::vector<std::vector<uint32_t>> &O_sparse);
 
@@ -357,6 +363,10 @@ class decoder
   /// @brief The decoder's D matrix in sparse format
   std::vector<std::vector<uint32_t>> D_sparse;
 
+  /// @brief CUDA device id consumed from the construction parameters by
+  /// decoder::get(); -1 = unpinned. See get_cuda_device_id().
+  int cuda_device_id_ = -1;
+
 private:
   decode_result_type result_type_ = decode_result_type::decode_to_errs;
 };

libs/qec/include/cudaq/qec/realtime/decoding_config.h

@@ -184,6 +184,11 @@ struct decoder_config {
   /// Defaults to cpu_roce.  Set to gpu_roce for decoders where syndrome bits
   /// are DMA'd directly to GPU VRAM (e.g. nv_qldpc_decoder with RelayBP).
   DecoderTransport transport = DecoderTransport::cpu_roce;
+  /// CUDA device this decoder is pinned to at construction (see the
+  /// "cuda_device_id" decoder parameter). Placement knob common to any
+  /// GPU-accelerated decoder, hence at this level rather than inside the
+  /// per-decoder custom args. Unset = unpinned.
+  std::optional<int> cuda_device_id;
   uint64_t block_size = 0;
   uint64_t syndrome_size = 0;
   std::vector<std::int64_t> H_sparse;

libs/qec/lib/CMakeLists.txt

@@ -150,6 +150,10 @@ target_link_libraries(${DECODERS_LIBRARY_NAME}
   PUBLIC
     $<LINK_LIBRARY:WHOLE_ARCHIVE,cudaqx-core>
     fmt::fmt-header-only
+  PRIVATE
+    # decoder::get() consumes the "cuda_device_id" construction parameter
+    # (validation via cudaGetDeviceCount + persistent cudaSetDevice pin).
+    CUDA::cudart
 )
 
 target_link_libraries(${LIBRARY_NAME}

libs/qec/lib/decoder.cpp

@@ -8,10 +8,12 @@
 
 #include "cudaq/qec/decoder.h"
 #include "cuda-qx/core/library_utils.h"
+#include "hardware_guards.h"
 #include "cudaq/qec/logger.h"
 #include "cudaq/qec/plugin_loader.h"
 #include "cudaq/qec/version.h"
 #include <cassert>
+#include <cuda_runtime_api.h>
 #include <dlfcn.h>
 #include <filesystem>
 #include <fmt/ranges.h>
@@ -123,8 +125,33 @@ std::string decoder::get_version() const {
 
 std::future<decoder_result>
 decoder::decode_async(const std::vector<float_t> &syndrome) {
-  return std::async(std::launch::async,
-                    [this, syndrome] { return this->decode(syndrome); });
+  // Captured by value: the worker must not dereference decoder members to
+  // find its device. The std::async thread is brand-new and unpinned, so it
+  // guards itself for the duration of the call (the one exception to the
+  // one-thread-owns-one-decoder persistent pin).
+  const int cuda_id = cuda_device_id_;
+  return std::async(std::launch::async, [this, syndrome, cuda_id] {
+    cudaq::qec::detail_affinity::CudaDeviceGuard dev(cuda_id);
+    return this->decode(syndrome);
+  });
+}
+
+/// Reads "cuda_device_id" from the construction parameters. Absent -> -1.
+/// Negative or >= cudaGetDeviceCount() -> std::runtime_error (fail fast:
+/// never silently decode on the wrong GPU).
+static int read_cuda_device_id(const cudaqx::heterogeneous_map &params) {
+  if (!params.contains("cuda_device_id"))
+    return -1;
+  const int value = params.get<int>("cuda_device_id");
+  if (value < 0)
+    throw std::runtime_error("cuda_device_id must be >= 0 (got " +
+                             std::to_string(value) + ")");
+  int count = 0;
+  if (cudaGetDeviceCount(&count) != cudaSuccess || value >= count)
+    throw std::runtime_error("cuda_device_id " + std::to_string(value) +
+                             " is out of range: " + std::to_string(count) +
+                             " CUDA device(s) visible");
+  return value;
 }
 
 std::unique_ptr<decoder>
@@ -138,7 +165,26 @@ decoder::get(const std::string &name, const decoder_init &init,
         "invalid decoder requested: " + name +
         ". Run with CUDAQ_LOG_LEVEL=info (environment variable) to see "
         "additional plugin diagnostics at startup.");
-  return iter->second(init, param_map);
+  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));
+  // The key is consumed here; strip it so plugins that strictly validate
+  // their parameter keys do not reject it.
+  cudaqx::heterogeneous_map plugin_params;
+  for (const auto &kv : param_map)
+    if (kv.first != "cuda_device_id")
+      plugin_params.insert(kv.first, kv.second);
+  auto d = iter->second(init, plugin_params);
+  d->cuda_device_id_ = cuda_device_id;
+  return d;
 }
 
 namespace details {

libs/qec/lib/hardware_guards.h

@@ -0,0 +1,73 @@
+/*******************************************************************************
+ * Copyright (c) 2022 - 2026 NVIDIA Corporation & Affiliates.                  *
+ * All rights reserved.                                                        *
+ *                                                                             *
+ * This source code and the accompanying materials are made available under    *
+ * the terms of the Apache License 2.0 which accompanies this distribution.    *
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda_runtime_api.h>
+#include <stdexcept>
+#include <string>
+
+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
+/// with NUMA guards; the nv-qldpc follow-up mirrors its use).
+///
+/// This guard is for threads that do NOT follow the one-thread-owns-one-
+/// decoder persistent pin (e.g. the fresh worker spawned by decode_async).
+class CudaDeviceGuard {
+public:
+  explicit CudaDeviceGuard(int target) {
+    if (target < 0)
+      return;
+    int count = 0;
+    if (cudaGetDeviceCount(&count) != cudaSuccess || target >= count)
+      throw std::runtime_error("cuda_device_id " + std::to_string(target) +
+                               " is out of range: " + std::to_string(count) +
+                               " CUDA device(s) visible");
+    // If the current device is unreadable, skip restoration rather than
+    // restore to a guessed device; the set below still applies.
+    if (cudaGetDevice(&prev_) != cudaSuccess)
+      prev_ = -1;
+    cudaError_t err = cudaSetDevice(target);
+    if (err != cudaSuccess)
+      throw std::runtime_error("CudaDeviceGuard: cudaSetDevice(" +
+                               std::to_string(target) +
+                               ") failed: " + cudaGetErrorString(err));
+    restore_ = (prev_ >= 0 && prev_ != target);
+  }
+  ~CudaDeviceGuard() {
+    if (restore_)
+      (void)cudaSetDevice(prev_);
+  }
+  CudaDeviceGuard(const CudaDeviceGuard &) = delete;
+  CudaDeviceGuard &operator=(const CudaDeviceGuard &) = delete;
+
+private:
+  int prev_ = -1;
+  bool restore_ = false;
+};
+
+} // namespace cudaq::qec::detail_affinity

libs/qec/lib/realtime/config.cpp

@@ -720,6 +720,7 @@ struct MappingTraits<cudaq::qec::decoding::config::decoder_config> {
     io.mapRequired("type", config.type);
     io.mapOptional("transport", config.transport,
                    cudaq::qec::decoding::config::DecoderTransport::cpu_roce);
+    io.mapOptional("cuda_device_id", config.cuda_device_id);
     io.mapRequired("block_size", config.block_size);
     io.mapRequired("syndrome_size", config.syndrome_size);
     io.mapRequired("H_sparse", config.H_sparse);

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,35 @@ 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 && 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 +88,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;

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,13 @@
 
 #include "DecodingSession.h"
 #include "RpcWireFormat.h"
+#include "../../hardware_guards.h"
 #include "cudaq/qec/logger.h"
 
 #include <chrono>
 #include <cstring>
+#include <cuda_runtime_api.h>
+#include <future>
 #include <stdexcept>
 #include <vector>
 
@@ -61,7 +64,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);