Fix block FP8 MTP in ModelOpt mixed checkpoints (#55513)
Signed-off-by: Shiyang Chen <shiychen@nvidia.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>
sychen52 committed
Sep 8, 2026 at 15:47 UTC
60ad959b6f1a5c8f602edbd608c8decbc0788c50
5 files changed
+119
-1
tests/models/qwen4_exp/test_config.py
+25
@@ -1,6 +1,7 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+from importlib import import_module
5
from types import SimpleNamespace
6
from unittest.mock import patch
7
@@ -19,6 +20,8 @@ from vllm.models.qwen4_exp.config import (
20
from vllm.models.qwen4_exp.nvidia.model_state import Qwen4ExpModelState
21
from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
22
23
+from ...utils import spawn_new_process_for_each_test
24
+
25
26
def _text_config(**kwargs) -> Qwen4ExpTextConfig:
27
values = {
@@ -89,6 +92,28 @@ def test_qwen4_exp_mtp_returns_sample_and_multi_streams() -> None:
92
assert returned_multi_hidden is multi_hidden
93
94
95
+@spawn_new_process_for_each_test
96
+@pytest.mark.parametrize("backend", ["amd", "nvidia"])
97
+def test_qwen4_exp_mtp_remaps_mixed_precision_layer_indices(backend: str) -> None:
98
+ mtp_module = import_module(f"vllm.models.qwen4_exp.{backend}.mtp")
99
+
100
+ quantized_layers = {
101
+ "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
102
+ "mtp.layers.0.mlp.experts": {
103
+ "quant_algo": "FP8_BLOCK_SCALES",
104
+ "group_size": 128,
105
+ },
106
+ }
107
+
108
+ assert mtp_module._remap_quantized_layers(quantized_layers, 48) == {
109
+ "model.language_model.layers.0.mlp.experts": {"quant_algo": "NVFP4"},
110
+ "mtp.layers.48.mlp.experts": {
111
+ "quant_algo": "FP8_BLOCK_SCALES",
112
+ "group_size": 128,
113
+ },
114
+ }
115
+
116
+
117
@pytest.mark.parametrize("wrapped_config", [False, True])
118
def test_qwen4_exp_mtp_override_sets_draft_config(
119
wrapped_config: bool,
tests/quantization/test_modelopt.py
+32
@@ -25,6 +25,7 @@ from vllm.model_executor.kernels.linear import (
25
MarlinNvFp4LinearKernel,
26
)
27
from vllm.model_executor.layers.attention import Attention
28
+from vllm.model_executor.layers.fused_moe import RoutedExperts
29
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
30
from vllm.model_executor.layers.quantization.modelopt import (
31
LINEAR_ALGOS,
@@ -224,6 +225,37 @@ def test_modelopt_mixed_precision_dispatches_every_linear_algo(algo):
225
assert isinstance(method, ModelOptLinearMethod), (algo, type(method).__name__)
226
227
228
+@pytest.mark.parametrize("algo", ["FP8_PB_WO", "FP8_BLOCK_SCALES"])
229
+def test_modelopt_mixed_precision_dispatches_block_fp8_moe(algo):
230
+ config = ModelOptMixedPrecisionConfig.from_config(
231
+ {
232
+ "quantization": {
233
+ "quant_algo": "MIXED_PRECISION",
234
+ "quantized_layers": {
235
+ "mtp.layers.48.mlp.experts": {
236
+ "quant_algo": algo,
237
+ "group_size": 128,
238
+ }
239
+ },
240
+ }
241
+ }
242
+ )
243
+ layer = MagicMock(spec=RoutedExperts)
244
+ expected = object()
245
+
246
+ with patch(
247
+ "vllm.model_executor.layers.quantization.modelopt.Fp8MoEMethod",
248
+ return_value=expected,
249
+ ) as method_cls:
250
+ method = config.get_quant_method(layer, "mtp.layers.48.mlp.experts")
251
+
252
+ assert method is expected
253
+ fp8_config, called_layer = method_cls.call_args.args
254
+ assert called_layer is layer
255
+ assert fp8_config.weight_block_size == [128, 128]
256
+ assert fp8_config.activation_scheme == "dynamic"
257
+
258
+
259
def test_modelopt_nvfp4_leaves_excluded_parallel_lm_head_unquantized():
260
config = ModelOptNvFp4Config(
261
is_checkpoint_nvfp4_serialized=True,
vllm/model_executor/layers/quantization/modelopt.py
+26
-1
@@ -60,6 +60,7 @@ from vllm.model_executor.layers.quantization.base_config import (
60
QuantizationConfig,
61
QuantizeMethodBase,
62
)
63
+from vllm.model_executor.layers.quantization.fp8 import Fp8Config, Fp8MoEMethod
64
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
65
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
66
FP8_SCALE_SENTINEL,
@@ -107,6 +108,11 @@ if TYPE_CHECKING:
108
109
logger = init_logger(__name__)
110
111
+# ``FP8_PB_WO`` is ModelOpt's canonical 2D block-FP8 name. Early composed
112
+# Qwen3.8-Flash-Next checkpoints used ``FP8_BLOCK_SCALES`` for the same tensor
113
+# layout, so retain it as a checkpoint-compatibility alias.
114
+_BLOCK_FP8_MOE_ALGOS = ("FP8_PB_WO", "FP8_BLOCK_SCALES")
115
+
116
# Single source of truth for the ModelOpt linear algos.
117
#
118
# Each entry is (owning config's name, mixed-precision sub-config attribute).
@@ -1491,11 +1497,28 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1497
self.w4a16_nvfp4_config = w4a16_nvfp4_config
1498
self.mxfp8_config = mxfp8_config
1499
1500
+ block_sizes = {
1501
+ int(layer_info.get("group_size", 128))
1502
+ for layer_info in quantized_layers.values()
1503
+ if layer_info.get("quant_algo", "").upper() in _BLOCK_FP8_MOE_ALGOS
1504
+ }
1505
+ if len(block_sizes) > 1:
1506
+ raise ValueError(
1507
+ "MIXED_PRECISION currently requires all block-FP8 MoE layers "
1508
+ f"to use one group_size, got {sorted(block_sizes)}."
1509
+ )
1510
+ block_size = next(iter(block_sizes), 128)
1511
+ self.fp8_block_config = Fp8Config(
1512
+ is_checkpoint_fp8_serialized=True,
1513
+ activation_scheme="dynamic",
1514
+ weight_block_size=[block_size, block_size],
1515
+ )
1516
+
1517
def has_blocked_weights(self) -> bool:
1518
# Same gate as ModelOptFp8Config.has_blocked_weights, resolved per
1519
# layer: "+quant_fp8" must be on as soon as any layer is block-scaled.
1520
return any(
1498
- info.get("quant_algo", "").upper() == "FP8_PB_WO"
1521
+ info.get("quant_algo", "").upper() in _BLOCK_FP8_MOE_ALGOS
1522
for info in self.quantized_layers.values()
1523
)
1524
@@ -1730,6 +1753,8 @@ class ModelOptMixedPrecisionConfig(ModelOptQuantConfigBase):
1753
return build_linear_method(getattr(self, subcfg_attr), quant_algo, prefix)
1754
1755
if isinstance(layer, RoutedExperts):
1756
+ if quant_algo in _BLOCK_FP8_MOE_ALGOS:
1757
+ return Fp8MoEMethod(self.fp8_block_config, layer)
1758
if quant_algo == "FP8":
1759
return ModelOptFp8MoEMethod(
1760
quant_config=self.fp8_config,
vllm/models/qwen4_exp/amd/mtp.py
+18
@@ -78,6 +78,17 @@ def _remap_ignored_layers(
78
return remapped
79
80
81
+def _remap_quantized_layers(
82
+ quantized_layers: dict[str, dict],
83
+ mtp_start_layer_idx: int,
84
+) -> dict[str, dict]:
85
+ """Map checkpoint MTP layer indices to standalone draft indices."""
86
+ return {
87
+ _remap_ignored_layers([name], mtp_start_layer_idx)[0]: layer_info
88
+ for name, layer_info in quantized_layers.items()
89
+ }
90
+
91
+
92
def _remap_mtp_weight_name(name: str) -> str | None:
93
"""Map Qwen4Exp checkpoint paths into the standalone draft model."""
94
@@ -136,6 +147,13 @@ def _make_draft_vllm_config(
147
"exclude_modules",
148
_remap_ignored_layers(exclude_modules, mtp_start_layer_idx),
149
)
150
+ quantized_layers = getattr(draft_quant_config, "quantized_layers", None)
151
+ if quantized_layers:
152
+ setattr( # noqa: B010
153
+ draft_quant_config,
154
+ "quantized_layers",
155
+ _remap_quantized_layers(quantized_layers, mtp_start_layer_idx),
156
+ )
157
158
draft_vllm_config = replace(
159
vllm_config,
vllm/models/qwen4_exp/nvidia/mtp.py
+18
@@ -77,6 +77,17 @@ def _remap_ignored_layers(
77
return remapped
78
79
80
+def _remap_quantized_layers(
81
+ quantized_layers: dict[str, dict],
82
+ mtp_start_layer_idx: int,
83
+) -> dict[str, dict]:
84
+ """Map checkpoint MTP layer indices to standalone draft indices."""
85
+ return {
86
+ _remap_ignored_layers([name], mtp_start_layer_idx)[0]: layer_info
87
+ for name, layer_info in quantized_layers.items()
88
+ }
89
+
90
+
91
def _remap_mtp_weight_name(name: str) -> str | None:
92
"""Map Qwen4Exp checkpoint paths into the standalone draft model."""
93
@@ -135,6 +146,13 @@ def _make_draft_vllm_config(
146
"exclude_modules",
147
_remap_ignored_layers(exclude_modules, mtp_start_layer_idx),
148
)
149
+ quantized_layers = getattr(draft_quant_config, "quantized_layers", None)
150
+ if quantized_layers:
151
+ setattr( # noqa: B010
152
+ draft_quant_config,
153
+ "quantized_layers",
154
+ _remap_quantized_layers(quantized_layers, mtp_start_layer_idx),
155
+ )
156
157
draft_vllm_config = replace(
158
vllm_config,