[Quant][Feature] Support online MXFP8 quantization for MoE and dense models

tests/models/quantization/test_mxfp8.py

@@ -0,0 +1,104 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""E2E tests for online MXFP8 quantization.
+
+Loads a BF16 model with ``--quantization mxfp8`` (online quantization) and
+compares log-probabilities against the same model served in BF16 without
+quantization.  This exercises the full pipeline: config parsing,
+``Mxfp8OnlineLinearMethod``, ``Mxfp8OnlineMoEMethod``, weight loading,
+online quantization / shuffling, and inference through ``apply_monolithic``.
+
+Layer skipping (``modules_to_not_convert``) is configured in the model's
+``config.json`` under ``quantization_config`` and is not tested here.
+
+``example_prompts`` is a pytest fixture (from conftest.py) that loads 8
+diverse prompts from ``tests/prompts/example.txt``.
+"""
+
+import pytest
+
+from tests.quantization.utils import is_quant_method_supported
+
+from ..utils import check_logprobs_close
+
+# A small MoE model that fits on a single GPU and has both linear + MoE layers.
+MOE_MODEL = "Qwen/Qwen3-30B-A3B"
+# A small dense model (no MoE) to validate the linear-only path.
+DENSE_MODEL = "Qwen/Qwen3-0.6B"
+
+MAX_MODEL_LEN = 1024
+MAX_TOKENS = 4
+NUM_LOG_PROBS = 8
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("mxfp8"),
+    reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
+)
+@pytest.mark.quant_model
+@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
+def test_mxfp8_logprobs(
+    vllm_runner,
+    example_prompts,
+    model: str,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """Compare BF16 baseline logprobs against online MXFP8-quantized model.
+
+    Runs the same model twice -- once in BF16 (baseline) and once with
+    online MXFP8 quantization -- then checks that the top log-probabilities
+    are close.  Only 4 tokens are generated to keep the test fast while
+    still catching numerical divergence.
+    """
+    with monkeypatch.context() as m:
+        m.setenv("TOKENIZERS_PARALLELISM", "true")
+
+        with vllm_runner(
+            model,
+            max_model_len=MAX_MODEL_LEN,
+            enforce_eager=True,
+        ) as vllm_model:
+            baseline_outputs = vllm_model.generate_greedy_logprobs(
+                example_prompts, MAX_TOKENS, NUM_LOG_PROBS
+            )
+
+        with vllm_runner(
+            model,
+            max_model_len=MAX_MODEL_LEN,
+            enforce_eager=True,
+            quantization="mxfp8",
+        ) as vllm_model:
+            test_outputs = vllm_model.generate_greedy_logprobs(
+                example_prompts, MAX_TOKENS, NUM_LOG_PROBS
+            )
+
+        check_logprobs_close(
+            outputs_0_lst=baseline_outputs,
+            outputs_1_lst=test_outputs,
+            name_0="bf16",
+            name_1="mxfp8",
+        )
+
+
+@pytest.mark.skipif(
+    not is_quant_method_supported("mxfp8"),
+    reason="mxfp8 is not supported on this GPU type (requires sm_100+).",
+)
+@pytest.mark.quant_model
+@pytest.mark.parametrize("model", [DENSE_MODEL, MOE_MODEL], ids=["dense", "moe"])
+def test_mxfp8_generation(vllm_runner, model: str) -> None:
+    """Smoke test: verify online MXFP8 model generates coherent text."""
+    prompt = "1 2 3 4 5"
+    with vllm_runner(
+        model,
+        enforce_eager=True,
+        quantization="mxfp8",
+        max_model_len=MAX_MODEL_LEN,
+    ) as vllm_model:
+        output = vllm_model.generate_greedy([prompt], max_tokens=5)
+
+    generated = output[0][1]
+    assert len(generated) > len(prompt), (
+        f"MXFP8 model produced no new tokens. Output: {generated!r}"
+    )

vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py

@@ -23,6 +23,8 @@
     kFp8Dynamic128Sym,
     kFp8Static128BlockSym,
     kFp8StaticTensorSym,
+    kMxfp8Dynamic,
+    kMxfp8Static,
 )
 from vllm.platforms import current_platform
 
@@ -67,11 +69,54 @@ def _supports_no_act_and_mul() -> bool:
         """Does not support non-gated MoE (i.e. Nanotron-3-Nano)."""
         return True
 
+    @staticmethod
+    def _supports_quant_scheme(
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        """Supports Fp8 per-tensor, Fp8 block, and MXFP8."""
+        SUPPORTED_W_A = [
+            (kFp8Static128BlockSym, kFp8Dynamic128Sym),
+            (kFp8StaticTensorSym, kFp8StaticTensorSym),
+            (kMxfp8Static, kMxfp8Dynamic),
+        ]
+        return (weight_key, activation_key) in SUPPORTED_W_A
+
     @staticmethod
     def _supports_activation(activation: MoEActivation) -> bool:
         """Supports only SiLU and RELU^2 non-gated activation."""
         return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL]
 
+    @staticmethod
+    def _supports_routing_method(
+        routing_method: RoutingMethodType,
+        weight_key: QuantKey | None,
+        activation_key: QuantKey | None,
+    ) -> bool:
+        """Monolithic kernels need to express router support."""
+        # NOTE(dbari): TopK routing could also be enabled, but need to validate models
+        # NOTE(dbari): Default is not implemented and should not be enabled until it is
+        if (weight_key, activation_key) in [
+            (kFp8Static128BlockSym, kFp8Dynamic128Sym),
+            (kMxfp8Static, kMxfp8Dynamic),
+        ]:
+            # NOTE(rob): potentially allow others here. This is a conservative list.
+            return routing_method in [
+                RoutingMethodType.DeepSeekV3,
+                RoutingMethodType.Renormalize,
+                RoutingMethodType.RenormalizeNaive,
+            ]
+        elif (weight_key, activation_key) == (kFp8StaticTensorSym, kFp8StaticTensorSym):
+            # NOTE(dbari): as above, potentially allow others here.
+            return routing_method in [
+                RoutingMethodType.DeepSeekV3,
+                RoutingMethodType.Llama4,
+                RoutingMethodType.Renormalize,
+                RoutingMethodType.RenormalizeNaive,
+            ]
+        else:
+            raise ValueError("Unsupported quantization scheme.")
+
     @staticmethod
     def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
         """Monolithic kernel so only use with naive DP/EP and TP."""
@@ -113,9 +158,10 @@ def _supports_quant_scheme(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        """Supports Fp8 block."""
+        """Supports Fp8 block and MXFP8."""
         SUPPORTED_W_A = [
             (kFp8Static128BlockSym, kFp8Dynamic128Sym),
+            (kMxfp8Static, kMxfp8Dynamic),
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A
 
@@ -159,6 +205,7 @@ def apply(
         apply_router_weight_on_input: bool,
     ):
         import flashinfer
+        from flashinfer.fused_moe import Fp8QuantizationType
 
         # Pack topk_ids and topk_weights into single tensor
         # Format: (expert_id << 16) | (weight_bf16.view(int16))
@@ -175,6 +222,16 @@ def apply(
 
         assert a1q_scale is not None
 
+        is_mxfp8 = self.quant_config.block_shape == [1, 32]
+        if is_mxfp8:
+            fp8_quant_type = Fp8QuantizationType.MxFp8
+            use_shuffled_weight = True
+            hidden_states_scale = a1q_scale
+        else:
+            fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
+            use_shuffled_weight = False
+            hidden_states_scale = a1q_scale.t().contiguous()
+
         # `trtllm_fp8_block_scale_routed_moe` has a bug and does not write to the
         # output tensor in-place so we need to manually copy the result to the
         # output tensor
@@ -183,7 +240,7 @@ def apply(
             topk_ids=packed_topk_ids,
             routing_bias=None,
             hidden_states=hidden_states,
-            hidden_states_scale=a1q_scale.t().contiguous(),  # type: ignore[union-attr]
+            hidden_states_scale=hidden_states_scale,
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale,
             gemm2_weights=w2,
@@ -197,8 +254,9 @@ def apply(
             local_num_experts=self.local_num_experts,
             routed_scaling_factor=None,
             routing_method_type=1,
-            use_shuffled_weight=False,
+            use_shuffled_weight=use_shuffled_weight,
             weight_layout=0,
+            fp8_quantization_type=fp8_quant_type,
             # output=output,
         )
         output.copy_(result)
@@ -240,10 +298,11 @@ def _supports_quant_scheme(
         weight_key: QuantKey | None,
         activation_key: QuantKey | None,
     ) -> bool:
-        """Supports Fp8 per-tensor and Fp8 block."""
+        """Supports Fp8 per-tensor, Fp8 block, and MXFP8."""
         SUPPORTED_W_A = [
             (kFp8Static128BlockSym, kFp8Dynamic128Sym),
             (kFp8StaticTensorSym, kFp8StaticTensorSym),
+            (kMxfp8Static, kMxfp8Dynamic),
         ]
         return (weight_key, activation_key) in SUPPORTED_W_A
 
@@ -256,7 +315,10 @@ def _supports_routing_method(
         """Monolithic kernels need to express router support."""
         # NOTE(dbari): TopK routing could also be enabled, but need to validate models
         # NOTE(dbari): Default is not implemented and should not be enabled until it is
-        if (weight_key, activation_key) == (kFp8Static128BlockSym, kFp8Dynamic128Sym):
+        if (weight_key, activation_key) in [
+            (kFp8Static128BlockSym, kFp8Dynamic128Sym),
+            (kMxfp8Static, kMxfp8Dynamic),
+        ]:
             # NOTE(rob): potentially allow others here. This is a conservative list.
             return routing_method in [
                 RoutingMethodType.DeepSeekV3,
@@ -274,7 +336,7 @@ def _supports_routing_method(
         else:
             raise ValueError("Unsupported quantization scheme.")
 
-    def _apply_per_block(
+    def _apply_block_scale(
         self,
         hidden_states: torch.Tensor,
         w1: torch.Tensor,
@@ -291,32 +353,38 @@ def _apply_per_block(
         routed_scaling_factor: float | None = None,
         topk_group: int | None = None,
     ) -> torch.Tensor:
-        # Delay import for non-CUDA.
         import flashinfer
+        from flashinfer.fused_moe import Fp8QuantizationType
 
         assert not apply_router_weight_on_input
         assert activation == MoEActivation.SILU
-
-        if self.routing_method_type == RoutingMethodType.DeepSeekV3:
-            router_logits = router_logits.to(torch.float32)
-
         assert self.topk <= global_num_experts
         assert self.topk <= 10
         assert global_num_experts % 4 == 0
-        assert self.quant_config.block_shape == [128, 128]
-        # Routing kernel expects #experts <= #threads 512
+        assert self.quant_config.block_shape in [[128, 128], [1, 32]]
+        # Kernel expects #experts <= #threads 512
         assert global_num_experts <= 512
-
-        # Kernel requires transposed hidden state scales
         # TODO: fuse into the quant kernel.
         assert a1q_scale is not None
-        a1q_scale_t = a1q_scale.t().contiguous()
+
+        if self.routing_method_type == RoutingMethodType.DeepSeekV3:
+            router_logits = router_logits.to(torch.float32)
+
+        is_mxfp8 = self.quant_config.block_shape == [1, 32]
+        if is_mxfp8:
+            fp8_quant_type = Fp8QuantizationType.MxFp8
+            use_shuffled_weight = True
+            hidden_states_scale = a1q_scale
+        else:
+            fp8_quant_type = Fp8QuantizationType.DeepSeekFp8
+            use_shuffled_weight = False
+            hidden_states_scale = a1q_scale.t().contiguous()
 
         return flashinfer.fused_moe.trtllm_fp8_block_scale_moe(
             routing_logits=router_logits,
             routing_bias=e_score_correction_bias,
             hidden_states=hidden_states,
-            hidden_states_scale=a1q_scale_t,
+            hidden_states_scale=hidden_states_scale,
             gemm1_weights=w1,
             gemm1_weights_scale=self.quant_config.w1_scale,
             gemm2_weights=w2,
@@ -330,7 +398,8 @@ def _apply_per_block(
             local_num_experts=self.local_num_experts,
             routed_scaling_factor=routed_scaling_factor,
             routing_method_type=self.routing_method_type,
-            use_shuffled_weight=False,
+            use_shuffled_weight=use_shuffled_weight,
+            fp8_quantization_type=fp8_quant_type,
         )
 
     def _apply_per_tensor(
@@ -409,7 +478,7 @@ def apply(
         topk_group: int | None = None,
     ) -> torch.Tensor:
         if self.quant_config.block_shape is not None:
-            return self._apply_per_block(
+            return self._apply_block_scale(
                 hidden_states,
                 w1,
                 w2,
@@ -441,6 +510,6 @@ def apply(
             )
         else:
             raise NotImplementedError(
-                "Only per-block and per-tensor quantization are supported in "
-                f"{self.__class__.__name__}."
+                "Only per-block, per-tensor, and MXFP8 quantization are "
+                f"supported in {self.__class__.__name__}."
             )

vllm/model_executor/layers/fused_moe/oracle/fp8.py

@@ -444,7 +444,7 @@ def convert_to_fp8_moe_kernel_format(
         Fp8MoeBackend.FLASHINFER_CUTLASS,
         Fp8MoeBackend.FLASHINFER_TRTLLM,
     ]:
-        w13, w2, w13_scale = prepare_fp8_moe_layer_for_fi(
+        w13, w2, w13_scale, w2_scale = prepare_fp8_moe_layer_for_fi(
             layer=layer,
             w13=w13,
             w2=w2,
@@ -512,6 +512,21 @@ def make_fp8_moe_quant_config(
             g1_alphas=(w1_scale * a1_scale).squeeze(),
             g2_alphas=(w2_scale * a2_scale).squeeze(),
         )
+    # MXFP8 uses "mxfp8" quant_dtype so the prepare step dispatches to
+    # _mxfp8_e4m3_quantize rather than standard FP8 block quantization.
+    # Non-swizzled layout is required since the TRTLLM kernel expects
+    # scales in (num_tokens, hidden_dim // 32) format.
+    if block_shape == [1, 32]:
+        return FusedMoEQuantConfig.make(
+            "mxfp8",
+            w1_scale=w1_scale,
+            w2_scale=w2_scale,
+            a1_scale=a1_scale,
+            a2_scale=a2_scale,
+            block_shape=block_shape,
+            is_nvfp4_scale_swizzled=False,
+        )
+
     # All other backends use normal config.
     return fp8_w8a8_moe_quant_config(
         w1_scale=w1_scale,