Address review: gate fix on is_input_quantized; expand tests

@meenchen: the unconditional input_quantizer.enable() in the
uncalibrated branch wrongly turned on input quantization for weight-only
AWQ configs (e.g. INT4_AWQ_CFG, where the user's config sets
*input_quantizer enable=False and setup() therefore never disabled it).
Gate the entire postprocess block — per-channel-amax collapse and
enable() — behind module.awq_lite.is_input_quantized so weight-only
configs are untouched.

@coderabbitai: strengthen the existing regression test to also assert
the export-critical scalar amax invariant (axis=None, numel==1) when
amax exists, and add a companion test on INT4_AWQ_CFG asserting the
uncalibrated linear's input_quantizer stays disabled.

The NVFP4 test now requires CUDA (dynamic block quantization is
CUDA-only), guarded with pytest.mark.skipif.

Signed-off-by: Chenjie Luo <chenjiel@nvidia.com>

Author: cjluo-nv

modelopt/torch/quantization/model_calib.py

@@ -1316,19 +1316,21 @@ def postprocess(module, name):
                     dtype=w_dtype,
                     device=w_device,
                 )
-                # Mirror the calibrated postprocess path: collapse any
-                # per-channel _amax left over from cache_mode max_calibrate
-                # into a per-tensor scalar (so preprocess_linear_fusion's
-                # numel==1 assertion passes), and re-enable the quantizer
-                # (awq_lite.setup disabled it). Without these, export drops
-                # input_scale and per-expert MoE loaders fail.
-                if module.input_quantizer.amax is not None:
-                    act_amax = module.input_quantizer.amax
-                    module.input_quantizer._amax_for_smoothing = act_amax.cpu()
-                    module.input_quantizer.reset_amax()
-                    module.input_quantizer.axis = None
-                    module.input_quantizer.amax = act_amax.amax()
-                module.input_quantizer.enable()
+                # Mirror the calibrated postprocess path, gated on
+                # is_input_quantized so weight-only AWQ configs (where
+                # setup() never disabled input_quantizer) stay untouched.
+                # Collapse any per-channel _amax left over from cache_mode
+                # max_calibrate into a per-tensor scalar so
+                # preprocess_linear_fusion's numel==1 assertion passes, and
+                # re-enable the quantizer (awq_lite.setup disabled it).
+                if module.awq_lite.is_input_quantized:
+                    if module.input_quantizer.amax is not None:
+                        act_amax = module.input_quantizer.amax
+                        module.input_quantizer._amax_for_smoothing = act_amax.cpu()
+                        module.input_quantizer.reset_amax()
+                        module.input_quantizer.axis = None
+                        module.input_quantizer.amax = act_amax.amax()
+                    module.input_quantizer.enable()
             else:
                 with enable_weight_access_and_writeback(module, model, name_to_module):
                     postprocess(module, name)

tests/unit/torch/quantization/test_calib.py

@@ -312,6 +312,21 @@ def test_padded_awq():
     model(torch.randn(2, 16, 16))
 
 
+class _TwoBranchModel(nn.Module):
+    """Two parallel linears; only the first is exercised by forward_loop."""
+
+    def __init__(self):
+        super().__init__()
+        self.calibrated = nn.Linear(16, 16, bias=False)
+        self.uncalibrated = nn.Linear(16, 16, bias=False)
+
+    def forward(self, x, branch="calibrated"):
+        if branch == "calibrated":
+            return self.calibrated(x)
+        return self.uncalibrated(x)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="NVFP4 dynamic block quant is CUDA-only")
 def test_awq_lite_uncalibrated_linear_keeps_input_quantizer_enabled():
     """Regression test for NVBug 6143871.
 
@@ -322,27 +337,16 @@ def test_awq_lite_uncalibrated_linear_keeps_input_quantizer_enabled():
     + _export_quantized_weight) drops the input_scale buffer and inference
     runtimes that read per-expert input_scale (e.g. TRT-LLM CutlassFusedMoE)
     crash with KeyError on '<idx>.w1.input_scale'.
-    """
-
-    class _TwoBranchModel(nn.Module):
-        """Two parallel linears; only the first is exercised by forward_loop."""
-
-        def __init__(self):
-            super().__init__()
-            self.calibrated = nn.Linear(16, 16, bias=False)
-            self.uncalibrated = nn.Linear(16, 16, bias=False)
-
-        def forward(self, x, branch="calibrated"):
-            if branch == "calibrated":
-                return self.calibrated(x)
-            return self.uncalibrated(x)
 
+    Also asserts the export-critical scalar amax invariant (axis=None,
+    numel==1) — preprocess_linear_fusion enforces it for fused-expert groups.
+    """
     torch.manual_seed(0)
-    model = _TwoBranchModel()
+    model = _TwoBranchModel().cuda()
 
     def _forward_loop(m):
         for _ in range(2):
-            m(torch.randn(2, 16, 16), branch="calibrated")
+            m(torch.randn(2, 16, 16, device="cuda"), branch="calibrated")
 
     mtq.quantize(model, mtq.NVFP4_AWQ_LITE_CFG, _forward_loop)
 
@@ -351,6 +355,39 @@ def _forward_loop(m):
         "Uncalibrated linear's input_quantizer must remain enabled after "
         "awq_lite postprocess so export emits input_scale (NVBug 6143871)."
     )
+    uncal_q = model.uncalibrated.input_quantizer
+    # When amax exists (cache-hit but search-miss path), it must be the
+    # scalar form export expects — preprocess_linear_fusion asserts numel==1.
+    # When it's None (truly never routed), set_expert_quantizer_amax will
+    # populate it during export.
+    if uncal_q.amax is not None:
+        assert uncal_q.axis is None
+        assert uncal_q.amax.numel() == 1
+
+
+def test_awq_lite_uncalibrated_weight_only_keeps_input_quantizer_disabled():
+    """Weight-only AWQ companion to NVBug 6143871.
+
+    For weight-only AWQ configs (input_quantizer disabled), awq_lite.setup()
+    never touches the input_quantizer, so the postprocess uncalibrated branch
+    must NOT enable it — doing so turns on quantization the user's config had
+    explicitly opted out of.
+    """
+    torch.manual_seed(0)
+    model = _TwoBranchModel()
+
+    def _forward_loop(m):
+        for _ in range(2):
+            m(torch.randn(2, 16, 16), branch="calibrated")
+
+    mtq.quantize(model, mtq.INT4_AWQ_CFG, _forward_loop)
+
+    assert not model.calibrated.input_quantizer.is_enabled
+    assert not model.uncalibrated.input_quantizer.is_enabled, (
+        "Weight-only AWQ must not flip on the input_quantizer for "
+        "uncalibrated layers — that would silently quantize activations "
+        "the user's config left in full precision."
+    )
 
 
 def test_smoothquant_enable_disable():