[Bugfix][Quantization] Fix fp32 weight scale for mxfp4 quantization and per-expert checkpoint mapping (#51419)
Signed-off-by: Isotr0py <Isotr0py@outlook.com>
Isotr0py committed
Aug 10, 2026 at 14:08 UTC
3b4c86e489c532f109e602c67b38ac94dc5b78fe
2 files changed
+67
vllm/model_executor/layers/fused_moe/routed_experts.py
+27
@@ -896,6 +896,14 @@ class RoutedExperts(PluggableLayer):
896
break
897
continue
898
matched = True
899
+ is_per_expert_fused_w13 = (
900
+ not is_fused
901
+ and shard_id in {"w1", "w3"}
902
+ and any(
903
+ f".{fused_name}." in qual_name
904
+ for fused_name in ("gate_up_proj", "w13")
905
+ )
906
+ )
907
weight_name = qual_name.replace(weight_name, param_name)
908
param_name = weight_name.removeprefix(f"{self.layer_name}.")
909
param = getattr(self, param_name, None)
@@ -919,6 +927,12 @@ class RoutedExperts(PluggableLayer):
927
else:
928
experts_shard = fused_weight
929
start = 0
930
+ elif is_per_expert_fused_w13:
931
+ shard_index = 0 if shard_id == "w1" else 1
932
+ experts_shard = loaded_weight.chunk(2, dim=0)[
933
+ shard_index
934
+ ].unsqueeze(0)
935
+ start = expert_id
936
else:
937
# loaded_weight is a single expert weight, so we add a dummy expert
938
# dimension to unify the loading logic with the fused case
@@ -1069,6 +1083,19 @@ class RoutedExperts(PluggableLayer):
1083
(f"{w13}weight", f"experts.{gate_up}", 1, "w3"),
1084
(f"{w2}weight", f"experts.{ckpt_down_proj_name}", 0, "w2"),
1085
]
1086
+ fused_mapping.extend(
1087
+ (
1088
+ w13,
1089
+ (
1090
+ f"experts.{physical_to_logical_map[expert_id]}."
1091
+ f"{gate_up}.{lora_base_layer_prefix}"
1092
+ ),
1093
+ expert_id,
1094
+ shard_id,
1095
+ )
1096
+ for expert_id in range(num_physical_experts)
1097
+ for shard_id in ("w1", "w3")
1098
+ )
1099
1100
per_expert_mapping = [
1101
# (param_name, weight_name, expert_id, shard_id)
vllm/model_executor/layers/quantization/mxfp4.py
+40
@@ -569,6 +569,38 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
569
activation=self.moe.activation,
570
)
571
572
+ @staticmethod
573
+ def _encode_mxfp4_weight_scale(loaded_weight: torch.Tensor) -> torch.Tensor:
574
+ if loaded_weight.dtype == torch.uint8:
575
+ return loaded_weight
576
+ if loaded_weight.dtype == torch.float8_e8m0fnu:
577
+ return loaded_weight.view(torch.uint8)
578
+ if loaded_weight.is_floating_point():
579
+ return loaded_weight.to(torch.float8_e8m0fnu).view(torch.uint8)
580
+ return loaded_weight
581
+
582
+ @staticmethod
583
+ def get_scale_weight_loader(weight_loader):
584
+ def mxfp4_weight_loader(
585
+ param: torch.nn.Parameter,
586
+ loaded_weight: torch.Tensor,
587
+ weight_name: str,
588
+ shard_id: str,
589
+ expert_id: int,
590
+ return_success: bool = False,
591
+ ) -> bool | None:
592
+ loaded_weight = Mxfp4MoEMethod._encode_mxfp4_weight_scale(loaded_weight)
593
+ return weight_loader(
594
+ param,
595
+ loaded_weight,
596
+ weight_name,
597
+ shard_id,
598
+ expert_id,
599
+ return_success=return_success,
600
+ )
601
+
602
+ return mxfp4_weight_loader
603
+
604
def create_weights(
605
self,
606
layer: RoutedExperts,
@@ -587,6 +619,8 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
619
layer.num_experts = num_experts
620
self.intermediate_size = intermediate_size_per_partition
621
self.hidden_size = hidden_size
622
+ weight_loader = extra_weight_attrs.pop("weight_loader")
623
+ scale_weight_loader = Mxfp4MoEMethod.get_scale_weight_loader(weight_loader)
624
625
# Fused gate_up_proj (column parallel)
626
w13_weight = torch.nn.Parameter(
@@ -600,6 +634,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
634
)
635
layer.register_parameter("w13_weight", w13_weight)
636
set_weight_attrs(w13_weight, extra_weight_attrs)
637
+ set_weight_attrs(w13_weight, {"weight_loader": weight_loader})
638
639
w13_weight_scale = torch.nn.Parameter(
640
torch.zeros(
@@ -612,6 +647,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
647
)
648
layer.register_parameter("w13_weight_scale", w13_weight_scale)
649
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
650
+ set_weight_attrs(w13_weight_scale, {"weight_loader": scale_weight_loader})
651
w13_weight_scale.quant_method = "block"
652
653
# down_proj (row parallel)
@@ -626,6 +662,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
662
)
663
layer.register_parameter("w2_weight", w2_weight)
664
set_weight_attrs(w2_weight, extra_weight_attrs)
665
+ set_weight_attrs(w2_weight, {"weight_loader": weight_loader})
666
667
w2_weight_scale = torch.nn.Parameter(
668
torch.zeros(
@@ -638,6 +675,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
675
)
676
layer.register_parameter("w2_weight_scale", w2_weight_scale)
677
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
678
+ set_weight_attrs(w2_weight_scale, {"weight_loader": scale_weight_loader})
679
w2_weight_scale.quant_method = "block"
680
681
if self.moe.has_bias:
@@ -651,6 +689,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
689
)
690
layer.register_parameter("w13_bias", w13_bias)
691
set_weight_attrs(w13_bias, extra_weight_attrs)
692
+ set_weight_attrs(w13_bias, {"weight_loader": weight_loader})
693
694
w2_bias = torch.nn.Parameter(
695
torch.zeros(
@@ -662,6 +701,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
701
)
702
layer.register_parameter("w2_bias", w2_bias)
703
set_weight_attrs(w2_bias, extra_weight_attrs)
704
+ set_weight_attrs(w2_bias, {"weight_loader": weight_loader})
705
706
def _setup_kernel(
707
self,