tensorrt_llm/_torch/pyexecutor/model_loader.py (2)

283-287: Initialize dummy weights under no_grad and avoid .data where possible.

Wrap ops in torch.no_grad() to avoid autograd tracking. Keep the existing logic but minimize .data usage.

Apply:

 def _initialize_dummy_weights(model: torch.nn.Module,
                               low: float = -1e-3,
                               high: float = 1e-3,
                               seed: int = 0) -> None:
     """Initializes model weights with random dummy values for testing purposes."""
 
     # This function's logic is copied directly from the original file
     def _get_random_min_max(dtype: torch.dtype) -> Tuple[int, int]:
@@
-        for param in model.state_dict().values():
-            generator = torch.Generator(device=param.data.device)
-            generator.manual_seed(seed)
-            dtype = param.data.dtype
-
-            if param.data.element_size() < 2:
-                tmp_param = torch.empty_like(param.data,
-                                             dtype=torch.float16,
-                                             device=param.data.device)
-                quant_min, quant_max = _get_random_min_max(dtype)
-                tmp_param.uniform_(quant_min, quant_max, generator=generator)
-                param.data.copy_(tmp_param.to(dtype))
-            elif torch.is_floating_point(param):
-                param.uniform_(low, high, generator=generator)
+        with torch.no_grad():
+            for tensor in model.state_dict().values():
+                generator = torch.Generator(device=tensor.device)
+                generator.manual_seed(seed)
+                dtype = tensor.dtype
+
+                if tensor.element_size() < 2:
+                    tmp_param = torch.empty_like(tensor,
+                                                 dtype=torch.float16,
+                                                 device=tensor.device)
+                    quant_min, quant_max = _get_random_min_max(dtype)
+                    tmp_param.uniform_(quant_min, quant_max, generator=generator)
+                    tensor.copy_(tmp_param.to(dtype))
+                elif torch.is_floating_point(tensor):
+                    tensor.uniform_(low, high, generator=generator)

Also applies to: 298-312

---

230-244: KV/mamba dtype validation usage looks consistent. Minor: log the resolved values.

In addition to mutating config, consider logging the final kv_cache_quant_algo and mamba_ssm_cache_dtype to aid debugging.

Example:

         validate_and_set_kv_cache_quant(
             config, self.pytorch_backend_config.kv_cache_dtype)
         validate_and_set_mamba_ssm_cache_dtype(
             config, self.pytorch_backend_config.mamba_ssm_cache_dtype)
+        logger.info(f"Resolved kv_cache_quant_algo={config.quant_config.kv_cache_quant_algo}, "
+                    f"mamba_ssm_cache_dtype={getattr(config.quant_config, 'mamba_ssm_cache_dtype', None)}")

tensorrt_llm/_torch/pyexecutor/model_engine.py (1)

38-41: Prefer module import over direct class import to match style guide.

Import the module and reference the class via the module to keep namespace clean.

-from ..models.modeling_utils import DecoderModelForCausalLM
+from ..models import modeling_utils

And later:

-                if isinstance(self.model, DecoderModelForCausalLM):
+                if isinstance(self.model, modeling_utils.DecoderModelForCausalLM):

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: jiaganc
PR: NVIDIA/TensorRT-LLM#7031
File: tensorrt_llm/bench/dataclasses/configuration.py:90-104
Timestamp: 2025-08-26T09:37:10.463Z
Learning: In TensorRT-LLM's bench configuration, the `get_pytorch_perf_config()` method returns `self.pytorch_config` which is a Dict[str, Any] that can contain default values including `cuda_graph_config`, making the fallback `llm_args["cuda_graph_config"]` safe to use.

Learnt from: amitz-nv
PR: NVIDIA/TensorRT-LLM#7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.