feat(quant): add constant_amax to pin activation input_scale (NVFP4 experts input_scale=1.0)

CHANGELOG.rst

@@ -44,6 +44,7 @@ Changelog
 - Add Torch-TensorRT FP8 deployment example for HuggingFace ViT (``examples/torch_trt/``): ``torch_tensorrt_ptq.py`` covers ``mtq.quantize`` → ``torch_tensorrt.compile(ir="dynamo")``, and ``torch_tensorrt_accuracy.py`` reports the compiled model's ImageNet-1k top-1/top-5 accuracy via the ``onnx_ptq`` ``evaluate`` harness (the unquantized baseline is Torch-TensorRT-compiled too, for an apples-to-apples comparison). Ships a ViT-tuned FP8 PTQ recipe under ``modelopt_recipes/huggingface/vit/ptq/`` (``fp8.yaml``) composed from the shared ``modelopt_recipes/configs/`` units: it quantizes the encoder Linears, patch-embed ``nn.Conv2d``, ``classifier``, and per-block LayerNorm inputs plus the attention Q/K/V BMMs and softmax. Verified on ``google/vit-base-patch16-224`` (ImageNet-1k 50k validation): FP8 stays within 0.13 pp Top-1 of the FP16 baseline.
 - Add **AutoQuantize recipe** support: ``mtq.auto_quantize`` can be driven declaratively from a YAML recipe (``RecipeType.AUTO_QUANTIZE`` / ``AutoQuantizeConfig``) specifying candidate formats, the ``effective_bits`` target, cost model (incl. ``active_moe`` and ``excluded_module_name_patterns``), scoring method, and disabled layers. Adds an ``effective_bits`` cost-model override on ``QuantizeConfig`` / ``QuantizerAttributeConfig`` (block-scale-accurate NVFP4 = 4.5 via ``configs/numerics/nvfp4``). Shipped recipes live under ``modelopt_recipes/general/auto_quantize/`` and model-specific ones under ``modelopt_recipes/huggingface/<model>/auto_quantize/``.
 - Add ``rotate.mode`` to torch quantizer configs. The default ``"rotate"`` keeps the existing rotate-before-quantize behavior; ``"rotate_back"`` enables fake-quant rotate → quantize → rotate-back for TensorQuantizer.
+- Add a ``constant_amax`` ``QuantizerAttributeConfig`` field that pins a quantizer's ``amax`` to a fixed value and skips activation calibration (no forward statistics collected). Unlike ``use_constant_amax`` (which hardcodes 448.0 for KV-cache cast math and registers no buffer), ``constant_amax`` stores the constant on the ``_amax`` buffer so it is used by both the fake-quant forward and the exported scaling factor. For NVFP4 activations the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX)``, so ``constant_amax: 2688.0`` yields ``input_scale == 1.0``. Ships a new recipe ``modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml`` that applies this to the MoE expert activation quantizers.
 
 **Bug Fixes**
 

modelopt/torch/quantization/config.py

@@ -671,6 +671,43 @@ def validate_calibrator(cls, v, info: ValidationInfo):
         """,
     )
 
+    constant_amax: float | None = ModeloptField(
+        default=None,
+        title="Pin the quantizer amax to a constant value and skip calibration.",
+        description="""If set, the quantizer ``amax`` is fixed to this constant value and no
+        activation calibration is performed (no forward statistics are collected). The constant
+        is stored on the ``_amax`` buffer, so it is used by both the fake-quant forward pass and
+        the exported scaling factor (e.g. ``input_scale``).
+
+        This differs from ``use_constant_amax``, which pins amax to the FP8 E4M3 range (448.0)
+        for KV-cache cast math and does not register an ``_amax`` buffer. For NVFP4 activations
+        the exported ``input_scale`` equals ``amax / (E2M1_MAX * E4M3_MAX) = amax / (6 * 448)``;
+        setting ``constant_amax`` to ``2688.0`` therefore yields an exported ``input_scale`` of
+        ``1.0``.
+        """,
+    )
+
+    @field_validator("constant_amax")
+    @classmethod
+    def validate_constant_amax(cls, v):
+        """Validate that constant_amax, when set, is a positive value."""
+        assert v is None or v > 0, "constant_amax must be a positive value."
+        return v
+
+    @model_validator(mode="after")
+    def validate_constant_amax_modes(self):
+        """Forbid combining ``use_constant_amax`` and ``constant_amax``.
+
+        Both pin the amax but disagree on the value: ``use_constant_amax`` uses the FP8 E4M3
+        range (448.0) for the fake-quant forward and registers no ``_amax`` buffer, while
+        ``constant_amax`` pins ``_amax`` to its configured value for both forward and export.
+        Setting both would make the simulated and exported scales silently diverge.
+        """
+        assert not (self.use_constant_amax and self.constant_amax is not None), (
+            "use_constant_amax and constant_amax are mutually exclusive; set only one."
+        )
+        return self
+
 
 class LayerwiseConfig(ModeloptBaseConfig):
     """Nested config for layer-by-layer calibration behavior."""

modelopt/torch/quantization/model_calib.py

@@ -920,8 +920,8 @@ def enable_stats_collection(model: nn.Module):
     """Enable stats collection for all quantizers in the model."""
     for name, module in model.named_modules():
         if isinstance(module, TensorQuantizer) and not module._disabled:
-            if module._use_constant_amax:
-                # use_constant_amax quantizers use a fixed amax and don't need calibration.
+            if module._use_constant_amax or module._constant_amax is not None:
+                # Quantizers with a constant amax use a fixed amax and don't need calibration.
                 # Disable quantization during calibration so it doesn't affect other quantizers.
                 module.disable_quant()
                 continue
@@ -938,8 +938,8 @@ def finish_stats_collection(model: nn.Module, method: str | None = None, **kwarg
         if not isinstance(module, TensorQuantizer) or module._disabled:
             continue
 
-        if module._use_constant_amax:
-            # Re-enable quantization for use_constant_amax quantizers disabled in enable_stats_collection.
+        if module._use_constant_amax or module._constant_amax is not None:
+            # Re-enable quantization for constant-amax quantizers disabled in enable_stats_collection.
             module.enable_quant()
             continue
 

modelopt/torch/quantization/nn/modules/tensor_quantizer.py

@@ -202,6 +202,7 @@ def __init__(
             self.amax = amax
 
         self._use_constant_amax = False
+        self._constant_amax = None
         self.set_from_attribute_config(quant_attribute_cfg)
 
         self._if_quant = if_quant
@@ -244,6 +245,13 @@ def _block_sizes_setter(val):
                     self._calibrator._axis = None
             return val
 
+        def _constant_amax_setter(val):
+            if val is not None:
+                # Pin amax to the constant on the _amax buffer so it is used by both the
+                # fake-quant forward pass and export; calibration is skipped in model_calib.
+                self.amax = float(val)
+            return val
+
         # Some attributes need custom handling.
         # By default, attributes from config are mapped to a name ``f"_{attribute}"``
         _custom_setters: dict[str, tuple[str, Callable]] = {
@@ -255,6 +263,7 @@ def _block_sizes_setter(val):
             "backend": ("backend", lambda val: val),
             "backend_extra_args": ("backend_extra_args", lambda val: val or {}),
             "use_constant_amax": ("_use_constant_amax", lambda val: val),
+            "constant_amax": ("_constant_amax", _constant_amax_setter),
         }
 
         for attribute, val in attribute_cfg.items():
@@ -722,6 +731,10 @@ def _get_amax(self, inputs):
             return torch.tensor(torch.finfo(torch.float8_e4m3fn).max, device=inputs.device)
         if hasattr(self, "_amax"):
             amax = self._amax
+            # A constant_amax buffer is registered at config time (on CPU) and may not have
+            # followed a later `model.to(device)`; align it with the input device on the fly.
+            if amax.device != inputs.device:
+                amax = amax.to(inputs.device)
         else:
             reduce_axis = quant_utils.convert_quantization_axis_to_reduce_axis(inputs, self._axis)
             amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach()

modelopt_recipes/general/ptq/nvfp4_experts_only_input_scale1-kv_fp8_cast.yaml

@@ -0,0 +1,63 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Composed PTQ recipe for expert-only NVFP4 quantization with a fixed activation input_scale of
+# 1.0 (no activation calibration), plus FP8 KV-cache cast mode.
+#
+# The expert weight quantizers use standard NVFP4 (per-tensor weight scale + dynamic block
+# scales, calibrated via max). The expert input (activation) quantizers pin the per-tensor amax
+# to a constant 2688.0 = E2M1_MAX * E4M3_MAX (6 * 448) so the exported ``input_scale`` is exactly
+# ``amax / (E2M1_MAX * E4M3_MAX) = 1.0``; the per-block E4M3 activation scales remain dynamic.
+# No activation statistics are collected for these quantizers.
+
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  nvfp4: configs/numerics/nvfp4
+  kv_fp8_cast: configs/ptq/units/kv_fp8_cast
+
+metadata:
+  recipe_type: ptq
+  description: >-
+    Applies NVFP4 only to expert-layer weight and input quantizers, pinning the expert activation
+    input_scale to 1.0 (constant amax = 2688, no activation calibration), plus FP8 KV-cache cast
+    mode using constant amax; uses max calibration for the remaining (weight) quantizers.
+quantize:
+  algorithm:
+    method: max
+    # Max calibration is fast and does not typically need checkpointing.
+    # layerwise=false required for VLMs where the decoder layers are nested under
+    # `model.language_model.layers` (layerwise_calibrate can't find them otherwise).
+    layerwise: false
+  quant_cfg:
+    - $import: base_disable_all
+    - quantizer_name: '*.experts.*weight_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*.experts.*input_quantizer'
+      cfg:
+        $import: nvfp4
+        # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration.
+        constant_amax: 2688.0
+    - quantizer_name: '*block_sparse_moe*weight_quantizer'
+      cfg:
+        $import: nvfp4
+    - quantizer_name: '*block_sparse_moe*input_quantizer'
+      cfg:
+        $import: nvfp4
+        # 2688 = E2M1_MAX * E4M3_MAX (6 * 448) -> exported input_scale == 1.0, no calibration.
+        constant_amax: 2688.0
+    - $import: kv_fp8_cast
+    - $import: default_disabled_quantizers

tests/_test_utils/torch/quantization/tensor_quantizer_common.py

@@ -28,6 +28,7 @@
     max_calibrate,
 )
 from modelopt.torch.quantization.nn import QuantLinear, SequentialQuantizer, TensorQuantizer
+from modelopt.torch.quantization.qtensor.nvfp4_tensor import NVFP4QTensor
 
 
 class TensorQuantizerTester:
@@ -267,6 +268,81 @@ def test_use_constant_amax_skips_calibration(self):
         assert not model["tq_const"]._disabled
         assert model["tq_const"]._if_quant
 
+    def test_constant_amax_registers_amax(self):
+        """constant_amax pins the _amax buffer to the configured value (used by fwd and export)."""
+        tq = TensorQuantizer(
+            QuantizerAttributeConfig(
+                num_bits=(2, 1),
+                block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                constant_amax=2688.0,
+            )
+        ).to(self.device)
+
+        assert tq._constant_amax == 2688.0
+        assert hasattr(tq, "_amax")
+        assert float(tq._amax) == 2688.0
+
+        # Default (unset) constant_amax must not register an _amax buffer.
+        tq_default = TensorQuantizer(
+            QuantizerAttributeConfig(
+                num_bits=(2, 1),
+                block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+            )
+        ).to(self.device)
+        assert tq_default._constant_amax is None
+        assert not hasattr(tq_default, "_amax")
+
+    def test_constant_amax_nvfp4_input_scale(self):
+        """For NVFP4, constant_amax=2688 (= E2M1_MAX * E4M3_MAX) exports input_scale == 1.0."""
+        tq = TensorQuantizer(
+            QuantizerAttributeConfig(
+                num_bits=(2, 1),
+                block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                constant_amax=2688.0,
+            )
+        ).to(self.device)
+
+        input_scale = NVFP4QTensor.get_activation_scaling_factor(tq)
+        assert torch.allclose(input_scale.float(), torch.ones_like(input_scale.float()))
+
+    def test_constant_amax_skips_calibration(self):
+        """constant_amax quantizers are excluded from calibration and keep their fixed amax."""
+        model = nn.ModuleDict(
+            {
+                "tq_const": TensorQuantizer(
+                    QuantizerAttributeConfig(
+                        num_bits=(2, 1),
+                        block_sizes={-1: 16, "type": "dynamic", "scale_bits": (4, 3)},
+                        constant_amax=2688.0,
+                    )
+                ),
+                "tq_calib": TensorQuantizer(QuantizerAttributeConfig(num_bits=8)),
+            }
+        ).to(self.device)
+
+        enable_stats_collection(model)
+        # constant_amax quantizer: quant disabled during calibration, not collecting stats.
+        assert not model["tq_const"]._if_calib
+        assert not model["tq_const"]._if_quant
+        assert float(model["tq_const"]._amax) == 2688.0
+
+        finish_stats_collection(model)
+        # Re-enabled and amax is NOT overwritten by calibration.
+        assert model["tq_const"]._if_quant
+        assert float(model["tq_const"]._amax) == 2688.0
+
+    def test_constant_amax_must_be_positive(self):
+        """constant_amax must be a positive value."""
+        with pytest.raises(Exception, match="constant_amax must be a positive value"):
+            QuantizerAttributeConfig(num_bits=(2, 1), constant_amax=-1.0)
+        with pytest.raises(Exception, match="constant_amax must be a positive value"):
+            QuantizerAttributeConfig(num_bits=(2, 1), constant_amax=0.0)
+
+    def test_constant_amax_mutually_exclusive_with_use_constant_amax(self):
+        """use_constant_amax and constant_amax cannot both be set."""
+        with pytest.raises(Exception, match="mutually exclusive"):
+            QuantizerAttributeConfig(num_bits=8, use_constant_amax=True, constant_amax=2688.0)
+
     def test_modelopt_state(self):
         # Test loading of amax from ref to test
         tensor_quantizer_ref = TensorQuantizer(QuantizerAttributeConfig(num_bits=4), amax=10.0)