From 40358750868669c52f0a0c709510af4f7196a2dd Mon Sep 17 00:00:00 2001 From: Chang Liu <9713593+chang-l@users.noreply.github.com> Date: Fri, 14 Aug 2026 02:11:35 -0700 Subject: [PATCH 1/5] [TRTLLM-15404][fix] VisualGen: refuse static quant recipes against unquantized checkpoints (silent weight corruption) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Static quant recipes (dynamic:false) against a BF16 checkpoint built quantized Linears but skipped load-time quantization, casting bf16 weights into fp8/fp4 buffers with default (FP8-QDQ: 1.0) or uninitialized (BLOCK_SCALES/NVFP4) scales — broken outputs with no error. Fail fast in load_linear_weights when a static recipe meets a high-precision weight without the expected scale tensors. Evidence: reproduced on Wan2.2-TI2V-5B (B200, 1.3.0rc24): dynamic:false vs BF16 ckpt = garbage output, no exception; dynamic:true works (fp8-bw -13.3%, LPIPS 0.125); ModelOpt FP8/NVFP4 static ckpts verified not to trip the guard. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> --- .../_torch/visual_gen/quantization/loader.py | 51 ++++++++++ .../visual_gen/test_quant_static_guard.py | 99 +++++++++++++++++++ 2 files changed, 150 insertions(+) create mode 100644 tests/unittest/_torch/visual_gen/test_quant_static_guard.py diff --git a/tensorrt_llm/_torch/visual_gen/quantization/loader.py b/tensorrt_llm/_torch/visual_gen/quantization/loader.py index 8b042ed25e38..b59045b09028 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/loader.py +++ b/tensorrt_llm/_torch/visual_gen/quantization/loader.py @@ -18,6 +18,15 @@ 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. +_STATIC_SCALE_KEYS = { + QuantAlgo.FP8: ("weight_scale",), + QuantAlgo.FP8_BLOCK_SCALES: ("weight_scale",), + QuantAlgo.NVFP4: ("weight_scale", "weight_scale_2"), +} + class DynamicLinearWeightLoader: """ @@ -135,6 +144,43 @@ 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. + """ + expected_scales = _STATIC_SCALE_KEYS.get(quant_algo) + if expected_scales 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 + + missing = [key for key in expected_scales if key not in weight_dict] + if missing: + raise ValueError( + f"Static quantization ({quant_algo.name}) is configured for module " + f"'{name}', but the checkpoint provides a {weight.dtype} weight " + f"without the expected scale tensor(s) {missing}. 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 +317,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 diff --git a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py new file mode 100644 index 000000000000..b103e4ca7ec5 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py @@ -0,0 +1,99 @@ +# 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 + +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 + + 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") From c9efdc39b66fb6abee18df2900af189620011fb6 Mon Sep 17 00:00:00 2001 From: Chang Liu <9713593+chang-l@users.noreply.github.com> Date: Tue, 18 Aug 2026 12:45:44 -0700 Subject: [PATCH 2/5] [TRTLLM-15404][fix] wire test_quant_static_guard into l0_b200 pre-merge visual_gen block Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index bc6a2427170f..344ffdea7d1a 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -213,6 +213,7 @@ l0_b200: - unittest/_torch/visual_gen/test_warmup.py - unittest/_torch/visual_gen/test_cache_dit.py - unittest/_torch/visual_gen/test_quant_ops.py + - unittest/_torch/visual_gen/test_quant_static_guard.py - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py From 938872c1ae2ab2268fb00d657e0b2375a8380466 Mon Sep 17 00:00:00 2001 From: Chang Liu <9713593+chang-l@users.noreply.github.com> Date: Tue, 18 Aug 2026 20:04:50 -0700 Subject: [PATCH 3/5] [TRTLLM-15404][fix] drop cpu_only marker so guard tests collect in the B200 unittest stage The L0 unittest wrapper always runs with -m 'not cpu_only' and no stage runs cpu_only tests, so the module was collected as 8 deselected / 0 selected -> pytest exit 5, reported as a failure (build 54726). Sibling visual_gen unittests listed in l0_b200.yml carry no cpu_only marker. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> --- tests/unittest/_torch/visual_gen/test_quant_static_guard.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py index b103e4ca7ec5..d9100046b0ad 100644 --- a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py +++ b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py @@ -17,8 +17,6 @@ from tensorrt_llm.models.modeling_utils import QuantConfig from tensorrt_llm.quantization.mode import QuantAlgo -pytestmark = pytest.mark.cpu_only - class _StubLinear: """Records load_weights calls; quant_algo resolves via the loader's global config.""" From c7a31125ccf0e799e034e8e4d69852b6df22f95b Mon Sep 17 00:00:00 2001 From: Chang Liu <9713593+chang-l@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:55:21 -0700 Subject: [PATCH 4/5] [TRTLLM-15404][fix] fail closed for static algos without a registered scale layout Review r3809965203: W4A16_AWQ / W4A8_AWQ (and W8A8_SQ_PER_CHANNEL) are accepted by the config algo_map but had no _STATIC_SCALE_KEYS entry, so the guard returned early and the uninitialized-scale corruption still applied (their LinearMethods allocate weight_scale with torch.empty). Restructure the guard so any static recipe seeing a high-precision weight on a non-excluded module raises: with the missing-scale detail when the algo's checkpoint layout is registered, or a fails-closed message when it is not. Verified: 11 unit tests green in the staging release container. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> --- .../_torch/visual_gen/quantization/loader.py | 41 +++++++++++++------ .../visual_gen/test_quant_static_guard.py | 14 +++++++ 2 files changed, 43 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/quantization/loader.py b/tensorrt_llm/_torch/visual_gen/quantization/loader.py index b59045b09028..8c1c33a59adc 100644 --- a/tensorrt_llm/_torch/visual_gen/quantization/loader.py +++ b/tensorrt_llm/_torch/visual_gen/quantization/loader.py @@ -21,6 +21,9 @@ # 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",), @@ -156,9 +159,13 @@ def _check_static_quant_scales( ``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. """ - expected_scales = _STATIC_SCALE_KEYS.get(quant_algo) - if expected_scales is None or self.dynamic_weight_quant: + if quant_algo is None or self.dynamic_weight_quant: return if self.quant_config is not None: @@ -169,17 +176,27 @@ def _check_static_quant_scales( if weight is None or weight.dtype not in (torch.bfloat16, torch.float16, torch.float32): return - missing = [key for key in expected_scales if key not in weight_dict] - if missing: - raise ValueError( - f"Static quantization ({quant_algo.name}) is configured for module " - f"'{name}', but the checkpoint provides a {weight.dtype} weight " - f"without the expected scale tensor(s) {missing}. 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." + 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 diff --git a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py index d9100046b0ad..7428fd16bd62 100644 --- a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py +++ b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py @@ -54,6 +54,20 @@ def test_static_recipe_vs_unquantized_checkpoint_raises(self, quant_algo): 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() From a4114618a360b3feeff9accbe00f0da10ce8e43e Mon Sep 17 00:00:00 2001 From: Chang Liu <9713593+chang-l@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:57:07 -0700 Subject: [PATCH 5/5] [TRTLLM-15404][test] move guard test to the CPU lane (l0_cpu.yml + cpu_only marker) Review r3809922300: the test uses a stub Linear and never touches a device, so run it in the CPU-Generic stages (which select with -m cpu_only) instead of spending B200 time. Restores the cpu_only pytestmark and moves the list entry from l0_b200.yml to l0_cpu.yml. Verified in the staging container: -m cpu_only selects and passes all 11 tests. Signed-off-by: Chang Liu <9713593+chang-l@users.noreply.github.com> --- tests/integration/test_lists/test-db/l0_b200.yml | 1 - tests/integration/test_lists/test-db/l0_cpu.yml | 1 + tests/unittest/_torch/visual_gen/test_quant_static_guard.py | 4 ++++ 3 files changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 4a9c515bb853..bfaa0c637875 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -223,7 +223,6 @@ l0_b200: - unittest/_torch/visual_gen/test_warmup.py - unittest/_torch/visual_gen/test_cache_dit.py - unittest/_torch/visual_gen/test_quant_ops.py - - unittest/_torch/visual_gen/test_quant_static_guard.py - unittest/_torch/visual_gen/test_attention_cute_dsl.py - unittest/_torch/visual_gen/test_attention_cute_dsl_vsa.py - unittest/_torch/visual_gen/test_attention_trtllm_sage.py diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index 4947141a35a9..d7c85557ce8d 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/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 diff --git a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py index 7428fd16bd62..af2a40e2bad9 100644 --- a/tests/unittest/_torch/visual_gen/test_quant_static_guard.py +++ b/tests/unittest/_torch/visual_gen/test_quant_static_guard.py @@ -17,6 +17,10 @@ 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."""