[TRTLLM-15404][fix] VisualGen: refuse static quant recipes against unquantized checkpoints (silent weight corruption)

tensorrt_llm/_torch/visual_gen/quantization/loader.py

@@ -18,6 +18,18 @@
 from tensorrt_llm.quantization.mode import QuantAlgo
 from tensorrt_llm.quantization.utils import fp4_utils
 
+# Checkpoint scale tensors each static (pre-quantized) recipe expects
+# alongside the quantized weight (e.g. as exported by ModelOpt). Used to
+# detect unquantized checkpoints when dynamic weight quantization is off.
+# Algos absent from this dict have no verified checkpoint layout in VisualGen;
+# the guard fails closed for them (a high-precision weight under any static
+# recipe is refused rather than silently cast over uninitialized scales).
+_STATIC_SCALE_KEYS = {
+    QuantAlgo.FP8: ("weight_scale",),
+    QuantAlgo.FP8_BLOCK_SCALES: ("weight_scale",),
+    QuantAlgo.NVFP4: ("weight_scale", "weight_scale_2"),
+}
+
 
 class DynamicLinearWeightLoader:
     """
@@ -135,6 +147,57 @@ def _get_quant_algo_for_layer(self, name: str) -> Optional[QuantAlgo]:
 
         return None
 
+    def _check_static_quant_scales(
+        self, weight_dict: Dict[str, torch.Tensor], quant_algo: Optional[QuantAlgo], name: str
+    ) -> None:
+        """Refuse static quant recipes against checkpoints without scales.
+
+        With ``dynamic_weight_quant=False`` the loader expects the checkpoint
+        to provide pre-quantized weights plus their scale tensors. If a module
+        was built for a quantized recipe but the checkpoint holds a
+        high-precision weight with no scales (an unquantized checkpoint),
+        ``Linear.load_weights`` would silently cast the weight into the
+        quantized buffer while the scales keep their default or uninitialized
+        values, corrupting the model without any error. Fail fast instead.
+
+        Algos without an entry in ``_STATIC_SCALE_KEYS`` (e.g. the AWQ
+        variants, which also allocate ``weight_scale`` uninitialized) fail
+        closed: a high-precision weight under any static recipe is refused
+        even when the expected scale names are unknown.
+        """
+        if quant_algo is None or self.dynamic_weight_quant:
+            return
+
+        if self.quant_config is not None:
+            if self.quant_config.is_module_excluded_from_quantization(name):
+                return
+
+        weight = weight_dict.get("weight")
+        if weight is None or weight.dtype not in (torch.bfloat16, torch.float16, torch.float32):
+            return
+
+        expected_scales = _STATIC_SCALE_KEYS.get(quant_algo)
+        if expected_scales is not None:
+            missing = [key for key in expected_scales if key not in weight_dict]
+            if not missing:
+                return
+            detail = f"without the expected scale tensor(s) {missing}"
+        else:
+            detail = (
+                "and no checkpoint scale layout is registered for this algo in "
+                "_STATIC_SCALE_KEYS, so the scales cannot be verified (the "
+                "guard fails closed for unverified algos)"
+            )
+        raise ValueError(
+            f"Static quantization ({quant_algo.name}) is configured for module "
+            f"'{name}', but the checkpoint provides a {weight.dtype} weight "
+            f"{detail}. The checkpoint appears to be unquantized; loading it "
+            "would silently corrupt the weights. Use a checkpoint that carries "
+            "quantized weights and scales (e.g. exported by ModelOpt), or "
+            "enable load-time quantization by setting 'dynamic': true in "
+            "quant_config."
+        )
+
     def _should_dynamic_quantize(
         self, weight_dict: Dict[str, torch.Tensor], quant_algo: Optional[QuantAlgo], name: str
     ) -> bool:
@@ -271,6 +334,11 @@ def load_linear_weights(
         else:
             quant_algo = self._get_quant_algo_for_layer(name)
 
+        # Static (pre-quantized) recipes must not be loaded from checkpoints
+        # that do not carry the expected scale tensors.
+        for weight_dict in weight_dicts:
+            self._check_static_quant_scales(weight_dict, quant_algo, name)
+
         # Special handling for fused NVFP4 dynamic quantization
         # Fused weights (Q,K,V or gate,up) must be quantized TOGETHER
         # to ensure consistent global scale

tests/integration/test_lists/test-db/l0_cpu.yml

@@ -47,6 +47,7 @@ l0_cpu:
   - unittest/_torch/visual_gen/test_flux_infer.py
   - unittest/_torch/visual_gen/test_ltx2_pipeline.py
   - unittest/_torch/visual_gen/test_ltx2_transformer.py
+  - unittest/_torch/visual_gen/test_quant_static_guard.py
   - unittest/_torch/visual_gen/test_teacache.py
   - unittest/_torch/visual_gen/test_tensor_payload.py
   - unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py

tests/unittest/_torch/visual_gen/test_quant_static_guard.py

@@ -0,0 +1,115 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+"""Static quant recipes must be refused against unquantized checkpoints.
+
+With ``dynamic_weight_quant=False`` the loader expects pre-quantized weights
+plus their scale tensors from the checkpoint. Loading a high-precision
+checkpoint through a static recipe silently corrupts the weights (scales stay
+at their default or uninitialized values), so the loader must raise instead.
+"""
+
+import pytest
+import torch
+
+from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig
+from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader
+from tensorrt_llm.models.modeling_utils import QuantConfig
+from tensorrt_llm.quantization.mode import QuantAlgo
+
+# Pure-CPU test (stub Linear, no device): runs in the CPU lane (l0_cpu.yml),
+# which selects with `-m cpu_only`; GPU stages deselect it via `not cpu_only`.
+pytestmark = pytest.mark.cpu_only
+
+
+class _StubLinear:
+    """Records load_weights calls; quant_algo resolves via the loader's global config."""
+
+    quant_config = None
+
+    def __init__(self):
+        self.loaded = None
+
+    def load_weights(self, weight_dicts):
+        self.loaded = weight_dicts
+
+
+def _make_loader(quant_algo, dynamic_weight_quant=False, exclude_modules=None):
+    model_config = DiffusionModelConfig(
+        quant_config=QuantConfig(quant_algo=quant_algo, exclude_modules=exclude_modules),
+        dynamic_weight_quant=dynamic_weight_quant,
+    )
+    return DynamicLinearWeightLoader(model_config)
+
+
+def _bf16_weights():
+    return {"weight": torch.zeros(8, 16, dtype=torch.bfloat16)}
+
+
+class TestStaticQuantGuard:
+    @pytest.mark.parametrize(
+        "quant_algo",
+        [QuantAlgo.FP8, QuantAlgo.FP8_BLOCK_SCALES, QuantAlgo.NVFP4],
+    )
+    def test_static_recipe_vs_unquantized_checkpoint_raises(self, quant_algo):
+        loader = _make_loader(quant_algo)
+        module = _StubLinear()
+        with pytest.raises(ValueError, match="appears to be unquantized"):
+            loader.load_linear_weights(module, "blocks.0.attn1.to_q", [_bf16_weights()])
+        assert module.loaded is None
+
+    @pytest.mark.parametrize(
+        "quant_algo",
+        [QuantAlgo.W4A16_AWQ, QuantAlgo.W4A8_AWQ, QuantAlgo.W8A8_SQ_PER_CHANNEL],
+    )
+    def test_unregistered_static_algo_fails_closed(self, quant_algo):
+        """Algos accepted by config parsing but absent from _STATIC_SCALE_KEYS
+        (their VisualGen checkpoint layout is unverified) must also refuse a
+        high-precision weight instead of silently skipping the check."""
+        loader = _make_loader(quant_algo)
+        module = _StubLinear()
+        with pytest.raises(ValueError, match="fails closed for unverified algos"):
+            loader.load_linear_weights(module, "blocks.0.attn1.to_q", [_bf16_weights()])
+        assert module.loaded is None
+
+    def test_static_fp8_checkpoint_with_scales_loads(self):
+        loader = _make_loader(QuantAlgo.FP8)
+        module = _StubLinear()
+        weight_dict = {
+            "weight": torch.zeros(8, 16, dtype=torch.float8_e4m3fn),
+            "weight_scale": torch.ones(1, dtype=torch.float32),
+            "input_scale": torch.ones(1, dtype=torch.float32),
+        }
+        loader.load_linear_weights(module, "blocks.0.attn1.to_q", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_static_nvfp4_checkpoint_with_scales_loads(self):
+        loader = _make_loader(QuantAlgo.NVFP4)
+        module = _StubLinear()
+        weight_dict = {
+            "weight": torch.zeros(8, 8, dtype=torch.uint8),
+            "weight_scale": torch.zeros(8, 1, dtype=torch.float8_e4m3fn),
+            "weight_scale_2": torch.ones(1, dtype=torch.float32),
+        }
+        loader.load_linear_weights(module, "blocks.0.attn1.to_q", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_excluded_module_keeps_high_precision_weights(self):
+        loader = _make_loader(QuantAlgo.FP8, exclude_modules=["proj_out"])
+        module = _StubLinear()
+        weight_dict = _bf16_weights()
+        loader.load_linear_weights(module, "proj_out", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_unquantized_recipe_is_unaffected(self):
+        loader = _make_loader(None)
+        module = _StubLinear()
+        weight_dict = _bf16_weights()
+        loader.load_linear_weights(module, "blocks.0.attn1.to_q", [weight_dict])
+        assert module.loaded == [weight_dict]
+
+    def test_dynamic_recipe_skips_the_guard(self):
+        loader = _make_loader(QuantAlgo.FP8, dynamic_weight_quant=True)
+        # The guard is a no-op for dynamic recipes; quantization happens later
+        # in _maybe_dynamic_quantize (GPU path, not exercised here).
+        loader._check_static_quant_scales(_bf16_weights(), QuantAlgo.FP8, "blocks.0.attn1.to_q")