diff --git a/tensorrt_llm/_torch/__init__.py b/tensorrt_llm/_torch/__init__.py
index 23257d91504..7d2de6d643c 100644
--- a/tensorrt_llm/_torch/__init__.py
+++ b/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"]
diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py
index a563bc98f28..6b4014dc333 100644
--- a/tensorrt_llm/llmapi/llm_args.py
+++ b/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]"
         },
     )
 
diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt
index f6a876ad01f..9ac39d159b4 100644
--- a/tests/integration/test_lists/waives.txt
+++ b/tests/integration/test_lists/waives.txt
@@ -424,7 +424,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)
diff --git a/tests/unittest/api_stability/api_stability_core.py b/tests/unittest/api_stability/api_stability_core.py
index 1014f1a22fa..3dc33819e0b 100644
--- a/tests/unittest/api_stability/api_stability_core.py
+++ b/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,
diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml
index 8284cc0a0db..cf39b325458 100644
--- a/tests/unittest/api_stability/references/llm.yaml
+++ b/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]