[TRTLLM-15151][chore] load models lazily - #17281
Open
QiJune wants to merge 7 commits into
Open
Conversation
QiJune
requested review from
JunyiXu-nv,
Tabrizian,
brnguyen2,
reasonsolo,
schetlur-nv and
zhaoyangwang-nvidia
August 5, 2026 04:04
Contributor
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughChangesThe package now loads public modules and symbols lazily. Static indexes support on-demand provider imports and registration. Executor and serving paths defer model-related imports. Shared steady-clock conversion moved to Lazy loading and registration
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant ModelLoader
participant MODEL_ARCH_TO_MODULE
participant ProviderModule
participant MODEL_CLASS_MAPPING
ModelLoader->>MODEL_ARCH_TO_MODULE: resolve checkpoint architecture
MODEL_ARCH_TO_MODULE->>ProviderModule: import provider module
ProviderModule->>MODEL_CLASS_MAPPING: register model class
ModelLoader->>MODEL_CLASS_MAPPING: resolve registered class
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 1⚔️ Resolve merge conflicts 💡
🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/models/modeling_auto.py (1)
36-40: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick winAlign EAGLE3 architecture naming. When
architectures[0]isDeepseekV3ForCausalLM, this code generatesEAGLE3DeepseekV3ForCausalLM. NeitherMODEL_ARCH_TO_MODULEnorregister_auto_modeldefines that exact key; the registered key isEagle3DeepSeekV3ForCausalLM.ensure_model_registeredtherefore skips the lazy import, and model resolution returnsNone. Align the rewrite, registration, and index entries, and add a regression test.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_auto.py` around lines 36 - 40, Update the EAGLE3 architecture rewrite in the model-architecture resolution flow around ensure_model_registered so DeepseekV3ForCausalLM produces the exact registered key Eagle3DeepSeekV3ForCausalLM, using consistent spelling and casing across MODEL_ARCH_TO_MODULE and register_auto_model. Add a regression test covering this architecture and verifying successful model resolution.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/models/__init__.py (1)
111-112: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueInclude the lazy names in
__dir__.
__dir__reports__all__and the already-resolved globals. It omits anyMODEL_CLASS_TO_MODULEkey that is not in__all__, so those names are invisible todir()and to tab completion until first access. The top-level package uses the union with its lazy table; match that here.♻️ Proposed fix
-def __dir__(): - return sorted(set(__all__) | set(globals())) +def __dir__() -> list[str]: + return sorted(set(__all__) | set(globals()) | set(MODEL_CLASS_TO_MODULE))🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/__init__.py` around lines 111 - 112, Update __dir__ to include the keys from MODEL_CLASS_TO_MODULE alongside __all__ and existing globals, matching the top-level package’s lazy-name discovery behavior so unresolved model names appear in dir() and tab completion.tensorrt_llm/__init__.py (2)
195-199: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winDo not convert unrelated import failures into
AttributeError.The fallback catches every
ModuleNotFoundError. A submodule that exists but imports a missing optional dependency also raisesModuleNotFoundError. The user then seesmodule 'tensorrt_llm' has no attribute 'X'instead of the real missing dependency. Check that the failed module is the requested one before converting.♻️ Proposed fix
try: return importlib.import_module(f'.{name}', __name__) - except ModuleNotFoundError: - raise AttributeError( - f"module {__name__!r} has no attribute {name!r}") from None + except ModuleNotFoundError as e: + if e.name != f'{__name__}.{name}': + raise # a real missing dependency inside the submodule + raise AttributeError( + f"module {__name__!r} has no attribute {name!r}") from None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/__init__.py` around lines 195 - 199, Update the import fallback around importlib.import_module in __getattr__ so it converts ModuleNotFoundError to AttributeError only when the missing module is the requested submodule; re-raise unrelated dependency import failures unchanged.
185-185: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAdd type annotations to the module hooks.
The coding guidelines require an annotation on every function. Use
def __getattr__(name: str) -> object:anddef __dir__() -> list[str]:.As per coding guidelines: "Annotate every function, use
Nonefor procedures".Also applies to: 202-202
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/__init__.py` at line 185, Update the module hooks __getattr__ and __dir__ with the required annotations: use a str parameter and object return type for __getattr__, and a list[str] return type for __dir__. Keep their existing behavior unchanged.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/__init__.py`:
- Around line 97-108: The __getattr__ fallback in __init__ must also resolve
direct submodule attributes such as checkpoints and hf_parameter_utils, not only
names beginning with modeling_. Attempt importing the requested submodule for
other names as well, convert only genuine module-not-found misses to
AttributeError, and preserve the existing modeling_multimodal_utils coverage and
model-class resolution behavior.
In `@tensorrt_llm/_torch/models/_arch_index.py`:
- Line 93: Remove the "SomeVLModel" mapping from MODEL_ARCH_TO_MODULE, leaving
the surrounding architecture registrations unchanged.
- Around line 18-99: Add a CI consistency check for MODEL_ARCH_TO_MODULE that
discovers every architecture registered by `@register_auto_model`, including
qualified and attribute-based decorator forms, and compares those registrations
against the mapping in both directions. Fail the check when an architecture is
missing from either set, so ensure_model_registered can import and register
every decorated model.
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 883-887: Record caught ImportError instances in a module-level
dictionary keyed by architecture within the lazy import path, while retaining
the existing warning. Update get_model_architecture to append the stored import
failure to its unknown-architecture error when the architecture is indexed but
its import failed. Keep non-ImportError exceptions propagating unchanged.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Line 4884: Update the comment near the token-range checks to replace the
Unicode hyphen in “token‐range” with an ASCII hyphen, yielding “token-range” and
satisfying Ruff RUF003.
---
Outside diff comments:
In `@tensorrt_llm/_torch/models/modeling_auto.py`:
- Around line 36-40: Update the EAGLE3 architecture rewrite in the
model-architecture resolution flow around ensure_model_registered so
DeepseekV3ForCausalLM produces the exact registered key
Eagle3DeepSeekV3ForCausalLM, using consistent spelling and casing across
MODEL_ARCH_TO_MODULE and register_auto_model. Add a regression test covering
this architecture and verifying successful model resolution.
---
Nitpick comments:
In `@tensorrt_llm/__init__.py`:
- Around line 195-199: Update the import fallback around importlib.import_module
in __getattr__ so it converts ModuleNotFoundError to AttributeError only when
the missing module is the requested submodule; re-raise unrelated dependency
import failures unchanged.
- Line 185: Update the module hooks __getattr__ and __dir__ with the required
annotations: use a str parameter and object return type for __getattr__, and a
list[str] return type for __dir__. Keep their existing behavior unchanged.
In `@tensorrt_llm/_torch/models/__init__.py`:
- Around line 111-112: Update __dir__ to include the keys from
MODEL_CLASS_TO_MODULE alongside __all__ and existing globals, matching the
top-level package’s lazy-name discovery behavior so unresolved model names
appear in dir() and tab completion.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 5305a36f-c8f2-44e1-bb21-0350d829c389
📒 Files selected for processing (12)
tensorrt_llm/__init__.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_llama.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/perf_metrics_manager.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_utils.pytensorrt_llm/commands/utils.pytensorrt_llm/serve/responses_utils.py
Comment on lines
+18
to
+99
| 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", | ||
| } |
Contributor
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Compare `@register_auto_model` architectures with MODEL_ARCH_TO_MODULE.
set -uo pipefail
python - <<'PY'
import pathlib, re, ast
root = pathlib.Path("tensorrt_llm/_torch/models")
index_src = (root / "_arch_index.py").read_text()
tree = ast.parse(index_src)
tables = {}
for node in tree.body:
if isinstance(node, ast.Assign) and isinstance(node.targets[0], ast.Name):
try:
tables[node.targets[0].id] = ast.literal_eval(node.value)
except ValueError:
pass
arch_index = tables.get("MODEL_ARCH_TO_MODULE", {})
decorated = {}
pat = re.compile(r'`@register_auto_model`\(\s*["\']([^"\']+)["\']')
for path in sorted(root.rglob("modeling_*.py")):
for arch in pat.findall(path.read_text()):
decorated.setdefault(arch, set()).add(path.stem)
missing = sorted(a for a in decorated if a not in arch_index)
extra = sorted(a for a in arch_index if a not in decorated)
wrong = sorted(
(a, arch_index[a], sorted(mods))
for a, mods in decorated.items()
if a in arch_index and arch_index[a] not in mods
)
print("registered but MISSING from MODEL_ARCH_TO_MODULE:", missing)
print("in MODEL_ARCH_TO_MODULE but NOT registered:", extra)
print("mapped to the WRONG module:", wrong)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 327
🏁 Script executed:
#!/bin/bash
set -euo pipefail
root="tensorrt_llm/_torch/models"
printf '%s\n' '--- index and registration references ---'
rg -n -C 4 'MODEL_ARCH_TO_MODULE|ensure_model_registered|register_auto_model|LlavaLlamaModel|PixtralVisionModel' "$root" \
-g '*.py'
printf '%s\n' '--- candidate module outlines ---'
for file in "$root"/_arch_index.py "$root"/modeling_vila.py "$root"/modeling_pixtral.py; do
if [ -f "$file" ]; then
echo "### $file"
ast-grep outline "$file"
fi
done
printf '%s\n' '--- index module existence ---'
python3 - <<'PY'
import ast
import pathlib
root = pathlib.Path("tensorrt_llm/_torch/models")
tree = ast.parse((root / "_arch_index.py").read_text())
table = next(
node.value
for node in tree.body
if isinstance(node, ast.Assign)
and any(isinstance(target, ast.Name) and target.id == "MODEL_ARCH_TO_MODULE"
for target in node.targets)
)
mapping = ast.literal_eval(table)
for architecture, module in mapping.items():
path = root / f"{module}.py"
print(f"{architecture}: {module}.py {'OK' if path.is_file() else 'MISSING'}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
root="tensorrt_llm/_torch/models"
printf '%s\n' '--- relevant lazy-loading code ---'
rg -n -C 8 'def ensure_model_registered|def register_auto_model|MODEL_ARCH_TO_MODULE|import_module' \
"$root/_arch_index.py" "$root/modeling_auto.py" "$root/modeling_utils.py"
printf '%s\n' '--- non-literal registration sites for indexed architectures ---'
rg -n -C 3 'LlavaLlamaModel|PixtralVisionModel|model_architecture|register_auto_model' \
"$root/modeling_vila.py" "$root/modeling_pixtral.py" "$root/modeling_mistral.py"
printf '%s\n' '--- AST comparison, including constant expressions ---'
python3 - <<'PY'
import ast
import pathlib
root = pathlib.Path("tensorrt_llm/_torch/models")
index_tree = ast.parse((root / "_arch_index.py").read_text())
index_node = next(
node.value for node in index_tree.body
if isinstance(node, ast.Assign)
and any(isinstance(target, ast.Name)
and target.id == "MODEL_ARCH_TO_MODULE"
for target in node.targets)
)
index = ast.literal_eval(index_node)
registrations = {}
for path in sorted(root.glob("modeling_*.py")):
tree = ast.parse(path.read_text())
constants = {
node.targets[0].id: ast.literal_eval(node.value)
for node in tree.body
if isinstance(node, ast.Assign)
and len(node.targets) == 1
and isinstance(node.targets[0], ast.Name)
and isinstance(node.value, ast.Constant)
and isinstance(node.value.value, str)
}
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
if not (isinstance(func, ast.Name) and func.id == "register_auto_model"):
continue
if len(node.args) != 1:
continue
arg = node.args[0]
value = None
if isinstance(arg, ast.Constant) and isinstance(arg.value, str):
value = arg.value
elif (isinstance(arg, ast.Attribute)
and isinstance(arg.value, ast.Name)
and arg.value.id in constants
and arg.attr == "model_architecture"):
value = constants[arg.value.id]
if value is not None:
registrations.setdefault(value, set()).add(path.stem)
missing = sorted(set(registrations) - set(index))
extra = sorted(set(index) - set(registrations))
wrong = sorted(
(arch, index[arch], sorted(registrations[arch]))
for arch in sorted(set(index) & set(registrations))
if index[arch] not in registrations[arch]
)
print("registered but missing:", missing)
print("indexed but not found by AST scan:", extra)
print("mapped to a different module:", wrong)
for arch in ("LlavaLlamaModel", "PixtralVisionModel"):
print(f"{arch}: registered in {sorted(registrations.get(arch, ()))}, "
f"indexed as {index.get(arch)!r}")
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 13967
🏁 Script executed:
#!/bin/bash
set -euo pipefail
root="tensorrt_llm/_torch/models"
python3 - <<'PY'
import pathlib
import re
root = pathlib.Path("tensorrt_llm/_torch/models")
index = {}
entry = re.compile(r'^\s*"([^"]+)":\s*"([^"]+)",\s*$')
for line in (root / "_arch_index.py").read_text().splitlines():
match = entry.match(line)
if match:
index[match.group(1)] = match.group(2)
registrations = {}
unresolved = []
decorator = re.compile(
r'^\s*@(?:[\w.]+\.)?register_auto_model\(\s*([^)]*?)\s*\)\s*$'
)
literal = re.compile(r"""^["']([^"']+)["']$""")
attribute = re.compile(r'^([A-Za-z_]\w*)\.([A-Za-z_]\w*)$')
for path in sorted(root.glob("modeling_*.py")):
lines = path.read_text().splitlines()
constants = {}
for line in lines:
match = re.match(
r'^\s*(\w+)\s*=\s*["\']([^"\']+)["\']\s*$', line
)
if match:
constants[match.group(1)] = match.group(2)
match = re.match(
r'^\s*model_architecture\s*=\s*["\']([^"\']+)["\']\s*$', line
)
if match:
constants["model_architecture"] = match.group(1)
for line_number, line in enumerate(lines, 1):
match = decorator.match(line)
if not match:
continue
argument = match.group(1).strip()
literal_match = literal.match(argument)
if literal_match:
architecture = literal_match.group(1)
else:
attribute_match = attribute.match(argument)
architecture = (
constants.get(attribute_match.group(2))
if attribute_match and attribute_match.group(1) == "VilaConfig"
else None
)
if architecture is None:
unresolved.append(f"{path}:{line_number}: {argument}")
else:
registrations.setdefault(architecture, set()).add(path.stem)
missing = sorted(set(registrations) - set(index))
extra = sorted(set(index) - set(registrations))
wrong = sorted(
(architecture, index[architecture], sorted(registrations[architecture]))
for architecture in sorted(set(index) & set(registrations))
if index[architecture] not in registrations[architecture]
)
print("registered but missing:", missing)
print("indexed but not registered:", extra)
print("mapped to a different module:", wrong)
print("unresolved decorators:", unresolved)
print("LlavaLlamaModel:", sorted(registrations.get("LlavaLlamaModel", ())))
print("PixtralVisionModel:", sorted(registrations.get("PixtralVisionModel", ())))
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 747
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tensorrt_llm/_torch/models/_arch_index.py"
sed -n '1,220p' "$file"
python3 - <<'PY'
import pathlib
import re
root = pathlib.Path("tensorrt_llm/_torch/models")
text = (root / "_arch_index.py").read_text()
match = re.search(
r"MODEL_ARCH_TO_MODULE\s*=\s*\{(?P<body>.*?)\n\}",
text,
re.S,
)
if not match:
raise SystemExit("MODEL_ARCH_TO_MODULE block not found")
index = dict(re.findall(
r'^\s*"([^"]+)":\s*"([^"]+)",\s*$',
match.group("body"),
re.M,
))
registrations = {}
unresolved = []
decorator = re.compile(
r'^\s*@(?:[\w.]+\.)?register_auto_model\(\s*([^)]*?)\s*\)\s*$'
)
literal = re.compile(r"""^["']([^"']+)["']$""")
attribute = re.compile(r'^([A-Za-z_]\w*)\.([A-Za-z_]\w*)$')
for path in sorted(root.glob("modeling_*.py")):
lines = path.read_text().splitlines()
class_attributes = {}
current_class = None
for line in lines:
class_match = re.match(r'^\s*class\s+(\w+)', line)
if class_match:
current_class = class_match.group(1)
attr_match = re.match(
r'^\s*model_architecture\s*=\s*["\']([^"\']+)["\']\s*$',
line,
)
if attr_match and current_class:
class_attributes[current_class] = attr_match.group(1)
for line_number, line in enumerate(lines, 1):
match = decorator.match(line)
if not match:
continue
argument = match.group(1).strip()
literal_match = literal.match(argument)
if literal_match:
architecture = literal_match.group(1)
else:
attribute_match = attribute.match(argument)
architecture = (
class_attributes.get(attribute_match.group(1))
if attribute_match else None
)
if architecture is None:
unresolved.append(f"{path}:{line_number}: {argument}")
else:
registrations.setdefault(architecture, set()).add(path.stem)
print("registered but missing:", sorted(set(registrations) - set(index)))
print("indexed but not registered:", sorted(set(index) - set(registrations)))
print("mapped to a different module:", sorted(
(architecture, index[architecture], sorted(registrations[architecture]))
for architecture in sorted(set(index) & set(registrations))
if index[architecture] not in registrations[architecture]
))
print("unresolved decorators:", unresolved)
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 8682
Add an automated consistency check for MODEL_ARCH_TO_MODULE.
If an architecture is absent from the index, ensure_model_registered skips its import and runtime registration fails. Add a CI check that compares @register_auto_model registrations with MODEL_ARCH_TO_MODULE, including qualified and attribute-based decorators.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/_arch_index.py` around lines 18 - 99, Add a CI
consistency check for MODEL_ARCH_TO_MODULE that discovers every architecture
registered by `@register_auto_model`, including qualified and attribute-based
decorator forms, and compares those registrations against the mapping in both
directions. Fail the check when an architecture is missing from either set, so
ensure_model_registered can import and register every decorated model.
Comment on lines
+883
to
+887
| 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}") |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Preserve the import failure so callers can report it.
The architecture is present in the static index, so a failed import is a real defect, not an unknown architecture. Today the traceback is dropped into a warning, and get_model_architecture then raises Unknown model architecture: <arch>. That message hides the root cause (a missing optional dependency, an incompatible transformers version, or a broken modeling_* module). Under the previous eager import, the original ImportError reached the user.
Record the failure and let the caller include it. For example, keep a module-level dict[str, ImportError] of failed architectures and append the stored error to the Unknown model architecture message.
Also note that only ImportError is caught. A modeling_* module that raises any other exception at import time still propagates, so the two failure modes behave differently.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/models/modeling_utils.py` around lines 883 - 887, Record
caught ImportError instances in a module-level dictionary keyed by architecture
within the lazy import path, while retaining the existing warning. Update
get_model_architecture to append the stored import failure to its
unknown-architecture error when the architecture is indexed but its import
failed. Keep non-ImportError exceptions propagating unchanged.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py (1)
203-234: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick winAdd a focused model-lookup contract test.
Test coverage summary: No test function covers
_lookup_model_clsorget_registered_model_class. The test file is scheduled intests/integration/test_lists/test-db/l0_a10.yml. Add tests for architecture lookup and the no-architecture fallback.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py` around lines 203 - 234, Add focused tests in the precheck test module covering the model-lookup contract: verify _lookup_model_cls resolves the expected class for a declared architecture and that get_registered_model_class handles the no-architecture case with its documented fallback. Reuse the existing _INTERNAL_APIS setup and registered model symbols, and keep the tests limited to these two lookup behaviors.Sources: Coding guidelines, Path instructions
tensorrt_llm/serve/openai_video_routes.py (1)
1-2: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winAdd the required NVIDIA copyright header to both modified Python files.
Both files begin with
from __future__ import annotationsand have no header before the first import. Add or update the NVIDIA header and set the year to 2026.
tensorrt_llm/serve/openai_video_routes.py#L1-L2: add the header before the future import.tensorrt_llm/serve/visual_gen_utils.py#L1-L2: add the header before the future import.As per coding guidelines, modified source files must contain the NVIDIA copyright header with the year of the latest meaningful modification.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_video_routes.py` around lines 1 - 2, Add the required NVIDIA copyright and SPDX license header with year 2026 before the future import in tensorrt_llm/serve/openai_video_routes.py at lines 1-2 and tensorrt_llm/serve/visual_gen_utils.py at lines 1-2; apply the same header format consistently in both files.Source: Coding guidelines
🧹 Nitpick comments (4)
tests/unittest/others/test_lazy_model_zoo.py (1)
117-143: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd return annotations to all functions.
Add
-> Noneto the ninetest_*functions. Add-> tuple[dict[str, set[str]], dict[str, set[str]]]to_decorated_registrations.Test coverage: nine test functions are present, with no matching entries in
tests/integration/test_lists/. No CBTS artifact is available; coverage verdict: needs follow-up.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/unittest/others/test_lazy_model_zoo.py` around lines 117 - 143, Add return annotations to every function in the test module: annotate all nine test_* functions with -> None, and annotate _decorated_registrations with -> tuple[dict[str, set[str]], dict[str, set[str]]].Source: Coding guidelines
tensorrt_llm/_torch/models/_arch_index.py (1)
11-15: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueReplace the reference to "the lazy-import PR".
The regeneration instruction points to a pull request instead of a checked-in artifact. A maintainer who adds a model cannot find that script later. Name the script path, or state that entries are added by hand.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/_arch_index.py` around lines 11 - 15, Update the regeneration comment in _arch_index.py to remove the reference to the lazy-import PR, and instead state that entries should be added by hand or identify the checked-in script path used to regenerate them.tensorrt_llm/_torch/models/modeling_utils.py (1)
882-889: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAnnotate the decorator signatures.
decorator(cls)has no parameter or return annotation, andregister_auto_modelreturns an unannotated callable. Add annotations, as the coding guidelines require.🔧 Proposed fix
-def register_auto_model(name: str): - - def decorator(cls): +def register_auto_model(name: str) -> Callable[[type], type]: + + def decorator(cls: type) -> type:As per coding guidelines: "Annotate every function, use
Nonefor procedures".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/models/modeling_utils.py` around lines 882 - 889, Update register_auto_model and its nested decorator function to include parameter and return type annotations, including the decorator’s cls parameter and register_auto_model’s callable return type. Ensure the decorator’s return annotation matches the class it returns, and annotate any procedure return as None where applicable.Source: Coding guidelines
tensorrt_llm/serve/openai_server.py (1)
116-124: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAnnotate
objin_is_visual_gen_instance.Add
obj: objector another precise type. The function already annotates its return value.As per coding guidelines, every function must have annotations.
Proposed fix
-def _is_visual_gen_instance(obj) -> bool: +def _is_visual_gen_instance(obj: object) -> bool:🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_server.py` around lines 116 - 124, Annotate the obj parameter in _is_visual_gen_instance with object (or another precise compatible type), preserving the existing return annotation and visual-gen detection logic.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 869-871: The built-in module checks use an overly broad prefix
match, misclassifying sibling packages as built-in. Update
_is_builtin_model_class in tensorrt_llm/_torch/models/modeling_utils.py (lines
869-871) and the registrant_module check in tensorrt_llm/inputs/registry.py
(lines 764-766) to accept only the exact built-in package name or names
beginning with that package followed by a dot.
In `@tensorrt_llm/inputs/registry.py`:
- Around line 872-894: Update get_registered_image_model_types,
get_registered_video_model_types, and get_registered_audio_model_types to wrap
each generator expression in tuple(), matching get_registered_model_types.
Preserve the existing provider initialization and modality filtering while
ensuring each method returns a reusable, indexable tuple.
In `@tests/unittest/llmapi/test_config_database.py`:
- Line 27: Add tests/unittest/llmapi/test_config_database.py to the test entries
in llm_config_database.yml, preserving the existing list structure. No test
implementation changes are needed; verify the updated test list with pytest
tests/unittest/.
In `@tests/unittest/others/test_lazy_model_zoo.py`:
- Around line 154-161: Extend the index validation around MODEL_ARCH_TO_MODULE
to detect stale architecture entries: identify keys present in
MODEL_ARCH_TO_MODULE but absent from arch_truth, and assert that this set is
empty with a diagnostic listing the stale entries. Preserve the existing missing
and wrong-module checks.
---
Outside diff comments:
In `@tensorrt_llm/serve/openai_video_routes.py`:
- Around line 1-2: Add the required NVIDIA copyright and SPDX license header
with year 2026 before the future import in
tensorrt_llm/serve/openai_video_routes.py at lines 1-2 and
tensorrt_llm/serve/visual_gen_utils.py at lines 1-2; apply the same header
format consistently in both files.
In `@tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py`:
- Around line 203-234: Add focused tests in the precheck test module covering
the model-lookup contract: verify _lookup_model_cls resolves the expected class
for a declared architecture and that get_registered_model_class handles the
no-architecture case with its documented fallback. Reuse the existing
_INTERNAL_APIS setup and registered model symbols, and keep the tests limited to
these two lookup behaviors.
---
Nitpick comments:
In `@tensorrt_llm/_torch/models/_arch_index.py`:
- Around line 11-15: Update the regeneration comment in _arch_index.py to remove
the reference to the lazy-import PR, and instead state that entries should be
added by hand or identify the checked-in script path used to regenerate them.
In `@tensorrt_llm/_torch/models/modeling_utils.py`:
- Around line 882-889: Update register_auto_model and its nested decorator
function to include parameter and return type annotations, including the
decorator’s cls parameter and register_auto_model’s callable return type. Ensure
the decorator’s return annotation matches the class it returns, and annotate any
procedure return as None where applicable.
In `@tensorrt_llm/serve/openai_server.py`:
- Around line 116-124: Annotate the obj parameter in _is_visual_gen_instance
with object (or another precise compatible type), preserving the existing return
annotation and visual-gen detection logic.
In `@tests/unittest/others/test_lazy_model_zoo.py`:
- Around line 117-143: Add return annotations to every function in the test
module: annotate all nine test_* functions with -> None, and annotate
_decorated_registrations with -> tuple[dict[str, set[str]], dict[str,
set[str]]].
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f97f4a63-e85b-46ec-a97d-69425f5083a5
📒 Files selected for processing (19)
tensorrt_llm/__init__.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_mistral.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/commands/serve.pytensorrt_llm/commands/utils.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/worker.pytensorrt_llm/inputs/registry.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/openai_video_routes.pytensorrt_llm/serve/visual_gen_utils.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/llmapi/test_config_database.pytests/unittest/others/test_lazy_model_zoo.py
🚧 Files skipped from review as they are similar to previous changes (3)
- tensorrt_llm/_torch/pyexecutor/model_loader.py
- tensorrt_llm/_torch/models/init.py
- tensorrt_llm/init.py
| 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 |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
ast-grep outline tests/unittest/llmapi/test_config_database.py \
--items all --type function
rg -n -C 3 \
'get_registered_model_class|_get_default_values_for_config|model_defaults' \
tests/unittest/llmapi/test_config_database.py
(fd -t f -e yml -e yaml . tests/integration/test_lists || true) |
while IFS= read -r list_file; do
rg -n -C 2 'test_config_database|llmapi' "$list_file" || true
doneRepository: NVIDIA/TensorRT-LLM
Length of output: 28378
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- diff ---'
git diff -- tests/unittest/llmapi/test_config_database.py
printf '%s\n' '--- file header and helper ---'
sed -n '1,125p' tests/unittest/llmapi/test_config_database.py
printf '%s\n' '--- matching test-list entries with file names ---'
while IFS= read -r list_file; do
matches=$(rg -n '(^|/)(test_config_database\.py|test_config_database)(:| |$)' "$list_file" || true)
if [ -n "$matches" ]; then
printf '%s\n%s\n' "$list_file" "$matches"
fi
done < <(fd -t f -e yml -e yaml . tests/integration/test_lists)Repository: NVIDIA/TensorRT-LLM
Length of output: 4838
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- matching test-list entries ---'
find tests/integration/test_lists -type f \( -name '*.yml' -o -name '*.yaml' \) -print0 |
while IFS= read -r -d '' list_file; do
if rg -n '(^|/)(test_config_database\.py|test_config_database)(:| |$)' "$list_file"; then
printf 'FILE: %s\n' "$list_file"
fi
done
printf '%s\n' '--- helper call sites ---'
rg -n -C 4 '_get_default_values_for_config|test_database_yaml_config_does_not_set_default_leaves' \
tests/unittest/llmapi/test_config_database.py
printf '%s\n' '--- registry definitions and call sites ---'
rg -n -C 5 'def get_registered_model_class|MODEL.*REGISTR|register.*model|import.*model' \
tensorrt_llm/_torch/models tensorrt_llm/_torch/models/modeling_utils.pyRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- custom model defaults ---'
rg -n -C 4 'def get_model_defaults|get_model_defaults\s*=|get_model_defaults\(' \
tensorrt_llm/_torch/models --glob '*.py'
printf '%s\n' '--- lookup architectures and models ---'
rg -n '^[[:space:]]*(model|arch|config_path):' examples/configs/database/lookup.yaml \
examples/configs/curated/lookup.yaml | head -160
printf '%s\n' '--- architecture index entries for custom-default classes ---'
rg -n -C 2 'MODEL_ARCH_TO_MODULE|MODEL_CLASS_TO_MODULE|Deepseek|Qwen|Llama|Whisper|Gemma|Phi|Mistral' \
tensorrt_llm/_torch/models/_arch_index.py | head -240
printf '%s\n' '--- test module imports of model providers ---'
rg -n '^import |^from ' tests/unittest/llmapi/test_config_database.py \
| rg 'model|torch|tensorrt_llm' || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 30942
🏁 Script executed:
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("tests/unittest/llmapi/test_config_database.py")
source = test_path.read_text(encoding="utf-8")
tree = ast.parse(source)
functions = {
node.name: node
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
helper = functions["_get_default_values_for_config"]
coverage_test = functions["test_database_yaml_config_does_not_set_default_leaves"]
helper_calls_resolver = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "get_registered_model_class"
for node in ast.walk(helper)
)
coverage_test_calls_helper = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "_get_default_values_for_config"
for node in ast.walk(coverage_test)
)
provider_imported_by_test = "modeling_deepseekv4" in source
lookup = Path("examples/configs/curated/lookup.yaml").read_text(encoding="utf-8")
provider = Path("tensorrt_llm/_torch/models/modeling_deepseekv4.py").read_text(
encoding="utf-8"
)
arch_index = Path("tensorrt_llm/_torch/models/_arch_index.py").read_text(
encoding="utf-8"
)
print(f"helper_calls_get_registered_model_class={helper_calls_resolver}")
print(f"parametrized_test_calls_helper={coverage_test_calls_helper}")
print(f"lookup_contains_DeepseekV4ForCausalLM={'DeepseekV4ForCausalLM' in lookup}")
print(f"provider_defines_get_model_defaults={'def get_model_defaults' in provider}")
print(
"arch_index_maps_DeepseekV4ForCausalLM="
f"{'\"DeepseekV4ForCausalLM\": \"modeling_deepseekv4\"' in arch_index}"
)
print(f"test_directly_imports_modeling_deepseekv4={provider_imported_by_test}")
assert helper_calls_resolver
assert coverage_test_calls_helper
assert "DeepseekV4ForCausalLM" in lookup
assert "def get_model_defaults" in provider
assert '"DeepseekV4ForCausalLM": "modeling_deepseekv4"' in arch_index
assert not provider_imported_by_test
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 260
🏁 Script executed:
python3 - <<'PY'
import ast
from pathlib import Path
test_path = Path("tests/unittest/llmapi/test_config_database.py")
source = test_path.read_text(encoding="utf-8")
tree = ast.parse(source)
functions = {
node.name: node
for node in ast.walk(tree)
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
}
helper = functions["_get_default_values_for_config"]
coverage_test = functions["test_database_yaml_config_does_not_set_default_leaves"]
helper_calls_resolver = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "get_registered_model_class"
for node in ast.walk(helper)
)
coverage_test_calls_helper = any(
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "_get_default_values_for_config"
for node in ast.walk(coverage_test)
)
provider_imported_by_test = "modeling_deepseekv4" in source
lookup = Path("examples/configs/curated/lookup.yaml").read_text(encoding="utf-8")
provider = Path("tensorrt_llm/_torch/models/modeling_deepseekv4.py").read_text(
encoding="utf-8"
)
arch_index = Path("tensorrt_llm/_torch/models/_arch_index.py").read_text(
encoding="utf-8"
)
deepseek_arch_entry = '"DeepseekV4ForCausalLM": "modeling_deepseekv4"'
print(f"helper_calls_get_registered_model_class={helper_calls_resolver}")
print(f"parametrized_test_calls_helper={coverage_test_calls_helper}")
print(f"lookup_contains_DeepseekV4ForCausalLM={'DeepseekV4ForCausalLM' in lookup}")
print(f"provider_defines_get_model_defaults={'def get_model_defaults' in provider}")
print(f"arch_index_maps_DeepseekV4ForCausalLM={deepseek_arch_entry in arch_index}")
print(f"test_directly_imports_modeling_deepseekv4={provider_imported_by_test}")
assert helper_calls_resolver
assert coverage_test_calls_helper
assert "DeepseekV4ForCausalLM" in lookup
assert "def get_model_defaults" in provider
assert deepseek_arch_entry in arch_index
assert not provider_imported_by_test
PYRepository: NVIDIA/TensorRT-LLM
Length of output: 413
🏁 Script executed:
#!/bin/bash
set -euo pipefail
printf '%s\n' '--- test-db list files ---'
find tests/integration/test_lists -type f -path '*test-db*' \
\( -name '*.yml' -o -name '*.yaml' \) -print | sort
printf '%s\n' '--- qa list files ---'
find tests/integration/test_lists -type f -path '*qa*' \
\( -name '*.yml' -o -name '*.yaml' \) -print | sort
printf '%s\n' '--- list files containing the llmapi unittest group ---'
find tests/integration/test_lists -type f \
\( -name '*.yml' -o -name '*.yaml' \) -print0 |
while IFS= read -r -d '' list_file; do
if rg -q 'unittest/llmapi/' "$list_file"; then
printf '%s\n' "$list_file"
fi
done | sortRepository: NVIDIA/TensorRT-LLM
Length of output: 4861
Add test_config_database.py to the test list.
Test coverage summary: No test function changed. The existing parametrized test covers _get_default_values_for_config, including lazy loading for DeepseekV4ForCausalLM. Add unittest/llmapi/test_config_database.py to tests/integration/test_lists/qa/llm_config_database.yml. Run pytest tests/unittest/ before merge.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/llmapi/test_config_database.py` at line 27, Add
tests/unittest/llmapi/test_config_database.py to the test entries in
llm_config_database.yml, preserving the existing list structure. No test
implementation changes are needed; verify the updated test list with pytest
tests/unittest/.
Sources: Coding guidelines, Path instructions
1 task
Contributor
|
Note GitHub couldn't provide a complete incremental comparison for this pull request, so CodeRabbit is performing a full review instead. This review may take a little longer. |
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Line 4885: Update the comment near the token-range check to replace the
Unicode hyphen in “token‐range” with the ASCII “-”, resolving Ruff RUF003
without changing the surrounding logic.
In `@tests/unittest/others/test_lazy_model_zoo.py`:
- Around line 183-213: Update test_resolver_keeps_existing_registration to
execute its registry assertions and cleanup through _run_fresh, isolating
MODEL_CLASS_MAPPING and imported provider state from the surrounding test
process. Preserve the existing-registration assertion and unknown-architecture
check within that fresh execution, and remove the unconditional restoration that
overwrites prior global state.
- Around line 64-343: The new tests in
tests/unittest/others/test_lazy_model_zoo.py are not included in any test-list
configuration, leaving their coverage insufficient. Add
tests/unittest/others/test_lazy_model_zoo.py to the appropriate test-db or qa
test list and run pytest tests/unittest/;
tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py requires no
direct change because its existing l0_b200.yml entry already covers the
_create_server-only modification.
🪄 Autofix
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 4c5ae1b9-dfe1-46f7-89d0-86cbba5f8fee
📒 Files selected for processing (24)
tensorrt_llm/__init__.pytensorrt_llm/_torch/models/__init__.pytensorrt_llm/_torch/models/_arch_index.pytensorrt_llm/_torch/models/modeling_auto.pytensorrt_llm/_torch/models/modeling_llama.pytensorrt_llm/_torch/models/modeling_mistral.pytensorrt_llm/_torch/models/modeling_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/perf_metrics_manager.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_utils.pytensorrt_llm/commands/serve.pytensorrt_llm/commands/utils.pytensorrt_llm/executor/proxy.pytensorrt_llm/executor/worker.pytensorrt_llm/inputs/registry.pytensorrt_llm/serve/openai_server.pytensorrt_llm/serve/openai_video_routes.pytensorrt_llm/serve/responses_utils.pytensorrt_llm/serve/visual_gen_utils.pytests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.pytests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.pytests/unittest/llmapi/test_config_database.pytests/unittest/others/test_lazy_model_zoo.py
🚧 Files skipped from review as they are similar to previous changes (21)
- tests/unittest/llmapi/test_config_database.py
- tensorrt_llm/serve/openai_video_routes.py
- tensorrt_llm/executor/proxy.py
- tensorrt_llm/_torch/models/modeling_mistral.py
- tensorrt_llm/serve/responses_utils.py
- tensorrt_llm/executor/worker.py
- tests/scripts/perf-sanity/cache_transceiver_precheck/run_precheck.py
- tensorrt_llm/_torch/models/modeling_llama.py
- tensorrt_llm/serve/visual_gen_utils.py
- tensorrt_llm/init.py
- tensorrt_llm/_torch/models/modeling_auto.py
- tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py
- tensorrt_llm/_utils.py
- tensorrt_llm/commands/serve.py
- tensorrt_llm/_torch/models/init.py
- tensorrt_llm/_torch/models/_arch_index.py
- tensorrt_llm/serve/openai_server.py
- tensorrt_llm/commands/utils.py
- tensorrt_llm/_torch/pyexecutor/model_loader.py
- tensorrt_llm/_torch/models/modeling_utils.py
- tensorrt_llm/inputs/registry.py
Comment on lines
+64
to
+343
| 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": "<custom_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") | ||
| """) | ||
| ) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
fd -t f . tests/integration/test_lists | sort
rg -n -C 2 'test_lazy_model_zoo|test_trtllm_serve_endpoints' \
tests/integration/test_lists || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 5373
🏁 Script executed:
git status --short
git diff --stat -- tests/unittest/others/test_lazy_model_zoo.py tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py
git diff --unified=3 -- tests/unittest/others/test_lazy_model_zoo.py tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py
rg -n 'unittest/others|test_lazy_model_zoo|test_trtllm_serve_endpoints' tests/integration/test_listsRepository: NVIDIA/TensorRT-LLM
Length of output: 2915
Add test_lazy_model_zoo.py to a test list.
- Added tests:
test_import_does_not_load_zoo_or_visual_gen,test_lazy_attribute_access_resolves_and_caches,test_placeholder_registry_resolves_in_fresh_process,test_arch_index_matches_decorators,test_models_package_missing_submodule_is_attribute_error,test_resolver_keeps_existing_registration,test_builtin_decorator_does_not_override_external_registration,test_custom_registration_survives_direct_provider_import,test_external_multimodal_override_keeps_provider_importable, andtest_external_sibling_registrations_not_clobbered. test_trtllm_serve_endpoints.pychanges only_create_server; the file is listed intests/integration/test_lists/test-db/l0_b200.yml.test_lazy_model_zoo.pyis absent from alltest-db/andqa/lists. Coverage verdict: insufficient.- Add the file to the appropriate test list and run
pytest tests/unittest/.
📍 Affects 2 files
tests/unittest/others/test_lazy_model_zoo.py#L64-L343(this comment)tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py#L310-L324
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unittest/others/test_lazy_model_zoo.py` around lines 64 - 343, The new
tests in tests/unittest/others/test_lazy_model_zoo.py are not included in any
test-list configuration, leaving their coverage insufficient. Add
tests/unittest/others/test_lazy_model_zoo.py to the appropriate test-db or qa
test list and run pytest tests/unittest/;
tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py requires no
direct change because its existing l0_b200.yml entry already covers the
_create_server-only modification.
Sources: Coding guidelines, Path instructions
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64049 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64049 [ run ] completed with state
|
Collaborator
Author
|
/bot run --stage-list "CPU-Generic-x86-1" |
Collaborator
|
PR_Github #64150 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64150 [ run ] completed with state |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64201 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64422 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64422 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64489 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64489 [ run ] completed with state
|
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64569 [ run ] triggered by Bot. Commit: |
zhaoyangwang-nvidia
approved these changes
Aug 7, 2026
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
Approve with nit
|
|
||
| # Trailing dot: package-boundary match, so "tensorrt_llm._torch.models_custom" | ||
| # counts as external. Registrants are always modeling_* submodules. | ||
| _BUILTIN_MODULE_PREFIX = "tensorrt_llm._torch.models." |
Collaborator
There was a problem hiding this comment.
The "built-in registrations only fill empty slots" rule is encoded twice: this _BUILTIN_MODULE_PREFIX and modeling_utils._is_builtin_model_class (line 863). Even the comments are near-duplicates — both explain the package boundary with the same tensorrt_llm._torch.models_custom example — and set_placeholder_metadata's docstring has to point at the other registry to explain itself. The two differ in signature (one takes a class, one a module-name string) but ultimately test the same thing.
Would a single predicate in _arch_index.py work?
Collaborator
|
PR_Github #64569 [ run ] completed with state
|
get_steady_clock_now_in_seconds lives in _utils so the executor and the serving frontends share one clock source without the pyexecutor importing the serve tree; responses_utils keeps a re-export for existing importers. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Importing tensorrt_llm no longer executes the whole product tree (previously ~99% of product modules ran at import time through the eager import chain). The public surface resolves on first attribute access via module __getattr__ backed by an explicit name -> (module, attr) table, with a plain-submodule fallback for names that used to be reachable as attributes of the eager chain. A TYPE_CHECKING block keeps the original imports visible to static tooling, and __dir__ covers discovery. Fallout fixed with the cut-over: - sampler.py's module-level edge into the speculative package is cut (function-local get_force_num_accepted_tokens, TYPE_CHECKING-only SpecTreeManager): with the top-level package lazy, a process whose first touch is pyexecutor.sampler.sampler entered the sampler <-> speculative cycle from the losing side. - test_finegrained_mixed_dtype_gemm.py imported only tensorrt_llm but used the tensorrt_llm.quantization.functional attribute chain, which only the eager import tree made reachable; import the submodule explicitly, like its sibling w4a8/w4a16 tests already do. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Importing tensorrt_llm._torch.models no longer imports every modeling_* module (each of which executes its model class bodies and @register_auto_model decorators at interpreter startup). Class access resolves via PEP 562 __getattr__ over a static class -> module index; architecture resolution goes through the new modeling_utils.get_registered_model_class / get_registered_vision_encoder, which import the providing module on demand via the MODEL_ARCH_TO_MODULE index in _arch_index.py. Because built-in modules may now run their decorators after external code (--custom_module_dirs) has claimed an architecture, every registry applies one priority rule: built-in registrations only fill empty slots, never overwrite; external registrations keep main's last-wins order. The multimodal placeholder registry resolves providers on demand the same way, and enumeration APIs import all indexed providers first. register_vision_encoder maps a class to the architectures it declared (per-class bookkeeping) instead of scanning the mapping by identity, so a built-in module still imports cleanly after losing its slot. Consumers that relied on import side effects are moved to the resolver API (model_loader, modeling_auto, executor proxy/worker, serve CLI arch probe, cache-transceiver precheck). py_executor probes modeling_llama via sys.modules instead of importing it at startup. The DeepSeek tokenizer wrappers import tensorrt_llm._torch.configs on demand so the AutoConfig registration for TRT-LLM-only model_types (an import side effect the eager chain used to guarantee) is in place before their bare AutoTokenizer.from_pretrained calls; function-local on purpose, since a module-level import would cycle through _torch/__init__ -> llmapi.llm -> llmapi.tokenizer. tests/unittest/others/test_lazy_model_zoo.py pins the contracts: import tensorrt_llm stays thin, lazy attribute access resolves and caches, the static index stays in sync with the decorators (AST scan), and on-demand registration never overrides external registrations. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Plain LLM serving never pays the visual_gen import cost: the serve CLI and the serving modules import the visual_gen tree only inside the VisualGen code paths (function-local imports), with TYPE_CHECKING blocks covering the annotations. OpenAIServer detects VisualGen generators via a sys.modules probe (if tensorrt_llm.visual_gen was never imported, the generator cannot be a VisualGen instance), exposed as a cached_property: __init__ pre-caches it, so a probe patched during construction sticks, while instances built without __init__ (the disaggregated unit tests use OpenAIServer.__new__) still resolve it on demand on the request path. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
The runtime tree (tensorrt_llm._torch.visual_gen) imports config leaf modules from the public package (.args, .sparse_attention, .output), and importing a leaf always executes the package __init__ first. The eager 'from .visual_gen import VisualGen' there re-enters the still-initializing runtime tree (visual_gen.py needs DiffusionRequest from it), so any process whose first import entered via tensorrt_llm._torch.visual_gen failed with a circular ImportError -- previously masked by the eager top-level import chain initializing the public package first (CI 52358: every visual_gen unit test failed at collection). Resolving names on attribute access keeps the package __init__ side-effect free, so both entry orders work; it also drops the eager tensorrt_llm.models.modeling_utils pull from the config path. The OpenAIServer probe now checks the module that defines VisualGen instead of the package: with the lazy __init__, a process that only imported config leaves has the package in sys.modules, and reading VisualGen off it would import the whole runtime tree via __getattr__. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
…y_pass torch.ops.trtllm.mla_custom_op_inplace is registered as an import side effect of modules/mla.py, which the eagerly loaded model zoo used to provide (models/__init__ -> modeling_deepseekv3 -> modules.mla). With the zoo lazy, the two tests referencing the op failed with '_OpNamespace trtllm has no attribute mla_custom_op_inplace' (CI 52358); a test must import the ops it uses itself. Production code is unaffected: the compilation passes go through get_optional_trtllm_op, which tolerates unregistered ops. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
- Single is_builtin_zoo_module predicate in _arch_index.py; the model class registry and the multimodal placeholder registry both delegate to it instead of encoding the package-boundary rule twice. - py_executor derives the Llama4 sys.modules probe target from MODEL_ARCH_TO_MODULE, so re-homing the class updates the probe (or fails the index sync test) instead of silently disabling the multimodal token-range exemption. - test_lazy_model_zoo covers MODEL_CLASS_TO_MODULE: every non-eager __all__ name has an index entry, no stale entries, and each mapped module defines its class (AST scan). - Drop the dead 'see the lazy-import PR' docstring reference in _arch_index.py; point at the sync test instead. - Comment in proxy.py that the MODEL_CLASS_MAPPING snapshot sent to workers is intentionally partial under lazy loading. Signed-off-by: junq <22017000+QiJune@users.noreply.github.com>
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64594 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64594 [ run ] completed with state
|
yibinl-nvidia
approved these changes
Aug 7, 2026
yibinl-nvidia
left a comment
yibinl-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
approved on visual gen changes.
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64731 [ run ] triggered by Bot. Commit: |
Collaborator
Author
|
/bot run |
Collaborator
|
PR_Github #64791 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #64731 [ run ] completed with state |
Collaborator
|
PR_Github #64791 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
tensorrt_llmimports.get_steady_clock_now_in_seconds()to_utils.test_config_database.pyandtest_lazy_model_zoo.pyto the CPU pre-merge test list.Dev Engineer Review
__getattr__and__dir__preserve public API behavior.QA Engineer Review
_is_visual_gen_instance.LlamaForCausalLMfrommodeling_llama.test_config_database.pyandtest_lazy_model_zoo.pytotests/integration/test_lists/test-db/l0_cpu.yml.Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.