[ROCm][Quant] Requantize serialized MXFP8 linears to FP8 PTPC

docs/features/quantization/online.md

@@ -84,6 +84,23 @@ vllm serve openai/gpt-oss-20b --quantization-config.moe.activation mxfp8
 
 Combine with `--moe-backend` to pin a specific kernel family.
 
+### Requantize ModelOpt MXFP8 linears on ROCm
+
+On ROCm, serialized ModelOpt MXFP8 linears can be converted to
+per-channel-weight/per-token-activation FP8 at load time. For example:
+
+```bash
+vllm serve MiniMaxAI/MiniMax-M3-MXFP8 \
+  --linear-backend auto \
+  --quantization-config.linear fp8_per_channel
+```
+
+The ModelOpt source method loads the serialized MXFP8 values and E8M0 block
+scales, reconstructs the weight in BF16, and passes it to the generic online
+PTPC method. Checkpoint-excluded and user-ignored linears retain their original
+scheme. MoE remains on checkpoint quantization. This path requires an AITER
+preshuffled per-token FP8 kernel.
+
 ### Separate Schemes for Dense and MoE Layers
 
 You can apply different quantization schemes to dense linear layers and MoE expert layers via the `linear` and `moe` fields. Each accepts either a full spec dict, or a bare string naming an online shorthand (e.g. `"fp8_per_block"`) or weight format (e.g. `"fp8_per_block_static"`); fields not set fall back to the shorthand defaults.

tests/quantization/test_modelopt.py

@@ -6,6 +6,7 @@
 """
 
 import os
+from types import SimpleNamespace
 from typing import Any, NoReturn
 from unittest.mock import MagicMock, Mock, patch
 
@@ -14,14 +15,24 @@
 
 from tests.quantization.utils import is_quant_method_supported
 from vllm.config.model import ModelConfig
-from vllm.model_executor.layers.linear import UnquantizedLinearMethod
+from vllm.config.quantization import QuantizationConfigArgs
+from vllm.model_executor.layers.linear import LinearBase, UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.modelopt import (
     ModelOptFp8Config,
     ModelOptMixedPrecisionConfig,
     ModelOptMxFp8Config,
+    ModelOptMxFp8LinearMethod,
     ModelOptNvFp4Config,
     ModelOptNvFp4LinearMethod,
 )
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PtpcOnlineLinearMethod,
+)
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
+    MXFP8_SCALE_DTYPE,
+    MXFP8_VALUE_DTYPE,
+    dequant_mxfp8_to_bf16,
+)
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -240,6 +251,196 @@ def test_modelopt_mixed_precision_does_not_infer_missing_sibling_linear(
     assert isinstance(method, UnquantizedLinearMethod)
 
 
+@pytest.mark.parametrize(
+    ("is_rocm", "online_args", "uses_ptpc"),
+    [
+        (True, QuantizationConfigArgs(linear="fp8_per_channel"), True),
+        (
+            True,
+            QuantizationConfigArgs(
+                linear="fp8_per_channel",
+                ignore=["model.layers.0.self_attn.qkv_proj"],
+            ),
+            False,
+        ),
+        (True, None, False),
+        (False, QuantizationConfigArgs(linear="fp8_per_channel"), False),
+    ],
+)
+def test_modelopt_mxfp8_uses_generic_online_ptpc_override(
+    is_rocm, online_args, uses_ptpc
+):
+    prefix = "model.layers.0.self_attn.qkv_proj"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            dtype=torch.bfloat16,
+            hf_config=SimpleNamespace(model_type="not_minimax"),
+            quantization_config=online_args,
+        ),
+        kernel_config=SimpleNamespace(linear_backend="auto"),
+    )
+    config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+    )
+    layer = MagicMock(spec=LinearBase)
+
+    with (
+        patch(
+            "vllm.model_executor.layers.quantization.modelopt.get_current_vllm_config",
+            return_value=vllm_config,
+        ),
+        patch(
+            "vllm.model_executor.layers.quantization.online.fp8."
+            "get_current_vllm_config",
+            return_value=vllm_config,
+        ),
+        patch.object(current_platform, "is_rocm", return_value=is_rocm),
+        patch(
+            "vllm.model_executor.layers.quantization.modelopt.init_mxfp8_linear_kernel"
+        ),
+    ):
+        method = config.get_quant_method(layer, prefix)
+
+    if uses_ptpc:
+        assert isinstance(method, Fp8PtpcOnlineLinearMethod)
+        assert isinstance(method.requantization_source, ModelOptMxFp8LinearMethod)
+    else:
+        assert isinstance(method, ModelOptMxFp8LinearMethod)
+
+
+def test_modelopt_mxfp8_ptpc_rejects_incompatible_rocm_backend():
+    prefix = "model.layers.0.self_attn.qkv_proj"
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            dtype=torch.bfloat16,
+            quantization_config=QuantizationConfigArgs(linear="fp8_per_channel"),
+        ),
+        kernel_config=SimpleNamespace(linear_backend="emulation"),
+    )
+    config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+    )
+    layer = MagicMock(spec=LinearBase)
+
+    with (
+        patch(
+            "vllm.model_executor.layers.quantization.modelopt.get_current_vllm_config",
+            return_value=vllm_config,
+        ),
+        patch(
+            "vllm.model_executor.layers.quantization.online.fp8."
+            "get_current_vllm_config",
+            return_value=vllm_config,
+        ),
+        patch.object(current_platform, "is_rocm", return_value=True),
+        pytest.raises(
+            ValueError,
+            match="ModelOpt MXFP8.*use --linear-backend=auto",
+        ),
+    ):
+        config.get_quant_method(layer, prefix)
+
+
+def test_modelopt_mxfp8_ptpc_loads_and_requantizes_source_weight(monkeypatch):
+    class FakeAiterKernel:
+        def __init__(self):
+            self.processed = False
+
+        def process_weights_after_loading(self, layer):
+            self.processed = True
+
+    fake_kernel = FakeAiterKernel()
+    vllm_config = SimpleNamespace(
+        model_config=SimpleNamespace(dtype=torch.bfloat16),
+        kernel_config=SimpleNamespace(linear_backend="auto"),
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.online.fp8.get_current_vllm_config",
+        lambda: vllm_config,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.layers.quantization.online.fp8.init_fp8_linear_kernel",
+        lambda **kwargs: fake_kernel,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.kernels.linear."
+        "AiterPreshuffledPerTokenFp8ScaledMMLinearKernel",
+        FakeAiterKernel,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_rank",
+        lambda: 0,
+    )
+    monkeypatch.setattr(
+        "vllm.model_executor.parameter.get_tensor_model_parallel_world_size",
+        lambda: 1,
+    )
+    monkeypatch.setattr(current_platform, "is_rocm", lambda: True)
+
+    captured: list[torch.Tensor] = []
+
+    def fake_scaled_fp8_quant(weight, *, scale=None, **kwargs):
+        assert scale is None
+        assert kwargs == {"use_per_token_if_dynamic": True}
+        captured.append(weight.clone())
+        return (
+            torch.empty_like(weight, dtype=MXFP8_VALUE_DTYPE),
+            torch.ones((weight.shape[0], 1), dtype=torch.float32),
+        )
+
+    monkeypatch.setattr(
+        "vllm._custom_ops.scaled_fp8_quant",
+        fake_scaled_fp8_quant,
+    )
+
+    config = ModelOptMxFp8Config(
+        is_checkpoint_mxfp8_serialized=True,
+        kv_cache_quant_algo=None,
+        exclude_modules=[],
+    )
+    source_method = ModelOptMxFp8LinearMethod(config, init_kernel=False)
+    method = Fp8PtpcOnlineLinearMethod()
+    method.set_requantization_source(source_method)
+    layer = torch.nn.Module()
+
+    method.create_weights(
+        layer,
+        input_size_per_partition=64,
+        output_partition_sizes=[16],
+        input_size=64,
+        output_size=16,
+        params_dtype=torch.bfloat16,
+        weight_loader=MagicMock(),
+    )
+
+    assert method.uses_meta_device is False
+    assert layer.weight.shape == (16, 64)
+    assert layer.weight.dtype == MXFP8_VALUE_DTYPE
+    assert layer.weight_scale.shape == (16, 2)
+    assert layer.weight_scale.dtype == MXFP8_SCALE_DTYPE
+    assert method.fp8_linear is fake_kernel
+
+    source_weight = torch.linspace(-1.5, 1.5, 16 * 64).reshape(16, 64)
+    source_weight = source_weight.to(MXFP8_VALUE_DTYPE)
+    source_scale = torch.arange(32, dtype=torch.uint8).reshape(16, 2)
+    source_scale = (source_scale % 3 + 126).to(MXFP8_SCALE_DTYPE)
+    layer.weight.data.copy_(source_weight)
+    layer.weight_scale.data.copy_(source_scale)
+    expected_bf16 = dequant_mxfp8_to_bf16(source_weight, source_scale)
+
+    method.process_weights_after_loading(layer)
+
+    assert len(captured) == 1
+    torch.testing.assert_close(captured[0], expected_bf16, rtol=0, atol=0)
+    assert layer.weight.shape == (64, 16)
+    assert layer.weight_scale.shape == (16, 1)
+    assert fake_kernel.processed
+
+
 def test_vocab_parallel_embedding_weight_loader_accepts_scalar_scale():
     holder = Mock()
     scale = torch.nn.Parameter(torch.empty(1))

vllm/model_executor/layers/quantization/modelopt.py

@@ -10,6 +10,7 @@
 import vllm.envs as envs
 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
 from vllm.config import get_current_vllm_config
+from vllm.config.quantization import QuantizationConfigArgs
 from vllm.logger import init_logger
 from vllm.model_executor.kernels.linear import (
     MarlinNvFp4LinearKernel,
@@ -58,6 +59,12 @@
     QuantizeMethodBase,
 )
 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
+from vllm.model_executor.layers.quantization.online.base import (
+    OnlineQuantizationConfig,
+)
+from vllm.model_executor.layers.quantization.online.fp8 import (
+    Fp8PtpcOnlineLinearMethod,
+)
 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
     process_fp8_input_tensor_strategy_moe,
     process_fp8_weight_channel_strategy,
@@ -70,6 +77,7 @@
     MXFP8_BLOCK_SIZE,
     MXFP8_SCALE_DTYPE,
     MXFP8_VALUE_DTYPE,
+    dequant_mxfp8_to_bf16,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
@@ -93,6 +101,7 @@
     PerTensorScaleParameter,
 )
 from vllm.model_executor.utils import replace_parameter, set_weight_attrs
+from vllm.platforms import current_platform
 
 if TYPE_CHECKING:
     from vllm.model_executor.models.utils import WeightsMapper
@@ -1714,6 +1723,40 @@ def get_name(self) -> QuantizationMethods:
     def get_supported_act_dtypes(self) -> list[torch.dtype]:
         return [torch.bfloat16]
 
+    def get_quant_method(
+        self, layer: torch.nn.Module, prefix: str
+    ) -> "QuantizeMethodBase | None":
+        if current_platform.is_rocm() and isinstance(layer, LinearBase):
+            model_config = get_current_vllm_config().model_config
+            args = model_config.quantization_config
+            if isinstance(args, QuantizationConfigArgs) and not self.is_layer_excluded(
+                prefix
+            ):
+                target_method = OnlineQuantizationConfig(args).get_quant_method(
+                    layer, prefix
+                )
+                if isinstance(target_method, Fp8PtpcOnlineLinearMethod):
+                    linear_backend = (
+                        get_current_vllm_config().kernel_config.linear_backend
+                    )
+                    if linear_backend not in {"auto", "aiter"}:
+                        raise ValueError(
+                            "ModelOpt MXFP8 to FP8 PTPC requantization requires "
+                            "the AITER linear kernel; use --linear-backend=auto "
+                            "or --linear-backend=aiter, got "
+                            f"--linear-backend={linear_backend}."
+                        )
+                    source_method = ModelOptMxFp8LinearMethod(self, init_kernel=False)
+                    target_method.set_requantization_source(source_method)
+                    logger.info_once(
+                        "ModelOpt MXFP8 linear override: checkpoint MXFP8 + "
+                        "E8M0 -> BF16 -> FP8 PTPC; unspecified linears and "
+                        "MoE retain checkpoint quantization.",
+                        scope="global",
+                    )
+                    return target_method
+        return super().get_quant_method(layer, prefix)
+
     @classmethod
     def get_min_capability(cls) -> int:
         # Marlin kernel supports MXFP8 on SM80+
@@ -1779,7 +1822,9 @@ def _from_config(
 class ModelOptMxFp8LinearMethod(LinearMethodBase):
     """Linear method for ModelOpt MXFP8 quantization."""
 
-    def __init__(self, quant_config: ModelOptMxFp8Config) -> None:
+    def __init__(
+        self, quant_config: ModelOptMxFp8Config, *, init_kernel: bool = True
+    ) -> None:
         self.quant_config = quant_config
 
         if not self.quant_config.is_checkpoint_mxfp8_serialized:
@@ -1788,7 +1833,7 @@ def __init__(self, quant_config: ModelOptMxFp8Config) -> None:
                 "Dynamic quantization is not supported."
             )
 
-        self.kernel = init_mxfp8_linear_kernel()
+        self.kernel = init_mxfp8_linear_kernel() if init_kernel else None
 
     def create_weights(
         self,
@@ -1853,6 +1898,12 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
         if layer.weight.element_size() >= 2:
             return
 
+        self._validate_serialized_weight(layer)
+        assert self.kernel is not None
+        self.kernel.process_weights_after_loading(layer)
+
+    @staticmethod
+    def _validate_serialized_weight(layer: torch.nn.Module) -> None:
         # Validate weight tensor
         if layer.weight.ndim != 2:
             raise ValueError(
@@ -1876,14 +1927,20 @@ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
             f" got {layer.weight_scale.dtype}"
         )
 
-        self.kernel.process_weights_after_loading(layer)
+    def dequantize_weight(self, layer: torch.nn.Module) -> torch.Tensor:
+        """Reconstruct the serialized MXFP8 weight for online requantization."""
+        self._validate_serialized_weight(layer)
+        return dequant_mxfp8_to_bf16(
+            layer.weight.contiguous(), layer.weight_scale.contiguous()
+        )
 
     def apply(
         self,
         layer: torch.nn.Module,
         x: torch.Tensor,
         bias: torch.Tensor | None = None,
     ) -> torch.Tensor:
+        assert self.kernel is not None
         return self.kernel.apply_weights(layer, x, bias)