[Quantization] Support compressed-tensors MXFP4 MoE with FP8 activations via Mxfp4MoEMethod (#59815)

Signed-off-by: mgoin <mgoin64@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
Co-authored-by: Kevin H. Luu <khluu000@gmail.com>

tests/quantization/test_compressed_tensors.py

@@ -18,7 +18,7 @@ from compressed_tensors.quantization import (
 
 from tests.models.utils import check_logprobs_close
 from tests.quantization.utils import load_model_without_vllm_runner
-from vllm.config import set_current_vllm_config
+from vllm.config import VllmConfig, set_current_vllm_config
 from vllm.config.kernel import KernelConfig
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.kernels.linear import (
@@ -54,6 +54,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
     GroupShape,
     QuantKey,
     ScaleDesc,
+    kFp8Dynamic128Sym,
 )
 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
 from vllm.platforms import current_platform
@@ -1020,3 +1021,144 @@ def test_compressed_tensors_mxfp4(vllm_runner):
         llm.apply_model(check_model)
         output = llm.generate_greedy("Hello my name is", max_tokens=4)
         assert output
+
+
+def _mxfp4_moe_layer(monkeypatch, input_activations):
+    from types import SimpleNamespace
+
+    from tests.kernels.moe.utils import make_dummy_moe_config
+    from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend
+    from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
+    from vllm.model_executor.layers.quantization import mxfp4
+
+    # Exercise scheme routing and allocation without selecting a GPU kernel.
+    monkeypatch.setattr(
+        mxfp4,
+        "select_deepseek_v4_mxfp4_moe_backend",
+        lambda moe: (Mxfp4MoeBackend.NONE, None),
+    )
+    weights = {
+        "num_bits": 4,
+        "type": "float",
+        "strategy": "group",
+        "group_size": 32,
+        "symmetric": True,
+        "dynamic": False,
+    }
+    quant_config = CompressedTensorsConfig.from_config(
+        {
+            "format": "mxfp4-pack-quantized",
+            "config_groups": {
+                "group_0": {
+                    # Name targets only match if the layer's own checkpoint
+                    # projection names are used for the lookup.
+                    "targets": [r"re:.*experts\.\d+\.w\d+$"],
+                    "weights": weights,
+                    "input_activations": input_activations,
+                }
+            },
+        }
+    )
+    return RoutedExperts(
+        "model.layers.0.ffn.experts",
+        torch.bfloat16,
+        make_dummy_moe_config(num_experts=2, hidden_dim=64, intermediate_size=32),
+        quant_config,
+        expert_map_manager=SimpleNamespace(
+            local_num_experts=2,
+            placement_strategy="linear",
+            expert_map=None,
+            expert_mask=None,
+            routing_tables=None,
+            map_global_to_local=lambda index: index,
+        ),
+        ckpt_gate_proj_name="w1",
+        ckpt_down_proj_name="w2",
+        ckpt_up_proj_name="w3",
+    )
+
+
+def test_mxfp4_moe_with_fp8_inputs_shares_native_mxfp4_method(monkeypatch):
+    """MXFP4 weights with dynamic FP8 inputs load compressed-tensors names but
+    otherwise behave as the native MXFP4 method, so its backends, padding and
+    monolithic support are not reimplemented."""
+    from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
+
+    layer = _mxfp4_moe_layer(
+        monkeypatch,
+        {
+            "num_bits": 8,
+            "type": "float",
+            "strategy": "group",
+            "group_size": 128,
+            "symmetric": True,
+            "dynamic": True,
+        },
+    )
+
+    assert isinstance(layer.quant_method, Mxfp4MoEMethod)
+    assert not layer.quant_method.supports_pre_processed_weights
+    # The checkpoint's FP8 activations reach backends such as Humming that
+    # would otherwise default to BF16 activations.
+    assert layer.quant_method.activation_key == kFp8Dynamic128Sym
+    assert {"w13_weight_packed", "w2_weight_packed"} <= dict(
+        layer.named_parameters()
+    ).keys()
+    w13 = layer.w13_weight_packed
+
+    layer.quant_method.process_weights_after_loading(layer)
+
+    assert layer.w13_weight is w13
+    assert not hasattr(layer, "w13_weight_packed")
+
+
+@pytest.mark.skipif(
+    not current_platform.is_cuda(), reason="FP8 x MXFP4 linear kernels need CUDA"
+)
+def test_mxfp4_linear_with_fp8_inputs_requests_fp8_activations():
+    """Dense MXFP4 layers with dynamic FP8 inputs (e.g. GLM-5.3 dense MLPs)
+    must not be routed as W4A4."""
+    from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
+        CompressedTensorsW4A8Mxfp4,
+    )
+
+    weights = QuantizationArgs(
+        num_bits=4, type="float", strategy="group", group_size=32, symmetric=True
+    )
+    inputs = QuantizationArgs(
+        num_bits=8,
+        type="float",
+        strategy="group",
+        group_size=128,
+        symmetric=True,
+        dynamic=True,
+    )
+    config = CompressedTensorsConfig.from_config(
+        {"format": "mxfp4-pack-quantized", "config_groups": {}}
+    )
+    with set_current_vllm_config(VllmConfig()):
+        scheme = config._get_scheme_from_parts(weights, inputs)
+
+    assert isinstance(scheme, CompressedTensorsW4A8Mxfp4)
+    assert scheme.kernel.config.activation_quant_key == kFp8Dynamic128Sym
+
+
+def test_mxfp4_moe_with_mxfp4_inputs_keeps_w4a4_method(monkeypatch):
+    """Only FP8 inputs take the shared path; W4A4 keeps its own kernels."""
+    from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_mxfp4 import (  # noqa: E501
+        CompressedTensorsW4A4Mxfp4MoEMethod,
+    )
+
+    layer = _mxfp4_moe_layer(
+        monkeypatch,
+        {
+            "num_bits": 4,
+            "type": "float",
+            "strategy": "group",
+            "group_size": 32,
+            "symmetric": True,
+            "dynamic": True,
+        },
+    )
+
+    assert isinstance(layer.quant_method, CompressedTensorsW4A4Mxfp4MoEMethod)

vllm/model_executor/kernels/linear/mxfp4/humming.py

@@ -56,9 +56,17 @@ class HummingMxFp4LinearKernel(MxFp4LinearKernel):
         }
 
         convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
-        input_schema = quant_key_to_input_schema(self.config.activation_quant_key)
+        act_key = self.config.activation_quant_key
+        input_schema = quant_key_to_input_schema(act_key)
+        # Keep checkpoint-declared FP8 activations instead of falling back to
+        # A16 for group-scaled weights.
         self.layer_config = prepare_humming_linear_layer_config(
-            layer, quant_config, input_schema=input_schema
+            layer,
+            quant_config,
+            input_schema=input_schema,
+            allow_input_schema_fallback=(
+                act_key is None or act_key.dtype != current_platform.fp8_dtype()
+            ),
         )
         self.compute_config = get_humming_linear_compute_config()
         self.locks = create_humming_locks(layer.weight.device)

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

@@ -1480,6 +1480,7 @@ def convert_weight_to_mxfp4_moe_kernel_format(
     _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None,
     activation: MoEActivation | None = None,
     use_separated_a4w4: bool = False,
+    activation_key: QuantKey | None = None,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -1541,8 +1542,20 @@ def convert_weight_to_mxfp4_moe_kernel_format(
             convert_to_humming_moe_kernel_format,
         )
 
+        # The env override wins; otherwise use the checkpoint's activation
+        # scheme instead of Humming's default BF16 activations.
+        input_schema = None
+        if activation_key is not None and not envs.VLLM_HUMMING_INPUT_QUANT_CONFIG:
+            from vllm.model_executor.layers.quantization.utils.humming import (
+                quant_key_to_input_schema,
+            )
+
+            input_schema = quant_key_to_input_schema(activation_key)
         convert_to_humming_moe_kernel_format(
-            layer, quant_config={"quant_method": "mxfp4"}
+            layer,
+            quant_config={"quant_method": "mxfp4"},
+            input_schema=input_schema,
+            allow_input_schema_fallback=False,
         )
         return (
             layer.w13_weight,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py

@@ -44,6 +44,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsW4A4Mxfp4,
     CompressedTensorsW4A8Fp8,
     CompressedTensorsW4A8Int,
+    CompressedTensorsW4A8Mxfp4,
     CompressedTensorsW8A8Fp8,
     CompressedTensorsW8A8Int8,
     CompressedTensorsW8A8Mxfp8,
@@ -63,6 +64,11 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
     should_ignore_layer,
 )
 from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    GroupShape,
+    QuantKey,
+    create_fp8_quant_key,
+)
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     ParallelLMHead,
     VocabParallelEmbedding,
@@ -449,6 +455,39 @@ class CompressedTensorsConfig(QuantizationConfig):
             and is_symmetric
         )
 
+    @classmethod
+    def _is_mxfp4_w4a8_fp8(
+        cls, weight_quant: QuantizationArgs, input_quant: QuantizationArgs
+    ) -> bool:
+        if not cls._is_mxfp4(weight_quant) or input_quant is None:
+            return False
+
+        is_float_type = input_quant.type == QuantizationType.FLOAT
+        is_8_bits = input_quant.num_bits == 8
+        is_static_weight = not weight_quant.dynamic
+        is_dynamic_input = bool(input_quant.dynamic)
+        is_token_or_group = input_quant.strategy in (
+            QuantizationStrategy.TOKEN,
+            QuantizationStrategy.GROUP,
+        )
+
+        return (
+            is_float_type
+            and is_8_bits
+            and is_static_weight
+            and is_dynamic_input
+            and is_token_or_group
+            and input_quant.symmetric
+        )
+
+    @staticmethod
+    def _fp8_dynamic_input_key(input_quant: QuantizationArgs) -> QuantKey:
+        if input_quant.strategy == QuantizationStrategy.GROUP:
+            group_shape = GroupShape(1, input_quant.group_size)
+        else:
+            group_shape = GroupShape.PER_TOKEN
+        return create_fp8_quant_key(static=False, group_shape=group_shape)
+
     @staticmethod
     def _is_mxfp8(quant_args: QuantizationArgs) -> bool:
         if quant_args is None:
@@ -764,6 +803,11 @@ class CompressedTensorsConfig(QuantizationConfig):
                 )
             return CompressedTensorsW4A4Fp4()
 
+        if self._is_mxfp4_w4a8_fp8(weight_quant, input_quant):
+            return CompressedTensorsW4A8Mxfp4(
+                activation_quant_key=self._fp8_dynamic_input_key(input_quant)
+            )
+
         if self._is_mxfp4(weight_quant):
             return CompressedTensorsW4A4Mxfp4()
 

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py

@@ -34,8 +34,12 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
         # make sure quantization config for Linear can target it
         quant_config._add_fused_moe_to_target_scheme_map()
         unfused_names = [
-            layer_name + proj_name
-            for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
+            f"{layer_name}.0.{proj_name}"
+            for proj_name in (
+                layer.ckpt_gate_proj_name,
+                layer.ckpt_up_proj_name,
+                layer.ckpt_down_proj_name,
+            )
         ]
         # TODO: refactor this to use expert_mapping and check all layer numbers
         all_scheme_dicts = [
@@ -59,6 +63,16 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
         input_quant = scheme_dict.get("input_activations")
         format = scheme_dict.get("format")
 
+        if quant_config._is_mxfp4_w4a8_fp8(weight_quant, input_quant):
+            from .compressed_tensors_moe_w4a8_mxfp4 import (
+                CompressedTensorsW4A8Mxfp4MoEMethod,
+            )
+
+            return CompressedTensorsW4A8Mxfp4MoEMethod(
+                layer.moe_config,
+                activation_key=quant_config._fp8_dynamic_input_key(input_quant),
+            )
+
         if quant_config._is_mxfp4(weight_quant):
             from .compressed_tensors_moe_w4a4_mxfp4 import (
                 CompressedTensorsW4A4Mxfp4MoEMethod,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py

@@ -0,0 +1,34 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.model_executor.layers.fused_moe import RoutedExperts
+from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
+
+_PACKED_WEIGHTS = ("w13_weight", "w2_weight")
+
+
+class CompressedTensorsW4A8Mxfp4MoEMethod(Mxfp4MoEMethod):
+    """MXFP4 weights with dynamic FP8 activations in compressed-tensors format.
+
+    Only the checkpoint naming and the declared activation scheme differ from
+    the native MXFP4 path, so backend selection, weight conversion and
+    execution are inherited.
+    """
+
+    # The packed -> unpacked rename means a processed state dict no longer
+    # matches the parameters registered by create_weights.
+    supports_pre_processed_weights = False
+
+    def create_weights(self, layer: RoutedExperts, *args, **kwargs):
+        super().create_weights(layer, *args, **kwargs)
+        for name in _PACKED_WEIGHTS:
+            param = getattr(layer, name)
+            delattr(layer, name)
+            layer.register_parameter(f"{name}_packed", param)
+
+    def process_weights_after_loading(self, layer: RoutedExperts) -> None:
+        for name in _PACKED_WEIGHTS:
+            param = getattr(layer, f"{name}_packed")
+            delattr(layer, f"{name}_packed")
+            layer.register_parameter(name, param)
+        super().process_weights_after_loading(layer)

vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py

@@ -6,6 +6,7 @@ from .compressed_tensors_w4a4_mxfp4 import CompressedTensorsW4A4Mxfp4
 from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
 from .compressed_tensors_w4a8_fp8 import CompressedTensorsW4A8Fp8
 from .compressed_tensors_w4a8_int import CompressedTensorsW4A8Int
+from .compressed_tensors_w4a8_mxfp4 import CompressedTensorsW4A8Mxfp4
 from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
 from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
 from .compressed_tensors_w8a8_mxfp8 import CompressedTensorsW8A8Mxfp8
@@ -23,6 +24,7 @@ __all__ = [
     "CompressedTensorsW8A8Int8",
     "CompressedTensorsW8A8Fp8",
     "CompressedTensorsW4A4Mxfp4",
+    "CompressedTensorsW4A8Mxfp4",
     "CompressedTensorsW4A4Fp4",
     "CompressedTensorsW4A8Int",
     "CompressedTensorsW4A8Fp8",

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py

@@ -10,6 +10,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
     CompressedTensorsScheme,
 )
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
     kMxfp4Dynamic,
 )
 from vllm.model_executor.parameter import (
@@ -35,10 +36,10 @@ class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme):
     Otherwise: W4A16 weight-only via Marlin.
     """
 
-    def __init__(self):
+    def __init__(self, activation_quant_key: QuantKey = kMxfp4Dynamic):
         self.group_size = 32
         self.kernel = init_mxfp4_linear_kernel(
-            activation_quant_key=kMxfp4Dynamic,
+            activation_quant_key=activation_quant_key,
         )
 
     @classmethod

vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_mxfp4.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_w4a4_mxfp4 import (  # noqa: E501
+    CompressedTensorsW4A4Mxfp4,
+)
+from vllm.model_executor.layers.quantization.utils.quant_utils import QuantKey
+
+logger = init_logger(__name__)
+
+__all__ = ["CompressedTensorsW4A8Mxfp4"]
+
+
+class CompressedTensorsW4A8Mxfp4(CompressedTensorsW4A4Mxfp4):
+    """MXFP4 weights with dynamic FP8 activations in compressed-tensors format.
+
+    The weight layout matches W4A4; only the activation quantization differs.
+    The checkpoint's FP8 activations are used when a kernel supports them,
+    otherwise this falls back to the W4A4 activation format.
+    """
+
+    def __init__(self, activation_quant_key: QuantKey):
+        try:
+            super().__init__(activation_quant_key=activation_quant_key)
+        except ValueError as e:
+            logger.warning_once(
+                "No MXFP4 linear kernel supports %s activations on this "
+                "platform; falling back to MXFP4 activations. %s",
+                activation_quant_key,
+                e,
+            )
+            super().__init__()

vllm/model_executor/layers/quantization/mxfp4.py

@@ -38,7 +38,10 @@ from vllm.model_executor.layers.quantization.base_config import (
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    QuantKey,
+    is_layer_skipped,
+)
 from vllm.model_executor.utils import (
     is_weights_pre_processed,
     replace_parameter,
@@ -541,9 +544,12 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
 
     supports_pre_processed_weights = True
 
-    def __init__(self, moe: FusedMoEConfig):
+    def __init__(self, moe: FusedMoEConfig, activation_key: QuantKey | None = None):
         super().__init__(moe)
         self.weight_dtype = "mxfp4"
+        # Checkpoint-declared activation quantization. Only backends that pick
+        # their activation format from the layer (Humming) consume it.
+        self.activation_key = activation_key
         self.is_k3_situ_aiter = _use_k3_situ_aiter(moe)
         self.is_k3_situ_int4_gfx942 = _use_k3_situ_int4_gfx942(moe)
         self.experts_cls: type[mk.FusedMoEExperts] | None
@@ -822,6 +828,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
                 _cache_permute_indices=self._cache_permute_indices,
                 activation=self.moe.activation,
                 use_separated_a4w4=self.moe.use_mxfp4_w4a4_dsv4,
+                activation_key=self.activation_key,
             )
         )
 

vllm/model_executor/layers/quantization/utils/humming/linear.py

@@ -59,6 +59,7 @@ def prepare_humming_linear_layer_config(
     input_quant_config: dict | None = None,
     *,
     input_schema: "BaseInputSchema | None" = None,
+    allow_input_schema_fallback: bool = True,
 ) -> "LayerConfig":
     from vllm.utils.humming import (
         BaseInputSchema,
@@ -119,6 +120,7 @@ def prepare_humming_linear_layer_config(
         weight_schema=weight_schema,
         input_schema=input_schema,
         param_dtype=layer.params_dtype,
+        allow_fallback=allow_input_schema_fallback,
     )
     for name in ("input_scale", "input_scale_2"):
         if name != input_schema.static_tensor_scale_name:

vllm/models/deepseek_v4/nvidia/model.py

@@ -752,6 +752,7 @@ class DeepseekV4MoE(nn.Module):
             renormalize=config.norm_topk_prob,
             quant_config=quant_config,
             prefix=f"{prefix}.experts",
+            ckpt_names=("w1", "w2", "w3"),
             scoring_func=self.scoring_func,
             routed_scaling_factor=self.routed_scaling_factor,
             e_score_correction_bias=self.gate.e_score_correction_bias,