From b217d4fd040d8e3e63509f8cfe3ef5c3cf46d716 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:47:58 +0800 Subject: [PATCH 1/6] [None][chore] Move steady-clock helper out of serve into _utils get_steady_clock_now_in_seconds was defined in serve/responses_utils, so the executor imported the serve package (openai types, openai_harmony, transformers processors) at module level just for a one-line clock wrapper. Move it to tensorrt_llm._utils next to the other bindings helpers; responses_utils re-exports it for its serve-side importers. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- .../_torch/pyexecutor/perf_metrics_manager.py | 2 +- tensorrt_llm/_utils.py | 11 ++++++++++- tensorrt_llm/serve/responses_utils.py | 7 ++----- 3 files changed, 13 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py index cb13757956a0..ca0cb4a5c0c7 100644 --- a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py @@ -11,8 +11,8 @@ import torch +from tensorrt_llm._utils import get_steady_clock_now_in_seconds from tensorrt_llm.logger import logger -from tensorrt_llm.serve.responses_utils import get_steady_clock_now_in_seconds from .llm_request import PerfTimingInfo diff --git a/tensorrt_llm/_utils.py b/tensorrt_llm/_utils.py index 71574f3e1aa5..6f4e32812cfc 100644 --- a/tensorrt_llm/_utils.py +++ b/tensorrt_llm/_utils.py @@ -63,7 +63,7 @@ has_nvml = False # isort: on -from tensorrt_llm.bindings import DataType, LayerType +from tensorrt_llm.bindings import DataType, LayerType, steady_clock_now from tensorrt_llm.bindings.BuildInfo import ENABLE_MULTI_DEVICE from tensorrt_llm.logger import logger @@ -109,6 +109,15 @@ def numpy_to_torch(x): return torch.from_numpy(x) +def get_steady_clock_now_in_seconds() -> float: + """Time from the C++ runtime's steady clock, in seconds. + + Shared by the executor and the serving frontends so their perf-metric + timestamps are directly comparable. + """ + return steady_clock_now().total_seconds() + + def CUASSERT(cuda_ret): err = cuda_ret[0] if err != cudart.cudaError_t.cudaSuccess: diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index edd36f8624f9..00b568a9cdb8 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -40,7 +40,8 @@ ToolDescription, load_harmony_encoding) from transformers import AutoProcessor, PretrainedConfig -from tensorrt_llm.bindings import steady_clock_now +from tensorrt_llm._utils import \ + get_steady_clock_now_in_seconds # noqa: F401 (re-export) from tensorrt_llm.executor import GenerationResult from tensorrt_llm.inputs.utils import async_apply_chat_template from tensorrt_llm.llmapi import SamplingParams @@ -113,10 +114,6 @@ def _decode_tokens( return _get_encoding().decode(tokens) -def get_steady_clock_now_in_seconds() -> float: - return steady_clock_now().total_seconds() - - def _parse_response_input( input_msg: ResponseInputOutputItem, prev_responses: list[Union[ResponseOutputItem, ResponseReasoningItem]] From 22833f24557491292e5120c9b31af6944251896c Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:49:06 +0800 Subject: [PATCH 2/6] [None][feat] Load the PyTorch model zoo lazily _torch/models/__init__.py eagerly imported every modeling_* module, so any process importing tensorrt_llm executed ~50 model files' class bodies and registration decorators at startup. Replace this with: - _arch_index.py: static tables mapping architecture name and public class name to the providing module (generated from the @register_auto_model decorators and the previous eager import list); - PEP 562 __getattr__ on the package: attribute access imports just the providing module; bare modeling_* submodule access keeps working; - ensure_model_registered() in modeling_utils: architecture-based resolution (AutoModelForCausalLM, get_model_architecture, model_loader preference lookup) imports the providing module on demand so its decorators run; - commands/utils checks the static index instead of importing the zoo; - py_executor matches Llama4 by class name and model_loader builds its post-transform profile registry on first use, so neither imports a model-zoo module at module level; - drop the shadowed legacy MistralForCausalLM registration in modeling_llama (modeling_mistral always won under eager import order; under lazy loading the duplicate would be order-dependent). Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- tensorrt_llm/_torch/models/__init__.py | 142 ++++++--------- tensorrt_llm/_torch/models/_arch_index.py | 167 ++++++++++++++++++ tensorrt_llm/_torch/models/modeling_auto.py | 9 +- tensorrt_llm/_torch/models/modeling_llama.py | 6 +- tensorrt_llm/_torch/models/modeling_utils.py | 24 +++ .../_torch/pyexecutor/model_loader.py | 41 +++-- tensorrt_llm/_torch/pyexecutor/py_executor.py | 12 +- tensorrt_llm/commands/utils.py | 13 +- 8 files changed, 300 insertions(+), 114 deletions(-) create mode 100644 tensorrt_llm/_torch/models/_arch_index.py diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 4f4905d83086..9a0ecd5e546e 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -1,81 +1,39 @@ -import transformers +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""PyTorch-backend model zoo, loaded lazily. + +Importing this package no longer imports every ``modeling_*`` module (which +executes all model class bodies and their ``@register_auto_model`` decorators +at interpreter startup). Instead: + +- attribute access (``models.LlamaForCausalLM`` or + ``from tensorrt_llm._torch.models import LlamaForCausalLM``) imports just the + providing module, via PEP 562 ``__getattr__`` and ``MODEL_CLASS_TO_MODULE``; +- architecture-based resolution (``AutoModelForCausalLM``) imports on demand + via ``modeling_utils.ensure_model_registered`` and ``MODEL_ARCH_TO_MODULE``. +""" +import importlib # Importing _torch.configs triggers AutoConfig registration for TRT-LLM-only # model_types (deepseek_v32, kimi_k2, gemma4_unified) so AutoConfig / -# AutoTokenizer.from_pretrained work under transformers >= 5.5; see -# _torch/configs/__init__.py. +# AutoTokenizer.from_pretrained work under transformers >= 5.5; this must stay +# eager — see _torch/configs/__init__.py. import tensorrt_llm._torch.configs # noqa: F401 -from .modeling_afmoe import AfmoeForCausalLM +from ._arch_index import MODEL_CLASS_TO_MODULE from .modeling_auto import AutoModelForCausalLM -from .modeling_bart import (BartForConditionalGeneration, - MBartForConditionalGeneration) -from .modeling_bert import BertForSequenceClassification -from .modeling_clip import CLIPVisionModel -from .modeling_cohere2 import Cohere2ForCausalLM -from .modeling_cosmos3 import Cosmos3Model -from .modeling_deepseekv3 import DeepseekV3ForCausalLM -from .modeling_deepseekv4 import DeepseekV4ForCausalLM -from .modeling_exaone4 import Exaone4ForCausalLM -from .modeling_exaone4_5 import Exaone4_5_ForConditionalGeneration -from .modeling_exaone_moe import ExaoneMoeForCausalLM -from .modeling_gemma3 import Gemma3ForCausalLM -from .modeling_gemma3vl import Gemma3VLM -from .modeling_gemma4 import Gemma4ForCausalLM -from .modeling_gemma4_unified import Gemma4UnifiedForConditionalGeneration -from .modeling_gemma4mm import Gemma4ForConditionalGeneration -from .modeling_glm import Glm4MoeForCausalLM -from .modeling_gpt_oss import GptOssForCausalLM -from .modeling_hunyuan_dense import HunYuanDenseV1ForCausalLM -from .modeling_hunyuan_moe import HunYuanMoEV1ForCausalLM -from .modeling_hyperclovax import HCXVisionForCausalLM -from .modeling_kimi_k25 import KimiK25ForConditionalGeneration -from .modeling_laguna import LagunaForCausalLM -from .modeling_llama import LlamaForCausalLM -from .modeling_llava_next import LlavaNextModel -from .modeling_minicpmv4_6 import MiniCPMV4_6Model -from .modeling_minimaxm2 import MiniMaxM2ForCausalLM -from .modeling_minimaxm3 import (MiniMaxM3ForCausalLM, - MiniMaxM3VLForConditionalGeneration) -from .modeling_mistral import Mistral3VLM, MistralForCausalLM -from .modeling_mixtral import MixtralForCausalLM -from .modeling_nemotron import NemotronForCausalLM -from .modeling_nemotron_h import NemotronHForCausalLM -from .modeling_nemotron_nano import NemotronH_Nano_VL_V2 -from .modeling_nemotron_nas import NemotronNASForCausalLM -from .modeling_phi3 import Phi3ForCausalLM -from .modeling_phi4mm import Phi4MMForCausalLM -from .modeling_qwen import (Qwen2ForCausalLM, Qwen2ForProcessRewardModel, - Qwen2ForRewardModel) -from .modeling_qwen2vl import Qwen2_5_VLModel, Qwen2VLModel -from .modeling_qwen3 import Qwen3ForCausalLM -from .modeling_qwen3_5 import (Qwen3_5ForCausalLM, Qwen3_5MoeForCausalLM, - Qwen3_5MoeVLModel, Qwen3_5VLModel) -from .modeling_qwen3_moe import Qwen3MoeForCausalLM -from .modeling_qwen3_next import Qwen3NextForCausalLM -from .modeling_qwen3vl import Qwen3VLModel -from .modeling_qwen3vl_moe import Qwen3MoeVLModel -from .modeling_qwen_image_bench import QwenImageBenchModel -from .modeling_qwen_moe import Qwen2MoeForCausalLM -from .modeling_seedoss import SeedOssForCausalLM -from .modeling_siglip import SiglipVisionModel -from .modeling_starcoder2 import Starcoder2ForCausalLM -from .modeling_step3p7 import Step3p7ForCausalLM -from .modeling_step3p7vl import Step3p7VLForConditionalGeneration -from .modeling_t5 import T5ForConditionalGeneration from .modeling_utils import get_model_architecture -from .modeling_vila import VilaModel -from .modeling_whisper import WhisperForConditionalGeneration -# Note: for better readiblity, this should have same order as imports above __all__ = [ "AfmoeForCausalLM", "AutoModelForCausalLM", "BartForConditionalGeneration", "BertForSequenceClassification", "CLIPVisionModel", + "Cohere2ForCausalLM", "Cosmos3Model", "DeepseekV3ForCausalLM", + "DeepseekV4ForCausalLM", "Exaone4ForCausalLM", "Exaone4_5_ForConditionalGeneration", "ExaoneMoeForCausalLM", @@ -84,21 +42,27 @@ "Gemma4ForCausalLM", "Gemma4ForConditionalGeneration", "Gemma4UnifiedForConditionalGeneration", + "Glm4MoeForCausalLM", + "GptOssForCausalLM", "HCXVisionForCausalLM", - "LagunaForCausalLM", "HunYuanDenseV1ForCausalLM", "HunYuanMoEV1ForCausalLM", "KimiK25ForConditionalGeneration", + "LagunaForCausalLM", "LlamaForCausalLM", "LlavaNextModel", + "MBartForConditionalGeneration", "MiniCPMV4_6Model", + "MiniMaxM2ForCausalLM", + "MiniMaxM3ForCausalLM", + "MiniMaxM3VLForConditionalGeneration", "Mistral3VLM", "MistralForCausalLM", "MixtralForCausalLM", - "DeepseekV4ForCausalLM", - "NemotronH_Nano_VL_V2", + "MllamaForConditionalGeneration", "NemotronForCausalLM", "NemotronHForCausalLM", + "NemotronH_Nano_VL_V2", "NemotronNASForCausalLM", "Phi3ForCausalLM", "Phi4MMForCausalLM", @@ -106,41 +70,43 @@ "Qwen2ForProcessRewardModel", "Qwen2ForRewardModel", "Qwen2MoeForCausalLM", - "SiglipVisionModel", - "Starcoder2ForCausalLM", - "T5ForConditionalGeneration", - "MBartForConditionalGeneration", - "get_model_architecture", - "VilaModel", "Qwen2VLModel", "Qwen2_5_VLModel", "Qwen3ForCausalLM", "Qwen3MoeForCausalLM", + "Qwen3MoeVLModel", + "Qwen3NextForCausalLM", + "Qwen3VLModel", "Qwen3_5ForCausalLM", "Qwen3_5MoeForCausalLM", - "QwenImageBenchModel", "Qwen3_5MoeVLModel", "Qwen3_5VLModel", - "Qwen3NextForCausalLM", - "Qwen3MoeVLModel", - "GptOssForCausalLM", + "QwenImageBenchModel", "SeedOssForCausalLM", - "Glm4MoeForCausalLM", - "Qwen3VLModel", - "MiniMaxM2ForCausalLM", - "MiniMaxM3ForCausalLM", - "MiniMaxM3VLForConditionalGeneration", - "Cohere2ForCausalLM", + "SiglipVisionModel", + "Starcoder2ForCausalLM", "Step3p7ForCausalLM", "Step3p7VLForConditionalGeneration", + "T5ForConditionalGeneration", + "VilaModel", "WhisperForConditionalGeneration", + "get_model_architecture", ] -if transformers.__version__ >= "4.45.1": - from .modeling_mllama import MllamaForConditionalGeneration # noqa - __all__.append("MllamaForConditionalGeneration") -else: - print( - f"Failed to import MllamaForConditionalGeneration as transformers.__version__ {transformers.__version__} < 4.45.1" - ) +def __getattr__(name: str): + module_name = MODEL_CLASS_TO_MODULE.get(name) + if module_name is None: + # Also resolve bare submodule access (models.modeling_llama) so callers + # that relied on the previously-eager submodule attributes keep working. + if name.startswith("modeling_"): + return importlib.import_module(f".{name}", __name__) + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + module = importlib.import_module(f".{module_name}", __name__) + attr = getattr(module, name) + globals()[name] = attr # cache: subsequent access skips __getattr__ + return attr + + +def __dir__(): + return sorted(set(__all__) | set(globals())) diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py new file mode 100644 index 000000000000..b62600adfedd --- /dev/null +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -0,0 +1,167 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Static index of the PyTorch-backend model zoo, used for lazy loading. + +Model implementations register themselves via ``@register_auto_model`` as an +import side effect. The zoo is imported lazily, so these tables record, without +importing anything, which ``modeling_*`` module provides which architecture +(``MODEL_ARCH_TO_MODULE``) and which public class (``MODEL_CLASS_TO_MODULE``). + +Regenerate after adding/moving a model: + scan ``@register_auto_model("")`` decorators and the public class + names under ``tensorrt_llm/_torch/models`` (see the lazy-import PR for the + one-off script), or add the new entry by hand next to its neighbors. +""" + +# Architecture name (HF ``config.architectures[0]``, possibly rewritten by +# ``AutoModelForCausalLM._resolve_class``) -> providing module. +MODEL_ARCH_TO_MODULE = { + "AfmoeForCausalLM": "modeling_afmoe", + "BartForConditionalGeneration": "modeling_bart", + "BertForSequenceClassification": "modeling_bert", + "CLIPVisionModel": "modeling_clip", + "Cohere2ForCausalLM": "modeling_cohere2", + "Cosmos3ForConditionalGeneration": "modeling_cosmos3", + "DeciLMForCausalLM": "modeling_nemotron_nas", + "DeepseekV32ForCausalLM": "modeling_deepseekv3", + "DeepseekV3ForCausalLM": "modeling_deepseekv3", + "DeepseekV4ForCausalLM": "modeling_deepseekv4", + "EAGLE3LlamaForCausalLM": "modeling_speculative", + "Eagle3DeepSeekV3ForCausalLM": "modeling_speculative", + "Exaone4ForCausalLM": "modeling_exaone4", + "Exaone4_5_ForConditionalGeneration": "modeling_exaone4_5", + "ExaoneMoEForCausalLM": "modeling_exaone_moe", + "Gemma3ForCausalLM": "modeling_gemma3", + "Gemma3ForConditionalGeneration": "modeling_gemma3vl", + "Gemma4AssistantForCausalLM": "modeling_gemma4", + "Gemma4ForCausalLM": "modeling_gemma4", + "Gemma4ForConditionalGeneration": "modeling_gemma4mm", + "Gemma4UnifiedForConditionalGeneration": "modeling_gemma4_unified", + "Glm4MoeForCausalLM": "modeling_glm", + "GlmMoeDsaForCausalLM": "modeling_deepseekv3", + "GptOssForCausalLM": "modeling_gpt_oss", + "HCXVisionForCausalLM": "modeling_hyperclovax", + "HCXVisionModel": "modeling_hyperclovax", + "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", + "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", + "KimiK25ForConditionalGeneration": "modeling_kimi_k25", + "LagunaForCausalLM": "modeling_laguna", + "Llama4ForConditionalGeneration": "modeling_llama", + "LlamaForCausalLM": "modeling_llama", + "LlavaLlamaModel": "modeling_vila", + "LlavaNextForConditionalGeneration": "modeling_llava_next", + "MBartForConditionalGeneration": "modeling_bart", + "MTPDraftModelForCausalLM": "modeling_speculative", + "MiniCPMV4_6ForConditionalGeneration": "modeling_minicpmv4_6", + "MiniMaxM2ForCausalLM": "modeling_minimaxm2", + "MiniMaxM3SparseForCausalLM": "modeling_minimaxm3", + "MiniMaxM3SparseForConditionalGeneration": "modeling_minimaxm3", + "Mistral3ForConditionalGeneration": "modeling_mistral", + "MistralForCausalLM": "modeling_mistral", + "MistralLarge3EagleForCausalLM": "modeling_speculative", + "MistralLarge3ForCausalLM": "modeling_mistral_large3", + "MixtralForCausalLM": "modeling_mixtral", + "MllamaForConditionalGeneration": "modeling_mllama", + "NemotronForCausalLM": "modeling_nemotron", + "NemotronHForCausalLM": "modeling_nemotron_h", + "NemotronHPuzzleForCausalLM": "modeling_nemotron_h", + "NemotronH_Nano_Omni_Reasoning_V3": "modeling_nemotron_nano", + "NemotronH_Nano_VL_V2": "modeling_nemotron_nano", + "Phi3ForCausalLM": "modeling_phi3", + "Phi4MMForCausalLM": "modeling_phi4mm", + "PixtralForConditionalGeneration": "modeling_mistral", + "PixtralVisionModel": "modeling_pixtral", + "Qwen2ForCausalLM": "modeling_qwen", + "Qwen2ForProcessRewardModel": "modeling_qwen", + "Qwen2ForRewardModel": "modeling_qwen", + "Qwen2MoeForCausalLM": "modeling_qwen_moe", + "Qwen2VLForConditionalGeneration": "modeling_qwen2vl", + "Qwen2_5_VLForConditionalGeneration": "modeling_qwen2vl", + "Qwen3ForCausalLM": "modeling_qwen3", + "Qwen3ForTextEmbedding": "modeling_qwen3", + "Qwen3MoeForCausalLM": "modeling_qwen3_moe", + "Qwen3NextForCausalLM": "modeling_qwen3_next", + "Qwen3VLForConditionalGeneration": "modeling_qwen3vl", + "Qwen3VLMoeForConditionalGeneration": "modeling_qwen3vl_moe", + "Qwen3_5ForCausalLM": "modeling_qwen3_5", + "Qwen3_5ForConditionalGeneration": "modeling_qwen3_5", + "Qwen3_5MoeForCausalLM": "modeling_qwen3_5", + "Qwen3_5MoeForConditionalGeneration": "modeling_qwen3_5", + "QwenImageBenchForConditionalGeneration": "modeling_qwen_image_bench", + "SeedOssForCausalLM": "modeling_seedoss", + "SiglipVisionModel": "modeling_siglip", + "SomeVLModel": "modeling_utils", + "Starcoder2ForCausalLM": "modeling_starcoder2", + "Step3p5ForCausalLM": "modeling_step3p7", + "Step3p7ForConditionalGeneration": "modeling_step3p7vl", + "T5ForConditionalGeneration": "modeling_t5", + "WhisperForConditionalGeneration": "modeling_whisper", +} + +# Public class name exported by ``tensorrt_llm._torch.models`` -> providing module. +MODEL_CLASS_TO_MODULE = { + "AfmoeForCausalLM": "modeling_afmoe", + "BartForConditionalGeneration": "modeling_bart", + "BertForSequenceClassification": "modeling_bert", + "CLIPVisionModel": "modeling_clip", + "Cohere2ForCausalLM": "modeling_cohere2", + "Cosmos3Model": "modeling_cosmos3", + "DeepseekV3ForCausalLM": "modeling_deepseekv3", + "DeepseekV4ForCausalLM": "modeling_deepseekv4", + "Exaone4ForCausalLM": "modeling_exaone4", + "Exaone4_5_ForConditionalGeneration": "modeling_exaone4_5", + "ExaoneMoeForCausalLM": "modeling_exaone_moe", + "Gemma3ForCausalLM": "modeling_gemma3", + "Gemma3VLM": "modeling_gemma3vl", + "Gemma4ForCausalLM": "modeling_gemma4", + "Gemma4ForConditionalGeneration": "modeling_gemma4mm", + "Gemma4UnifiedForConditionalGeneration": "modeling_gemma4_unified", + "Glm4MoeForCausalLM": "modeling_glm", + "GptOssForCausalLM": "modeling_gpt_oss", + "HCXVisionForCausalLM": "modeling_hyperclovax", + "HunYuanDenseV1ForCausalLM": "modeling_hunyuan_dense", + "HunYuanMoEV1ForCausalLM": "modeling_hunyuan_moe", + "KimiK25ForConditionalGeneration": "modeling_kimi_k25", + "LagunaForCausalLM": "modeling_laguna", + "LlamaForCausalLM": "modeling_llama", + "LlavaNextModel": "modeling_llava_next", + "MBartForConditionalGeneration": "modeling_bart", + "MiniCPMV4_6Model": "modeling_minicpmv4_6", + "MiniMaxM2ForCausalLM": "modeling_minimaxm2", + "MiniMaxM3ForCausalLM": "modeling_minimaxm3", + "MiniMaxM3VLForConditionalGeneration": "modeling_minimaxm3", + "Mistral3VLM": "modeling_mistral", + "MistralForCausalLM": "modeling_mistral", + "MixtralForCausalLM": "modeling_mixtral", + "MllamaForConditionalGeneration": "modeling_mllama", + "NemotronForCausalLM": "modeling_nemotron", + "NemotronHForCausalLM": "modeling_nemotron_h", + "NemotronH_Nano_VL_V2": "modeling_nemotron_nano", + "NemotronNASForCausalLM": "modeling_nemotron_nas", + "Phi3ForCausalLM": "modeling_phi3", + "Phi4MMForCausalLM": "modeling_phi4mm", + "Qwen2ForCausalLM": "modeling_qwen", + "Qwen2ForProcessRewardModel": "modeling_qwen", + "Qwen2ForRewardModel": "modeling_qwen", + "Qwen2MoeForCausalLM": "modeling_qwen_moe", + "Qwen2VLModel": "modeling_qwen2vl", + "Qwen2_5_VLModel": "modeling_qwen2vl", + "Qwen3ForCausalLM": "modeling_qwen3", + "Qwen3MoeForCausalLM": "modeling_qwen3_moe", + "Qwen3MoeVLModel": "modeling_qwen3vl_moe", + "Qwen3NextForCausalLM": "modeling_qwen3_next", + "Qwen3VLModel": "modeling_qwen3vl", + "Qwen3_5ForCausalLM": "modeling_qwen3_5", + "Qwen3_5MoeForCausalLM": "modeling_qwen3_5", + "Qwen3_5MoeVLModel": "modeling_qwen3_5", + "Qwen3_5VLModel": "modeling_qwen3_5", + "QwenImageBenchModel": "modeling_qwen_image_bench", + "SeedOssForCausalLM": "modeling_seedoss", + "SiglipVisionModel": "modeling_siglip", + "Starcoder2ForCausalLM": "modeling_starcoder2", + "Step3p7ForCausalLM": "modeling_step3p7", + "Step3p7VLForConditionalGeneration": "modeling_step3p7vl", + "T5ForConditionalGeneration": "modeling_t5", + "VilaModel": "modeling_vila", + "WhisperForConditionalGeneration": "modeling_whisper", +} diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 2802ff59fd53..695c6261d9b2 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -4,7 +4,8 @@ from ..utils import model_extra_attrs from .modeling_utils import (MODEL_CLASS_MAPPING, MODEL_CLASS_VISION_ENCODER_MAPPING, - DecoderModelForCausalLM, TConfig, TModel) + DecoderModelForCausalLM, TConfig, TModel, + ensure_model_registered) class AutoModelForCausalLM(Generic[TModel, TConfig]): @@ -17,6 +18,9 @@ def _resolve_class(config: ModelConfig) -> Optional[Type]: return None model_arch = pretrained_config.architectures[0] + # The model zoo is imported lazily: pull in the module providing this + # architecture so its registration decorators have run. + ensure_model_registered(model_arch) if config.mm_encoder_only: vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get( @@ -33,12 +37,14 @@ def _resolve_class(config: ModelConfig) -> Optional[Type]: model_arch = model_arch.replace("Eagle3", "") # Strip the appended EAGLE3 model_arch = "EAGLE3" + model_arch + ensure_model_registered(model_arch) if model_arch in ( "DeepseekV3ForCausalLM", "Glm4MoeForCausalLM", "ExaoneMoEForCausalLM" ) and config.spec_config is not None and config.spec_config.max_draft_len == 0: model_arch = "MTPDraftModelForCausalLM" + ensure_model_registered(model_arch) return MODEL_CLASS_MAPPING.get(model_arch) @@ -48,6 +54,7 @@ def from_config( ) -> DecoderModelForCausalLM[TModel, TConfig]: if config.mm_encoder_only: model_arch = config.pretrained_config.architectures[0] + ensure_model_registered(model_arch) vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get( model_arch) if vision_encoder_info is None: diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py index 78851f9de78f..6be6c944ba77 100644 --- a/tensorrt_llm/_torch/models/modeling_llama.py +++ b/tensorrt_llm/_torch/models/modeling_llama.py @@ -1606,7 +1606,11 @@ def setup_aliases(self) -> None: layer.next_attn = self.model.layers[idx + 1].self_attn -@register_auto_model("MistralForCausalLM") +# NOTE: deliberately NOT decorated with @register_auto_model: the +# "MistralForCausalLM" architecture is owned by modeling_mistral (which, under +# the previous eager import order, always overwrote this legacy registration). +# With the model zoo imported lazily, a duplicate registration here would make +# the resolved class depend on import order. class MistralForCausalLM(DecoderModelForCausalLM[LlamaModel, LlamaConfig]): def __init__( diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index eb91a6456209..57e714dceff1 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -2,6 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 import contextlib +import importlib import inspect import math import os @@ -32,6 +33,7 @@ from ..modules.logits_processor import LogitsProcessor from ..modules.rms_norm import RMSNorm from ..speculative import SpecMetadata +from ._arch_index import MODEL_ARCH_TO_MODULE @contextlib.contextmanager @@ -864,6 +866,27 @@ def decorator(cls): return decorator +def ensure_model_registered(model_arch: str) -> None: + """Import the module that provides ``model_arch``, if it isn't loaded yet. + + Model implementations register themselves in ``MODEL_CLASS_MAPPING`` (and + the sibling registries) as an import side effect. With the model zoo + imported lazily, this is the hook that turns an architecture name into + "the decorators have run". Architectures missing from the static index + (e.g. registered dynamically by user code) and modules that fail to import + are left to the caller's normal missing-architecture handling; the warning + keeps the root cause visible. + """ + module_name = MODEL_ARCH_TO_MODULE.get(model_arch) + if module_name is None: + return + try: + importlib.import_module(f"tensorrt_llm._torch.models.{module_name}") + except ImportError as e: + logger.warning(f"Lazy import of {module_name} for architecture " + f"{model_arch} failed: {e!r}") + + def register_vision_encoder( vision_encoder_cls: Type[nn.Module], vlm_base_model: Optional[Type[nn.Module]] = None, @@ -962,6 +985,7 @@ def get_model_architecture( cls = None if model_config.architectures is not None and len( model_config.architectures) > 0: + ensure_model_registered(model_config.architectures[0]) cls = MODEL_CLASS_MAPPING.get(model_config.architectures[0]) else: raise RuntimeError("Model architecture is not provided.") diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 2a6ea23076aa..8bfb748093e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -34,11 +34,11 @@ from ...llmapi.llm_args import LoadFormat from ..model_config import ModelConfig -from ..models import AutoModelForCausalLM, LlamaForCausalLM +from ..models import AutoModelForCausalLM from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader from ..models.modeling_utils import (MODEL_CLASS_MAPPING, DecoderModelForCausalLM, MetaInitMode, - timing) + ensure_model_registered, timing) from ..modules.fused_moe.moe_load_balancer import ( MoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode @@ -342,17 +342,29 @@ class ModelLoader: This class isolates model loading logic from the main execution engine. """ _MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION = 1 - _POST_TRANSFORM_PROFILE_REGISTRY = PostTransformProfileRegistry( - profiles=(PostTransformProfile( - profile_id="llama-for-causal-lm-target-v1", - root_model_class=LlamaForCausalLM, - architecture="LlamaForCausalLM", - model_type="llama", - speculative_mode=None, - protocol_version=_MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, - transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, - transfer_scope=PostTransformTransferScope.TARGET_MODEL, - ), )) + _POST_TRANSFORM_PROFILE_REGISTRY: Optional[ + PostTransformProfileRegistry] = None + + @classmethod + def _post_transform_profile_registry(cls) -> PostTransformProfileRegistry: + # Built on first use: the profile references a model-zoo class, and + # importing it here (not at module level) keeps the zoo lazy for + # processes that import model_loader but never qualify a model. + if cls._POST_TRANSFORM_PROFILE_REGISTRY is None: + from ..models.modeling_llama import LlamaForCausalLM + cls._POST_TRANSFORM_PROFILE_REGISTRY = PostTransformProfileRegistry( + profiles=(PostTransformProfile( + profile_id="llama-for-causal-lm-target-v1", + root_model_class=LlamaForCausalLM, + architecture="LlamaForCausalLM", + model_type="llama", + speculative_mode=None, + protocol_version=cls. + _MX_STAGED_RECEIVER_TRANSFORM_PROTOCOL_VERSION, + transform_abi_id=LLAMA_POST_TRANSFORM_LAYOUT_ABI_V1, + transfer_scope=PostTransformTransferScope.TARGET_MODEL, + ), )) + return cls._POST_TRANSFORM_PROFILE_REGISTRY def __init__(self, llm_args: TorchLlmArgs, @@ -448,6 +460,7 @@ def load_config_and_apply_defaults( preference_cls = model_cls architectures = getattr(config.pretrained_config, 'architectures', None) if architectures: + ensure_model_registered(architectures[0]) preference_cls = MODEL_CLASS_MAPPING.get(architectures[0], model_cls) @@ -1131,7 +1144,7 @@ def _qualify_post_transform_profile( enabled_features = set() if loads_draft_weights: enabled_features.add(PostTransformFeature.SEPARATE_DRAFT_MODEL) - return cls._POST_TRANSFORM_PROFILE_REGISTRY.qualify( + return cls._post_transform_profile_registry().qualify( root_model_class=type(model), architecture=config_identity.architecture, model_type=config_identity.model_type, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index 6e3252b48688..c0c353c171f4 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -18,7 +18,6 @@ from strenum import StrEnum from tensorrt_llm.llmapi import DisaggScheduleStyle -from tensorrt_llm.serve.responses_utils import get_steady_clock_now_in_seconds try: from cuda.bindings import runtime as cudart @@ -26,6 +25,7 @@ from cuda import cudart from tensorrt_llm._utils import (CUASSERT, customized_gc_thresholds, + get_steady_clock_now_in_seconds, is_trace_enabled, mpi_comm, mpi_disabled, nvtx_range, set_thread_local_mpi_comm, trace_func) @@ -50,7 +50,6 @@ from ..distributed import Distributed from ..distributed.communicator import ReduceOp from ..expert_statistic import ExpertStatistic -from ..models.modeling_llama import Llama4ForConditionalGeneration from ..models.modeling_multimodal_mixin import \ maybe_prefetch_mm_encoder_for_next_iter from ..models.modeling_utils import DecoderModelForCausalLM @@ -4882,9 +4881,12 @@ def _forward_step_inter_pp(self, def _validate_token_id_range(self, request: LlmRequest) -> None: if isinstance(self.model_engine.model, DecoderModelForCausalLM): - # Only skip token‐range checks for Llama4 when the request has multimodal data - if isinstance(self.model_engine.model, - Llama4ForConditionalGeneration): + # Only skip token‐range checks for Llama4 when the request has multimodal data. + # Matched by class name (equivalent to isinstance incl. subclasses) + # so this module does not import a model-zoo module at startup, + # which would defeat the zoo's lazy loading. + if any(c.__name__ == "Llama4ForConditionalGeneration" + for c in type(self.model_engine.model).__mro__): has_mm = bool(request.py_multimodal_data) if has_mm: logger.debug( diff --git a/tensorrt_llm/commands/utils.py b/tensorrt_llm/commands/utils.py index d33282ffe85e..1ab0afbe7a5f 100644 --- a/tensorrt_llm/commands/utils.py +++ b/tensorrt_llm/commands/utils.py @@ -92,13 +92,16 @@ def has_registered_llm_architecture(model_path: str) -> bool: with open(config_path) as f: architectures = json.load(f).get("architectures") or [] - # Importing the models package runs the @register_auto_model decorators - # that populate MODEL_CLASS_MAPPING. Done lazily to avoid a heavy import at - # module load time and to sidestep circular imports via commands.serve. - import tensorrt_llm._torch.models # noqa: F401 + # The model zoo is imported lazily, so MODEL_CLASS_MAPPING only holds what + # has been resolved so far; the static index knows every built-in + # architecture without importing anything. Imported here (not at module + # load) to sidestep circular imports via commands.serve. + from tensorrt_llm._torch.models._arch_index import MODEL_ARCH_TO_MODULE from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING - return any(arch in MODEL_CLASS_MAPPING for arch in architectures) + return any( + arch in MODEL_CLASS_MAPPING or arch in MODEL_ARCH_TO_MODULE for arch in architectures + ) def get_is_diffusion_only_model(model_path: str): From e8430e506ea051872bcc18c5dcb68be797f02676 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Wed, 5 Aug 2026 11:49:44 +0800 Subject: [PATCH 3/6] [None][feat] Make the tensorrt_llm top-level namespace lazy (PEP 562) 'import tensorrt_llm' eagerly pulled the whole public surface (both model zoos, quantization, runtime, tools, llmapi, visual_gen), which executed ~99% of product modules in every process at startup. Keep the environment setup, _common._init(), logger and version banner eager; resolve every public name on first attribute access instead, with a plain-submodule fallback so previously-reachable module attributes keep working. A TYPE_CHECKING block preserves the original imports for static tooling. Also maps KvCacheConfig, which __all__ listed but the eager chain never actually imported. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- tensorrt_llm/__init__.py | 115 ++++++++++++++++++++++++++++++++------- 1 file changed, 95 insertions(+), 20 deletions(-) diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index d707ca0e0694..c73f7e628f26 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -99,34 +99,109 @@ def _setup_vendored_triton_kernels(): _setup_vendored_triton_kernels() +# The package's public surface is loaded lazily (PEP 562): importing +# tensorrt_llm no longer executes the whole product tree (previously ~99% of +# product modules ran at import time through the eager chain below). Accessing +# any public name (tensorrt_llm.LLM, tensorrt_llm.models, ...) imports just +# what that name needs. The TYPE_CHECKING block keeps the original imports +# visible to static tooling. +import importlib +from typing import TYPE_CHECKING + # Need to import torch before tensorrt_llm library, otherwise some shared binary files # cannot be found for the public PyTorch, raising errors like: # ImportError: libc10.so: cannot open shared object file: No such file or directory import torch # noqa -import tensorrt_llm._torch.models as torch_models -import tensorrt_llm.math_utils as math_utils -import tensorrt_llm.models as models -import tensorrt_llm.quantization as quantization -import tensorrt_llm.runtime as runtime -import tensorrt_llm.tools as tools - from ._common import _init -from ._mnnvl_utils import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo -from ._utils import (default_gpus_per_node, local_mpi_rank, local_mpi_size, - mpi_barrier, mpi_comm, mpi_rank, mpi_world_size, - set_mpi_comm, str_dtype_to_torch) -from .disaggregated_params import DisaggregatedParams -from .llmapi import LLM, AsyncLLM, MultimodalEncoder -from .llmapi.llm_args import LlmArgs, TorchLlmArgs from .logger import logger -from .mapping import Mapping -from .models.automodel import AutoConfig, AutoModelForCausalLM -from .sampling_params import SamplingParams from .version import __version__ -from .visual_gen import (ExtraParamSchema, VisualGen, VisualGenArgs, - VisualGenMetrics, VisualGenOutput, VisualGenParams, - VisualGenResult) + +if TYPE_CHECKING: + import tensorrt_llm._torch.models as torch_models + import tensorrt_llm.math_utils as math_utils + import tensorrt_llm.models as models + import tensorrt_llm.quantization as quantization + import tensorrt_llm.runtime as runtime + import tensorrt_llm.tools as tools + + from ._mnnvl_utils import MnnvlMemory, MnnvlMoe, MoEAlltoallInfo + from ._utils import (default_gpus_per_node, local_mpi_rank, local_mpi_size, + mpi_barrier, mpi_comm, mpi_rank, mpi_world_size, + set_mpi_comm, str_dtype_to_torch) + from .disaggregated_params import DisaggregatedParams + from .llmapi import LLM, AsyncLLM, KvCacheConfig, MultimodalEncoder + from .llmapi.llm_args import LlmArgs, TorchLlmArgs + from .mapping import Mapping + from .models.automodel import AutoConfig, AutoModelForCausalLM + from .sampling_params import SamplingParams + from .visual_gen import (ExtraParamSchema, VisualGen, VisualGenArgs, + VisualGenMetrics, VisualGenOutput, VisualGenParams, + VisualGenResult) + +# Public name -> (source module, attribute); attribute None = the module itself. +_LAZY_ATTRS = { + 'torch_models': ('tensorrt_llm._torch.models', None), + 'math_utils': ('tensorrt_llm.math_utils', None), + 'models': ('tensorrt_llm.models', None), + 'quantization': ('tensorrt_llm.quantization', None), + 'runtime': ('tensorrt_llm.runtime', None), + 'tools': ('tensorrt_llm.tools', None), + 'MnnvlMemory': ('tensorrt_llm._mnnvl_utils', 'MnnvlMemory'), + 'MnnvlMoe': ('tensorrt_llm._mnnvl_utils', 'MnnvlMoe'), + 'MoEAlltoallInfo': ('tensorrt_llm._mnnvl_utils', 'MoEAlltoallInfo'), + 'default_gpus_per_node': ('tensorrt_llm._utils', 'default_gpus_per_node'), + 'local_mpi_rank': ('tensorrt_llm._utils', 'local_mpi_rank'), + 'local_mpi_size': ('tensorrt_llm._utils', 'local_mpi_size'), + 'mpi_barrier': ('tensorrt_llm._utils', 'mpi_barrier'), + 'mpi_comm': ('tensorrt_llm._utils', 'mpi_comm'), + 'mpi_rank': ('tensorrt_llm._utils', 'mpi_rank'), + 'mpi_world_size': ('tensorrt_llm._utils', 'mpi_world_size'), + 'set_mpi_comm': ('tensorrt_llm._utils', 'set_mpi_comm'), + 'str_dtype_to_torch': ('tensorrt_llm._utils', 'str_dtype_to_torch'), + 'DisaggregatedParams': + ('tensorrt_llm.disaggregated_params', 'DisaggregatedParams'), + 'LLM': ('tensorrt_llm.llmapi', 'LLM'), + 'AsyncLLM': ('tensorrt_llm.llmapi', 'AsyncLLM'), + 'MultimodalEncoder': ('tensorrt_llm.llmapi', 'MultimodalEncoder'), + 'KvCacheConfig': ('tensorrt_llm.llmapi', 'KvCacheConfig'), + 'LlmArgs': ('tensorrt_llm.llmapi.llm_args', 'LlmArgs'), + 'TorchLlmArgs': ('tensorrt_llm.llmapi.llm_args', 'TorchLlmArgs'), + 'Mapping': ('tensorrt_llm.mapping', 'Mapping'), + 'AutoConfig': ('tensorrt_llm.models.automodel', 'AutoConfig'), + 'AutoModelForCausalLM': + ('tensorrt_llm.models.automodel', 'AutoModelForCausalLM'), + 'SamplingParams': ('tensorrt_llm.sampling_params', 'SamplingParams'), + 'ExtraParamSchema': ('tensorrt_llm.visual_gen', 'ExtraParamSchema'), + 'VisualGen': ('tensorrt_llm.visual_gen', 'VisualGen'), + 'VisualGenArgs': ('tensorrt_llm.visual_gen', 'VisualGenArgs'), + 'VisualGenMetrics': ('tensorrt_llm.visual_gen', 'VisualGenMetrics'), + 'VisualGenOutput': ('tensorrt_llm.visual_gen', 'VisualGenOutput'), + 'VisualGenParams': ('tensorrt_llm.visual_gen', 'VisualGenParams'), + 'VisualGenResult': ('tensorrt_llm.visual_gen', 'VisualGenResult'), +} + + +def __getattr__(name): + entry = _LAZY_ATTRS.get(name) + if entry is not None: + module_name, attr = entry + module = importlib.import_module(module_name) + value = module if attr is None else getattr(module, attr) + globals()[name] = value # cache: subsequent access skips __getattr__ + return value + # Fall back to plain submodules (tensorrt_llm.functional, .profiler, ...) + # that used to be reachable as attributes via the eager import chain. + try: + return importlib.import_module(f'.{name}', __name__) + except ModuleNotFoundError: + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None + + +def __dir__(): + return sorted(set(__all__) | set(globals()) | set(_LAZY_ATTRS)) + __all__ = [ 'AutoConfig', From 635f9fb43bc4a7e8777a970523ea054abe875226 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:54:19 +0800 Subject: [PATCH 4/6] [None][feat] Keep visual_gen imports out of LLM serving processes The OpenAI server and the serve/bench CLIs imported the visual_gen tree at module level, so every plain LLM serving process paid the import cost of the whole visual_gen package. Type-only uses move under TYPE_CHECKING, runtime uses import locally inside the VisualGen code paths, and the isinstance checks go through a sys.modules probe: if visual_gen was never imported, the generator cannot be a VisualGen instance. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- tensorrt_llm/commands/serve.py | 14 +++++++--- tensorrt_llm/commands/utils.py | 5 +++- tensorrt_llm/serve/openai_server.py | 33 ++++++++++++++++------- tensorrt_llm/serve/openai_video_routes.py | 9 ++++++- tensorrt_llm/serve/visual_gen_utils.py | 10 +++++-- 5 files changed, 53 insertions(+), 18 deletions(-) diff --git a/tensorrt_llm/commands/serve.py b/tensorrt_llm/commands/serve.py index 65691acbf91c..13a81831dfb0 100644 --- a/tensorrt_llm/commands/serve.py +++ b/tensorrt_llm/commands/serve.py @@ -14,7 +14,7 @@ import time import uuid from pathlib import Path -from typing import Any, Dict, NamedTuple, Optional, Sequence, Set +from typing import TYPE_CHECKING, Any, Dict, NamedTuple, Optional, Sequence, Set import click import torch @@ -51,8 +51,11 @@ MODEL_TYPE_TO_TOOL_PARSER, resolve_auto_tool_parser) from tensorrt_llm.tools.importlib_utils import import_custom_module_from_dir from tensorrt_llm.usage import config as _telemetry_config -from tensorrt_llm.visual_gen import VisualGen -from tensorrt_llm.visual_gen.args import VisualGenArgs + +if TYPE_CHECKING: + # Type-only: the visual_gen tree is imported lazily inside the VisualGen + # code paths so plain LLM serving never pays its import cost. + from tensorrt_llm.visual_gen.args import VisualGenArgs # Global variable to store the Popen object of the child process _child_p_global: Optional[subprocess.Popen] = None @@ -880,7 +883,7 @@ def launch_visual_gen_server( host: str, port: int, model: str, - visual_gen_args: Optional[VisualGenArgs] = None, + visual_gen_args: Optional["VisualGenArgs"] = None, metadata_server_cfg: Optional[MetadataServerConfig] = None, middleware: Sequence[str] = (), ): @@ -898,6 +901,7 @@ def launch_visual_gen_server( # races the same port and all but one die EADDRINUSE. VisualGen() on a # worker rank never returns (sys.exit in __init__). from tensorrt_llm._torch.visual_gen.executor import _detect_external_launch + from tensorrt_llm.visual_gen import VisualGen ext = _detect_external_launch() if ext is not None and ext[0] != 0: VisualGen(model=model, args=visual_gen_args) @@ -1502,6 +1506,8 @@ def _serve_llm(): internal_disagg_auth_key=internal_disagg_auth_key) def _serve_visual_gen(): + from tensorrt_llm.visual_gen.args import VisualGenArgs + parsed_visual_gen_args = (VisualGenArgs.from_yaml(visual_gen_args) if visual_gen_args is not None else None) diff --git a/tensorrt_llm/commands/utils.py b/tensorrt_llm/commands/utils.py index 1ab0afbe7a5f..011f772d9145 100644 --- a/tensorrt_llm/commands/utils.py +++ b/tensorrt_llm/commands/utils.py @@ -8,7 +8,6 @@ from click.core import ParameterSource from tensorrt_llm.llmapi.utils import download_hf_partial -from tensorrt_llm.visual_gen.args import ParallelConfig logger = logging.getLogger(__name__) @@ -200,6 +199,10 @@ def get_visual_gen_num_gpus(diffusion_config: dict) -> int: Uses ParallelConfig.model_construct (skips env validators) so this is safe to call from non-worker processes. """ + # Imported here (not at module load) so LLM-only CLI processes never pull + # the visual_gen tree. + from tensorrt_llm.visual_gen.args import ParallelConfig + parallel = diffusion_config.get("parallel_config", {}) if isinstance(parallel, dict): parallel = ParallelConfig.model_construct(**parallel) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 367caab10ec1..7f46a768a1db 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -8,6 +8,7 @@ import re import signal import socket +import sys import time import traceback import uuid @@ -17,8 +18,8 @@ from datetime import datetime from http import HTTPStatus from pathlib import Path -from typing import (Annotated, Any, AsyncGenerator, AsyncIterator, List, - Optional, Union) +from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, + AsyncIterator, List, Optional, Union) import uvicorn from fastapi import Body, FastAPI, Request @@ -104,11 +105,24 @@ build_visual_gen_timing_headers from tensorrt_llm.serve.visual_gen_utils import parse_visual_gen_params from tensorrt_llm.version import __version__ as VERSION -from tensorrt_llm.visual_gen import VisualGen from .._utils import nvtx_mark, set_prometheus_multiproc_dir from .harmony_adapter import HarmonyAdapter, get_harmony_adapter +if TYPE_CHECKING: + from tensorrt_llm.visual_gen import VisualGen + + +def _is_visual_gen_instance(obj) -> bool: + """isinstance(obj, VisualGen) without importing the visual_gen tree. + + If tensorrt_llm.visual_gen was never imported in this process, obj + cannot be a VisualGen instance, so the LLM serving path never pays + the visual_gen import cost. + """ + visual_gen = sys.modules.get("tensorrt_llm.visual_gen") + return visual_gen is not None and isinstance(obj, visual_gen.VisualGen) + # yapf: enable # msgspec msgpack is an opt-in transport for the disagg orchestrator->worker @@ -259,7 +273,7 @@ def _iteration_stats_buffer_maxlen( def __init__( self, - generator: Union[LLM, MultimodalEncoder, VisualGen], + generator: Union[LLM, MultimodalEncoder, "VisualGen"], model: str, tool_parser: Optional[str], server_role: Optional[ServerRole], @@ -274,7 +288,7 @@ def __init__( media_load_workers: int = 8, internal_disagg_auth_key: Optional[str] = None): self.generator = generator - self._is_visual_gen = isinstance(generator, VisualGen) + self._is_visual_gen = _is_visual_gen_instance(generator) self._embedding_max_queue_delay = embedding_max_queue_delay self._embedding_max_queue_size = embedding_max_queue_size self.embedding_batcher: Optional[EncodeBatcher] = None @@ -391,7 +405,7 @@ async def lifespan(app: FastAPI): self.disagg_cluster_config, self.disagg_cluster_storage) # VisualGen has no args - if not isinstance(self.generator, VisualGen): + if not self._is_visual_gen: # Start energy monitoring if enabled if getattr(self.generator.args, "enable_energy_metrics", False): try: @@ -478,9 +492,8 @@ async def validation_exception_handler(_, exc): return JSONResponse(status_code=400, content={"error": str(exc)}) if self.server_role is ServerRole.VISUAL_GEN: - assert isinstance( - self.generator, VisualGen - ), "generator must be a VisualGen for VISUAL_GEN server" + assert self._is_visual_gen, \ + "generator must be a VisualGen for VISUAL_GEN server" self.register_visual_gen_routes() elif self.server_role is ServerRole.MM_ENCODER: assert isinstance( @@ -754,7 +767,7 @@ async def await_disconnected(self, raw_request: Request, promise): @property def postproc_worker_enabled(self) -> bool: - if isinstance(self.generator, VisualGen): + if self._is_visual_gen: return False return True if self.generator.args.num_postprocess_workers > 0 else False diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 3601c6f9624b..a409e8f701ff 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -10,6 +10,8 @@ is strictly unchanged from the inlined version. """ +from __future__ import annotations + import asyncio import base64 import json @@ -19,6 +21,7 @@ import uuid from http import HTTPStatus from pathlib import Path +from typing import TYPE_CHECKING from fastapi import Request from fastapi.responses import FileResponse, JSONResponse, Response @@ -30,7 +33,11 @@ from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest, VideoJob, VideoJobList from tensorrt_llm.serve.visual_gen_metrics import build_visual_gen_timing_headers from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE, parse_visual_gen_params -from tensorrt_llm.visual_gen.params import VisualGenParams + +if TYPE_CHECKING: + # Type-only: importing tensorrt_llm.visual_gen at runtime would pull the + # whole visual_gen tree into every LLM serving process. + from tensorrt_llm.visual_gen.params import VisualGenParams def _video_content_type(suffix: str) -> str: diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 31dee1f07d6d..afd5d4d6e5a1 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -1,12 +1,18 @@ +from __future__ import annotations + import asyncio import base64 import os -from typing import Any, Dict, List, Optional +from typing import TYPE_CHECKING, Any, Dict, List, Optional from tensorrt_llm.inputs.media_io import is_isobmff_image_bytes, sniff_media_kind from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest -from tensorrt_llm.visual_gen import VisualGen, VisualGenParams + +if TYPE_CHECKING: + # Type-only: importing tensorrt_llm.visual_gen at runtime would pull the + # whole visual_gen tree into every LLM serving process. + from tensorrt_llm.visual_gen import VisualGen, VisualGenParams # Per-field warnings for OpenAI-shaped knobs that the engine has no # semantic for. Each entry maps the request attribute to the message From d553db65c32865adb8565031d1d430fed001d7cc Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Wed, 5 Aug 2026 14:58:34 +0800 Subject: [PATCH 5/6] [None][fix] Harden the lazy model zoo: resolver API, registration priority, multimodal registry Review follow-ups on the lazy-loading change, squashed. Resolver API. Architecture lookups collapse into a single entry point that imports the built-in provider on demand: get_registered_model_class and get_registered_vision_encoder (ensure_model_registered becomes the private _ensure_model_registered). Each resolver short-circuits on its own registry only, so an external registration satisfies the model-class lookup without importing the built-in module, while a vision-encoder or placeholder lookup with an empty slot still pulls the provider in. All consumers (modeling_auto, model_loader's transceiver preference, get_model_architecture, the cache-transceiver precheck, test_config_database) go through the resolvers instead of pairing a manual import with a raw MODEL_CLASS_MAPPING.get(). Registration priority. One rule, applied by every registry (model class, vision encoder, placeholder metadata): built-in registrations only fill empty slots and never overwrite, external registrations always win. This is equivalent to main's eager order -- built-ins ran first there, and no architecture is double-registered among built-ins -- and needs no registrant bookkeeping. register_auto_model records the declared architectures on the class itself so register_vision_encoder no longer scans the mapping by identity (an external override used to make the built-in provider's import raise). Multimodal registry. The placeholder registry resolves its provider module on demand through a new static MULTIMODAL_MODEL_TYPE_TO_MODULE index, so model_type-keyed queries (trtllm-bench dataset prep, quickstart) work in a fresh process that never loads a model; enumeration APIs import all indexed providers first. Error handling. Lazy-import failures only swallow "the requested module itself does not exist" (both the model zoo and the PEP 562 package fallbacks); a missing dependency inside an existing module propagates instead of surfacing as unknown architecture or a missing attribute. The Llama4 check in py_executor probes sys.modules and uses a real isinstance; the VisualGen endpoint tests patch the probe helper. Drops the SomeVLModel index entry picked up from a docstring example. Adds fresh-process tests pinning the lazy contracts: import stays thin, attribute access resolves and caches, the static index matches the registration decorators, external registrations survive built-in provider imports (including the Qwen3VL vision-encoder path), and placeholder lookups resolve without a loaded model. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- tensorrt_llm/__init__.py | 7 +- tensorrt_llm/_torch/models/__init__.py | 14 +- tensorrt_llm/_torch/models/_arch_index.py | 47 ++- tensorrt_llm/_torch/models/modeling_auto.py | 21 +- .../_torch/models/modeling_mistral.py | 3 +- tensorrt_llm/_torch/models/modeling_utils.py | 126 ++++++- .../_torch/pyexecutor/model_loader.py | 10 +- tensorrt_llm/_torch/pyexecutor/py_executor.py | 17 +- tensorrt_llm/executor/proxy.py | 3 +- tensorrt_llm/executor/worker.py | 3 +- tensorrt_llm/inputs/registry.py | 79 +++- .../run_precheck.py | 8 +- .../visual_gen/test_trtllm_serve_endpoints.py | 11 +- tests/unittest/llmapi/test_config_database.py | 4 +- tests/unittest/others/test_lazy_model_zoo.py | 343 ++++++++++++++++++ 15 files changed, 631 insertions(+), 65 deletions(-) create mode 100644 tests/unittest/others/test_lazy_model_zoo.py diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index c73f7e628f26..a1ed59a0dea1 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -194,7 +194,12 @@ def __getattr__(name): # that used to be reachable as attributes via the eager import chain. try: return importlib.import_module(f'.{name}', __name__) - except ModuleNotFoundError: + except ModuleNotFoundError as e: + # Only translate "no such submodule" into AttributeError. A + # ModuleNotFoundError raised *inside* an existing submodule (missing + # dependency) is a real error and must propagate unchanged. + if e.name != f'{__name__}.{name}': + raise raise AttributeError( f"module {__name__!r} has no attribute {name!r}") from None diff --git a/tensorrt_llm/_torch/models/__init__.py b/tensorrt_llm/_torch/models/__init__.py index 9a0ecd5e546e..4a8d5199b09f 100644 --- a/tensorrt_llm/_torch/models/__init__.py +++ b/tensorrt_llm/_torch/models/__init__.py @@ -10,7 +10,8 @@ ``from tensorrt_llm._torch.models import LlamaForCausalLM``) imports just the providing module, via PEP 562 ``__getattr__`` and ``MODEL_CLASS_TO_MODULE``; - architecture-based resolution (``AutoModelForCausalLM``) imports on demand - via ``modeling_utils.ensure_model_registered`` and ``MODEL_ARCH_TO_MODULE``. + via ``modeling_utils.get_registered_model_class`` and + ``MODEL_ARCH_TO_MODULE``. """ import importlib @@ -100,7 +101,16 @@ def __getattr__(name: str): # Also resolve bare submodule access (models.modeling_llama) so callers # that relied on the previously-eager submodule attributes keep working. if name.startswith("modeling_"): - return importlib.import_module(f".{name}", __name__) + try: + return importlib.import_module(f".{name}", __name__) + except ModuleNotFoundError as e: + # Only translate "no such submodule" into AttributeError; a + # ModuleNotFoundError raised *inside* an existing submodule + # (missing dependency) must propagate unchanged. + if e.name != f"{__name__}.{name}": + raise + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None raise AttributeError(f"module {__name__!r} has no attribute {name!r}") module = importlib.import_module(f".{module_name}", __name__) attr = getattr(module, name) diff --git a/tensorrt_llm/_torch/models/_arch_index.py b/tensorrt_llm/_torch/models/_arch_index.py index b62600adfedd..10eea7fc9f20 100644 --- a/tensorrt_llm/_torch/models/_arch_index.py +++ b/tensorrt_llm/_torch/models/_arch_index.py @@ -5,12 +5,14 @@ Model implementations register themselves via ``@register_auto_model`` as an import side effect. The zoo is imported lazily, so these tables record, without importing anything, which ``modeling_*`` module provides which architecture -(``MODEL_ARCH_TO_MODULE``) and which public class (``MODEL_CLASS_TO_MODULE``). +(``MODEL_ARCH_TO_MODULE``), which public class (``MODEL_CLASS_TO_MODULE``), +and which multimodal ``model_type`` (``MULTIMODAL_MODEL_TYPE_TO_MODULE``). Regenerate after adding/moving a model: - scan ``@register_auto_model("")`` decorators and the public class - names under ``tensorrt_llm/_torch/models`` (see the lazy-import PR for the - one-off script), or add the new entry by hand next to its neighbors. + scan ``@register_auto_model("")`` / ``@register_input_processor(..., + model_type="")`` decorators and the public class names under + ``tensorrt_llm/_torch/models`` (see the lazy-import PR for the one-off + script), or add the new entry by hand next to its neighbors. """ # Architecture name (HF ``config.architectures[0]``, possibly rewritten by @@ -90,7 +92,6 @@ "QwenImageBenchForConditionalGeneration": "modeling_qwen_image_bench", "SeedOssForCausalLM": "modeling_seedoss", "SiglipVisionModel": "modeling_siglip", - "SomeVLModel": "modeling_utils", "Starcoder2ForCausalLM": "modeling_starcoder2", "Step3p5ForCausalLM": "modeling_step3p7", "Step3p7ForConditionalGeneration": "modeling_step3p7vl", @@ -165,3 +166,39 @@ "VilaModel": "modeling_vila", "WhisperForConditionalGeneration": "modeling_whisper", } + +# Multimodal ``model_type`` (HF ``config.model_type``, as passed to +# ``register_input_processor`` / ``set_placeholder_metadata``) -> providing +# module. Lets the multimodal placeholder registry resolve a model type +# without eagerly importing the zoo. ``qwen3_5`` is also registered by +# ``modeling_qwen_image_bench`` with byte-identical metadata; the real model's +# module is indexed. +MULTIMODAL_MODEL_TYPE_TO_MODULE = { + "NemotronH_Nano_Omni_Reasoning_V3": "modeling_nemotron_nano", + "NemotronH_Nano_VL_V2": "modeling_nemotron_nano", + "cosmos3": "modeling_cosmos3", + "cosmos3_omni": "modeling_cosmos3", + "exaone4_5": "modeling_exaone4_5", + "gemma3": "modeling_gemma3vl", + "gemma4": "modeling_gemma4mm", + "gemma4_unified": "modeling_gemma4_unified", + "hyperclovax_vlm": "modeling_hyperclovax", + "kimi_k25": "modeling_kimi_k25", + "llama4": "modeling_llama", + "llava_llama": "modeling_vila", + "llava_next": "modeling_llava_next", + "minicpmv4_6": "modeling_minicpmv4_6", + "minimax_m3_vl": "modeling_minimaxm3", + "mistral3": "modeling_mistral", + "mistral_common": "modeling_mistral", + "mistral_large_3": "modeling_mistral", + "phi4mm": "modeling_phi4mm", + "qwen2_5_vl": "modeling_qwen2vl", + "qwen2_vl": "modeling_qwen2vl", + "qwen3_5": "modeling_qwen3_5", + "qwen3_5_moe": "modeling_qwen3_5", + "qwen3_vl": "modeling_qwen3vl", + "qwen3_vl_moe": "modeling_qwen3vl_moe", + "step3p7": "modeling_step3p7vl", + "whisper": "modeling_whisper", +} diff --git a/tensorrt_llm/_torch/models/modeling_auto.py b/tensorrt_llm/_torch/models/modeling_auto.py index 695c6261d9b2..14fd37043614 100644 --- a/tensorrt_llm/_torch/models/modeling_auto.py +++ b/tensorrt_llm/_torch/models/modeling_auto.py @@ -2,10 +2,9 @@ from ..model_config import ModelConfig from ..utils import model_extra_attrs -from .modeling_utils import (MODEL_CLASS_MAPPING, - MODEL_CLASS_VISION_ENCODER_MAPPING, - DecoderModelForCausalLM, TConfig, TModel, - ensure_model_registered) +from .modeling_utils import (DecoderModelForCausalLM, TConfig, TModel, + get_registered_model_class, + get_registered_vision_encoder) class AutoModelForCausalLM(Generic[TModel, TConfig]): @@ -18,13 +17,9 @@ def _resolve_class(config: ModelConfig) -> Optional[Type]: return None model_arch = pretrained_config.architectures[0] - # The model zoo is imported lazily: pull in the module providing this - # architecture so its registration decorators have run. - ensure_model_registered(model_arch) if config.mm_encoder_only: - vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get( - model_arch) + vision_encoder_info = get_registered_vision_encoder(model_arch) if vision_encoder_info is None: return None vision_encoder_cls, _ = vision_encoder_info @@ -37,16 +32,14 @@ def _resolve_class(config: ModelConfig) -> Optional[Type]: model_arch = model_arch.replace("Eagle3", "") # Strip the appended EAGLE3 model_arch = "EAGLE3" + model_arch - ensure_model_registered(model_arch) if model_arch in ( "DeepseekV3ForCausalLM", "Glm4MoeForCausalLM", "ExaoneMoEForCausalLM" ) and config.spec_config is not None and config.spec_config.max_draft_len == 0: model_arch = "MTPDraftModelForCausalLM" - ensure_model_registered(model_arch) - return MODEL_CLASS_MAPPING.get(model_arch) + return get_registered_model_class(model_arch) @staticmethod def from_config( @@ -54,9 +47,7 @@ def from_config( ) -> DecoderModelForCausalLM[TModel, TConfig]: if config.mm_encoder_only: model_arch = config.pretrained_config.architectures[0] - ensure_model_registered(model_arch) - vision_encoder_info = MODEL_CLASS_VISION_ENCODER_MAPPING.get( - model_arch) + vision_encoder_info = get_registered_vision_encoder(model_arch) if vision_encoder_info is None: raise ValueError( f"Unknown architecture for AutoModelForMultimodalEncoder: {model_arch}" diff --git a/tensorrt_llm/_torch/models/modeling_mistral.py b/tensorrt_llm/_torch/models/modeling_mistral.py index 55c4633cc4bf..78c7ec0561df 100644 --- a/tensorrt_llm/_torch/models/modeling_mistral.py +++ b/tensorrt_llm/_torch/models/modeling_mistral.py @@ -675,7 +675,8 @@ def call_with_text_prompt( placeholder_map={"image": "[IMG]"}, placeholder_placement=MultimodalPlaceholderPlacement.BEFORE_TEXT, content_format=ContentFormat.PASSTHROUGH, - )) + ), + registrant_module=__name__) MistralNativeInputProcessor._registered_model_type = "mistral_common" diff --git a/tensorrt_llm/_torch/models/modeling_utils.py b/tensorrt_llm/_torch/models/modeling_utils.py index 57e714dceff1..4acfdc386e04 100755 --- a/tensorrt_llm/_torch/models/modeling_utils.py +++ b/tensorrt_llm/_torch/models/modeling_utils.py @@ -857,36 +857,116 @@ def infer_max_seq_len(self) -> int: CHECKPOINT_LOADER_FORMAT_DEFAULT_MAPPING = {} +# Registration priority under lazy loading: on main the built-in zoo imported +# first and external code (--custom_module_dirs, user modules) overrode it +# later. With the zoo imported lazily, built-in modules may run their +# decorators *after* an external registration, so every registry applies one +# rule: built-in registrations only fill empty slots, never overwrite. +# Anything already present outranks a built-in — it is either an external +# registration (which must keep its main-order priority) or another built-in +# (no architecture is double-registered among built-ins, so filling empty +# slots is equivalent to main). External registrations always overwrite. +def _is_builtin_model_class(cls) -> bool: + return getattr(cls, "__module__", + "").startswith("tensorrt_llm._torch.models") + + +# Architecture names each decorated class declared via ``register_auto_model``, +# kept per class (``cls.__dict__``, never inherited). Recorded even when the +# class loses the ``MODEL_CLASS_MAPPING`` slot to an external registration, so +# stacked decorators (``register_vision_encoder``) can still map the class to +# its architectures instead of scanning the mapping by identity. +_REGISTERED_ARCHS_ATTR = "_registered_architectures" + + def register_auto_model(name: str): def decorator(cls): + archs = cls.__dict__.get(_REGISTERED_ARCHS_ATTR) + if archs is None: + archs = set() + setattr(cls, _REGISTERED_ARCHS_ATTR, archs) + archs.add(name) + + existing = MODEL_CLASS_MAPPING.get(name) + if (existing is not None and existing is not cls + and _is_builtin_model_class(cls)): + logger.info( + f"Keeping existing registration " + f"{existing.__module__}.{existing.__name__} for architecture " + f"{name}; built-in {cls.__module__}.{cls.__name__} not " + f"registered.") + return cls MODEL_CLASS_MAPPING[name] = cls return cls return decorator -def ensure_model_registered(model_arch: str) -> None: +def _ensure_model_registered(model_arch: str) -> None: """Import the module that provides ``model_arch``, if it isn't loaded yet. Model implementations register themselves in ``MODEL_CLASS_MAPPING`` (and the sibling registries) as an import side effect. With the model zoo imported lazily, this is the hook that turns an architecture name into "the decorators have run". Architectures missing from the static index - (e.g. registered dynamically by user code) and modules that fail to import - are left to the caller's normal missing-architecture handling; the warning - keeps the root cause visible. + (e.g. registered dynamically by user code) are left to the caller's + normal missing-architecture handling. + + Internal: consumers go through ``get_registered_model_class`` / + ``get_registered_vision_encoder`` instead of pairing this with a raw + registry read. Each resolver short-circuits on *its own* registry only: + an external registration satisfies the model-class lookup without + importing the built-in provider, but a lookup in a sibling registry + (vision encoder, placeholder metadata) still triggers the import when + its slot is empty. Priority on that import is enforced inside each + registration decorator; the import itself is idempotent via + ``sys.modules``. """ module_name = MODEL_ARCH_TO_MODULE.get(model_arch) if module_name is None: return + full_name = f"tensorrt_llm._torch.models.{module_name}" try: - importlib.import_module(f"tensorrt_llm._torch.models.{module_name}") - except ImportError as e: + importlib.import_module(full_name) + except ModuleNotFoundError as e: + # Only swallow "the providing module itself is missing" (stale index + # entry); a missing dependency *inside* the module is a real error + # and must not be masked as "unknown architecture". + if e.name != full_name: + raise logger.warning(f"Lazy import of {module_name} for architecture " f"{model_arch} failed: {e!r}") +def get_registered_model_class(model_arch: str) -> Optional[Type[nn.Module]]: + """Resolve ``model_arch`` to its registered model class, or ``None``. + + The single entry point for architecture lookups: the model zoo is + imported lazily, so this resolves the built-in provider on demand before + reading the registry. Do not read ``MODEL_CLASS_MAPPING`` directly for + lookups — a raw ``.get()`` silently misses every not-yet-imported + built-in model. + """ + if model_arch not in MODEL_CLASS_MAPPING: + _ensure_model_registered(model_arch) + return MODEL_CLASS_MAPPING.get(model_arch) + + +def get_registered_vision_encoder( + model_arch: str) -> Optional[Tuple[Type[nn.Module], Optional[Type]]]: + """Resolve ``model_arch`` to its ``(vision_encoder_cls, vlm_base_model)``. + + Same on-demand resolution as ``get_registered_model_class``, for the + vision-encoder sibling registry: the provider import triggers when + *this* registry misses, even if an external class holds the + model-class slot. + """ + if model_arch not in MODEL_CLASS_VISION_ENCODER_MAPPING: + _ensure_model_registered(model_arch) + return MODEL_CLASS_VISION_ENCODER_MAPPING.get(model_arch) + + def register_vision_encoder( vision_encoder_cls: Type[nn.Module], vlm_base_model: Optional[Type[nn.Module]] = None, @@ -903,17 +983,34 @@ class SomeVLModel(...): """ def wrapper(model_cls: Type[nn.Module]) -> Type[nn.Module]: - registered = False - for arch_name, registered_cls in MODEL_CLASS_MAPPING.items(): - if registered_cls is model_cls: - MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = ( - vision_encoder_cls, vlm_base_model) - registered = True - if not registered: + # The architectures this class declared via register_auto_model. Do + # not scan MODEL_CLASS_MAPPING by identity: a built-in class may have + # lost its mapping slot to an external registration, and its module + # must still import cleanly. + archs = model_cls.__dict__.get(_REGISTERED_ARCHS_ATTR) + if not archs: + # Fallback for classes placed into the mapping directly instead + # of via the register_auto_model decorator. + archs = { + arch_name + for arch_name, registered_cls in MODEL_CLASS_MAPPING.items() + if registered_cls is model_cls + } + if not archs: raise ValueError( f"register_vision_encoder: model class {model_cls.__name__} is not registered " f"via register_auto_model; decorator order must ensure registration occurs first." ) + for arch_name in archs: + if (arch_name in MODEL_CLASS_VISION_ENCODER_MAPPING + and _is_builtin_model_class(model_cls)): + # Built-in registrations only fill empty slots (see the + # priority rule above register_auto_model). + logger.info(f"Keeping existing vision encoder registration for " + f"architecture {arch_name}.") + continue + MODEL_CLASS_VISION_ENCODER_MAPPING[arch_name] = (vision_encoder_cls, + vlm_base_model) return model_cls @@ -985,8 +1082,7 @@ def get_model_architecture( cls = None if model_config.architectures is not None and len( model_config.architectures) > 0: - ensure_model_registered(model_config.architectures[0]) - cls = MODEL_CLASS_MAPPING.get(model_config.architectures[0]) + cls = get_registered_model_class(model_config.architectures[0]) else: raise RuntimeError("Model architecture is not provided.") diff --git a/tensorrt_llm/_torch/pyexecutor/model_loader.py b/tensorrt_llm/_torch/pyexecutor/model_loader.py index 8bfb748093e0..7a852bf9bfd1 100644 --- a/tensorrt_llm/_torch/pyexecutor/model_loader.py +++ b/tensorrt_llm/_torch/pyexecutor/model_loader.py @@ -36,9 +36,8 @@ from ..model_config import ModelConfig from ..models import AutoModelForCausalLM from ..models.checkpoints.base_checkpoint_loader import BaseCheckpointLoader -from ..models.modeling_utils import (MODEL_CLASS_MAPPING, - DecoderModelForCausalLM, MetaInitMode, - ensure_model_registered, timing) +from ..models.modeling_utils import (DecoderModelForCausalLM, MetaInitMode, + get_registered_model_class, timing) from ..modules.fused_moe.moe_load_balancer import ( MoeLoadBalancer, maybe_create_moe_load_balancer) from ..virtual_memory import RestoreMode @@ -460,9 +459,8 @@ def load_config_and_apply_defaults( preference_cls = model_cls architectures = getattr(config.pretrained_config, 'architectures', None) if architectures: - ensure_model_registered(architectures[0]) - preference_cls = MODEL_CLASS_MAPPING.get(architectures[0], - model_cls) + preference_cls = get_registered_model_class( + architectures[0]) or model_cls # Resolve "auto" sentinel values after model defaults are applied. _resolve_transceiver_runtime_auto(llm_args, preference_cls, diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index c0c353c171f4..b928bc309834 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -5,6 +5,7 @@ import datetime import functools import os +import sys import threading import time import traceback @@ -4881,12 +4882,16 @@ def _forward_step_inter_pp(self, def _validate_token_id_range(self, request: LlmRequest) -> None: if isinstance(self.model_engine.model, DecoderModelForCausalLM): - # Only skip token‐range checks for Llama4 when the request has multimodal data. - # Matched by class name (equivalent to isinstance incl. subclasses) - # so this module does not import a model-zoo module at startup, - # which would defeat the zoo's lazy loading. - if any(c.__name__ == "Llama4ForConditionalGeneration" - for c in type(self.model_engine.model).__mro__): + # Only skip token‐range checks for Llama4 when the request has + # multimodal data. Probed via sys.modules so this module does not + # import a model-zoo module at startup (which would defeat the + # zoo's lazy loading): if modeling_llama was never imported, the + # engine's model cannot be a Llama4 instance. + modeling_llama = sys.modules.get( + "tensorrt_llm._torch.models.modeling_llama") + if modeling_llama is not None and isinstance( + self.model_engine.model, + modeling_llama.Llama4ForConditionalGeneration): has_mm = bool(request.py_multimodal_data) if has_mm: logger.debug( diff --git a/tensorrt_llm/executor/proxy.py b/tensorrt_llm/executor/proxy.py index 32afb4780ed1..eed0ce6ccccb 100644 --- a/tensorrt_llm/executor/proxy.py +++ b/tensorrt_llm/executor/proxy.py @@ -616,7 +616,8 @@ def mpi_done_callback(future: concurrent.futures.Future): tracer_init_kwargs = get_tracer().init_kwargs if enable_llm_tracer( ) else None - from tensorrt_llm._torch.models.modeling_auto import MODEL_CLASS_MAPPING + from tensorrt_llm._torch.models.modeling_utils import \ + MODEL_CLASS_MAPPING torch.cuda.Stream() # Strip the tokenizer from worker_kwargs to avoid MPI pickle failures. diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py index f173aabc224a..5c0b5758073f 100644 --- a/tensorrt_llm/executor/worker.py +++ b/tensorrt_llm/executor/worker.py @@ -227,7 +227,8 @@ def _print_stacks(): set_global_tracer(tracer) if _torch_model_class_mapping is not None: - from tensorrt_llm._torch.models.modeling_auto import MODEL_CLASS_MAPPING + from tensorrt_llm._torch.models.modeling_utils import \ + MODEL_CLASS_MAPPING MODEL_CLASS_MAPPING.update(**_torch_model_class_mapping) set_mpi_session_cpp(mpi_comm()) diff --git a/tensorrt_llm/inputs/registry.py b/tensorrt_llm/inputs/registry.py index d4f2f96a4561..eca8fdddeedd 100644 --- a/tensorrt_llm/inputs/registry.py +++ b/tensorrt_llm/inputs/registry.py @@ -15,6 +15,7 @@ # SPDX-License-Identifier: Apache-2.0 import enum +import importlib import traceback from abc import ABC, abstractmethod from dataclasses import dataclass, field @@ -728,11 +729,14 @@ class MultimodalPlaceholderRegistry: Registry for the multimodal models to keep track of the placeholder information. """ + _BUILTIN_MODULE_PREFIX = "tensorrt_llm._torch.models" + def __init__(self) -> None: self._multimodal_placeholder_by_model_type: Dict[ str, MultimodalPlaceholderMetadata] = {} def __str__(self) -> str: + self._ensure_all_providers_imported() s = "" for model_type, placeholder_metadata in self._multimodal_placeholder_by_model_type.items( ): @@ -745,22 +749,81 @@ def __str__(self) -> str: return s def set_placeholder_metadata( - self, model_type: str, - placeholder_metadata: MultimodalPlaceholderMetadata): + self, + model_type: str, + placeholder_metadata: MultimodalPlaceholderMetadata, + registrant_module: Optional[str] = None): + """Register placeholder metadata for ``model_type``. + + ``registrant_module`` identifies the registering module. Built-in + registrations (under the lazily imported model zoo) only fill empty + slots: they may run *after* an external implementation claimed the + model type and must not clobber it (same priority rule as the model + class registry). + """ + if (model_type in self._multimodal_placeholder_by_model_type + and registrant_module is not None + and registrant_module.startswith(self._BUILTIN_MODULE_PREFIX)): + logger.info(f"Keeping existing placeholder metadata for model " + f"type {model_type}.") + return self._multimodal_placeholder_by_model_type[ model_type] = placeholder_metadata + def _ensure_provider_imported(self, model_type: str) -> None: + """Import the modeling module that registers ``model_type``, if needed. + + Registration is an import side effect of the (lazily imported) model + zoo, so point lookups resolve their provider on demand; model types + registered dynamically (or already imported) are returned as-is. + Imports are function-local: modeling modules import this module, so a + top-level import of the zoo index would be circular. + """ + if model_type in self._multimodal_placeholder_by_model_type: + return + from tensorrt_llm._torch.models._arch_index import \ + MULTIMODAL_MODEL_TYPE_TO_MODULE + module_name = MULTIMODAL_MODEL_TYPE_TO_MODULE.get(model_type) + if module_name is None: + return + full_name = f"tensorrt_llm._torch.models.{module_name}" + try: + importlib.import_module(full_name) + except ModuleNotFoundError as e: + # Only swallow "the providing module itself is missing" (stale + # index entry); a missing dependency inside the module is a real + # error and must not be masked as an unregistered model type. + if e.name != full_name: + raise + logger.warning(f"Lazy import of {module_name} for model type " + f"{model_type} failed: {e!r}") + + def _ensure_all_providers_imported(self) -> None: + """Import every indexed multimodal provider (for enumeration APIs). + + Listing registered model types is only meaningful once all providers + have run their registration side effects; with the zoo lazy this + imports just the multimodal modeling modules, on first enumeration. + """ + from tensorrt_llm._torch.models._arch_index import \ + MULTIMODAL_MODEL_TYPE_TO_MODULE + for model_type in MULTIMODAL_MODEL_TYPE_TO_MODULE: + self._ensure_provider_imported(model_type) + def remove_placeholder_metadata(self, model_type: str): + self._ensure_provider_imported(model_type) if model_type not in self._multimodal_placeholder_by_model_type: raise ValueError(f"Model type '{model_type}' is not registered") del self._multimodal_placeholder_by_model_type[model_type] def is_valid(self, model_type: str, modality: str) -> bool: + self._ensure_provider_imported(model_type) return model_type in self._multimodal_placeholder_by_model_type and \ modality in self._multimodal_placeholder_by_model_type[model_type].placeholder_map def get_placeholder_metadata( self, model_type: str) -> MultimodalPlaceholderMetadata: + self._ensure_provider_imported(model_type) if model_type not in self._multimodal_placeholder_by_model_type: raise ValueError( f"Model type {model_type} is not registered in MultimodalPlaceholderRegistry" @@ -777,12 +840,14 @@ def get_placeholder(self, model_type: str, modality: str) -> str: def get_placeholder_placement( self, model_type: str) -> MultimodalPlaceholderPlacement: + self._ensure_provider_imported(model_type) if model_type not in self._multimodal_placeholder_by_model_type: raise ValueError(f"Model type '{model_type}' is not registered") return self._multimodal_placeholder_by_model_type[ model_type].placeholder_placement def get_placeholders_separator(self, model_type: str) -> str: + self._ensure_provider_imported(model_type) if model_type not in self._multimodal_placeholder_by_model_type: raise ValueError(f"Model type '{model_type}' is not registered") return self._multimodal_placeholder_by_model_type[ @@ -790,6 +855,7 @@ def get_placeholders_separator(self, model_type: str) -> str: def get_interleave_placeholders(self, model_type: str) -> bool: """Return whether the model opts in to interleaved placeholder insertion.""" + self._ensure_provider_imported(model_type) if model_type not in self._multimodal_placeholder_by_model_type: return False return self._multimodal_placeholder_by_model_type[ @@ -797,12 +863,14 @@ def get_interleave_placeholders(self, model_type: str) -> bool: def get_content_format(self, model_type: str) -> Optional[ContentFormat]: """Get the content format override for a model type, or None for auto-detect.""" + self._ensure_provider_imported(model_type) if model_type not in self._multimodal_placeholder_by_model_type: return None return self._multimodal_placeholder_by_model_type[ model_type].content_format def get_registered_image_model_types(self) -> Tuple[str, ...]: + self._ensure_all_providers_imported() return ( model_type for model_type in self._multimodal_placeholder_by_model_type @@ -810,6 +878,7 @@ def get_registered_image_model_types(self) -> Tuple[str, ...]: _multimodal_placeholder_by_model_type[model_type].placeholder_map) def get_registered_video_model_types(self) -> Tuple[str, ...]: + self._ensure_all_providers_imported() return ( model_type for model_type in self._multimodal_placeholder_by_model_type @@ -817,6 +886,7 @@ def get_registered_video_model_types(self) -> Tuple[str, ...]: _multimodal_placeholder_by_model_type[model_type].placeholder_map) def get_registered_audio_model_types(self) -> Tuple[str, ...]: + self._ensure_all_providers_imported() return ( model_type for model_type in self._multimodal_placeholder_by_model_type @@ -824,6 +894,7 @@ def get_registered_audio_model_types(self) -> Tuple[str, ...]: _multimodal_placeholder_by_model_type[model_type].placeholder_map) def get_registered_model_types(self) -> Tuple[str, ...]: + self._ensure_all_providers_imported() return tuple(self._multimodal_placeholder_by_model_type.keys()) @@ -894,7 +965,9 @@ def wrapper(model_cls: N) -> N: ) MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( - model_type, placeholder_metadata) + model_type, + placeholder_metadata, + registrant_module=model_cls.__module__) # Expose the registered model_type on the processor class so callers # can look it up without re-deriving it from the HF config. diff --git a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py index 19c6bf70c856..146c40173cc3 100644 --- a/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py +++ b/tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py @@ -196,12 +196,11 @@ def load_internal_apis(): import types import tensorrt_llm - import tensorrt_llm._torch.models # noqa: F401 - populates the model registry import tensorrt_llm.bindings import tensorrt_llm.bindings.executor as trtllm_executor from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.distributed import Distributed - from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING + from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm._torch.pyexecutor.hang_detector import HangDetector from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver @@ -232,7 +231,7 @@ def load_internal_apis(): KvCacheConfigCpp=trtllm_executor.KvCacheConfig, DisaggregatedParams=DisaggregatedParams, Distributed=Distributed, - MODEL_CLASS_MAPPING=MODEL_CLASS_MAPPING, + get_registered_model_class=get_registered_model_class, HangDetector=HangDetector, KVCacheManager=KVCacheManager, KVCacheManagerV2=KVCacheManagerV2, @@ -314,7 +313,8 @@ def _lookup_model_cls(model_dir): hf_view = type("HFConfigView", (), hf_cfg) # attribute access for the pref hook if not archs: return None, hf_view - return load_internal_apis().MODEL_CLASS_MAPPING.get(archs[0]), hf_view + # Resolves the lazily imported provider on demand, like serving does. + return load_internal_apis().get_registered_model_class(archs[0]), hf_view def resolve_model_prefs(model_dir, side, cache_cfg): diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index 69d17ea523ef..ae2b2410d75e 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -310,13 +310,18 @@ def result(self, timeout=None): def _create_server(generator: MockVisualGen, model_name: str = "test-model") -> TestClient: """Instantiate an OpenAIServer for VISUAL_GEN with a mocked generator. - We patch the ``VisualGen`` name inside the ``openai_server`` module so that - ``isinstance(generator, VisualGen)`` returns True for our mock. + The server detects VisualGen generators via ``_is_visual_gen_instance`` + (a sys.modules probe, so plain LLM serving never imports visual_gen) and + caches the result in ``__init__``; patching the probe during construction + makes it recognize our mock. """ from tensorrt_llm.llmapi.disagg_utils import ServerRole from tensorrt_llm.serve.openai_server import OpenAIServer - with patch("tensorrt_llm.serve.openai_server.VisualGen", MockVisualGen): + with patch( + "tensorrt_llm.serve.openai_server._is_visual_gen_instance", + return_value=True, + ): server = OpenAIServer( generator=generator, model=model_name, diff --git a/tests/unittest/llmapi/test_config_database.py b/tests/unittest/llmapi/test_config_database.py index 3c15db7e410d..b443c231e15e 100644 --- a/tests/unittest/llmapi/test_config_database.py +++ b/tests/unittest/llmapi/test_config_database.py @@ -24,7 +24,7 @@ from fastapi.testclient import TestClient from starlette.middleware.base import BaseHTTPMiddleware -from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING +from tensorrt_llm._torch.models.modeling_utils import get_registered_model_class from tensorrt_llm.commands.serve import _apply_fastapi_middlewares from tensorrt_llm.commands.serve import main as serve_main from tensorrt_llm.llmapi import llm_args as llm_args_module @@ -104,7 +104,7 @@ def _get_default_values_for_config(config_path: Path) -> dict: model = entry["model"] arch = entry["arch"] - model_cls = MODEL_CLASS_MAPPING.get(arch) + model_cls = get_registered_model_class(arch) if not model_cls or not hasattr(model_cls, "get_model_defaults"): return global_default diff --git a/tests/unittest/others/test_lazy_model_zoo.py b/tests/unittest/others/test_lazy_model_zoo.py new file mode 100644 index 000000000000..56408b3fee1a --- /dev/null +++ b/tests/unittest/others/test_lazy_model_zoo.py @@ -0,0 +1,343 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Guards for the lazily imported model zoo and top-level namespace. + +The package surface is loaded via PEP 562 and the PyTorch model zoo resolves +through the static tables in ``_arch_index.py``; registration is an import +side effect that now runs on demand instead of at ``import tensorrt_llm`` +time. These tests pin the three contracts that keep that scheme correct: + +- ``import tensorrt_llm`` stays thin (no model zoo / visual_gen in a fresh + process) while first attribute access still resolves and caches. +- The static index stays in sync with the ``@register_auto_model`` / + ``register_input_processor`` decorators it mirrors (a new model that + forgets its index entry fails here, not at model-load time in production). +- On-demand registration never overrides an existing (e.g. custom, via + ``--custom_module_dirs``) registration, and multimodal placeholder lookups + resolve their provider in a process that never loaded a model. +""" + +import ast +import os +import subprocess +import sys +import textwrap +from pathlib import Path + +_MODELS_DIR = Path(__file__).parents[3] / "tensorrt_llm" / "_torch" / "models" + +_SENTINEL = "LAZY-OK" + + +def _run_fresh(body: str, timeout: int = 300) -> None: + """Run ``body`` in a fresh interpreter and assert it prints the sentinel. + + Fresh subprocess on purpose: the pytest process has long since imported + tensorrt_llm (and, through other tests, parts of the model zoo), so + lazy-import claims are only observable in a new process. + """ + script = body + f"\nprint({_SENTINEL!r})\n" + result = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + timeout=timeout, + env=os.environ.copy(), + ) + assert result.returncode == 0 and _SENTINEL in result.stdout, ( + f"fresh-process check failed\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}" + ) + + +def test_import_does_not_load_zoo_or_visual_gen(): + _run_fresh( + textwrap.dedent("""\ + import sys + import tensorrt_llm + + loaded = [ + m for m in sys.modules + if m.startswith("tensorrt_llm._torch.models.modeling_") + or m == "tensorrt_llm.visual_gen" + or m.startswith("tensorrt_llm.visual_gen.") + ] + assert not loaded, f"import tensorrt_llm eagerly loaded: {loaded}" + """) + ) + + +def test_lazy_attribute_access_resolves_and_caches(): + _run_fresh( + textwrap.dedent("""\ + import tensorrt_llm + + sp = tensorrt_llm.SamplingParams + assert sp is tensorrt_llm.SamplingParams # cached in globals() + assert "SamplingParams" in vars(tensorrt_llm) + assert "SamplingParams" in dir(tensorrt_llm) + + try: + tensorrt_llm.definitely_not_an_attribute + except AttributeError: + pass + else: + raise AssertionError("missing attribute did not raise") + """) + ) + + +def test_placeholder_registry_resolves_in_fresh_process(): + # trtllm-bench dataset prep queries the placeholder registry by + # model_type in a process that never loads a model; the registry must + # import the provider on demand. + _run_fresh( + textwrap.dedent("""\ + from tensorrt_llm.inputs.registry import MULTIMODAL_PLACEHOLDER_REGISTRY + + assert MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid("llama4", "image") + assert MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder( + "llama4", "image") + assert "qwen2_vl" in MULTIMODAL_PLACEHOLDER_REGISTRY.get_registered_model_types() + """) + ) + + +def _decorated_registrations(): + """AST-scan the modeling files for the registrations the index mirrors.""" + arch_to_modules = {} + model_type_to_modules = {} + for path in sorted(_MODELS_DIR.glob("*.py")): + tree = ast.parse(path.read_text()) + for node in ast.walk(tree): + if not isinstance(node, ast.Call): + continue + func = node.func + name = getattr(func, "id", None) or getattr(func, "attr", None) + if name == "register_auto_model": + if node.args and isinstance(node.args[0], ast.Constant): + arch_to_modules.setdefault(node.args[0].value, set()).add(path.stem) + elif name in ("register_input_processor", "set_placeholder_metadata"): + model_type = None + if name == "register_input_processor": + if len(node.args) >= 2 and isinstance(node.args[1], ast.Constant): + model_type = node.args[1].value + elif node.args and isinstance(node.args[0], ast.Constant): + model_type = node.args[0].value + for kw in node.keywords: + if kw.arg == "model_type" and isinstance(kw.value, ast.Constant): + model_type = kw.value.value + if model_type is not None: + model_type_to_modules.setdefault(model_type, set()).add(path.stem) + return arch_to_modules, model_type_to_modules + + +def test_arch_index_matches_decorators(): + from tensorrt_llm._torch.models._arch_index import ( + MODEL_ARCH_TO_MODULE, + MULTIMODAL_MODEL_TYPE_TO_MODULE, + ) + + arch_truth, model_type_truth = _decorated_registrations() + + missing = set(arch_truth) - set(MODEL_ARCH_TO_MODULE) + assert not missing, f"architectures missing from _arch_index: {missing}" + wrong = { + arch: (MODEL_ARCH_TO_MODULE[arch], arch_truth[arch]) + for arch in MODEL_ARCH_TO_MODULE + if arch in arch_truth and MODEL_ARCH_TO_MODULE[arch] not in arch_truth[arch] + } + assert not wrong, f"index points at the wrong module: {wrong}" + + missing = set(model_type_truth) - set(MULTIMODAL_MODEL_TYPE_TO_MODULE) + assert not missing, f"model types missing from _arch_index: {missing}" + stale = set(MULTIMODAL_MODEL_TYPE_TO_MODULE) - set(model_type_truth) + assert not stale, f"stale model types in _arch_index: {stale}" + wrong = { + mt: (MULTIMODAL_MODEL_TYPE_TO_MODULE[mt], model_type_truth[mt]) + for mt in MULTIMODAL_MODEL_TYPE_TO_MODULE + if MULTIMODAL_MODEL_TYPE_TO_MODULE[mt] not in model_type_truth[mt] + } + assert not wrong, f"index points at the wrong module: {wrong}" + + +def test_models_package_missing_submodule_is_attribute_error(): + # The PEP 562 fallback must translate only "no such submodule" into + # AttributeError (so hasattr works), same as the top-level package. + import tensorrt_llm._torch.models as torch_models + + assert not hasattr(torch_models, "modeling_definitely_not_a_model") + + +def test_resolver_keeps_existing_registration(): + # A registration made by user code (--custom_module_dirs) must win over + # the built-in module, exactly as it does with the eager import order on + # main where the zoo loads first and custom code overrides it. + from tensorrt_llm._torch.models.modeling_utils import ( + MODEL_CLASS_MAPPING, + get_registered_model_class, + ) + + arch = "MistralForCausalLM" + + class _CustomStub: + pass + + MODEL_CLASS_MAPPING[arch] = _CustomStub + try: + assert get_registered_model_class(arch) is _CustomStub, ( + "resolving an architecture overrode its existing registration" + ) + finally: + # Restore the real class rather than deleting the entry: the + # built-in decorator skips occupied slots, so if the provider was + # (or gets) imported while the stub held the slot, a bare delete + # would leave the architecture unresolvable for the rest of the + # process (module imports are cached). + from tensorrt_llm._torch.models.modeling_mistral import MistralForCausalLM + + MODEL_CLASS_MAPPING[arch] = MistralForCausalLM + + # Unknown architectures resolve to None, left to the caller's handling. + assert get_registered_model_class("DefinitelyNotARegisteredArch") is None + + +def test_builtin_decorator_does_not_override_external_registration(): + # The priority guarantee lives in the registry itself: a built-in module + # may run its decorators after an external registration (direct imports, + # sibling architectures from the same module), and must not clobber it. + from tensorrt_llm._torch.models.modeling_utils import MODEL_CLASS_MAPPING, register_auto_model + + arch = "LazyZooTestOnlyArch" + assert arch not in MODEL_CLASS_MAPPING + + class _External: + pass + + class _Builtin: + pass + + _Builtin.__module__ = "tensorrt_llm._torch.models.modeling_fake" + + try: + register_auto_model(arch)(_External) + register_auto_model(arch)(_Builtin) + assert MODEL_CLASS_MAPPING[arch] is _External, ( + "built-in decorator overrode an external registration" + ) + + # The other direction stays last-wins: external code registering + # over a built-in is exactly the --custom_module_dirs use case. + del MODEL_CLASS_MAPPING[arch] + register_auto_model(arch)(_Builtin) + register_auto_model(arch)(_External) + assert MODEL_CLASS_MAPPING[arch] is _External + finally: + MODEL_CLASS_MAPPING.pop(arch, None) + + +def test_custom_registration_survives_direct_provider_import(): + # Regression for the full production path: a custom implementation is + # registered, then some other code path imports the built-in provider + # directly (e.g. model_loader's post-transform profile registry imports + # modeling_llama) -- the custom registration must survive the built-in + # module's decorators. Fresh process so modeling_llama is genuinely not + # imported yet when the custom registration happens. + _run_fresh( + textwrap.dedent("""\ + import importlib + from tensorrt_llm._torch.models.modeling_utils import ( + MODEL_CLASS_MAPPING, register_auto_model) + + @register_auto_model("LlamaForCausalLM") + class CustomLlama: + pass + + importlib.import_module( + "tensorrt_llm._torch.models.modeling_llama") + assert MODEL_CLASS_MAPPING["LlamaForCausalLM"] is CustomLlama, ( + "direct import of the built-in provider overrode the custom " + "registration") + """) + ) + + +def test_external_multimodal_override_keeps_provider_importable(): + # An external override of a multimodal architecture must not break the + # built-in provider's import: register_vision_encoder used to locate the + # freshly decorated class in MODEL_CLASS_MAPPING by identity and raise + # when the external registration had won the slot. The built-in vision + # encoder still fills the empty sibling slot, like the eager import + # order on main. + _run_fresh( + textwrap.dedent("""\ + import importlib + from tensorrt_llm._torch.models.modeling_utils import ( + MODEL_CLASS_MAPPING, MODEL_CLASS_VISION_ENCODER_MAPPING, + register_auto_model) + + arch = "Qwen3VLForConditionalGeneration" + + @register_auto_model(arch) + class CustomQwen3VL: + pass + + importlib.import_module( + "tensorrt_llm._torch.models.modeling_qwen3vl") + + assert MODEL_CLASS_MAPPING[arch] is CustomQwen3VL, ( + "built-in provider import overrode the external registration") + assert MODEL_CLASS_VISION_ENCODER_MAPPING.get(arch) is not None, ( + "built-in vision encoder did not fill the empty sibling slot") + """) + ) + + +def test_external_sibling_registrations_not_clobbered(): + # When the external implementation brings its own vision encoder and + # placeholder metadata, a later built-in import must not overwrite them. + _run_fresh( + textwrap.dedent("""\ + import importlib + from tensorrt_llm._torch.models.modeling_utils import ( + MODEL_CLASS_VISION_ENCODER_MAPPING, register_auto_model, + register_vision_encoder) + from tensorrt_llm.inputs.registry import ( + MULTIMODAL_PLACEHOLDER_REGISTRY, MultimodalPlaceholderMetadata) + + arch = "Qwen3VLForConditionalGeneration" + + class CustomEncoder: + pass + + @register_vision_encoder(CustomEncoder) + @register_auto_model(arch) + class CustomQwen3VL: + pass + + custom_metadata = MultimodalPlaceholderMetadata( + placeholder_map={"image": ""}) + MULTIMODAL_PLACEHOLDER_REGISTRY.set_placeholder_metadata( + "qwen3_vl", custom_metadata, registrant_module=__name__) + + importlib.import_module( + "tensorrt_llm._torch.models.modeling_qwen3vl") + + assert MODEL_CLASS_VISION_ENCODER_MAPPING[arch][0] is CustomEncoder, ( + "built-in import clobbered the external vision encoder") + assert MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder_metadata( + "qwen3_vl") is custom_metadata, ( + "built-in import clobbered the external placeholder metadata") + """) + ) From eb14eb84acbece72995778e3089dd829aed0ce48 Mon Sep 17 00:00:00 2001 From: junq <22017000+QiJune@users.noreply.github.com> Date: Thu, 6 Aug 2026 08:59:35 +0800 Subject: [PATCH 6/6] [None][fix] Import LlamaForCausalLM from its home module in the MX loader test test_model_loader_mx.py reached for LlamaForCausalLM through the model_loader module namespace at collection time; that attribute was only ever there as a side effect of a module-level import, which the lazy model zoo moved into _post_transform_profile_registry. Use the modeling_llama module (already imported by this test) directly. This was the persistent CPU-stage failure in CI pipeline 51978 (collection error in unittest/_torch/executor). Signed-off-by: junq <22017000+QiJune@users.noreply.github.com> --- tests/unittest/_torch/executor/test_model_loader_mx.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/executor/test_model_loader_mx.py b/tests/unittest/_torch/executor/test_model_loader_mx.py index 25686a15063e..7df5d5c055ef 100644 --- a/tests/unittest/_torch/executor/test_model_loader_mx.py +++ b/tests/unittest/_torch/executor/test_model_loader_mx.py @@ -124,14 +124,14 @@ def _moe_context(config, mapping): yield None -class _UnqualifiedLlamaForCausalLM(model_loader_mod.LlamaForCausalLM): +class _UnqualifiedLlamaForCausalLM(modeling_llama_mod.LlamaForCausalLM): pass def _tiny_llama_model( monkeypatch: pytest.MonkeyPatch, *, - model_class: type[nn.Module] = model_loader_mod.LlamaForCausalLM, + model_class: type[nn.Module] = modeling_llama_mod.LlamaForCausalLM, ) -> nn.Module: monkeypatch.setattr(modeling_llama_mod, "get_sm_version", lambda: 90) llama_config = LlamaConfig(