tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py (1)

1660-1675: Add strict=True to zip calls in recursive arg merge.

This merge logic assumes aligned lengths; explicit strictness makes mismatches fail fast and addresses the Ruff B905 findings in this block.

♻️ Proposed patch

-                return tuple(_merge_arg(cur, old) for cur, old in zip(current_arg, stored_arg))
+                return tuple(
+                    _merge_arg(cur, old)
+                    for cur, old in zip(current_arg, stored_arg, strict=True)
+                )
@@
-                return [_merge_arg(cur, old) for cur, old in zip(current_arg, stored_arg)]
+                return [
+                    _merge_arg(cur, old)
+                    for cur, old in zip(current_arg, stored_arg, strict=True)
+                ]
@@
-    new_args = [
-        _merge_arg(current_arg, stored_arg) for current_arg, stored_arg in zip(node.args, args)
-    ]
+    new_args = [
+        _merge_arg(current_arg, stored_arg)
+        for current_arg, stored_arg in zip(node.args, args, strict=True)
+    ]

Based on learnings: In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features (e.g., PEP 585 generics), so zip(..., strict=True) is available.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tensorrt_llm/_torch/auto_deploy/transform/library/sharding.py` around lines
1660 - 1675, The zip usage in the recursive merge (_merge_arg) and when building
new_args assumes equal-length iterables but doesn't fail on mismatch; update the
zip(...) calls in the tuple/list handling branches and the final new_args
construction to use zip(..., strict=True) so mismatched lengths raise
immediately. Locate the _merge_arg function and replace the two zip(...) usages
inside the tuple and list branches and the zip(...) in the new_args list
comprehension with zip(..., strict=True) to enforce strict alignment.

tests/unittest/_torch/auto_deploy/unit/singlegpu/test_hf_export_info.py (1)

7-10: Use module-namespace imports in this test.

This file introduces direct symbol imports; repository guidance prefers module imports with namespaced usage.

♻️ Proposed patch

 import operator

 import torch
-from torch import nn
-from torch.fx import symbolic_trace
+import torch.fx as fx
+import torch.nn as nn

-from tensorrt_llm._torch.auto_deploy.models.hf import TextModelExportInfo
+import tensorrt_llm._torch.auto_deploy.models.hf as hf_models
@@
-    gm = symbolic_trace(model)
+    gm = fx.symbolic_trace(model)
@@
-    export_info = TextModelExportInfo("dummy")
+    export_info = hf_models.TextModelExportInfo("dummy")

As per coding guidelines: "When importing in Python, always maintain the namespace. Import the module, not individual classes or functions..."

Also applies to: 27-30

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unittest/_torch/auto_deploy/unit/singlegpu/test_hf_export_info.py`
around lines 7 - 10, The test file uses direct symbol imports (nn,
symbolic_trace, TextModelExportInfo); change these to module-namespace imports
and update usages accordingly—e.g., replace "from torch import nn" with "import
torch.nn as nn" or "import torch" and use "torch.nn", replace "from torch.fx
import symbolic_trace" with "import torch.fx as fx" and use "fx.symbolic_trace",
and replace "from tensorrt_llm._torch.auto_deploy.models.hf import
TextModelExportInfo" with "import tensorrt_llm._torch.auto_deploy.models.hf as
hf" and use "hf.TextModelExportInfo"; also update the other occurrences noted
around lines 27-30 similarly.

tests/unittest/_torch/auto_deploy/unit/singlegpu/test_flashinfer_mla_cache_append.py (1)

47-54: Consider adding strict=True to zip() for defensive programming.

While batch_indices and positions are constructed with the same length in _make_append_metadata, adding strict=True would catch any future mismatches early.

♻️ Suggested fix

     for token_idx, (batch_idx, position) in enumerate(
-        zip(batch_indices.tolist(), positions.tolist())
+        zip(batch_indices.tolist(), positions.tolist(), strict=True)
     ):

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tests/unittest/_torch/auto_deploy/unit/singlegpu/test_flashinfer_mla_cache_append.py`
around lines 47 - 54, The loop in test_flashinfer_mla_cache_append.py that
iterates "for token_idx, (batch_idx, position) in
enumerate(zip(batch_indices.tolist(), positions.tolist()))" should use zip(...,
strict=True) to defensively ensure batch_indices and positions are the same
length; update that zip call to zip(batch_indices.tolist(), positions.tolist(),
strict=True) so any future mismatch (despite _make_append_metadata currently
producing equal lengths) raises immediately and surfaces the bug during tests.

tensorrt_llm/_torch/auto_deploy/tokenizers/mistral_small_4_119b/processing_mistral_small_4.py (2)

148-149: Same assertion pattern in processor - consider consistent error handling.

Apply the same improvement as suggested for the tokenizer class for consistency.

♻️ Suggested improvement

-        assert source_processor_config_path is not None
+        if source_processor_config_path is None:
+            raise FileNotFoundError(
+                f"Could not find {_PROCESSOR_CONFIG_FILE} for {source_model_name_or_path}"
+            )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tensorrt_llm/_torch/auto_deploy/tokenizers/mistral_small_4_119b/processing_mistral_small_4.py`
around lines 148 - 149, Replace the bare assert in processing_mistral_small_4.py
that checks source_processor_config_path with explicit error handling similar to
the tokenizer class: check if source_processor_config_path is None and raise a
clear ValueError (or RuntimeError) with a descriptive message identifying
source_processor_config_path and the expected config; then call
_load_json(Path(source_processor_config_path)) as before. This change should be
applied around the source_processor_config_path usage in the processing logic to
ensure consistent, explicit error reporting.

---

94-98: Consider replacing assertions with informative exceptions for user-facing errors.

The assert statements will raise AssertionError without context if the config files are missing. For better user experience, consider raising FileNotFoundError or ValueError with a descriptive message.

♻️ Suggested improvement

-        assert source_tokenizer_config_path is not None
+        if source_tokenizer_config_path is None:
+            raise FileNotFoundError(
+                f"Could not find {_TOKENIZER_CONFIG_FILE} for {source_model_name_or_path}"
+            )
         source_tokenizer_config = _load_json(Path(source_tokenizer_config_path))
 
         tokenizer_file = cached_file(source_model_name_or_path, _TOKENIZER_FILE, **kwargs)
-        assert tokenizer_file is not None
+        if tokenizer_file is None:
+            raise FileNotFoundError(
+                f"Could not find {_TOKENIZER_FILE} for {source_model_name_or_path}"
+            )

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In
`@tensorrt_llm/_torch/auto_deploy/tokenizers/mistral_small_4_119b/processing_mistral_small_4.py`
around lines 94 - 98, Replace the bare assert checks for missing tokenizer
files/configs with explicit exceptions: instead of "assert
source_tokenizer_config_path is not None" raise a FileNotFoundError or
ValueError that includes the variable value (source_tokenizer_config_path) and a
clear message before calling _load_json; likewise, after obtaining
tokenizer_file from cached_file(source_model_name_or_path, _TOKENIZER_FILE,
**kwargs) replace "assert tokenizer_file is not None" with a FileNotFoundError
that includes source_model_name_or_path and _TOKENIZER_FILE so callers get a
descriptive error rather than an AssertionError.