Enforce decoder GPU ownership in the realtime server

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/core/include/cuda-qx/core/kwargs_utils.h

@@ -1,5 +1,5 @@
 /****************************************************************-*- C++ -*-****
- * Copyright (c) 2024 - 2025 NVIDIA Corporation & Affiliates.                  *
+ * Copyright (c) 2024 - 2026 NVIDIA Corporation & Affiliates.                  *
  * All rights reserved.                                                        *
  *                                                                             *
  * This source code and the accompanying materials are made available under    *
@@ -13,6 +13,7 @@
 #include <nanobind/ndarray.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/vector.h>
+#include <stdexcept>
 
 namespace nb = nanobind;
 
@@ -41,7 +42,14 @@ inline heterogeneous_map hetMapFromKwargs(const nb::kwargs &kwargs) {
     if (nb::isinstance<nb::bool_>(value)) {
       result.insert(key, nb::cast<bool>(value));
     } else if (nb::isinstance<nb::int_>(value)) {
-      result.insert(key, nb::cast<std::size_t>(value));
+      std::size_t integer_value = 0;
+      try {
+        integer_value = nb::cast<std::size_t>(value);
+      } catch (...) {
+        throw std::runtime_error("Integer keyword argument '" + key +
+                                 "' must be non-negative and fit in size_t");
+      }
+      result.insert(key, integer_value);
     } else if (nb::isinstance<nb::float_>(value)) {
       result.insert(key, nb::cast<double>(value));
     } else if (nb::isinstance<nb::str>(value)) {

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;
+
   /// @brief Set the observable matrix.
   void set_O_sparse(const std::vector<std::vector<uint32_t>> &O_sparse);
 

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

@@ -184,6 +184,9 @@ 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 that owns this decoder's construction, execution, captured
+  /// graph resources, and teardown. Unset leaves device selection unchanged.
+  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,13 +8,19 @@
 
 #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 <any>
 #include <cassert>
+#include <cstdint>
+#include <cuda_runtime_api.h>
 #include <dlfcn.h>
 #include <filesystem>
 #include <fmt/ranges.h>
+#include <limits>
+#include <type_traits>
 #include <vector>
 
 INSTANTIATE_REGISTRY(cudaq::qec::decoder, const cudaq::qec::decoder_init &,
@@ -54,6 +60,11 @@ struct decoder::rt_impl {
   /// The id of the decoder (for instrumentation)
   uint32_t decoder_id = 0;
 
+  /// CUDA device selected by decoder::get(), or -1 when no placement was
+  /// requested. Keep this in the existing pimpl so adding placement does not
+  /// change the public decoder object layout used by external plugins.
+  int cuda_device_id = -1;
+
   bool is_sliding_window = false;
 
   /// The number of syndromes per round.  Only used for sliding window decoder.
@@ -121,12 +132,139 @@ std::string decoder::get_version() const {
   return ss.str();
 }
 
+int decoder::get_cuda_device_id() const { return pimpl->cuda_device_id; }
+
 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 = get_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).
+template <typename Integer>
+static bool try_read_cuda_device_id(const std::any &raw_value, int &result) {
+  const auto *value = std::any_cast<Integer>(&raw_value);
+  if (!value)
+    return false;
+
+  if constexpr (std::is_signed_v<Integer>) {
+    if (*value < 0)
+      throw std::runtime_error(
+          "cuda_device_id must be a non-negative integer (got " +
+          std::to_string(*value) + ")");
+  }
+
+  using unsigned_integer = std::make_unsigned_t<Integer>;
+  const auto magnitude =
+      static_cast<std::uintmax_t>(static_cast<unsigned_integer>(*value));
+  if (magnitude > static_cast<std::uintmax_t>(std::numeric_limits<int>::max()))
+    throw std::runtime_error("cuda_device_id is too large (got " +
+                             std::to_string(*value) +
+                             "); maximum supported value is " +
+                             std::to_string(std::numeric_limits<int>::max()));
+
+  result = static_cast<int>(magnitude);
+  return true;
+}
+
+static int read_cuda_device_id(const cudaqx::heterogeneous_map &params) {
+  if (!params.contains("cuda_device_id"))
+    return -1;
+
+  const std::any *raw_value = nullptr;
+  for (const auto &[key, value] : params) {
+    if (key == "cuda_device_id") {
+      raw_value = &value;
+      break;
+    }
+  }
+  if (!raw_value)
+    throw std::runtime_error("cuda_device_id is missing from parameter map");
+
+  int value = -1;
+  const bool is_integer =
+      try_read_cuda_device_id<signed char>(*raw_value, value) ||
+      try_read_cuda_device_id<unsigned char>(*raw_value, value) ||
+      try_read_cuda_device_id<short>(*raw_value, value) ||
+      try_read_cuda_device_id<unsigned short>(*raw_value, value) ||
+      try_read_cuda_device_id<int>(*raw_value, value) ||
+      try_read_cuda_device_id<unsigned int>(*raw_value, value) ||
+      try_read_cuda_device_id<long>(*raw_value, value) ||
+      try_read_cuda_device_id<unsigned long>(*raw_value, value) ||
+      try_read_cuda_device_id<long long>(*raw_value, value) ||
+      try_read_cuda_device_id<unsigned long long>(*raw_value, value);
+  if (!is_integer)
+    throw std::runtime_error("cuda_device_id must be an integer");
+
+  int count = 0;
+  const cudaError_t count_status = cudaGetDeviceCount(&count);
+  if (count_status != cudaSuccess)
+    throw std::runtime_error(
+        "cuda_device_id " + std::to_string(value) +
+        " could not be validated because cudaGetDeviceCount() failed: " +
+        cudaGetErrorString(count_status));
+  if (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;
+}
+
+/// Selects a device for decoder construction and restores the previous device
+/// unless commit() is called. This keeps the successful persistent-pin
+/// contract while making failed plugin construction transactional.
+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 for cuda_device_id: " + cudaGetErrorString(err));
+    selected_ = true;
+  }
+
+  ~ConstructionDevicePin() {
+    if (selected_ && !committed_ && previous_ != target_)
+      (void)cudaSetDevice(previous_);
+  }
+
+  ConstructionDevicePin(const ConstructionDevicePin &) = delete;
+  ConstructionDevicePin &operator=(const ConstructionDevicePin &) = delete;
+
+  void commit() {
+    const cudaError_t err = cudaSetDevice(target_);
+    if (err != cudaSuccess)
+      throw std::runtime_error(
+          "cudaSetDevice(" + std::to_string(target_) +
+          ") failed while finalizing cuda_device_id placement: " +
+          cudaGetErrorString(err));
+    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) {
@@ -138,7 +276,24 @@ 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.
+  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;
+  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->pimpl->cuda_device_id = cuda_device_id;
+  device_pin.commit();
+  return d;
 }
 
 namespace details {

libs/qec/lib/hardware_guards.h

@@ -0,0 +1,85 @@
+/*******************************************************************************
+ * 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>
+#include <string_view>
+
+namespace cudaq::qec::detail_affinity {
+
+/// Persistently select \p target for the calling thread. The device id must
+/// already have been validated by decoder::get(). This is the hot-path helper
+/// for long-lived decoder workers and shared realtime dispatcher threads.
+inline void select_cuda_device(int target, std::string_view context) {
+  if (target < 0)
+    return;
+
+  int current = -1;
+  cudaError_t err = cudaGetDevice(&current);
+  if (err != cudaSuccess)
+    throw std::runtime_error(
+        std::string(context) +
+        ": cudaGetDevice() failed: " + cudaGetErrorString(err));
+  if (current == target)
+    return;
+
+  err = cudaSetDevice(target);
+  if (err != cudaSuccess)
+    throw std::runtime_error(std::string(context) + ": 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 it can be
+/// reused by decoder implementation files without adding a public API.
+///
+/// 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;
+    const cudaError_t count_status = cudaGetDeviceCount(&count);
+    if (count_status != cudaSuccess)
+      throw std::runtime_error(
+          "CudaDeviceGuard: cudaGetDeviceCount() failed for cuda_device_id " +
+          std::to_string(target) + ": " + cudaGetErrorString(count_status));
+    if (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,8 @@ target_link_libraries(cudaq-qec-decoding-server
     # violates the realtime-server dependency-closure contract.
     cudaq-qec-decoders
     cudaq-qec-realtime-decoding
+  PRIVATE
+    CUDA::cudart
 )
 
 if(CUDAQ_GPU_ROCE_AVAILABLE)