From 48f6f82c7a93ec5b4afc35fd74b0695f3cdeb43d Mon Sep 17 00:00:00 2001 From: Sai Kiran Polisetty Date: Thu, 23 Jul 2026 20:46:49 +0530 Subject: [PATCH] Update --- src/instance_state.cc | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/instance_state.cc b/src/instance_state.cc index 415bbe1..2ba3206 100644 --- a/src/instance_state.cc +++ b/src/instance_state.cc @@ -4036,7 +4036,11 @@ TRTv3Interface::SetCudaGraphShape( TensorRTContext::CudaGraph* cuda_graph) { int batch_size = graph_spec.batch_size_; - *cuda_graph_key = std::vector{batch_size}; + // Non-batching models (max_batch_size == 0) use a batch size of 1 in the + // lookup key at inference time. Use the same value here so the captured + // graph can be found. + int64_t cuda_graph_key_batch_size = (batch_size == 0) ? 1 : batch_size; + *cuda_graph_key = std::vector{cuda_graph_key_batch_size}; auto& lower_bound_key = cuda_graph->lower_bound_key_; lower_bound_key.push_back( (graph_spec.lower_bound_batch_size_ == 0) @@ -4078,10 +4082,12 @@ TRTv3Interface::SetCudaGraphShape( cuda_graph->input_dims_.emplace_back(); } else { cuda_graph->input_dims_.emplace_back(); + // Prepend the batch dimension only for batching models, so the keys + // stay aligned with non-batching models. if (batch_size != 0) { cuda_graph->input_dims_.back().push_back(batch_size); + lower_bound_key.push_back(lower_bound_key[0]); } - lower_bound_key.push_back(lower_bound_key[0]); } auto& shape = cuda_graph->input_dims_.back(); shape.insert(shape.end(), it->second.begin(), it->second.end());