[TRTLLM-13250][fix] Address Wave 5 review findings

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Author: chienchunhung

tensorrt_llm/_torch/attention_backend/sparse/dsa.py

@@ -1518,14 +1518,14 @@ def __init__(self,
             # attribute queries do not end up frozen into a captured graph.
             warmup_heuristic_topk_decode(top_k=self.index_topk)
 
-    def cache_derived_state(self):
+    def cache_derived_state(self) -> None:
         """Fuse wk + weights_proj into single FP32 weight for F.linear GEMM under allow_tf32 (TF32 tensor cores on Ampere+)."""
         # wk: [head_dim, hidden_size] + weights_proj: [n_heads, hidden_size]
         # → fused: [head_dim + n_heads, hidden_size]
         self._fused_wk_wp_weight = torch.cat(
             [self.wk.weight.data, self.weights_proj.weight.data], dim=0)
 
-    def post_load_weights(self):
+    def post_load_weights(self) -> None:
         self.cache_derived_state()
 
     @staticmethod

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1921,7 +1921,7 @@ def load_weights(self, weights: ConsumableWeightsDict):
         weight_loader = DeepseekV3WeightLoader(self)
         weight_loader.load_weights(weights)
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_exaone_moe.py

@@ -725,7 +725,7 @@ def load_weights(
             allow_partial_loading=allow_partial_loading,
         )
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         # For the cross-layer residual+LN fusion.
         for idx, layer in enumerate(self.model.layers[: self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_glm.py

@@ -1074,7 +1074,7 @@ def load_weights(self, weights: ConsumableWeightsDict, allow_partial_loading: bo
         weight_loader = Glm4WeightLoader(self)
         weight_loader.load_weights(weights, allow_partial_loading=allow_partial_loading)
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(self.model.layers[: self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:
                 layer.next_layer_layernorm = self.model.norm

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -631,7 +631,7 @@ def load_weights(self, weights: Dict):
         else:
             self.load_hf_weights(weights)
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(
                 self.model.block[:self.config.num_hidden_layers]):
             if idx == 0:

tensorrt_llm/_torch/models/modeling_llama.py

@@ -1140,7 +1140,7 @@ def __init__(
     ):
         super().__init__(LlamaModel(model_config), model_config)
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:
@@ -1564,7 +1564,7 @@ def load_weights(self, weights: Dict, weight_mapper: BaseWeightMapper):
             if had_mm_encoder:
                 self.mm_encoder = saved_mm_encoder
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_llama_min_latency.py

@@ -323,7 +323,7 @@ def __init__(self,
 
     # After loading both gate_up_proj and down_proj, we need to set the scales needed by the special kernels and by
     # the trtllm-gen gemm+swiglu kernel.
-    def cache_derived_state(self):
+    def cache_derived_state(self) -> None:
         if self.gate_up_proj.has_fp8_qdq:
             # For the special gemm+swiglu kernel, we need to set the inverse of the output scale, which is the inverse
             # of down_proj's combined input scale.
@@ -332,7 +332,7 @@ def cache_derived_state(self):
             # combined input scale times inv_output_scale.
             self.gate_up_proj.trtllm_gen_global_scale = self.gate_up_proj.combined_scale * self.gate_up_proj.inv_output_scale
 
-    def post_load_weights(self):
+    def post_load_weights(self) -> None:
         self.cache_derived_state()
 
     def forward(
@@ -584,7 +584,7 @@ def __init__(
             dtype=model_config.pretrained_config.torch_dtype,
             quant_config=None)
 
-    def cache_derived_state(self):
+    def cache_derived_state(self) -> None:
         # Set min-latency quant scales for routed experts if we plan to use min-latency MoE kernels.
         # This is because the routed experts' input scale is after the score multiplication, so we must use the
         # pre-score scaling input scale, which happens to be shared expert's input scale.
@@ -600,7 +600,7 @@ def cache_derived_state(self):
                 fc1_input_dequant=pre_score_scaling_input_scale,
             )
 
-    def post_load_weights(self):
+    def post_load_weights(self) -> None:
         self.cache_derived_state()
 
     def compute_routed_output(

tensorrt_llm/_torch/models/modeling_qwen3_moe.py

@@ -417,7 +417,7 @@ def __init__(
         )
         self.preload_weight_modules = self.model.preload_weight_modules
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/models/modeling_qwen3_next.py

@@ -980,7 +980,7 @@ def load_weights(self, weights: dict, weight_mapper: BaseWeightMapper):
         new_weights = weight_mapper.preprocess_weights(weights)
         super().load_weights(new_weights, weight_mapper)
 
-    def setup_aliases(self):
+    def setup_aliases(self) -> None:
         for idx, layer in enumerate(
                 self.model.layers[:self.config.num_hidden_layers]):
             if idx == self.config.num_hidden_layers - 1:

tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py

@@ -650,9 +650,10 @@ def load_weights(self, weights: List[Dict], allow_partial_loading: bool = False)
         assert hasattr(self.backend, "load_weights"), (
             f"Backend {self.backend.__class__.__name__} must implement load_weights()"
         )
+        self._weights_transformed = False
         return self.backend.load_weights(weights, allow_partial_loading)
 
-    def transform_weights(self):
+    def transform_weights(self) -> None:
         """
         Transform weights - delegated to backend
 
@@ -665,17 +666,17 @@ def transform_weights(self):
         self.backend.transform_weights()
         self._weights_transformed = True
 
-    def cache_derived_state(self):
+    def cache_derived_state(self) -> None:
         """
         Cache derived state - delegated to backend
 
         """
         assert hasattr(self.backend, "cache_derived_state"), (
             f"Backend {self.backend.__class__.__name__} must implement cache_derived_state()"
         )
-        return self.backend.cache_derived_state()
+        self.backend.cache_derived_state()
 
-    def post_load_weights(self):
+    def post_load_weights(self) -> None:
         """
         Backward-compatible staged post-load processing - delegated to backend