diff --git a/tests/quantization/test_humming_mxfp4_block_fp8.py b/tests/quantization/test_humming_mxfp4_block_fp8.py index dbec989df3f9..115916e2baec 100644 --- a/tests/quantization/test_humming_mxfp4_block_fp8.py +++ b/tests/quantization/test_humming_mxfp4_block_fp8.py @@ -120,3 +120,92 @@ def test_default_activation_quant_config_defers_to_humming(): ) assert not qc.is_block_quantized assert qc.per_act_token_quant + + +def _ct_input_quant(**kwargs): + from compressed_tensors.quantization import QuantizationArgs + + return QuantizationArgs(num_bits=8, type="float", symmetric=True, **kwargs) + + +@pytest.mark.skipif(not has_humming(), reason="humming is not installed") +@pytest.mark.parametrize( + "input_quant_kwargs,expected_group_size", + [ + # The MXFP4xFP8_BLOCK recipe: dynamic FP8 activations, group 128. + (dict(dynamic=True, strategy="group", group_size=128), 128), + # Per-token dynamic FP8 activations keep Humming's per-token path. + (dict(dynamic=True, strategy="token"), 0), + ], +) +def test_checkpoint_activations_drive_humming_input_schema( + input_quant_kwargs, expected_group_size +): + """The checkpoint's ``input_activations`` pick the Humming activation + schema, keeping the group size that humming's own compressed-tensors + schema would drop for FP8.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_mxfp4 import ( # noqa: E501 + _humming_input_schema, + ) + from vllm.utils.humming import dtypes as humming_dtypes + + schema = _humming_input_schema(_ct_input_quant(**input_quant_kwargs)) + assert schema.a_dtype == humming_dtypes.float8e4m3 + assert schema.input_scale_group_size == expected_group_size + + +@pytest.mark.skipif(not has_humming(), reason="humming is not installed") +def test_weight_only_checkpoint_keeps_bf16_activations(): + """No ``input_activations`` means W4A16: Humming dequantizes the weights.""" + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_mxfp4 import ( # noqa: E501 + _humming_input_schema, + ) + + assert _humming_input_schema(None).a_dtype is None + + +@pytest.mark.skipif(not has_humming(), reason="humming is not installed") +def test_env_input_quant_config_overrides_checkpoint(monkeypatch: pytest.MonkeyPatch): + """VLLM_HUMMING_INPUT_QUANT_CONFIG wins: returning None leaves the + conversion on its env-driven path.""" + from vllm import envs + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_mxfp4 import ( # noqa: E501 + _humming_input_schema, + ) + + monkeypatch.setattr( + envs, "VLLM_HUMMING_INPUT_QUANT_CONFIG", {"dtype": "float8e4m3"} + ) + input_quant = _ct_input_quant(dynamic=True, strategy="group", group_size=128) + assert _humming_input_schema(input_quant) is None + + +def test_compressed_tensors_mxfp4_moe_backend_selects_humming( + monkeypatch: pytest.MonkeyPatch, +): + """``--moe-backend humming`` routes the compressed-tensors MXFP4 MoE + through the oracle instead of the Cutlass/Marlin device probe.""" + from types import SimpleNamespace + + from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import Mxfp4MoeBackend + from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 + compressed_tensors_moe_w4a4_mxfp4 as ct_mxfp4, + ) + + sentinel_experts = type("SentinelExperts", (), {}) + monkeypatch.setattr( + ct_mxfp4.CutlassExpertsMxfp4, "_supports_current_device", lambda: True + ) + monkeypatch.setattr( + ct_mxfp4, + "select_mxfp4_moe_backend", + lambda moe: (Mxfp4MoeBackend.HUMMING, sentinel_experts), + ) + + method = ct_mxfp4.CompressedTensorsW4A4Mxfp4MoEMethod( + SimpleNamespace(w13_num_shards=2, moe_backend="humming") + ) + + assert method.mxfp4_backend == Mxfp4MoeBackend.HUMMING + assert method.experts_cls is sentinel_experts + assert not method.use_cutlass_mxfp4 diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 4a8c72327862..fe569362a341 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from enum import Enum -from typing import TYPE_CHECKING, Literal, Union +from typing import TYPE_CHECKING, Any, Literal, Union import torch @@ -1307,6 +1307,7 @@ def convert_weight_to_mxfp4_moe_kernel_format( w2_bias: torch.Tensor | None = None, _cache_permute_indices: dict[torch.Size, torch.Tensor] | None = None, activation: MoEActivation | None = None, + humming_input_schema: Any | None = None, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -1357,8 +1358,13 @@ def convert_weight_to_mxfp4_moe_kernel_format( convert_to_humming_moe_kernel_format, ) + # A caller that knows the checkpoint's activation scheme passes a + # HummingInputSchema; otherwise the conversion falls back to + # VLLM_HUMMING_INPUT_QUANT_CONFIG. convert_to_humming_moe_kernel_format( - layer, quant_config={"quant_method": "mxfp4"} + layer, + quant_config={"quant_method": "mxfp4"}, + input_schema=humming_input_schema, ) return ( layer.w13_weight, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py index bcc930471da7..763e1f52be6a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe.py @@ -64,7 +64,9 @@ def get_moe_method( CompressedTensorsW4A4Mxfp4MoEMethod, ) - return CompressedTensorsW4A4Mxfp4MoEMethod(layer.moe_config) + return CompressedTensorsW4A4Mxfp4MoEMethod( + layer.moe_config, input_quant=input_quant + ) if quant_config._is_mxfp8(weight_quant): from .compressed_tensors_moe_w8a8_mxfp8 import ( diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py index f7d12c863593..1facbab1eb40 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py @@ -2,9 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import TYPE_CHECKING + import torch +from compressed_tensors.quantization import ( + QuantizationArgs, + QuantizationStrategy, + QuantizationType, +) import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm import envs from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import ( FusedMoeWeightScaleSupported, @@ -27,6 +35,7 @@ from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( B12X_BACKENDS, Mxfp4MoeBackend, + convert_weight_to_mxfp4_moe_kernel_format, make_mxfp4_moe_kernel, make_mxfp4_moe_quant_config, select_mxfp4_moe_backend, @@ -40,20 +49,82 @@ from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform +if TYPE_CHECKING: + from vllm.utils.humming import HummingInputSchema + logger = init_logger(__name__) +# moe_backend values whose experts class comes from the MXFP4 oracle instead of +# this method's own device probe. +ORACLE_MOE_BACKENDS = ("b12x", "humming") + +# Activation dtypes an MXFP4 checkpoint can pair with that Humming can consume. +_HUMMING_ACTIVATION_DTYPES: dict[tuple[str, int], str] = { + (QuantizationType.FLOAT, 8): "float8e4m3", + (QuantizationType.FLOAT, 4): "float4e2m1", +} + + +def _humming_input_schema( + input_quant: QuantizationArgs | None, +) -> "HummingInputSchema | None": + """Build the Humming activation schema from the checkpoint. + + Returns None when ``VLLM_HUMMING_INPUT_QUANT_CONFIG`` is set, so that the + env override keeps precedence over the checkpoint. + + This deliberately does not route through humming's + ``CompressedTensorsInputSchema``, which keeps the group size only for FP4 + activations and would silently turn a block-FP8 (group-128) activation into + a per-token one. + """ + if envs.VLLM_HUMMING_INPUT_QUANT_CONFIG: + return None + + from vllm.utils.humming import HummingInputSchema + + if input_quant is None or not input_quant.dynamic: + # W4A16: Humming dequantizes the MXFP4 weights, activations stay bf16. + # Static activation scales have no place to come from -- this method + # loads no input_scale -- so they fall back to weight-only as well. + return HummingInputSchema() + + a_dtype = _HUMMING_ACTIVATION_DTYPES.get((input_quant.type, input_quant.num_bits)) + if a_dtype is None: + raise NotImplementedError( + f"Humming MXFP4 MoE does not support " + f"{input_quant.type}{input_quant.num_bits} activations." + ) + + if input_quant.strategy == QuantizationStrategy.TOKEN: + group_size = 0 + elif input_quant.strategy in ( + QuantizationStrategy.GROUP, + QuantizationStrategy.TENSOR_GROUP, + ): + assert input_quant.group_size is not None + group_size = input_quant.group_size + else: + raise NotImplementedError( + f"Humming MXFP4 MoE does not support " + f"{input_quant.strategy} activation scales." + ) + + return HummingInputSchema(a_dtype=a_dtype, input_scale_group_size=group_size) + class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod): - def __init__(self, moe): + def __init__(self, moe, input_quant: QuantizationArgs | None = None): super().__init__(moe) self.group_size = 32 + self.input_quant = input_quant self.mxfp4_backend = Mxfp4MoeBackend.MARLIN # Backend selection must match the weight preparation below: CUTLASS - # swizzles scales, b12x and XPU consume checkpoint packing, and Marlin - # repacks weights and scales. + # swizzles scales, b12x and XPU consume checkpoint packing, Humming + # repacks into its own layout, and Marlin repacks weights and scales. self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device() self.experts_cls: type[mk.FusedMoEExperts] - if moe.moe_backend == "b12x": + if moe.moe_backend in ORACLE_MOE_BACKENDS: self.mxfp4_backend, experts_cls = select_mxfp4_moe_backend(moe) assert experts_cls is not None self.experts_cls = experts_cls @@ -147,7 +218,8 @@ def get_fused_moe_quant_config( w2_scale=layer.w2_weight_scale, ) else: - # b12x selects W4A8 or W4A16; XPU uses W4A4; Marlin uses W4A16. + # b12x selects W4A8 or W4A16; XPU uses W4A4; Marlin uses W4A16; + # Humming reads the schemas recorded by the weight conversion. return make_mxfp4_moe_quant_config( mxfp4_backend=self.mxfp4_backend, w1_scale=layer.w13_weight_scale, @@ -197,6 +269,19 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.w2_weight_scale = torch.nn.Parameter( torch.stack(swizzled_w2), requires_grad=False ) + elif self.mxfp4_backend == Mxfp4MoeBackend.HUMMING: + # Repacks the checkpoint's MXFP4 weights and e8m0 scales into the + # Humming kernel layout and records the per-sublayer LayerConfigs + # (and weight/input schemas) that the quant config reads back. + convert_weight_to_mxfp4_moe_kernel_format( + mxfp4_backend=self.mxfp4_backend, + layer=layer, + w13_weight=layer.w13_weight, + w2_weight=layer.w2_weight, + w13_weight_scale=layer.w13_weight_scale, + w2_weight_scale=layer.w2_weight_scale, + humming_input_schema=_humming_input_schema(self.input_quant), + ) elif self.mxfp4_backend in B12X_BACKENDS or current_platform.is_xpu(): pass else: