[Model] [Quantization] Add Ling hybrid MXFP4 routed experts support (#52114)
Signed-off-by: zexplorerhj <zhjoneson@163.com>
zexplorerhj committed
Aug 13, 2026 at 21:46 UTC
8e1131e1d1d47ae06acecb820b4c4a734697c0b2
2 files changed
+48
-4
vllm/model_executor/models/bailing_moe_v3.py
+45
-3
@@ -13,6 +13,7 @@ from collections.abc import Callable, Iterable
13
from math import lcm
14
from typing import TypeGuard
15
16
+import regex as re
17
import torch
18
import torch.nn as nn
19
import torch.nn.functional as F
@@ -214,9 +215,49 @@ def _is_block_fp8_config(
215
216
def _configure_ling_fp8_quant_config(
217
quant_config: QuantizationConfig | None,
218
+ config: PretrainedConfig,
219
) -> None:
218
- if _is_block_fp8_config(quant_config):
219
- quant_config.ignored_layers_match_mode = "suffix"
220
+ if not _is_block_fp8_config(quant_config):
221
+ return
222
+
223
+ quant_config.ignored_layers_match_mode = "suffix"
224
+ hf_quant_config = getattr(config, "quantization_config", None)
225
+ if not isinstance(hf_quant_config, dict):
226
+ return
227
+
228
+ quant_config.is_scale_e8m0 = ( # type: ignore[attr-defined]
229
+ hf_quant_config.get("scale_fmt") == "ue8m0"
230
+ )
231
+
232
+ routed_quant_method = hf_quant_config.get("routed_experts_quant_method")
233
+ if routed_quant_method is None:
234
+ return
235
+ if routed_quant_method != "mxfp4":
236
+ raise ValueError(
237
+ f"Unsupported routed experts quantization: {routed_quant_method!r}"
238
+ )
239
+
240
+ quant_config.store_dtype = "mxfp4"
241
+
242
+
243
+_LING_MXFP4_WEIGHTS_MAPPER = WeightsMapper(
244
+ orig_to_new_regex={
245
+ re.compile(
246
+ r"(\.mlp\.experts\.\d+\."
247
+ r"(?:gate_proj|up_proj|down_proj)\.weight_scale)_inv$"
248
+ ): r"\1"
249
+ }
250
+)
251
+
252
+
253
+def _maybe_remap_ling_mxfp4_weight_names(
254
+ weights: Iterable[tuple[str, torch.Tensor]],
255
+ quant_config: QuantizationConfig | None,
256
+) -> Iterable[tuple[str, torch.Tensor]]:
257
+ """Map Ling's MXFP4 expert scales to Mxfp4MoEMethod parameters."""
258
+ if isinstance(quant_config, Fp8Config) and quant_config.store_dtype == "mxfp4":
259
+ return _LING_MXFP4_WEIGHTS_MAPPER.apply(weights)
260
+ return weights
261
262
263
def _is_fp8_module_excluded(
@@ -1351,7 +1392,7 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):
1392
super().__init__()
1393
config = vllm_config.model_config.hf_config
1394
quant_config = vllm_config.quant_config
1354
- _configure_ling_fp8_quant_config(quant_config)
1395
+ _configure_ling_fp8_quant_config(quant_config, config)
1396
self.config = config
1397
self.quant_config = quant_config
1398
self.model = BailingMoeV3Model(
@@ -1421,6 +1462,7 @@ class BailingMoeV3ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsPP):
1462
1463
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1464
weights = self.hf_to_vllm_mapper.apply(weights)
1465
+ weights = _maybe_remap_ling_mxfp4_weight_names(weights, self.quant_config)
1466
params_dict = dict(self.named_parameters(remove_duplicate=False))
1467
loaded_params: set[str] = set()
1468
stacked_mappings = [
vllm/model_executor/models/bailing_moe_v3_mtp.py
+3
-1
@@ -31,6 +31,7 @@ from vllm.model_executor.models.bailing_moe_v3 import (
31
BailingMoeV3MoE,
32
_configure_ling_fp8_quant_config,
33
_maybe_pad_block_fp8_shared_expert_checkpoint_tensor,
34
+ _maybe_remap_ling_mxfp4_weight_names,
35
)
36
from vllm.model_executor.models.interfaces import SupportsPP
37
from vllm.model_executor.models.utils import (
@@ -231,7 +232,7 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):
232
super().__init__()
233
self.config = _get_draft_hf_config(vllm_config)
234
self.quant_config = vllm_config.quant_config
234
- _configure_ling_fp8_quant_config(self.quant_config)
235
+ _configure_ling_fp8_quant_config(self.quant_config, self.config)
236
self.model = BailingMoeV3MultiTokenPredictor(
237
vllm_config=vllm_config,
238
prefix=maybe_prefix(prefix, "model"),
@@ -281,6 +282,7 @@ class BailingMoeV3MTPModel(nn.Module, SupportsPP):
282
283
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
284
weights = self.hf_to_vllm_mapper.apply(weights)
285
+ weights = _maybe_remap_ling_mxfp4_weight_names(weights, self.quant_config)
286
stacked_params_mapping = [
287
(".fused_qkv_a_proj", ".q_a_proj", 0),
288
(".fused_qkv_a_proj", ".kv_a_proj_with_mqa", 1),