From c6d1fcf8632fd8ada96eec0a02057831907bd806 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Mon, 15 Sep 2025 18:00:19 +0800 Subject: [PATCH] Fix error when running trtllm-bench without cuda graph. Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> --- .../bench/dataclasses/configuration.py | 20 ++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/bench/dataclasses/configuration.py b/tensorrt_llm/bench/dataclasses/configuration.py index 6d8e703ee490..db6f2a9fafcf 100755 --- a/tensorrt_llm/bench/dataclasses/configuration.py +++ b/tensorrt_llm/bench/dataclasses/configuration.py @@ -90,15 +90,17 @@ def get_llm_args(self) -> Dict: if self.backend == "pytorch": cuda_graph_config = updated_llm_args.pop( "cuda_graph_config", llm_args["cuda_graph_config"]) - # Use runtime max_batch_size as cuda_graph_config.max_batch_size - # if both max_batch_size and batch_sizes are not set. - batch_sizes_set = cuda_graph_config.get("batch_sizes", - None) is not None - max_batch_size_set = cuda_graph_config.get("max_batch_size", - None) is not None - if not batch_sizes_set and not max_batch_size_set: - cuda_graph_config[ - "max_batch_size"] = self.settings_config.max_batch_size + if cuda_graph_config: + # Use runtime max_batch_size as cuda_graph_config.max_batch_size + # if both max_batch_size and batch_sizes are not set. + batch_sizes_set = cuda_graph_config.get("batch_sizes", + None) is not None + max_batch_size_set = cuda_graph_config.get( + "max_batch_size", None) is not None + if not batch_sizes_set and not max_batch_size_set: + cuda_graph_config[ + "max_batch_size"] = self.settings_config.max_batch_size + updated_llm_args["cuda_graph_config"] = cuda_graph_config return updated_llm_args