Add cuda_device_id placement knob for GPU decoders

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

@@ -183,6 +183,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

@@ -11,7 +11,9 @@
 #include "cudaq/qec/logger.h"
 #include "cudaq/qec/plugin_loader.h"
 #include "cudaq/qec/version.h"
+#include "hardware_guards.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,27 @@ 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,57 @@
+/*******************************************************************************
+ * 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 {
+
+/// 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

@@ -717,6 +717,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/qec_realtime_session.cpp

@@ -19,6 +19,7 @@
 #include <atomic>
 #include <cstdlib>
 #include <cstring>
+#include <cuda_runtime_api.h>
 #include <dlfcn.h>
 #include <stdexcept>
 #include <string>
@@ -126,6 +127,26 @@ cudaq::qec::decoder *get_decoder_or_throw(std::int64_t decoder_id) {
   return (*decoders)[static_cast<std::size_t>(decoder_id)].get();
 }
 
+// Point the calling thread at the decoder's pinned CUDA device before work
+// that allocates or launches on it. Set-and-leave (no restore): one
+// dispatcher thread serves all decoders, so the thread simply converges to
+// the device of the decoder it is currently serving; cudaSetDevice on an
+// already-current device is a cheap no-op.
+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));
+}
+
 // Two-ring response writer: the request stays in `rx_slot` (read-only); the
 // response is written into the distinct `tx_slot`.  The preserved header fields
 // (request_id, ptp_timestamp) must be echoed explicitly from rx to tx.  The
@@ -178,6 +199,7 @@ void enqueue_syndromes_host(const void *rx_slot, void *tx_slot,
     }
 
     auto *decoder = get_decoder_or_throw(body->decoder_id);
+    apply_decoder_cuda_device(decoder);
     // Reject requests larger than this decoder's per-decode window.  The slot
     // is sized for the largest decoder in the session, so an oversized request
     // for a smaller decoder can still fit the slot; without this guard it would
@@ -537,6 +559,7 @@ void qec_realtime_session::capture_decoder_graphs() {
           "qec_realtime_session::initialize: decoder " + std::to_string(i) +
           " does not support graph dispatch in DEVICE mode.");
 
+    apply_decoder_cuda_device(dec);
     // reserved_sms = 0 is intentional for the inproc_rpc desktop / CI path.
     void *raw = dec->capture_decode_graph(/*reserved_sms=*/0);
     if (!raw)

libs/qec/lib/realtime/realtime_decoding.cpp

@@ -156,6 +156,10 @@ namespace cudaq::qec::decoding::host {
 cudaqx::heterogeneous_map prepare_decoder_params(
     const cudaq::qec::decoding::config::decoder_config &decoder_config) {
   auto params = decoder_config.decoder_custom_args_to_heterogeneous_map();
+  // Placement knob: surfaced for every decoder type (deliberately before the
+  // trt-only early return below); consumed by decoder::get() at construction.
+  if (decoder_config.cuda_device_id.has_value())
+    params.insert("cuda_device_id", decoder_config.cuda_device_id.value());
   if (decoder_config.type != "trt_decoder")
     return params;
 

libs/qec/python/bindings/py_decoding_config.cpp

@@ -279,6 +279,7 @@ void bindDecodingConfig(nb::module_ &mod) {
       .def(nb::init<>())
       .def_rw("id", &decoder_config::id)
       .def_rw("type", &decoder_config::type)
+      .def_rw("cuda_device_id", &decoder_config::cuda_device_id)
       .def_rw("block_size", &decoder_config::block_size)
       .def_rw("syndrome_size", &decoder_config::syndrome_size)
       .def_rw("H_sparse", &decoder_config::H_sparse)

libs/qec/python/tests/test_decoder.py

@@ -1002,6 +1002,32 @@ def test_get_decoder_rejects_unknown_decoder_for_stim_dem_text():
         qec.get_decoder("__no_such_decoder__", "error(0.1) D0 L0\n")
 
 
+def test_decoder_cuda_device_id_invalid_raises():
+    H = create_test_matrix()
+    # A negative id is rejected before reaching the C++ guard: the kwargs
+    # marshalling layer stores Python ints as size_t, so nanobind refuses the
+    # negative value with a bare RuntimeError ("std::bad_cast").
+    with pytest.raises(RuntimeError):
+        qec.get_decoder("single_error_lut", H, cuda_device_id=-2)
+    # An out-of-range id flows through kwargs to decoder::get(), which raises
+    # a runtime_error naming the offending parameter.
+    with pytest.raises(RuntimeError, match="cuda_device_id"):
+        qec.get_decoder("single_error_lut", H, cuda_device_id=1 << 20)
+
+
+def test_decoder_cuda_device_id_valid():
+    H = create_test_matrix()
+    try:
+        d = qec.get_decoder("single_error_lut", H, cuda_device_id=0)
+    except RuntimeError as e:
+        if "out of range" in str(e):
+            pytest.skip("no CUDA device visible")
+        raise
+    syndrome = create_test_syndrome()
+    result = d.decode(syndrome)
+    assert len(result.result) == H.shape[1]
+
+
 def test_get_decoder_user_O_wins_over_dem_derived():
     dem_text = ("error(0.1) D0 L0\n"
                 "error(0.1) D1 L0\n"

libs/qec/unittests/CMakeLists.txt

@@ -35,7 +35,7 @@ find_package(CUDAToolkit REQUIRED)
 add_compile_options(-Wno-attributes) 
 
 add_executable(test_decoders test_decoders.cpp decoders/sample_decoder.cpp)
-target_link_libraries(test_decoders PRIVATE GTest::gtest_main cudaq-qec-decoders libstim)
+target_link_libraries(test_decoders PRIVATE GTest::gtest_main cudaq-qec-decoders libstim CUDA::cudart)
 add_dependencies(CUDAQXQECUnitTests test_decoders)
 gtest_discover_tests(test_decoders)