fix: To validate AI decoder RPC slot capacity to avoid gateway output overwrite

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

@@ -93,13 +93,15 @@ class ai_decoder_service {
   ai_decoder_service(const std::string &model_path, void **device_mailbox_slot,
                      const std::string &engine_save_path = "",
                      network_typing_override typing_override =
-                         network_typing_override::automatic);
+                         network_typing_override::automatic,
+                     size_t rpc_slot_size_bytes = 0);
 
   /// Create a passthrough (identity copy) instance for testing without TRT.
   static std::unique_ptr<ai_decoder_service>
   create_passthrough(void **device_mailbox_slot,
                      size_t input_bytes = 1600 * sizeof(float),
-                     size_t output_bytes = 1600 * sizeof(float));
+                     size_t output_bytes = 1600 * sizeof(float),
+                     size_t rpc_slot_size_bytes = 0);
 
   virtual ~ai_decoder_service();
 
@@ -113,6 +115,16 @@ class ai_decoder_service {
   /// @brief Size of the primary output tensor in bytes (forwarded to CPU)
   size_t get_output_size() const { return output_size_; }
 
+  /// @brief Set the total RPC slot size, including the RPC header/response.
+  ///
+  /// If left as zero, the service assumes slots are sized for the input
+  /// request frame (RPCHeader + input bytes). Graph capture rejects models
+  /// whose response frame would not fit.
+  void set_rpc_slot_size_bytes(size_t rpc_slot_size_bytes);
+
+  /// @brief Total RPC slot size in bytes, including the RPC header/response.
+  size_t get_rpc_slot_size_bytes() const { return rpc_slot_size_bytes_; }
+
   /// @brief Logical element count (detector count) of the primary input
   /// tensor.  Computed from the engine's tensor volume; independent of
   /// the IO dtype so FP16 / INT32 / FP8 inputs all report the same
@@ -132,14 +144,16 @@ class ai_decoder_service {
 protected:
   /// Passthrough constructor (no TRT, identity copy kernel only).
   ai_decoder_service(void **device_mailbox_slot, size_t input_bytes,
-                     size_t output_bytes);
+                     size_t output_bytes, size_t rpc_slot_size_bytes = 0);
 
   void load_engine(const std::string &path);
   void build_engine_from_onnx(const std::string &onnx_path,
                               const std::string &engine_save_path,
                               network_typing_override typing_override);
   void setup_bindings();
   void allocate_resources();
+  void set_default_rpc_slot_size();
+  void validate_rpc_slot_size() const;
 
   std::unique_ptr<nvinfer1::IRuntime> runtime_;
   std::unique_ptr<nvinfer1::ICudaEngine> engine_;
@@ -164,6 +178,7 @@ class ai_decoder_service {
   size_t output_size_ = 0;
   size_t input_num_elements_ = 0;
   size_t output_num_elements_ = 0;
+  size_t rpc_slot_size_bytes_ = 0;
 
   onnx_quant_info quant_info_;
 };

libs/qec/lib/realtime/ai_decoder_service.cu

@@ -13,6 +13,7 @@
 #include <cstdlib>
 #include <fstream>
 #include <iostream>
+#include <limits>
 #include <stdexcept>
 #include <string>
 
@@ -34,32 +35,40 @@ namespace cudaq::qec::realtime::experimental {
 
 __global__ void gateway_input_kernel(void **mailbox_slot_ptr,
                                      void *trt_fixed_input,
-                                     size_t copy_size_bytes) {
+                                     size_t copy_size_bytes,
+                                     size_t payload_capacity_bytes) {
   void *ring_buffer_data = *mailbox_slot_ptr;
   if (ring_buffer_data == nullptr)
     return;
 
   const char *src =
       (const char *)ring_buffer_data + sizeof(cudaq::realtime::RPCHeader);
   char *dst = (char *)trt_fixed_input;
+  size_t safe_copy_size = copy_size_bytes < payload_capacity_bytes
+                              ? copy_size_bytes
+                              : payload_capacity_bytes;
 
-  for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < copy_size_bytes;
+  for (size_t i = threadIdx.x + blockIdx.x * blockDim.x; i < safe_copy_size;
        i += blockDim.x * gridDim.x) {
     dst[i] = src[i];
   }
 }
 
 __global__ void gateway_output_kernel(void **mailbox_slot_ptr,
                                       const void *trt_fixed_output,
-                                      size_t result_size_bytes) {
+                                      size_t result_size_bytes,
+                                      size_t payload_capacity_bytes) {
   void *ring_buffer_data = *mailbox_slot_ptr;
   if (ring_buffer_data == nullptr)
     return;
 
   char *dst = (char *)ring_buffer_data + sizeof(cudaq::realtime::RPCHeader);
   const char *src = (const char *)trt_fixed_output;
+  size_t safe_result_size = result_size_bytes < payload_capacity_bytes
+                                ? result_size_bytes
+                                : payload_capacity_bytes;
 
-  for (int i = threadIdx.x + blockIdx.x * blockDim.x; i < result_size_bytes;
+  for (size_t i = threadIdx.x + blockIdx.x * blockDim.x; i < safe_result_size;
        i += blockDim.x * gridDim.x) {
     dst[i] = src[i];
   }
@@ -73,8 +82,8 @@ __global__ void gateway_output_kernel(void **mailbox_slot_ptr,
 
     auto *response = (cudaq::realtime::RPCResponse *)ring_buffer_data;
     response->magic = cudaq::realtime::RPC_MAGIC_RESPONSE;
-    response->status = 0;
-    response->result_len = static_cast<uint32_t>(result_size_bytes);
+    response->status = (safe_result_size == result_size_bytes) ? 0 : 1;
+    response->result_len = static_cast<uint32_t>(safe_result_size);
     response->request_id = rid;
     response->ptp_timestamp = pts;
     __threadfence_system();
@@ -147,29 +156,36 @@ void ai_decoder_service::Logger::log(Severity severity,
 ai_decoder_service::ai_decoder_service(const std::string &model_path,
                                        void **device_mailbox_slot,
                                        const std::string &engine_save_path,
-                                       network_typing_override typing_override)
-    : device_mailbox_slot_(device_mailbox_slot) {
+                                       network_typing_override typing_override,
+                                       size_t rpc_slot_size_bytes)
+    : device_mailbox_slot_(device_mailbox_slot),
+      rpc_slot_size_bytes_(rpc_slot_size_bytes) {
   std::string ext = model_path.substr(model_path.find_last_of('.'));
   if (ext == ".onnx") {
     build_engine_from_onnx(model_path, engine_save_path, typing_override);
   } else {
     load_engine(model_path);
   }
   setup_bindings();
+  set_default_rpc_slot_size();
   allocate_resources();
 }
 
 ai_decoder_service::ai_decoder_service(void **device_mailbox_slot,
-                                       size_t input_bytes, size_t output_bytes)
+                                       size_t input_bytes, size_t output_bytes,
+                                       size_t rpc_slot_size_bytes)
     : device_mailbox_slot_(device_mailbox_slot), input_size_(input_bytes),
-      output_size_(output_bytes) {
+      output_size_(output_bytes), rpc_slot_size_bytes_(rpc_slot_size_bytes) {
+  set_default_rpc_slot_size();
   allocate_resources();
 }
 
-std::unique_ptr<ai_decoder_service> ai_decoder_service::create_passthrough(
-    void **device_mailbox_slot, size_t input_bytes, size_t output_bytes) {
-  return std::unique_ptr<ai_decoder_service>(
-      new ai_decoder_service(device_mailbox_slot, input_bytes, output_bytes));
+std::unique_ptr<ai_decoder_service>
+ai_decoder_service::create_passthrough(void **device_mailbox_slot,
+                                       size_t input_bytes, size_t output_bytes,
+                                       size_t rpc_slot_size_bytes) {
+  return std::unique_ptr<ai_decoder_service>(new ai_decoder_service(
+      device_mailbox_slot, input_bytes, output_bytes, rpc_slot_size_bytes));
 }
 
 ai_decoder_service::~ai_decoder_service() {
@@ -485,7 +501,39 @@ void ai_decoder_service::allocate_resources() {
   }
 }
 
+void ai_decoder_service::set_rpc_slot_size_bytes(size_t rpc_slot_size_bytes) {
+  rpc_slot_size_bytes_ = rpc_slot_size_bytes;
+  set_default_rpc_slot_size();
+}
+
+void ai_decoder_service::set_default_rpc_slot_size() {
+  if (rpc_slot_size_bytes_ == 0)
+    rpc_slot_size_bytes_ = sizeof(cudaq::realtime::RPCHeader) + input_size_;
+}
+
+void ai_decoder_service::validate_rpc_slot_size() const {
+  const size_t request_frame_size =
+      sizeof(cudaq::realtime::RPCHeader) + input_size_;
+  const size_t response_frame_size =
+      sizeof(cudaq::realtime::RPCResponse) + output_size_;
+
+  if (output_size_ > std::numeric_limits<uint32_t>::max())
+    throw std::length_error(
+        "ai_decoder_service output exceeds RPCResponse::result_len capacity");
+
+  if (rpc_slot_size_bytes_ < request_frame_size) {
+    throw std::length_error(
+        "ai_decoder_service RPC slot is too small for the input payload");
+  }
+  if (rpc_slot_size_bytes_ < response_frame_size) {
+    throw std::length_error(
+        "ai_decoder_service RPC slot is too small for the output payload");
+  }
+}
+
 void ai_decoder_service::capture_graph(cudaStream_t stream) {
+  validate_rpc_slot_size();
+
   for (auto &b : all_bindings_) {
     context_->setTensorAddress(b.name.c_str(), b.d_buffer);
   }
@@ -499,13 +547,19 @@ void ai_decoder_service::capture_graph(cudaStream_t stream) {
   DECODER_CUDA_CHECK(
       cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
 
-  gateway_input_kernel<<<1, 128, 0, stream>>>(device_mailbox_slot_,
-                                              d_trt_input_, input_size_);
+  const size_t input_payload_capacity =
+      rpc_slot_size_bytes_ - sizeof(cudaq::realtime::RPCHeader);
+  const size_t output_payload_capacity =
+      rpc_slot_size_bytes_ - sizeof(cudaq::realtime::RPCResponse);
+
+  gateway_input_kernel<<<1, 128, 0, stream>>>(
+      device_mailbox_slot_, d_trt_input_, input_size_, input_payload_capacity);
   if (!context_->enqueueV3(stream))
     throw std::runtime_error(
         "TRT enqueueV3 failed during graph capture in ai_decoder_service");
   gateway_output_kernel<<<1, 128, 0, stream>>>(device_mailbox_slot_,
-                                               d_trt_output_, output_size_);
+                                               d_trt_output_, output_size_,
+                                               output_payload_capacity);
 
   DECODER_CUDA_CHECK(cudaStreamEndCapture(stream, &graph));
 

libs/qec/unittests/test_ai_decoder_quantized_onnx.cpp

@@ -14,7 +14,11 @@
 #include <cstdint>
 #include <cstring>
 #include <cuda_runtime.h>
+#include <filesystem>
+#include <fstream>
 #include <gtest/gtest.h>
+#include <memory>
+#include <stdexcept>
 #include <string>
 #include <vector>
 
@@ -55,6 +59,61 @@ bool isFp8HardwareAvailable() {
   return prop.major >= 9;
 }
 
+std::string build_output_larger_engine() {
+  auto engine_path = std::filesystem::path(::testing::TempDir()) /
+                     "ai_decoder_output_larger_than_input.engine";
+
+  auto builder = std::unique_ptr<nvinfer1::IBuilder>(
+      nvinfer1::createInferBuilder(ai_decoder_service::gLogger));
+  if (!builder)
+    throw std::runtime_error("Failed to create TensorRT builder");
+
+  auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(
+      builder->createNetworkV2(0));
+  auto config =
+      std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
+  if (!network || !config)
+    throw std::runtime_error("Failed to create TensorRT network/config");
+
+  nvinfer1::Dims input_dims{};
+  input_dims.nbDims = 1;
+  input_dims.d[0] = static_cast<int64_t>(kNumElements);
+  auto *input =
+      network->addInput("input", nvinfer1::DataType::kFLOAT, input_dims);
+  if (!input)
+    throw std::runtime_error("Failed to add TensorRT input");
+
+  // Build a valid asymmetric model: output is [0.0, input...], so the
+  // output payload is exactly one float larger than the input payload.
+  float zero = 0.0f;
+  nvinfer1::Weights zero_weights{nvinfer1::DataType::kFLOAT, &zero, 1};
+  nvinfer1::Dims zero_dims{};
+  zero_dims.nbDims = 1;
+  zero_dims.d[0] = 1;
+  auto *constant = network->addConstant(zero_dims, zero_weights);
+  if (!constant)
+    throw std::runtime_error("Failed to add TensorRT constant");
+
+  nvinfer1::ITensor *concat_inputs[] = {constant->getOutput(0), input};
+  auto *concat = network->addConcatenation(concat_inputs, 2);
+  if (!concat)
+    throw std::runtime_error("Failed to add TensorRT concatenation");
+  concat->setAxis(0);
+  concat->getOutput(0)->setName("output");
+  network->markOutput(*concat->getOutput(0));
+
+  auto plan = std::unique_ptr<nvinfer1::IHostMemory>(
+      builder->buildSerializedNetwork(*network, *config));
+  if (!plan)
+    throw std::runtime_error("Failed to build asymmetric TensorRT engine");
+
+  std::ofstream out(engine_path, std::ios::binary);
+  if (!out.good())
+    throw std::runtime_error("Failed to open temporary engine path");
+  out.write(static_cast<const char *>(plan->data()), plan->size());
+  return engine_path.string();
+}
+
 void write_rpc_slot(uint8_t *slot_host, const std::vector<float> &input) {
   std::memset(slot_host, 0, kSlotSize);
   rt_sdk::RPCHeader hdr{};
@@ -172,4 +231,67 @@ TEST_F(AiDecoderQuantizedOnnxSmokeTest, Fp8QdqRunsWithExpectedNumerics) {
   expect_identity_qdq(output, 1.0e-3f);
 }
 
+TEST_F(AiDecoderQuantizedOnnxSmokeTest, RejectsOversizedGatewayOutput) {
+  constexpr size_t kOutputBytes = (kNumElements + 1) * sizeof(float);
+
+  auto engine_path = build_output_larger_engine();
+  ai_decoder_service service(engine_path,
+                             reinterpret_cast<void **>(mailbox_dev_), "",
+                             network_typing_override::automatic);
+  ASSERT_EQ(service.get_input_size(), kPayloadBytes);
+  ASSERT_EQ(service.get_output_size(), kOutputBytes);
+
+  // By default the service assumes an RPC slot sized from the request payload.
+  // A larger output must be rejected before graph capture can write past it.
+  EXPECT_THROW(service.capture_graph(stream_), std::length_error);
+}
+
+TEST_F(AiDecoderQuantizedOnnxSmokeTest, WritesGatewayOutputWithinSlot) {
+  constexpr size_t kOutputBytes = (kNumElements + 1) * sizeof(float);
+  constexpr size_t kLargeSlotSize = CUDAQ_RPC_HEADER_SIZE + kOutputBytes;
+
+  uint8_t *slots_host = nullptr;
+  uint8_t *slots_dev = nullptr;
+  CUDA_ASSERT_OK(cudaHostAlloc(reinterpret_cast<void **>(&slots_host),
+                               2 * kLargeSlotSize, cudaHostAllocMapped));
+  auto slots_cleanup = std::unique_ptr<uint8_t, decltype(&cudaFreeHost)>(
+      slots_host, cudaFreeHost);
+  CUDA_ASSERT_OK(cudaHostGetDevicePointer(reinterpret_cast<void **>(&slots_dev),
+                                          slots_host, 0));
+
+  write_rpc_slot(slots_host, kInputs);
+  uint8_t *adjacent_slot = slots_host + kLargeSlotSize;
+  std::memset(adjacent_slot, 0xA5, kLargeSlotSize);
+  mailbox_host_[0] = slots_dev;
+
+  auto engine_path = build_output_larger_engine();
+  ai_decoder_service service(
+      engine_path, reinterpret_cast<void **>(mailbox_dev_), "",
+      network_typing_override::automatic, kLargeSlotSize);
+  ASSERT_EQ(service.get_input_size(), kPayloadBytes);
+  ASSERT_EQ(service.get_output_size(), kOutputBytes);
+
+  service.capture_graph(stream_);
+  ASSERT_NE(service.get_executable_graph(), nullptr);
+  CUDA_ASSERT_OK(cudaGraphLaunch(service.get_executable_graph(), stream_));
+  CUDA_ASSERT_OK(cudaStreamSynchronize(stream_));
+
+  auto *response = reinterpret_cast<rt_sdk::RPCResponse *>(slots_host);
+  EXPECT_EQ(response->magic, rt_sdk::RPC_MAGIC_RESPONSE);
+  EXPECT_EQ(response->status, 0u);
+  EXPECT_EQ(response->result_len, kOutputBytes);
+
+  const auto *output =
+      reinterpret_cast<const float *>(slots_host + CUDAQ_RPC_HEADER_SIZE);
+  EXPECT_FLOAT_EQ(output[0], 0.0f);
+  for (size_t i = 0; i < kInputs.size(); ++i)
+    EXPECT_FLOAT_EQ(output[i + 1], kInputs[i]);
+
+  // ASSERT that the larger, explicitly declared slot contains the whole
+  // response and that the next slot remains untouched.
+  for (size_t i = 0; i < kLargeSlotSize; ++i)
+    EXPECT_EQ(adjacent_slot[i], 0xA5) << "gateway_output_kernel overwrote byte "
+                                      << i << " of the adjacent RPC slot";
+}
+
 } // namespace