[NVBug 6143871] Fix awq_lite uncalibrated branch leaving input_quantizer disabled

modelopt/torch/quantization/model_calib.py

@@ -1316,6 +1316,21 @@ def postprocess(module, name):
                     dtype=w_dtype,
                     device=w_device,
                 )
+                # 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,84 @@ 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.
+
+    awq_lite.setup() disables the input_quantizer at the start of search. The
+    calibrated branch re-enables it inside postprocess(); the uncalibrated
+    branch (no cache-pass tokens, e.g. an MoE expert that never gets routed)
+    must do the same — otherwise downstream export (set_expert_quantizer_amax
+    + _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'.
+
+    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().cuda()
+
+    def _forward_loop(m):
+        for _ in range(2):
+            m(torch.randn(2, 16, 16, device="cuda"), branch="calibrated")
+
+    mtq.quantize(model, mtq.NVFP4_AWQ_LITE_CFG, _forward_loop)
+
+    assert model.calibrated.input_quantizer.is_enabled
+    assert model.uncalibrated.input_quantizer.is_enabled, (
+        "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():
     torch.manual_seed(1234)
     model = _SimpleMLP()