perf: Add completion queue sharding

src/command_line_parser.cc

@@ -307,6 +307,7 @@ enum TritonOptionId {
   OPTION_GRPC_ADDRESS,
   OPTION_GRPC_HEADER_FORWARD_PATTERN,
   OPTION_GRPC_INFER_THREAD_COUNT,
+  OPTION_GRPC_INFER_CQ_COUNT,
   OPTION_GRPC_INFER_ALLOCATION_POOL_SIZE,
   OPTION_GRPC_MAX_RESPONSE_POOL_SIZE,
   OPTION_GRPC_USE_SSL,
@@ -542,6 +543,10 @@ TritonParser::SetupOptions()
       {OPTION_GRPC_INFER_THREAD_COUNT, "grpc-infer-thread-count",
        Option::ArgInt,
        "The number of gRPC inference handler threads. Default is 2."});
+  grpc_options_.push_back(
+      {OPTION_GRPC_INFER_CQ_COUNT, "grpc-infer-cq-count", Option::ArgInt,
+       "The number of gRPC inference completion queues. Default is 0 "
+       "(one CQ per handler thread). Use 1 for legacy single-CQ behavior."});
   grpc_options_.push_back(
       {OPTION_GRPC_INFER_ALLOCATION_POOL_SIZE,
        "grpc-infer-allocation-pool-size", Option::ArgInt,
@@ -1479,6 +1484,15 @@ TritonParser::Parse(int argc, char** argv)
                 "the range 2 to 128.");
           }
           break;
+        case OPTION_GRPC_INFER_CQ_COUNT:
+          lgrpc_options.infer_cq_count_ = ParseOption<int>(optarg);
+          if (lgrpc_options.infer_cq_count_ < 0 ||
+              lgrpc_options.infer_cq_count_ > 128) {
+            throw ParseException(
+                "invalid argument for --grpc_infer_cq_count. Must be in "
+                "the range 0 to 128.");
+          }
+          break;
         case OPTION_GRPC_INFER_ALLOCATION_POOL_SIZE:
           lgrpc_options.infer_allocation_pool_size_ = ParseOption<int>(optarg);
           break;

src/grpc/grpc_server.cc

@@ -2386,7 +2386,14 @@ Server::Server(
   }
 
   common_cq_ = builder_.AddCompletionQueue();
-  model_infer_cq_ = builder_.AddCompletionQueue();
+
+  int cq_count = (options.infer_cq_count_ > 0)
+      ? std::min(options.infer_cq_count_, options.infer_thread_count_)
+      : options.infer_thread_count_;
+  for (int i = 0; i < cq_count; ++i) {
+    model_infer_cqs_.emplace_back(builder_.AddCompletionQueue());
+  }
+
   model_stream_infer_cq_ = builder_.AddCompletionQueue();
 
   // For testing purposes only, add artificial delay in grpc responses.
@@ -2408,7 +2415,7 @@ Server::Server(
   for (int i = 0; i < options.infer_thread_count_; ++i) {
     model_infer_handlers_.emplace_back(new ModelInferHandler(
         "ModelInferHandler", tritonserver_, trace_manager_, shm_manager_,
-        &service_, model_infer_cq_.get(),
+        &service_, model_infer_cqs_[i % cq_count].get(),
         options.infer_allocation_pool_size_ /* max_state_bucket_count */,
         options.max_response_pool_size_, options.infer_compression_level_,
         restricted_kv, options.forward_header_pattern_, &conn_mtx_, &conn_cnt_,
@@ -2488,6 +2495,8 @@ Server::GetOptions(Options& options, UnorderedMapType& options_map)
 
   RETURN_IF_ERR(GetValue(
       options_map, "infer_thread_count", &options.infer_thread_count_));
+  RETURN_IF_ERR(GetValue(
+      options_map, "infer_cq_count", &options.infer_cq_count_));
   RETURN_IF_ERR(GetValue(
       options_map, "infer_allocation_pool_size",
       &options.infer_allocation_pool_size_));
@@ -2622,7 +2631,9 @@ Server::Stop()
 
   // Shutdown completion queues
   common_cq_->Shutdown();
-  model_infer_cq_->Shutdown();
+  for (auto& cq : model_infer_cqs_) {
+    cq->Shutdown();
+  }
   model_stream_infer_cq_->Shutdown();
 
   // Must stop all handlers explicitly to wait for all the handler

src/grpc/grpc_server.h

@@ -88,6 +88,10 @@ struct Options {
   // The number of gRPC inference handler threads. Useful for
   // throughput tuning of models that are request handling bounded.
   int infer_thread_count_{2};
+  // The number of gRPC inference completion queues. Default is 1
+  // (single shared CQ, legacy behavior). Set to 0 for one CQ per
+  // handler thread, or N>1 for N sharded CQs.
+  int infer_cq_count_{1};
   // The maximum number of inference request/response objects that
   // remain allocated for reuse. As long as the number of in-flight
   // requests doesn't exceed this value there will be no
@@ -154,7 +158,7 @@ class Server {
   std::unique_ptr<::grpc::Server> server_;
 
   std::unique_ptr<::grpc::ServerCompletionQueue> common_cq_;
-  std::unique_ptr<::grpc::ServerCompletionQueue> model_infer_cq_;
+  std::vector<std::unique_ptr<::grpc::ServerCompletionQueue>> model_infer_cqs_;
   std::unique_ptr<::grpc::ServerCompletionQueue> model_stream_infer_cq_;
 
   std::unique_ptr<HandlerBase> common_handler_;