[nvbugs/5401156][fix] Avoid import all models when import trtllm._common (#6266)

Author: chang-l

tensorrt_llm/_torch/__init__.py

@@ -1,5 +1,4 @@
 from .llm import LLM
 from .model_config import MoeLoadBalancerConfig
-from .models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader
 
-__all__ = ["LLM", "MoeLoadBalancerConfig", "BaseCheckpointLoader"]
+__all__ = ["LLM", "MoeLoadBalancerConfig"]

tensorrt_llm/llmapi/llm_args.py

@@ -1894,7 +1894,8 @@ class TorchLlmArgs(BaseLlmArgs):
         default=None,
         description="The checkpoint loader to use for this LLM instance.",
         json_schema_extra={
-            "type": "Optional[tensorrt_llm._torch.BaseCheckpointLoader]"
+            "type":
+            "Optional[tensorrt_llm._torch.models.checkpoints.BaseCheckpointLoader]"
         },
     )
 

tests/integration/test_lists/waives.txt

@@ -423,7 +423,6 @@ test_e2e.py::test_ptp_quickstart_multimodal[gemma-3-27b-it-gemma/gemma-3-27b-it-
 test_e2e.py::test_ptp_quickstart_multimodal[gemma-3-27b-it-gemma/gemma-3-27b-it-image-False] SKIP (https://nvbgus/5401114)
 examples/test_recurrentgemma.py::test_llm_recurrentgemma_1gpu[use_cpp_session-recurrentgemma-2b-use_paged_cache-int4_awq-float16-enable_attn_plugin-enable_gemm_plugin] SKIP (https://nvbugs/5401233)
 examples/test_recurrentgemma.py::test_llm_recurrentgemma_2gpu[recurrentgemma-2b] SKIP (https://nvbugs/5401233)
-examples/test_multimodal.py::test_llm_multimodal_general[VILA1.5-3b-pp:1-tp:1-float16-bs:1-cpp_e2e:False-nb:1] SKIP (https://nvbugs/5401156)
 test_e2e.py::test_ptp_quickstart_multimodal[mistral-small-3.1-24b-instruct-Mistral-Small-3.1-24B-Instruct-2503-image-True] SKIP (https://nvbugs/5404005)
 accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_eagle3 SKIP (https://nvbugs/5409414)
 accuracy/test_llm_api_pytorch.py::TestLlama3_1_8BInstruct::test_fp8_beam_search SKIP (https://nvbugs/5409415)

tests/unittest/api_stability/api_stability_core.py

@@ -18,6 +18,9 @@
 
 import tensorrt_llm
 from tensorrt_llm import LLM
+# Import BaseCheckpointLoader for YAML processing
+from tensorrt_llm._torch.models.checkpoints.base_checkpoint_loader import \
+    BaseCheckpointLoader
 from tensorrt_llm.executor import GenerationResult
 from tensorrt_llm.executor.result import TokenLogprobs
 from tensorrt_llm.llmapi import (CalibConfig, CompletionOutput,

tests/unittest/api_stability/references/llm.yaml

@@ -71,7 +71,7 @@ methods:
         annotation: Optional[tensorrt_llm.llmapi.llm_args.CudaGraphConfig]
         default: null
       checkpoint_loader:
-        annotation: Optional[tensorrt_llm._torch.BaseCheckpointLoader]
+        annotation: Optional[tensorrt_llm._torch.models.checkpoints.BaseCheckpointLoader]
         default: null
       checkpoint_format:
         annotation: Optional[str]