Support FP8 per block (weight + dynamic per token activation) export

modelopt/torch/export/convert_hf_config.py

@@ -104,6 +104,19 @@ def _quant_algo_to_group_config(quant_algo: str, group_size: int | None = None)
             },
             "weights": {"dynamic": False, "num_bits": 8, "type": "float", "group_size": gs},
         }
+    elif quant_algo == "FP8_PB":
+        # Block-wise FP8 (W8A8). Weights in gsxgs blocks; activations quantized
+        # dynamically per-token at runtime (no input_activations entry).
+        gs = group_size or 128
+        return {
+            "weights": {
+                "dynamic": False,
+                "num_bits": 8,
+                "type": "float",
+                "strategy": "block",
+                "block_structure": [gs, gs],
+            },
+        }
     else:
         warnings.warn(
             f"Unsupported quantization algorithm '{quant_algo}' in "
@@ -166,6 +179,29 @@ def convert_hf_quant_config_format(input_config: dict[str, Any]) -> dict[str, An
     original_quantization_details = input_config.get("quantization", {})
     quant_algo_value = original_quantization_details.get("quant_algo")
 
+    # FP8_PB (block-wise FP8, W8A8): emit the native ``quant_method: fp8`` config
+    # vLLM/SGLang expect (weight_scale_inv + dynamic activations), matching the
+    # official Qwen3.5 FP8 checkpoint.
+    if quant_algo_value == "FP8_PB":
+        kv_cache_quant_algo = original_quantization_details.get("kv_cache_quant_algo")
+        assert not kv_cache_quant_algo, (
+            "FP8_PB export does not support kv_cache quantization yet "
+            f"(got kv_cache_quant_algo={kv_cache_quant_algo!r})."
+        )
+        group_size = original_quantization_details.get("group_size") or 128
+        exclude_modules = original_quantization_details.get("exclude_modules") or []
+        fp8_config: dict[str, Any] = {
+            "quant_method": "fp8",
+            "fmt": "e4m3",
+            "activation_scheme": "dynamic",
+            "weight_block_size": [group_size, group_size],
+            "modules_to_not_convert": exclude_modules,
+        }
+        producer_info = input_config.get("producer")
+        if producer_info:
+            fp8_config["producer"] = producer_info
+        return fp8_config
+
     # This structure is derived based on the example for "FP8" and "NVFP4"
     # TODO: Handle other quantization algorithms
     if quant_algo_value == "FP8":

modelopt/torch/export/model_config.py

@@ -42,6 +42,7 @@
 QUANTIZATION_NVFP4_AWQ = "nvfp4_awq"
 QUANTIZATION_FP8_PB_REAL = "fp8_pb_real"
 QUANTIZATION_FP8_PB_WO = "fp8_pb_wo"
+QUANTIZATION_FP8_PB_W8A8 = "fp8_pb_w8a8"
 QUANTIZATION_FP8_PC_PT = "fp8_pc_pt"
 
 KV_CACHE_FP8 = "FP8"

modelopt/torch/export/quant_utils.py

@@ -54,6 +54,7 @@
     KV_CACHE_NVFP4_AFFINE,
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
+    QUANTIZATION_FP8_PB_W8A8,
     QUANTIZATION_FP8_PB_WO,
     QUANTIZATION_FP8_PC_PT,
     QUANTIZATION_INT4_AWQ,
@@ -537,10 +538,12 @@ def _get_quantization_from_layer(layer, quantizer_attr_names: QuantizerAttrNames
                     and block_sizes.get("scale_bits") == (8, 0)
                 ):
                     return QUANTIZATION_MXFP8
-                if weight_quantizer.fake_quant:
-                    return QUANTIZATION_FP8_PB_WO
-                else:
+                # Block FP8: input quantizer enabled -> W8A8, else weight-only.
+                if not weight_quantizer.fake_quant:
                     return QUANTIZATION_FP8_PB_REAL
+                if input_quantizer is not None and input_quantizer.is_enabled:
+                    return QUANTIZATION_FP8_PB_W8A8
+                return QUANTIZATION_FP8_PB_WO
             if weight_quantizer.axis == 0:
                 return QUANTIZATION_FP8_PC_PT
             return QUANTIZATION_FP8
@@ -758,6 +761,12 @@ def process_layer_quant_config(layer_config_dict):
                 "quant_algo": "MXFP8",
                 "group_size": block_size_value,
             }
+        elif v == "fp8_pb_w8a8":
+            # Block-wise FP8, W8A8 at serve time.
+            layer_config = {
+                "quant_algo": "FP8_PB",
+                "group_size": block_size_value,
+            }
         else:
             layer_config = {"quant_algo": v}
 
@@ -865,7 +874,7 @@ def to_quantized_weight(
     if quantization == QUANTIZATION_MXFP8:
         return MXFP8QTensor.quantize_with_scale(weight, weights_scaling_factor)
 
-    if quantization == QUANTIZATION_FP8_PB_WO:
+    if quantization in (QUANTIZATION_FP8_PB_WO, QUANTIZATION_FP8_PB_W8A8):
         return FP8QTensor.quantize(
             weight, weights_scaling_factor.squeeze(), block_sizes={-1: block_size, -2: block_size}
         )[0]._quantized_data

modelopt/torch/export/unified_export_hf.py

@@ -78,6 +78,7 @@
 from .model_config import (
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
+    QUANTIZATION_FP8_PB_W8A8,
     QUANTIZATION_FP8_PC_PT,
     QUANTIZATION_MXFP8,
     QUANTIZATION_NONE,
@@ -741,9 +742,22 @@ def _export_quantized_weight(
 
     setattr(sub_module, weight_name, nn.Parameter(quantized_weight, requires_grad=False))
 
-    # Register the corrected weight_scale as a buffer
+    # Register the corrected weight scale as a buffer.
     if weight_scale is not None:
-        sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale)
+        if quantization_format == QUANTIZATION_FP8_PB_W8A8:
+            # Store per-block scale as 2-D weight_scale_inv (amax/448, not
+            # inverted); squeeze the keepdim block-amax [out, 1, in, 1].
+            if (
+                weight_scale.dim() == 4
+                and weight_scale.shape[1] == 1
+                and weight_scale.shape[3] == 1
+            ):
+                weight_scale = weight_scale.squeeze(3).squeeze(1)
+            sub_module.register_buffer(quantizer_attrs.weight_scale_inv, weight_scale)
+            if quantizer_attrs.weight_scale in sub_module._buffers:
+                del sub_module._buffers[quantizer_attrs.weight_scale]
+        else:
+            sub_module.register_buffer(quantizer_attrs.weight_scale, weight_scale)
 
     # Tied-weight dedup: if a previously-processed module shared the same
     # source weight memory, alias the packed weight + scale buffers so the
@@ -758,6 +772,7 @@ def _export_quantized_weight(
                 setattr(sub_module, weight_name, getattr(_prior, weight_name))
             for _attr in (
                 quantizer_attrs.weight_scale,
+                quantizer_attrs.weight_scale_inv,
                 quantizer_attrs.weight_scale_2,
                 quantizer_attrs.input_scale,
             ):

modelopt/torch/export/unified_export_megatron.py

@@ -41,6 +41,7 @@
     KV_CACHE_NVFP4,
     QUANTIZATION_FP8,
     QUANTIZATION_FP8_PB_REAL,
+    QUANTIZATION_FP8_PB_W8A8,
     QUANTIZATION_FP8_PB_WO,
     QUANTIZATION_NONE,
     QUANTIZATION_NVFP4,
@@ -296,6 +297,11 @@ def save_pretrained(
             quantization = "NVFP4"
         elif quantization_format == QUANTIZATION_W4A16_NVFP4:
             quantization = "W4A16_NVFP4"
+        elif quantization_format == QUANTIZATION_FP8_PB_W8A8:
+            raise NotImplementedError(
+                "Block-wise FP8 W8A8 (FP8_PB_W8A8) export is not supported on the "
+                "Megatron path; export via the HF path (unified_export_hf)."
+            )
 
         # We use the last PP rank and the 1st EP rank to write the config because
         # medusa_heads and eagle_module only exist in the last stage.
@@ -860,6 +866,11 @@ def _get_quantized_state(
         """
         name_to_value = {}
         qformat: str = self._get_quantization_format(module)
+        if qformat == QUANTIZATION_FP8_PB_W8A8:
+            raise NotImplementedError(
+                "Block-wise FP8 W8A8 (FP8_PB_W8A8) export is not supported on the "
+                "Megatron path; export via the HF path (unified_export_hf)."
+            )
         if qformat is None and "norm" not in prefix:
             self._record_excluded_module(prefix)
         block_size = get_weight_block_size(module)

modelopt/torch/quantization/utils/core_utils.py

@@ -272,6 +272,7 @@ def weight_attr_names(module: nn.Module) -> "Generator[str, None, None]":
         "input_quantizer",
         "output_quantizer",
         "weight_scale",
+        "weight_scale_inv",
         "weight_scale_2",
         "input_scale",
         "output_scale",
@@ -287,6 +288,7 @@ def quantizer_attr_names(weight_name: str = "weight") -> QuantizerAttrNames:
         input_quantizer=f"{prefix}input_quantizer",
         output_quantizer=f"{prefix}output_quantizer",
         weight_scale=f"{prefix}weight_scale",
+        weight_scale_inv=f"{prefix}weight_scale_inv",
         weight_scale_2=f"{prefix}weight_scale_2",
         input_scale=f"{prefix}input_scale",
         output_scale=f"{prefix}output_scale",

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

@@ -0,0 +1,41 @@
+# 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 FP8 E4M3 2D blockwise weights + dynamic per-token FP8
+# activations (W8A8). The dynamic input quantizer makes PTQ calibrate as W8A8;
+# it stores no input_scale, so exported weights match the weight-only preset.
+
+# modelopt-schema: modelopt.torch.quantization.config.QuantizeConfig
+imports:
+  base_disable_all: configs/ptq/units/base_disable_all
+  default_disabled_quantizers: configs/ptq/units/default_disabled_quantizers
+  fp8: configs/numerics/fp8
+
+algorithm: max
+quant_cfg:
+  - $import: base_disable_all
+  - quantizer_name: '*weight_quantizer'
+    cfg:
+      $import: fp8
+      block_sizes:
+        -1: 128
+        -2: 128
+  - quantizer_name: '*input_quantizer'
+    cfg:
+      $import: fp8
+      block_sizes:
+        -1: 128
+        type: dynamic
+  - $import: default_disabled_quantizers

tests/gpu/torch/export/test_export.py

@@ -141,6 +141,21 @@ def test_get_quantization_format(config, expected):
                 "exclude_modules": ["layer8"],
             },
         ),
+        (
+            {
+                "layer1.quantization": "fp8_pb_w8a8",  # 128x128 block-wise FP8 (W8A8)
+                "layer1.awq_block_size": 128,
+                "layer2.quantization": "fp8_pb_w8a8",
+                "layer2.awq_block_size": 128,
+                "layer8.quantization": None,
+            },
+            {
+                "quant_algo": "FP8_PB",
+                "kv_cache_quant_algo": None,
+                "group_size": 128,
+                "exclude_modules": ["layer8"],
+            },
+        ),
     ],
 )
 def test_process_layer_quant_config(layer_config_dict, expected_processed_dict):
@@ -520,3 +535,66 @@ def is_excluded(module_name: str) -> bool:
     assert not is_excluded("model.layers.0.mlp.experts.0.down_proj"), (
         f"Routed experts should not be excluded, got patterns: {exclude_modules}"
     )
+
+
+def test_fp8_pb_w8a8_export_uses_weight_scale_inv(tmp_path):
+    """W8A8 block-FP8 (FP8_PB) export stores per-block scales as weight_scale_inv
+    (DeepSeek/Qwen convention) and drops the plain weight_scale; activations are
+    dynamic so no input_scale is stored."""
+    from safetensors import safe_open
+
+    model = get_tiny_qwen3_moe().to("cuda")
+    model.config.architectures = ["Qwen3MoeForCausalLM"]
+
+    cfg = {
+        "quant_cfg": [
+            {"quantizer_name": "*", "enable": False},
+            {
+                "quantizer_name": "*weight_quantizer",
+                "cfg": {"num_bits": (4, 3), "block_sizes": {-1: 128, -2: 128}, "axis": None},
+                "enable": True,
+            },
+            {
+                "quantizer_name": "*input_quantizer",
+                "cfg": {
+                    "num_bits": (4, 3),
+                    "block_sizes": {-1: 128, "type": "dynamic"},
+                    "axis": None,
+                },
+                "enable": True,
+            },
+            {"quantizer_name": "*lm_head*", "enable": False},
+        ],
+        "algorithm": "max",
+    }
+    dummy_inputs = {k: v.to("cuda") for k, v in model.dummy_inputs.items()}
+    mtq.quantize(model, cfg, lambda m: m(**dummy_inputs))
+
+    export_dir = tmp_path / "fp8_pb_w8a8"
+    export_hf_checkpoint(model, export_dir=export_dir)
+
+    keys = set()
+    for st in export_dir.glob("*.safetensors"):
+        with safe_open(st, framework="pt") as f:
+            keys.update(f.keys())
+
+    assert any(k.endswith(".weight_scale_inv") for k in keys), (
+        "block FP8 must store weight_scale_inv"
+    )
+    assert not any(k.endswith(".weight_scale") for k in keys), "plain weight_scale must be dropped"
+    assert not any(k.endswith(".input_scale") for k in keys), (
+        "dynamic activations store no input_scale"
+    )
+
+    # Per-block scales must be 2-D [out_blocks, in_blocks] (DeepSeek/Qwen
+    # convention) -- not the 4-D [out_blocks, 1, in_blocks, 1] block-amax shape.
+    # This is what ModelOpt's _QuantFP8Linear reload path and vLLM/SGLang's stock
+    # block-FP8 loader expect.
+    for st in export_dir.glob("*.safetensors"):
+        with safe_open(st, framework="pt") as f:
+            for k in list(f.keys()):
+                if k.endswith(".weight_scale_inv"):
+                    scale_inv = f.get_tensor(k)
+                    assert scale_inv.ndim == 2, (
+                        f"{k} must be 2-D [out_blocks, in_blocks], got {tuple(scale_inv.shape)}"
+                    )

tests/unit/torch/export/test_get_quantization.py

@@ -26,16 +26,39 @@
 from modelopt.torch.export.layer_utils import get_quantization_format
 from modelopt.torch.export.model_config import (
     QUANTIZATION_FP8,
+    QUANTIZATION_FP8_PB_W8A8,
     QUANTIZATION_NVFP4,
     QUANTIZATION_W4A8_AWQ,
 )
 from modelopt.torch.export.quant_utils import get_quant_config
 from modelopt.torch.quantization.nn import NVFP4StaticQuantizer
 
+# Block-wise FP8 W8A8 on the ".1" linear: weight block + enabled dynamic input quantizer.
+_fp8_pb_w8a8_config = {
+    "quant_cfg": [
+        {"quantizer_name": "*", "enable": False},
+        {
+            "quantizer_name": "*.1.weight_quantizer",
+            "cfg": {"num_bits": (4, 3), "block_sizes": {-1: 128, -2: 128}, "axis": None},
+            "enable": True,
+        },
+        {
+            "quantizer_name": "*.1.input_quantizer",
+            "cfg": {"num_bits": (4, 3), "block_sizes": {-1: 128, "type": "dynamic"}, "axis": None},
+            "enable": True,
+        },
+    ],
+    "algorithm": "max",
+}
+
 
 @pytest.mark.parametrize(
     ("config", "expected"),
-    [(partial_fp8_config, QUANTIZATION_FP8), (partial_w4a8_config, QUANTIZATION_W4A8_AWQ)],
+    [
+        (partial_fp8_config, QUANTIZATION_FP8),
+        (partial_w4a8_config, QUANTIZATION_W4A8_AWQ),
+        (_fp8_pb_w8a8_config, QUANTIZATION_FP8_PB_W8A8),
+    ],
 )
 def test_get_quantization_format(config, expected):
     model = ToyModel()