[TRTLLM-13247][fix] Address CodeRabbit review comments

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

Author: chienchunhung

tensorrt_llm/_torch/modules/attention.py

@@ -3046,5 +3046,8 @@ def transform_weights(self) -> None:
                 self.v_b_proj, self.v_b_proj_scale, recipe=(1, 128, 128))
         self._weights_transformed = True
 
+    def cache_derived_state(self) -> None:
+        self._weights_transformed = True
+
     def post_load_weights(self) -> None:
         self.transform_weights()

tensorrt_llm/_torch/modules/linear.py

@@ -381,7 +381,7 @@ def load_weights(self,
             self.process_weights_after_loading(module)
 
     def transform_weights(self, module: Linear) -> None:
-        pass
+        ...
 
     def post_load_weights(self, module: Linear) -> None:
         self.transform_weights(module)
@@ -3148,6 +3148,9 @@ def transform_weights(self) -> None:
         self.quant_method.transform_weights(self)
         self._weights_transformed = True
 
+    def cache_derived_state(self) -> None:
+        self._weights_transformed = True
+
     def post_load_weights(self) -> None:
         self.transform_weights()
 

tests/unittest/_torch/pyexecutor/test_model_loader_mx.py

@@ -307,6 +307,10 @@ def test_linear_transform_weights_is_idempotent():
     linear.post_load_weights()
     assert linear.quant_method.transform_weights.call_count == 2
 
+    linear._weights_transformed = False
+    linear.cache_derived_state()
+    assert linear._weights_transformed is True
+
 
 def test_mla_transform_weights_is_idempotent(monkeypatch):
     monkeypatch.setattr(attention_mod, "get_sm_version", lambda: 120)
@@ -338,3 +342,7 @@ def fake_resmooth(weight, scale, recipe):
     assert mla.v_b_proj == "v_weight_transformed"
     assert mla.v_b_proj_scale == "v_scale_transformed"
     assert mla._weights_transformed is True
+
+    mla._weights_transformed = False
+    MLA.cache_derived_state(mla)
+    assert mla._weights_transformed is True