Enforce decoder GPU ownership in the realtime server

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/lib/realtime/decoding-server-cqr/SessionRegistry.cpp

@@ -7,6 +7,7 @@
  ******************************************************************************/
 
 #include "SessionRegistry.h"
+#include "../../hardware_guards.h"
 #include "../realtime_decoding.h"
 #include "cudaq/qec/logger.h"
 #include "cudaq/qec/realtime/decoding_config.h"
@@ -71,6 +72,11 @@ void SessionRegistry::load_from_config(const multi_decoder_config &config,
     CUDA_QEC_INFO("SessionRegistry: creating decoder id={} type={}", dc.id,
                   dc.type);
 
+    // Keep each construction transaction on its configured device, then
+    // restore the registry thread. Runtime ownership transfers to the
+    // session's dedicated worker in start_worker().
+    cudaq::qec::detail_affinity::CudaDeviceGuard construction_device(
+        dc.cuda_device_id.value_or(-1));
     auto decoder = cudaq::qec::decoding::host::create_realtime_decoder(dc);
     auto session = DecodingSession::create(std::move(decoder),
                                            make_default_mapping_table());

libs/qec/lib/realtime/qec_realtime_session.h

@@ -187,6 +187,10 @@ class __attribute__((visibility("default"))) qec_realtime_session {
   // ---- Lifetime / mode ----
   bool initialized_ = false;
   bool device_mode_ = false;
+  // Every DEVICE-mode decoder must share one CUDA owner, so each scheduler
+  // allocation, graph launch, and cleanup operation belongs to this device.
+  // -1 means none of the decoders requested explicit placement.
+  int device_mode_cuda_device_id_ = -1;
 
   // ---- Ring buffer (raw pointers; _dev aliases _host in HOST mode) ----
   static constexpr std::size_t kDefaultNumSlots = 8;

libs/qec/python/tests/test_decoders_yaml.py

@@ -128,6 +128,19 @@ def create_test_decoder_config_nv_qldpc(decoder_id):
     return config
 
 
+def test_cuda_device_id_yaml_roundtrip():
+    multi_config = qec.multi_decoder_config()
+    config = create_test_empty_decoder_config(0)
+    assert config.cuda_device_id is None
+    config.cuda_device_id = 2
+    multi_config.decoders = [config]
+
+    check_decoder_yaml_roundtrip(multi_config)
+    parsed = qec.multi_decoder_config.from_yaml_str(
+        multi_config.to_yaml_str(200))
+    assert parsed.decoders[0].cuda_device_id == 2
+
+
 def test_single_decoder():
     """
     Test YAML serialization/deserialization and creation of a single NV-QLDPC decoder.

libs/qec/unittests/CMakeLists.txt

@@ -48,7 +48,8 @@ target_link_libraries(test_decoders_yaml PRIVATE
   cudaq-qec-decoding-server
   cudaq-qec-realtime-decoding
   cudaq-qec-realtime-decoding-simulation
-  cudaq::cudaq)
+  cudaq::cudaq
+  CUDA::cudart)
 add_dependencies(CUDAQXQECUnitTests test_decoders_yaml)
 gtest_discover_tests(test_decoders_yaml)
 

libs/qec/unittests/decoders/pymatching/CMakeLists.txt

@@ -42,6 +42,8 @@ if(CUDAQ_REALTIME_ROOT AND CUDAQ_REALTIME_INCLUDE_DIR)
     GTest::gtest_main
     cudaq-qec-decoders
     cudaq-qec-realtime-decoding
+    cudaq::cudaq
+    CUDA::cudart
   )
   set_target_properties(test_pymatching_realtime PROPERTIES
     BUILD_RPATH "${CMAKE_BINARY_DIR}/lib;${CMAKE_BINARY_DIR}/lib/decoder-plugins"

libs/qec/unittests/decoders/pymatching/test_pymatching_realtime.cpp

@@ -12,8 +12,10 @@
 #include "cudaq/qec/decoder.h"
 #include "cudaq/qec/realtime/decoding_config.h"
 
+#include <cuda_runtime_api.h>
 #include <gtest/gtest.h>
 
+#include <atomic>
 #include <cstdint>
 #include <memory>
 #include <span>
@@ -24,6 +26,53 @@ namespace {
 
 using DecoderVec = std::vector<std::unique_ptr<cudaq::qec::decoder>>;
 
+class RealtimeSessionDeviceRecordingDecoder final : public cudaq::qec::decoder {
+public:
+  std::atomic<int> last_decode_device{-2};
+
+  RealtimeSessionDeviceRecordingDecoder(
+      const cudaq::qec::sparse_binary_matrix &H,
+      const cudaqx::heterogeneous_map &)
+      : decoder(H) {}
+
+  cudaq::qec::decoder_result
+  decode(const std::vector<cudaq::qec::float_t> &syndrome) override {
+    int device = -1;
+    if (cudaGetDevice(&device) != cudaSuccess)
+      device = -1;
+    last_decode_device.store(device, std::memory_order_release);
+    cudaq::qec::decoder_result result;
+    result.converged = true;
+    result.result = {syndrome.at(0)};
+    return result;
+  }
+
+  CUDAQ_EXTENSION_CUSTOM_CREATOR_FUNCTION(
+      RealtimeSessionDeviceRecordingDecoder,
+      static std::unique_ptr<cudaq::qec::decoder> create(
+          const cudaq::qec::decoder_init &init,
+          const cudaqx::heterogeneous_map &params) {
+        return cudaq::qec::make_pcm_decoder<
+            RealtimeSessionDeviceRecordingDecoder>(init, params);
+      })
+};
+CUDAQ_EXT_PT_REGISTER_TYPE(RealtimeSessionDeviceRecordingDecoder)
+
+class ScopedDeviceRestore {
+public:
+  ScopedDeviceRestore() {
+    if (cudaGetDevice(&previous) != cudaSuccess)
+      previous = -1;
+  }
+  ~ScopedDeviceRestore() {
+    if (previous >= 0)
+      (void)cudaSetDevice(previous);
+  }
+
+private:
+  int previous = -1;
+};
+
 DecoderVec make_pymatching_decoders(const std::vector<std::uint8_t> &h_vec,
                                     std::size_t syndrome_size,
                                     std::size_t block_size) {
@@ -208,3 +257,38 @@ TEST(PyMatchingRealtime, ConfiguresViaRealtimeDecoderConfig) {
 
   config::finalize_decoders();
 }
+
+TEST(RealtimeSessionCudaDeviceId, HostDispatcherUsesConfiguredDevice) {
+  int device_count = 0;
+  if (cudaGetDeviceCount(&device_count) != cudaSuccess || device_count < 2)
+    GTEST_SKIP() << "needs >= 2 CUDA devices";
+
+  ScopedDeviceRestore restore;
+  cudaqx::tensor<std::uint8_t> H({std::size_t{1}, std::size_t{1}});
+  H.at({0, 0}) = 1;
+  cudaqx::heterogeneous_map params;
+  params.insert("cuda_device_id", 1);
+  auto decoder = cudaq::qec::decoder::get(
+      "RealtimeSessionDeviceRecordingDecoder", H, params);
+  auto *recording =
+      dynamic_cast<RealtimeSessionDeviceRecordingDecoder *>(decoder.get());
+  ASSERT_NE(recording, nullptr);
+  decoder->set_decoder_id(0);
+  decoder->set_D_sparse(std::vector<std::vector<std::uint32_t>>{{0}});
+  decoder->set_O_sparse(std::vector<std::vector<std::uint32_t>>{{0}});
+
+  DecoderVec decoders;
+  decoders.push_back(std::move(decoder));
+  ASSERT_EQ(cudaSetDevice(0), cudaSuccess);
+  cudaq::qec::realtime::qec_realtime_session session(decoders);
+  session.initialize();
+
+  expect_corrections(session, {1}, std::vector<std::uint8_t>{1},
+                     /*counter=*/1);
+  EXPECT_EQ(recording->last_decode_device.load(std::memory_order_acquire), 1);
+
+  int current = -1;
+  ASSERT_EQ(cudaGetDevice(&current), cudaSuccess);
+  EXPECT_EQ(current, 0);
+  session.finalize();
+}

libs/qec/unittests/realtime/qec_graph_decode_test/qldpc_config_loader.cpp

@@ -7,6 +7,7 @@
  ******************************************************************************/
 
 #include "qldpc_config_loader.h"
+#include "realtime_decoding.h"
 
 #include "cudaq/qec/decoder.h"
 #include "cudaq/qec/realtime/decoding_config.h"
@@ -64,7 +65,7 @@ LoadedDecoder load_decoder_from_yaml(const std::string &yaml_path) {
     for (std::uint32_t j = h_row_ptr[r]; j < h_row_ptr[r + 1]; ++j)
       H_tensor.at({r, static_cast<std::size_t>(h_col_idx[j])}) = 1;
 
-  auto params = dec.decoder_custom_args_to_heterogeneous_map();
+  auto params = decoding::host::prepare_decoder_params(dec);
   auto plugin = decoder::get("nv-qldpc-decoder", H_tensor, params);
   if (!plugin)
     throw std::runtime_error(

libs/qec/unittests/realtime/qec_graph_decode_test/test_realtime_qldpc_graph_decoding.cpp

@@ -254,6 +254,7 @@ class GraphDecodeTest : public ::testing::Test {
     num_measurements_ = loaded.num_measurements;
     num_observables_ = loaded.num_observables;
     ASSERT_NE(loaded.decoder, nullptr);
+    EXPECT_EQ(loaded.decoder->get_cuda_device_id(), 0);
     printf("Config: num_measurements=%zu, num_observables=%zu\n",
            num_measurements_, num_observables_);
 

libs/qec/unittests/realtime/qec_roce_decode_test/data/config_nv_qldpc_relay.yml

@@ -2,6 +2,7 @@
 decoders:
   - id:              0
     type:            nv-qldpc-decoder
+    cuda_device_id:  0
     block_size:      69
     syndrome_size:   24
     H_sparse:        [ 0, 1, 2, 3, 4, 5, 6, -1, 1, 4, 7, 8, -1, 2, 5, 9, 10, 11, 12, 13, -1, 10, 11, 14, 15, -1, 3, 6, 16, 17, -1, 4, 5, 8, 13, 18, 19, 20, -1, 5, 6, 11, 15, 17, 19, 21, -1, 12, 13, 20,