From b246d260ed25d1ce46f589610a49e756f01a636c Mon Sep 17 00:00:00 2001 From: geobeau Date: Wed, 3 Jun 2026 15:47:46 +0000 Subject: [PATCH] Add completion queue sharding for gRPC inference Add --grpc-infer-cq-count to control the number of inference completion queues. Default is 1 (single shared CQ, same behavior as before). Set to 0 for one CQ per handler thread, or N>1 for N sharded CQs, to reduce contention at high throughput. --- src/command_line_parser.cc | 14 ++++++++++++++ src/grpc/grpc_server.cc | 17 ++++++++++++++--- src/grpc/grpc_server.h | 6 +++++- 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/src/command_line_parser.cc b/src/command_line_parser.cc index f228bcaf77..61f48de747 100644 --- a/src/command_line_parser.cc +++ b/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(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(optarg); break; diff --git a/src/grpc/grpc_server.cc b/src/grpc/grpc_server.cc index 10da64776e..4b70ad821f 100644 --- a/src/grpc/grpc_server.cc +++ b/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 diff --git a/src/grpc/grpc_server.h b/src/grpc/grpc_server.h index 6020bff7ad..e51b41631e 100644 --- a/src/grpc/grpc_server.h +++ b/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> model_infer_cqs_; std::unique_ptr<::grpc::ServerCompletionQueue> model_stream_infer_cq_; std::unique_ptr common_handler_;