Update interface

Signed-off-by: Tailing Yuan <yuantailing@gmail.com>

Author: yuantailing

examples/layer_wise_benchmarks/run_single.py

@@ -8,8 +8,7 @@
 from tensorrt_llm._torch.autotuner import AutoTuner, autotune
 from tensorrt_llm._torch.modules.multi_stream_utils import with_multi_stream
 from tensorrt_llm._utils import local_mpi_rank, mpi_rank, mpi_world_size
-from tensorrt_llm.tools.layer_wise_benchmarks.runner_base import BalanceMethod
-from tensorrt_llm.tools.layer_wise_benchmarks.runner_factory import get_runner_cls
+from tensorrt_llm.tools.layer_wise_benchmarks import BalanceMethod, get_runner_cls
 
 
 def comma_separated_ints(s):

tensorrt_llm/tools/layer_wise_benchmarks/__init__.py

@@ -0,0 +1,7 @@
+from .runner_factory import get_runner_cls
+from .runner_interface import BalanceMethod
+
+__all__ = [
+    "BalanceMethod",
+    "get_runner_cls",
+]

tensorrt_llm/tools/layer_wise_benchmarks/deepseekv3_runner.py

@@ -12,7 +12,7 @@
 from tensorrt_llm.functional import AllReduceStrategy
 from tensorrt_llm.mapping import Mapping
 
-from .runner_base import BalanceMethod, RunnerBase
+from .runner_interface import BalanceMethod, RunnerBase
 from .runner_utils import RunnerMixin, ceil_div
 
 

tensorrt_llm/tools/layer_wise_benchmarks/qwen3_next_runner.py

@@ -8,7 +8,7 @@
 from tensorrt_llm.functional import AllReduceStrategy
 from tensorrt_llm.mapping import Mapping
 
-from .runner_base import RunnerBase
+from .runner_interface import RunnerBase
 from .runner_utils import RunnerMixin
 
 

tensorrt_llm/tools/layer_wise_benchmarks/runner_factory.py

@@ -4,7 +4,7 @@
 from .qwen3_next_runner import Qwen3NextRunner
 
 
-def get_runner_cls(pretrained_model_name_or_path: str):
+def get_runner_cls(pretrained_model_name_or_path: str) -> type:
     pretrained_config = load_pretrained_config(pretrained_model_name_or_path)
     return {
         "deepseek_v3": DeepSeekV3Runner,

tensorrt_llm/tools/layer_wise_benchmarks/runner_base.py renamed to tensorrt_llm/tools/layer_wise_benchmarks/runner_interface.py

@@ -22,7 +22,7 @@
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.models.modeling_utils import QuantConfig
 
-from .runner_base import BalanceMethod
+from .runner_interface import BalanceMethod
 
 
 def ceil_div(a, b):