diff --git a/tests/integration/defs/perf/model_yaml_config.py b/tests/integration/defs/perf/pytorch_model_config.py
similarity index 95%
rename from tests/integration/defs/perf/model_yaml_config.py
rename to tests/integration/defs/perf/pytorch_model_config.py
index c9ada66638e1..7792377649af 100644
--- a/tests/integration/defs/perf/model_yaml_config.py
+++ b/tests/integration/defs/perf/pytorch_model_config.py
@@ -14,11 +14,11 @@
 # limitations under the License.
 # -*- coding: utf-8 -*-
 """
-Model yaml config for trtllm-bench perf tests
+Model pytorch yaml config for trtllm-bench perf tests
 """
 
 
-def get_model_yaml_config(model_label: str) -> dict:
+def get_model_yaml_config(model_label: str, input_lens: list[int]) -> dict:
     """
         Return the yaml config corresponding to the model label.
         Args:
@@ -32,7 +32,6 @@ def get_model_yaml_config(model_label: str) -> dict:
             'print_iter_log': True,
             'use_cuda_graph': True,
             'cuda_graph_padding_enabled': True,
-            'cuda_graph_max_batch_size': 4096,
         }
     }
     model_configs = {
diff --git a/tests/integration/defs/perf/test_perf.py b/tests/integration/defs/perf/test_perf.py
index 5882fedb4807..5258f923f50e 100644
--- a/tests/integration/defs/perf/test_perf.py
+++ b/tests/integration/defs/perf/test_perf.py
@@ -27,7 +27,7 @@
                                        print_warning)
 
 from ..conftest import get_llm_root, llm_models_root, trt_environment
-from .model_yaml_config import get_model_yaml_config
+from .pytorch_model_config import get_model_yaml_config
 from .utils import (AbstractPerfScriptTestClass, PerfBenchScriptTestCmds,
                     PerfMetricType, PerfScriptTestCmds, generate_test_nodes)