diff --git a/tests/models/qwen4_exp/test_config.py b/tests/models/qwen4_exp/test_config.py index 45af884a86..9d30365b81 100644 --- a/tests/models/qwen4_exp/test_config.py +++ b/tests/models/qwen4_exp/test_config.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from importlib import import_module from types import SimpleNamespace from unittest.mock import patch @@ -19,6 +20,8 @@ from vllm.models.qwen4_exp.config import ( from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState +from ...utils import spawn_new_process_for_each_test + def _text_config(**kwargs) -> Qwen4ExpTextConfig: values = { @@ -89,6 +92,28 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None: assert returned_multi_hidden is multi_hidden +@spawn_new_process_for_each_test +@pytest.mark.parametrize("backend", ["amd", "nvidia"]) +def test_qwen4_exp_mtp_remaps_mixed_precision_layer_indices(backend: str) -> None: + mtp_module = import_module(f"vllm.models.qwen4_exp.{backend}.mtp") + + quantized_layers = { + "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}, + "mtp.layers.0.mlp.experts": { + "quant_algo": "FP8_BLOCK_SCALES", + "group_size": 128, + }, + } + + assert mtp_module._remap_quantized_layers(quantized_layers, 48) == { + "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"}, + "mtp.layers.48.mlp.experts": { + "quant_algo": "FP8_BLOCK_SCALES", + "group_size": 128, + }, + } + + @pytest.mark.parametrize("wrapped_config", [False, True]) def test_qwen4_exp_mtp_override_sets_draft_config( wrapped_config: bool, diff --git a/tests/quantization/test_modelopt.py b/tests/quantization/test_modelopt.py index 151fbbf840..59a333bfe4 100644 --- a/tests/quantization/test_modelopt.py +++ b/tests/quantization/test_modelopt.py @@ -25,6 +25,7 @@ from vllm.model_executor.kernels.linear import ( MarlinNvFp4LinearKernel, ) from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import RoutedExperts from vllm.model_executor.layers.linear import UnquantizedLinearMethod from vllm.model_executor.layers.quantization.modelopt import ( LINEAR_ALGOS, @@ -224,6 +225,37 @@ def test_modelopt_mixed_precision_dispatches_every_linear_algo(algo): assert isinstance(method, ModelOptLinearMethod), (algo, type(method).__name__) +@pytest.mark.parametrize("algo", ["FP8_PB_WO", "FP8_BLOCK_SCALES"]) +def test_modelopt_mixed_precision_dispatches_block_fp8_moe(algo): + config = ModelOptMixedPrecisionConfig.from_config( + { + "quantization": { + "quant_algo": "MIXED_PRECISION", + "quantized_layers": { + "mtp.layers.48.mlp.experts": { + "quant_algo": algo, + "group_size": 128, + } + }, + } + } + ) + layer = MagicMock(spec=RoutedExperts) + expected = object() + + with patch( + "vllm.model_executor.layers.quantization.modelopt.Fp8MoEMethod", + return_value=expected, + ) as method_cls: + method = config.get_quant_method(layer, "mtp.layers.48.mlp.experts") + + assert method is expected + fp8_config, called_layer = method_cls.call_args.args + assert called_layer is layer + assert fp8_config.weight_block_size == [128, 128] + assert fp8_config.activation_scheme == "dynamic" + + def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized(): config = ModelOptNvFp4Config( is_checkpoint_nvfp4_serialized=True, diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index e2ab74c369..2abe60d7ad 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -60,6 +60,7 @@ from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, QuantizeMethodBase, ) +from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.utils.fp8_utils import ( FP8_SCALE_SENTINEL, @@ -107,6 +108,11 @@ if TYPE_CHECKING: logger = init_logger(__name__) +# ``FP8_PB_WO`` is ModelOpt's canonical 2D block-FP8 name. Early composed +# Qwen3.8-Flash-Next checkpoints used ``FP8_BLOCK_SCALES`` for the same tensor +# layout, so retain it as a checkpoint-compatibility alias. +_BLOCK_FP8_MOE_ALGOS = ("FP8_PB_WO", "FP8_BLOCK_SCALES") + # Single source of truth for the ModelOpt linear algos. # # Each entry is (owning config's name, mixed-precision sub-config attribute). @@ -1491,11 +1497,28 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): self.w4a16_nvfp4_config = w4a16_nvfp4_config self.mxfp8_config = mxfp8_config + block_sizes = { + int(layer_info.get("group_size", 128)) + for layer_info in quantized_layers.values() + if layer_info.get("quant_algo", "").upper() in _BLOCK_FP8_MOE_ALGOS + } + if len(block_sizes) > 1: + raise ValueError( + "MIXED_PRECISION currently requires all block-FP8 MoE layers " + f"to use one group_size, got {sorted(block_sizes)}." + ) + block_size = next(iter(block_sizes), 128) + self.fp8_block_config = Fp8Config( + is_checkpoint_fp8_serialized=True, + activation_scheme="dynamic", + weight_block_size=[block_size, block_size], + ) + def has_blocked_weights(self) -> bool: # Same gate as ModelOptFp8Config.has_blocked_weights, resolved per # layer: "+quant_fp8" must be on as soon as any layer is block-scaled. return any( - info.get("quant_algo", "").upper() == "FP8_PB_WO" + info.get("quant_algo", "").upper() in _BLOCK_FP8_MOE_ALGOS for info in self.quantized_layers.values() ) @@ -1730,6 +1753,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase): return build_linear_method(getattr(self, subcfg_attr), quant_algo, prefix) if isinstance(layer, RoutedExperts): + if quant_algo in _BLOCK_FP8_MOE_ALGOS: + return Fp8MoEMethod(self.fp8_block_config, layer) if quant_algo == "FP8": return ModelOptFp8MoEMethod( quant_config=self.fp8_config, diff --git a/vllm/models/qwen4_exp/amd/mtp.py b/vllm/models/qwen4_exp/amd/mtp.py index 5a77448eca..40ee6a6d47 100644 --- a/vllm/models/qwen4_exp/amd/mtp.py +++ b/vllm/models/qwen4_exp/amd/mtp.py @@ -78,6 +78,17 @@ def _remap_ignored_layers( return remapped +def _remap_quantized_layers( + quantized_layers: dict[str, dict], + mtp_start_layer_idx: int, +) -> dict[str, dict]: + """Map checkpoint MTP layer indices to standalone draft indices.""" + return { + _remap_ignored_layers([name], mtp_start_layer_idx)[0]: layer_info + for name, layer_info in quantized_layers.items() + } + + def _remap_mtp_weight_name(name: str) -> str | None: """Map Qwen4Exp checkpoint paths into the standalone draft model.""" @@ -136,6 +147,13 @@ def _make_draft_vllm_config( "exclude_modules", _remap_ignored_layers(exclude_modules, mtp_start_layer_idx), ) + quantized_layers = getattr(draft_quant_config, "quantized_layers", None) + if quantized_layers: + setattr( # noqa: B010 + draft_quant_config, + "quantized_layers", + _remap_quantized_layers(quantized_layers, mtp_start_layer_idx), + ) draft_vllm_config = replace( vllm_config, diff --git a/vllm/models/qwen4_exp/nvidia/mtp.py b/vllm/models/qwen4_exp/nvidia/mtp.py index 1650f84a30..d12b851d68 100644 --- a/vllm/models/qwen4_exp/nvidia/mtp.py +++ b/vllm/models/qwen4_exp/nvidia/mtp.py @@ -77,6 +77,17 @@ def _remap_ignored_layers( return remapped +def _remap_quantized_layers( + quantized_layers: dict[str, dict], + mtp_start_layer_idx: int, +) -> dict[str, dict]: + """Map checkpoint MTP layer indices to standalone draft indices.""" + return { + _remap_ignored_layers([name], mtp_start_layer_idx)[0]: layer_info + for name, layer_info in quantized_layers.items() + } + + def _remap_mtp_weight_name(name: str) -> str | None: """Map Qwen4Exp checkpoint paths into the standalone draft model.""" @@ -135,6 +146,13 @@ def _make_draft_vllm_config( "exclude_modules", _remap_ignored_layers(exclude_modules, mtp_start_layer_idx), ) + quantized_layers = getattr(draft_quant_config, "quantized_layers", None) + if quantized_layers: + setattr( # noqa: B010 + draft_quant_config, + "quantized_layers", + _remap_quantized_layers(quantized_layers, mtp_start_layer_idx), + ) draft_vllm_config = replace( vllm_config,