Add unscaled E5M2 fake quantization and Slurm eval support

examples/vllm_serve/README.md

@@ -100,6 +100,15 @@ MODELOPT_STATE_PATH=<vllm_fq_modelopt_state.pth> python vllm_serve_fakequant.py
 QUANT_CFG=<quant_cfg> QUANT_FILE_PATH=<quantizer_state.pth> python vllm_serve_fakequant.py <model_path> -tp 8 --host 0.0.0.0 --port 8000
 ```
 
+For scale-free E5M2 fake quantization, use the built-in `E5M2_DEFAULT_CFG`:
+
+```bash
+QUANT_CFG=E5M2_DEFAULT_CFG python vllm_serve_fakequant.py Qwen/Qwen3-4B --host 0.0.0.0 --port 8000
+```
+
+This simulates E5M2 by casting weights and activations directly to `torch.float8_e5m2`
+and back to their original dtype, without calibration or per-tensor scaling.
+
 ## Serve a model with sparse attention in vLLM
 
 Apply ModelOpt sparse attention at serve time. Right after model load, the launcher replaces each native attention implementation with its matching ModelOpt adapter: `ModelOptSparseAttentionImpl` for FlashAttention or `ModelOptSparseFlashInferImpl` for FlashInfer. Both adapters use the same Triton kernel with paged KV cache support.

modelopt/torch/quantization/config.py

@@ -428,6 +428,11 @@ def validate_num_bits(self):
             raise ValueError(
                 "Supported FPx quantization formats: FP8 (E4M3, E5M2), FP6(E3M2, E2M3), FP4(E2M1)."
             )
+        elif num_bits == (5, 2) and block_sizes is None:
+            if self.type != "dynamic" or self.axis is not None:
+                raise ValueError(
+                    "Unscaled E5M2 quantization requires type='dynamic' and axis=None."
+                )
         elif num_bits not in [(4, 3), (2, 1)] and (
             block_sizes is None or block_sizes.get("type", None) != "dynamic"
         ):
@@ -1659,6 +1664,7 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]:
     "configs/ptq/presets/model/int8_weight_only"
 )
 FP8_DEFAULT_CFG: dict[str, Any] = _load_quantize_config_dict("configs/ptq/presets/model/fp8")
+E5M2_DEFAULT_CFG: dict[str, Any] = _load_quantize_config_dict("configs/ptq/presets/model/e5m2")
 MAMBA_MOE_FP8_AGGRESSIVE_CFG: dict[str, Any] = _load_quantize_config_dict(
     "configs/ptq/presets/model/mamba_moe_fp8_aggressive"
 )
@@ -1755,6 +1761,7 @@ def _load_quantizer_cfg_dict_list(config_path: str) -> list[dict[str, Any]]:
 # DO NOT ADD NEW CONFIGS HERE. If you want to add a new general recipe, add it to
 # modelopt_recipes/general/ptq/ as a yaml file
 choices: set[str] = {
+    "E5M2_DEFAULT_CFG",
     "FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG",
     "FP8_AFFINE_KV_CFG",
     "FP8_DEFAULT_CFG",

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

@@ -56,6 +56,7 @@
 )
 from ...tensor_quant import (
     dynamic_block_quant,
+    fake_fp8,
     fake_tensor_quant,
     fp4_cast_ste,
     int_cast_ste,
@@ -895,7 +896,7 @@ def _fake_quantize(self, inputs):
             return entrypoint(inputs, self)
 
         amax = None
-        if not self.is_mx_format:
+        if not self.is_mx_format and self._num_bits != (5, 2):
             amax = self._get_amax(inputs)
 
         if self.block_sizes is not None and self.block_sizes.get("type", "static") == "dynamic":
@@ -921,7 +922,7 @@ def _fake_quantize(self, inputs):
             # Float-point quantization, e.g., FP8
             E, M = self._num_bits  # noqa: N806
 
-            outputs = scaled_e4m3(
+            outputs = fake_fp8(
                 inputs,
                 amax,
                 self._get_bias(inputs),
@@ -949,6 +950,9 @@ def _fake_quantize(self, inputs):
 
     def _check_onnx_readiness(self, inputs):
         """Check if quantizer is ready for ONNX export."""
+        if self._num_bits == (5, 2):
+            raise NotImplementedError("ONNX export does not support E5M2 fake quantization.")
+
         if not self.block_sizes or self.block_sizes.get("scale_bits", None) != (8, 0):
             assert hasattr(self, "_amax"), (
                 "Quantizer has not been calibrated. ONNX export requires the quantizer to be"

modelopt/torch/quantization/tensor_quant.py

@@ -64,6 +64,11 @@ def fp8_eager(x, amax):
     return _fp8_eager(x, amax)
 
 
+def e5m2_eager(x):
+    """Eager mode implementation of unscaled E5M2 quantization."""
+    return x.to(torch.float8_e5m2).to(x.dtype)
+
+
 def scaled_e4m3_impl(
     inputs: torch.Tensor,
     amax: torch.Tensor | None = None,
@@ -114,14 +119,18 @@ def fake_quant_impl(
 
 def _quantize_impl(
     inputs: torch.Tensor,
-    amax: torch.Tensor,
+    amax: torch.Tensor | None,
     num_bits: int = 8,
     exponent_bits: int = 0,
     unsigned: bool = False,
     narrow_range: bool = True,
 ):
     if num_bits == 8 and exponent_bits == 4:
         return scaled_e4m3_impl(inputs=inputs, amax=amax)
+    elif num_bits == 8 and exponent_bits == 5:
+        if amax is not None:
+            raise ValueError("E5M2 fake quantization does not support scaling.")
+        return e5m2_eager(inputs)
     elif isinstance(num_bits, int):
         return fake_quant_impl(
             inputs=inputs,
@@ -138,7 +147,7 @@ def _quantize_impl(
 
 def _quantize_impl_abstract(
     input: torch.Tensor,
-    amax: torch.Tensor,
+    amax: torch.Tensor | None,
     num_bits: int = 8,
     exponent_bits: int = 0,
     unsigned: bool = False,
@@ -222,7 +231,7 @@ def _dynamic_block_quantize_impl_abstract(
 try:
     torch.library.define(
         "tensorrt::quantize_op",
-        "(Tensor input, Tensor amax, int num_bits, int exponent_bits, "
+        "(Tensor input, Tensor? amax, int num_bits, int exponent_bits, "
         "bool unsigned, bool narrow_range) -> Tensor",
     )
     torch.library.define(
@@ -399,8 +408,8 @@ def backward(ctx, grad_outputs):
         return _fake_quant_backward_function(ctx, grad_outputs, num_args=10)
 
 
-class ScaledE4M3Function(Function):
-    """E4M3fy input with scale."""
+class FakeFP8Function(Function):
+    """Fake quantize input to a supported FP8 format."""
 
     @staticmethod
     @symbolic_helper.parse_args("v", "t", "t", "i", "i", "s", "b")
@@ -415,6 +424,8 @@ def symbolic(
         pass_through_bwd=False,
     ):
         """ONNX symbolic function."""
+        if E != 4 or M != 3:
+            raise NotImplementedError("ONNX export only supports E4M3 FP8 quantization.")
         from .export_onnx import export_fp8
 
         return export_fp8(g, inputs, amax, trt_high_precision_dtype)
@@ -432,8 +443,8 @@ def forward(
         pass_through_bwd=False,
     ):
         """Forward method."""
-        if E != 4 or M != 3:
-            raise NotImplementedError("Only support E=4 & M=3 for now.")
+        if (E, M) not in ((4, 3), (5, 2)):
+            raise NotImplementedError("Only E4M3 and E5M2 FP8 formats are supported.")
 
         if bias is not None:
             inputs = inputs - bias
@@ -444,7 +455,7 @@ def forward(
             inputs,
             amax,
             num_bits=8,
-            exponent_bits=4,
+            exponent_bits=E,
             unsigned=False,
             narrow_range=False,
         )
@@ -700,7 +711,8 @@ def backward(ctx, grad_outputs):
 
 
 fake_tensor_quant = FakeTensorQuantFunction.apply
-scaled_e4m3 = ScaledE4M3Function.apply
+fake_fp8 = FakeFP8Function.apply
+scaled_e4m3 = fake_fp8
 dynamic_block_quant = DynamicBlockQuantizationFunction.apply
 static_blockwise_fp4_fake_quant = StaticBlockwiseFP4FakeQuantFunction.apply
 fp4_cast_ste = FP4CastSTEFunction.apply

modelopt_recipes/configs/numerics/e5m2.yaml

@@ -0,0 +1,21 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+# Per-tensor unscaled FP8 E5M2 quantizer attributes.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizerAttributeConfig
+num_bits: e5m2
+axis:
+type: dynamic

modelopt_recipes/configs/ptq/presets/model/e5m2.yaml

@@ -0,0 +1,28 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+# QuantizeConfig preset for W8A8 FP8 E5M2 without scaling.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  w8a8_e5m2_e5m2: configs/ptq/units/w8a8_e5m2_e5m2
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+
+algorithm: max
+quant_cfg:
+  - $import: base_disable_all
+  - $import: w8a8_e5m2_e5m2
+  - $import: default_disabled_quantizers

modelopt_recipes/configs/ptq/units/w8a8_e5m2_e5m2.yaml

@@ -0,0 +1,27 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 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.
+
+# QuantizerCfgList snippet that enables unscaled FP8 E5M2 on weight and input quantizers.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizerCfgListConfig
+imports:
+  e5m2: configs/numerics/e5m2
+---
+  - quantizer_name: '*weight_quantizer'
+    cfg:
+      $import: e5m2
+  - quantizer_name: '*input_quantizer'
+    cfg:
+      $import: e5m2

tests/gpu/torch/quantization/test_quantize_cuda.py

@@ -84,6 +84,7 @@
     "config",
     [
         mtq.INT8_DEFAULT_CFG,
+        mtq.E5M2_DEFAULT_CFG,
         mtq.FP8_DEFAULT_CFG,
         mtq.W4A8_AWQ_BETA_CFG,
         mtq.INT8_SMOOTHQUANT_CFG,

tests/gpu/torch/quantization/test_tensor_quant_cuda.py

@@ -179,6 +179,17 @@ def test_zero_amax_per_channel_is_finite(self):
         assert torch.isfinite(xq).all()
 
 
+class TestUnscaledE5M2:
+    @pytest.mark.parametrize("device", ["cuda", "cpu"])
+    def test_matches_native_cast(self, device):
+        x = torch.randn(4, 4, device=device, dtype=torch.float32)
+        expected = x.to(torch.float8_e5m2).to(x.dtype)
+
+        actual = tensor_quant.fake_fp8(x, None, None, 5, 2)
+
+        assert torch.equal(actual, expected)
+
+
 class Testfp4:
     @pytest.mark.skipif(not NATIVE_E4M3_AVAILABLE, reason="Native E4M3 requires compute >= 8.9")
     def test_native_block_scale_underflows_to_zero(self):

tests/gpu/torch/quantization/test_tensor_quantizer_cuda.py

@@ -41,19 +41,30 @@ class TestBlockQuantCuda(BlockQuantTester):
 
 class TestTensorQuantizerE4M3:
     @pytest.mark.parametrize(
-        ("E", "M", "axis"), [(5, 2, None), (4, 3, None), (4, 3, 1), (7, 3, None)]
+        ("E", "M", "axis", "quant_type"),
+        [
+            (5, 2, None, "dynamic"),
+            (5, 2, None, "static"),
+            (5, 2, 1, "dynamic"),
+            (4, 3, None, "static"),
+            (4, 3, 1, "static"),
+            (7, 3, None, "static"),
+        ],
     )
-    def test_e4m3(self, E, M, axis):  # noqa: N803
-        is_error_expected = E != 4 or M != 3
+    def test_fp8(self, E, M, axis, quant_type):  # noqa: N803
+        is_e4m3 = (E, M) == (4, 3)
+        is_e5m2 = (E, M) == (5, 2) and axis is None and quant_type == "dynamic"
+        is_error_expected = not (is_e4m3 or is_e5m2)
         with pytest.raises(ValidationError) if is_error_expected else contextlib.nullcontext():
-            e4m3_desc = QuantizerAttributeConfig(num_bits=(E, M), axis=axis)
-            e4m3_quantizer = tensor_quantizer.TensorQuantizer(e4m3_desc).cuda()
+            fp8_desc = QuantizerAttributeConfig(num_bits=(E, M), axis=axis, type=quant_type)
+            fp8_quantizer = tensor_quantizer.TensorQuantizer(fp8_desc).cuda()
 
             x = torch.rand(3, 6, 7, 7).cuda()
 
-            e4m3_x = e4m3_quantizer(x)
-            ref = tensor_quant.scaled_e4m3(x, e4m3_quantizer._get_amax(x), None, E, M)
-            assert torch.allclose(e4m3_x, ref)
+            fp8_x = fp8_quantizer(x)
+            amax = None if is_e5m2 else fp8_quantizer._get_amax(x)
+            ref = tensor_quant.fake_fp8(x, amax, None, E, M)
+            assert torch.allclose(fp8_x, ref)
 
     def test_non_current_gpu(self, need_2_gpus):
         x = torch.randn(3, 4)

tests/unit/torch/quantization/test_config_validation.py

@@ -23,6 +23,7 @@
 import modelopt.torch.quantization as mtq
 from modelopt.torch.quantization.algorithms import _match_quantizer_cfg
 from modelopt.torch.quantization.config import (
+    E5M2_DEFAULT_CFG,
     FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG,
     FP8_DEFAULT_CFG,
     FP8_PER_CHANNEL_PER_TOKEN_CFG,
@@ -45,6 +46,7 @@
 
 
 def test_need_calibration():
+    assert not need_calibration(E5M2_DEFAULT_CFG)
     assert need_calibration(FP8_DEFAULT_CFG)
     assert not need_calibration(FP8_PER_CHANNEL_PER_TOKEN_CFG)
     assert not need_calibration(FP8_2D_BLOCKWISE_WEIGHT_ONLY_CFG)

tests/unit/torch/quantization/test_tensor_quant_cpu.py

@@ -23,6 +23,7 @@
 
 import modelopt.torch.quantization as mtq
 import modelopt.torch.quantization.nn.modules.tensor_quantizer as tensor_quantizer_module
+import modelopt.torch.quantization.tensor_quant as tensor_quant
 from modelopt.torch.quantization import QuantModuleRegistry
 from modelopt.torch.quantization.config import QuantizerAttributeConfig, RotateConfig
 from modelopt.torch.quantization.nn import (
@@ -37,6 +38,15 @@ class TestFakeTensorQuantCPU(FakeTensorQuantTester):
     device = "cpu"
 
 
+def test_unscaled_e5m2_matches_native_cast():
+    inputs = torch.tensor([-60000.0, -1.3, 0.0, 1.3, 60000.0], dtype=torch.float32)
+    expected = inputs.to(torch.float8_e5m2).to(inputs.dtype)
+
+    actual = tensor_quant.fake_fp8(inputs, None, None, 5, 2)
+
+    assert torch.equal(actual, expected)
+
+
 class TestQuantizerAttributeConfig:
     def test_scaled_mode(self):
         num_bits = np.random.randint(1, 16)
@@ -110,6 +120,25 @@ def test_num_bits(self):
         ):
             QuantizerAttributeConfig(enable=True, num_bits=(-1, 2))
 
+    def test_unscaled_e5m2_config(self):
+        config = QuantizerAttributeConfig(num_bits=(5, 2), axis=None, type="dynamic")
+        assert config.num_bits == (5, 2)
+
+        with pytest.raises(ValueError, match="requires type='dynamic' and axis=None"):
+            QuantizerAttributeConfig(num_bits=(5, 2), axis=None)
+
+        with pytest.raises(ValueError, match="requires type='dynamic' and axis=None"):
+            QuantizerAttributeConfig(num_bits=(5, 2), axis=0, type="dynamic")
+
+
+def test_unscaled_e5m2_tensor_quantizer():
+    quantizer = TensorQuantizer(
+        QuantizerAttributeConfig(num_bits=(5, 2), axis=None, type="dynamic")
+    )
+    inputs = torch.tensor([-60000.0, -1.3, 0.0, 1.3, 60000.0], dtype=torch.float32)
+
+    assert torch.equal(quantizer(inputs), inputs.to(torch.float8_e5m2).to(inputs.dtype))
+
 
 def _run_rotated_backend(monkeypatch, rotate):
     calls = []