diff --git a/tests/quantization/test_compressed_tensors.py b/tests/quantization/test_compressed_tensors.py index 126ac92eef..8120cc35f1 100644 --- a/tests/quantization/test_compressed_tensors.py +++ b/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) diff --git a/vllm/model_executor/kernels/linear/mxfp4/humming.py b/vllm/model_executor/kernels/linear/mxfp4/humming.py index 4325b917db..0bfb060770 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/humming.py +++ b/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) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 88e80c755d..3abc0cbede 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/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, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py index a3fc11be37..f5b73ade53 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py +++ b/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() 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 d965127b81..be2eed2ce4 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 @@ -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, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a8_mxfp4.py new file mode 100644 index 0000000000..ae2c88044f --- /dev/null +++ b/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) diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py index 0f46d581f0..d10ec25107 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py +++ b/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", diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py index f3b02b0037..4a410bce6a 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py +++ b/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 diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_mxfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a8_mxfp4.py new file mode 100644 index 0000000000..c718645766 --- /dev/null +++ b/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__() diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 69ddd0045b..8f42f23970 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/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, ) ) diff --git a/vllm/model_executor/layers/quantization/utils/humming/linear.py b/vllm/model_executor/layers/quantization/utils/humming/linear.py index c82f4fd33f..e9e5d1b6bd 100644 --- a/vllm/model_executor/layers/quantization/utils/humming/linear.py +++ b/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: diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 275dc33491..c5d6a2d15a 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/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,