Add hardware pinning for QEC decoders

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

@@ -245,6 +245,11 @@ class decoder
   std::size_t get_block_size() { return block_size; }
   std::size_t get_syndrome_size() { return syndrome_size; }
 
+  /// @brief Store hardware affinity parameters read from the constructor params
+  /// map. Called by decoder::get() after the plugin constructor returns.
+  /// Not intended for direct use by plugin authors.
+  void set_hardware_params(const cudaqx::heterogeneous_map &params);
+
   // -- Begin realtime decoding API --
 
   // Note: all of the current realtime decoding API is designed to be used with
@@ -357,6 +362,11 @@ class decoder
   /// @brief The decoder's D matrix in sparse format
   std::vector<std::vector<uint32_t>> D_sparse;
 
+  /// Target CUDA device for this decoder. -1 = inherit caller's device (no-op).
+  int cuda_device_id_ = -1;
+  /// Target NUMA node for this decoder. -1 = no binding.
+  int numa_node_id_ = -1;
+
 private:
   decode_result_type result_type_ = decode_result_type::decode_to_errs;
 };

libs/qec/include/cudaq/qec/device_affinity.h

@@ -0,0 +1,42 @@
+/*******************************************************************************
+ * Copyright (c) 2025 - 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-qx/core/heterogeneous_map.h"
+#include <stdexcept>
+#include <string>
+
+namespace cudaq::qec {
+
+namespace detail {
+/// Returns -1 if key is absent; throws if stored value is negative.
+/// Handles both int (YAML path) and std::size_t (Python kwargs path) storage.
+inline int read_pin_key(const cudaqx::heterogeneous_map &params,
+                        const std::string &key) {
+  if (!params.contains(key))
+    return -1;
+  int value = params.get<int>(key);
+  if (value < 0)
+    throw std::runtime_error(key + " must be >= 0 (got " +
+                             std::to_string(value) + ")");
+  return value;
+}
+} // namespace detail
+
+/// @brief GPU device id for a decoder. -1 = inherit current device (no-op).
+inline int read_cuda_device_id(const cudaqx::heterogeneous_map &params) {
+  return detail::read_pin_key(params, "cuda_device_id");
+}
+
+/// @brief NUMA node id for a decoder. -1 = no binding.
+inline int read_numa_node_id(const cudaqx::heterogeneous_map &params) {
+  return detail::read_pin_key(params, "numa_node_id");
+}
+
+} // namespace cudaq::qec

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

@@ -166,6 +166,8 @@ struct decoder_config {
   std::vector<std::int64_t> H_sparse;
   std::vector<std::int64_t> O_sparse;
   std::vector<std::int64_t> D_sparse;
+  std::optional<int> cuda_device_id;
+  std::optional<int> numa_node_id;
   std::variant<single_error_lut_config, multi_error_lut_config,
                nv_qldpc_decoder_config, sliding_window_config,
                trt_decoder_config, pymatching_config>

libs/qec/lib/decoder.cpp

@@ -9,13 +9,143 @@
 #include "cudaq/qec/decoder.h"
 #include "common/FmtCore.h"
 #include "cuda-qx/core/library_utils.h"
+#include "cudaq/qec/device_affinity.h"
 #include "cudaq/qec/plugin_loader.h"
 #include "cudaq/qec/version.h"
 #include "cudaq/runtime/logger/logger.h"
 #include <cassert>
+#include <cuda_runtime.h>
 #include <dlfcn.h>
 #include <filesystem>
 #include <vector>
+#if defined(__linux__)
+#include <fstream>
+#include <linux/mempolicy.h>
+#include <sched.h>
+#include <sstream>
+#include <sys/syscall.h>
+#include <unistd.h>
+#endif
+
+namespace {
+
+// RAII: sets the calling thread's CUDA current device, restores on destruction.
+// target < 0 = no-op. Throws std::runtime_error if target is out of range.
+struct CudaDeviceGuard {
+  int prev_ = -1;
+  bool active_ = false;
+
+  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) +
+          " out of range (device_count=" + std::to_string(count) + ")");
+    cudaGetDevice(&prev_);
+    if (prev_ != target) {
+      cudaSetDevice(target);
+      active_ = true;
+    }
+  }
+  ~CudaDeviceGuard() {
+    if (active_)
+      cudaSetDevice(prev_);
+  }
+  CudaDeviceGuard(const CudaDeviceGuard &) = delete;
+  CudaDeviceGuard &operator=(const CudaDeviceGuard &) = delete;
+};
+
+#if defined(__linux__)
+// Parse /sys/devices/system/node/node<N>/cpulist e.g. "0-7,16-23".
+// Returns false if the file is missing or empty.
+static bool build_node_cpuset(int node, cpu_set_t &out) {
+  std::ifstream f("/sys/devices/system/node/node" + std::to_string(node) +
+                  "/cpulist");
+  if (!f.is_open())
+    return false;
+  std::string list;
+  std::getline(f, list);
+  if (list.empty())
+    return false;
+  std::stringstream ss(list);
+  std::string range;
+  bool any = false;
+  while (std::getline(ss, range, ',')) {
+    auto dash = range.find('-');
+    int lo = std::stoi(range.substr(0, dash));
+    int hi =
+        (dash == std::string::npos) ? lo : std::stoi(range.substr(dash + 1));
+    for (int c = lo; c <= hi; ++c) {
+      if (c < CPU_SETSIZE) { // guard against OOB on >1024-CPU machines
+        CPU_SET(c, &out);
+        any = true;
+      }
+    }
+  }
+  return any;
+}
+
+// RAII: binds the calling thread's CPU affinity and memory policy to a NUMA
+// node, restores on destruction. node < 0 = no-op.
+// Uses raw Linux syscalls -- no libnuma dependency.
+struct NumaGuard {
+  bool mempol_set_ = false;
+  bool affinity_set_ = false;
+  bool has_prev_affinity_ = false;
+  cpu_set_t prev_set_{};
+
+  explicit NumaGuard(int node) {
+    if (node < 0)
+      return;
+    CPU_ZERO(&prev_set_);
+    if (sched_getaffinity(0, sizeof(prev_set_), &prev_set_) == 0)
+      has_prev_affinity_ = true;
+
+    cpu_set_t node_set;
+    CPU_ZERO(&node_set);
+    // Only pin CPU affinity when we can restore it; avoids permanent pinning
+    // if sched_getaffinity failed (e.g., container with locked cpuset).
+    if (has_prev_affinity_ && build_node_cpuset(node, node_set)) {
+      if (sched_setaffinity(0, sizeof(node_set), &node_set) == 0)
+        affinity_set_ = true;
+    }
+
+    if (node < static_cast<int>(sizeof(unsigned long) * 8)) {
+      unsigned long nodemask = 1UL << node;
+      if (syscall(SYS_set_mempolicy, MPOL_BIND, &nodemask,
+                  static_cast<unsigned long>(sizeof(nodemask) * 8)) == 0)
+        mempol_set_ = true;
+    }
+  }
+
+  ~NumaGuard() {
+    if (mempol_set_)
+      syscall(SYS_set_mempolicy, MPOL_DEFAULT, nullptr, 0UL);
+    if (affinity_set_)
+      sched_setaffinity(0, sizeof(prev_set_), &prev_set_);
+  }
+
+  NumaGuard(const NumaGuard &) = delete;
+  NumaGuard &operator=(const NumaGuard &) = delete;
+};
+#else
+struct NumaGuard {
+  explicit NumaGuard(int node) {
+    if (node < 0)
+      return;
+    static bool warned = false;
+    if (!warned) {
+      std::cerr << "[cudaq-qec] numa_node_id ignored: NUMA binding is only "
+                   "supported on Linux.\n";
+      warned = true;
+    }
+  }
+};
+#endif
+
+} // anonymous namespace
 
 INSTANTIATE_REGISTRY(cudaq::qec::decoder, const cudaq::qec::decoder_init &,
                      const cudaqx::heterogeneous_map &)
@@ -87,6 +217,11 @@ decoder::decoder(cudaq::qec::sparse_binary_matrix H)
     pimpl->should_log = ch[0] == '1' || ch[0] == 'y' || ch[0] == 'Y';
 }
 
+void decoder::set_hardware_params(const cudaqx::heterogeneous_map &params) {
+  cuda_device_id_ = cudaq::qec::read_cuda_device_id(params);
+  numa_node_id_ = cudaq::qec::read_numa_node_id(params);
+}
+
 // Provide a trivial implementation of for tensor<uint8_t> decode call. Child
 // classes should override this if they never want to pass through floats.
 decoder_result decoder::decode(const cudaqx::tensor<uint8_t> &syndrome) {
@@ -107,6 +242,10 @@ decoder_result decoder::decode(const cudaqx::tensor<uint8_t> &syndrome) {
 // should override this if they can do it more efficiently than this.
 std::vector<decoder_result>
 decoder::decode_batch(const std::vector<std::vector<float_t>> &syndrome) {
+  // Apply affinity once for the whole batch (not per-syndrome) to avoid
+  // repeated syscall overhead on the hot path.
+  CudaDeviceGuard dev(cuda_device_id_);
+  NumaGuard numa(numa_node_id_);
   std::vector<decoder_result> result;
   result.reserve(syndrome.size());
   for (auto &s : syndrome)
@@ -123,8 +262,15 @@ 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); });
+  // Capture by value: avoids a data race if the decoder is destroyed before
+  // the future resolves.
+  const int cuda_id = cuda_device_id_;
+  const int numa_id = numa_node_id_;
+  return std::async(std::launch::async, [this, syndrome, cuda_id, numa_id] {
+    CudaDeviceGuard dev(cuda_id);
+    NumaGuard numa(numa_id);
+    return this->decode(syndrome);
+  });
 }
 
 std::unique_ptr<decoder>
@@ -138,7 +284,14 @@ 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);
+  // Guards during construction so allocations land on the right hardware.
+  // Restored before this function returns; decode-time affinity is re-applied
+  // per call in decode_batch() and decode_async().
+  CudaDeviceGuard ctor_dev(cudaq::qec::read_cuda_device_id(param_map));
+  NumaGuard ctor_numa(cudaq::qec::read_numa_node_id(param_map));
+  auto d = iter->second(init, param_map);
+  d->set_hardware_params(param_map);
+  return d;
 }
 
 namespace details {

libs/qec/lib/decoders/plugins/trt_decoder/trt_decoder.cpp

@@ -879,6 +879,26 @@ decoder_result trt_decoder::decode(const std::vector<float_t> &syndrome) {
 
 std::vector<decoder_result>
 trt_decoder::decode_batch(const std::vector<std::vector<float_t>> &syndromes) {
+  // This override bypasses decoder::decode_batch()'s CudaDeviceGuard; apply it
+  // here so TRT inference lands on the right device regardless of the caller.
+  int _prev_dev = -1;
+  bool _dev_switched = false;
+  if (cuda_device_id_ >= 0) {
+    cudaGetDevice(&_prev_dev);
+    if (_prev_dev != cuda_device_id_) {
+      cudaSetDevice(cuda_device_id_);
+      _dev_switched = true;
+    }
+  }
+  struct _RestoreDevice {
+    int prev;
+    bool active;
+    ~_RestoreDevice() {
+      if (active)
+        cudaSetDevice(prev);
+    }
+  } _dev_guard{_prev_dev, _dev_switched};
+
   // Validate that we have syndromes to decode
   if (syndromes.empty()) {
     return {};

libs/qec/lib/realtime/config.cpp

@@ -534,6 +534,8 @@ struct MappingTraits<cudaq::qec::decoding::config::decoder_config> {
     io.mapRequired("H_sparse", config.H_sparse);
     io.mapRequired("O_sparse", config.O_sparse);
     io.mapRequired("D_sparse", config.D_sparse);
+    io.mapOptional("cuda_device_id", config.cuda_device_id);
+    io.mapOptional("numa_node_id", config.numa_node_id);
 
     // Validate that the number of rows in the H_sparse vector is equal to
     // syndrome_size.

libs/qec/lib/realtime/realtime_decoding.cpp

@@ -183,6 +183,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();
+  if (decoder_config.cuda_device_id.has_value())
+    params.insert("cuda_device_id", decoder_config.cuda_device_id.value());
+  if (decoder_config.numa_node_id.has_value())
+    params.insert("numa_node_id", decoder_config.numa_node_id.value());
   if (decoder_config.type != "trt_decoder")
     return params;
 

libs/qec/python/bindings/py_decoding_config.cpp

@@ -241,6 +241,8 @@ void bindDecodingConfig(nb::module_ &mod) {
       .def_rw("H_sparse", &decoder_config::H_sparse)
       .def_rw("O_sparse", &decoder_config::O_sparse)
       .def_rw("D_sparse", &decoder_config::D_sparse)
+      .def_rw("cuda_device_id", &decoder_config::cuda_device_id)
+      .def_rw("numa_node_id", &decoder_config::numa_node_id)
       .def_rw("decoder_custom_args", &decoder_config::decoder_custom_args)
       .def(
           "set_decoder_custom_args",

libs/qec/unittests/CMakeLists.txt

@@ -35,7 +35,8 @@ 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 cudaq-qec-realtime-decoding cudaq::cudaq libstim)
+target_link_libraries(test_decoders PRIVATE GTest::gtest_main cudaq-qec cudaq-qec-realtime-decoding cudaq::cudaq libstim CUDA::cudart)
+target_include_directories(test_decoders PRIVATE ${CUDAToolkit_INCLUDE_DIRS})
 add_dependencies(CUDAQXQECUnitTests test_decoders)
 gtest_discover_tests(test_decoders)
 
@@ -49,6 +50,11 @@ target_link_libraries(test_decoders_yaml PRIVATE
 add_dependencies(CUDAQXQECUnitTests test_decoders_yaml)
 gtest_discover_tests(test_decoders_yaml)
 
+add_executable(test_device_affinity test_device_affinity.cpp)
+target_link_libraries(test_device_affinity PRIVATE GTest::gtest_main cudaq-qec cudaq::cudaq)
+add_dependencies(CUDAQXQECUnitTests test_device_affinity)
+gtest_discover_tests(test_device_affinity)
+
 add_executable(test_qec test_qec.cpp)
 target_link_libraries(test_qec PRIVATE GTest::gtest_main cudaq-qec cudaq::cudaq-stim-target)
 add_dependencies(CUDAQXQECUnitTests test_qec)