A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 8, 19:00 UTC
[XPU][4/N] add mxfp4 moe model support (#33679)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Feb 6, 2026 at 13:03 UTC
7439e4f41b1f877e088d1b8c9ce4f2847c59423f
1 file changed
+53
-31
vllm/model_executor/layers/quantization/mxfp4.py
+53
-31
@@ -215,7 +215,7 @@ class Mxfp4Config(QuantizationConfig):
215
return UnquantizedLinearMethod()
216
elif isinstance(layer, FusedMoE):
217
if current_platform.is_xpu():
218
- return IpexMxfp4MoEMethod(layer.moe_config)
218
+ return XpuMxfp4MoEMethod(layer.moe_config)
219
else:
220
quant_method = Mxfp4MoEMethod(layer.moe_config)
221
quant_method.marlin_input_dtype = get_marlin_input_dtype(prefix)
@@ -1096,7 +1096,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
1096
raise ValueError(f"Unsupported backend: {self.mxfp4_backend}")
1097
1098
1099
-class IpexMxfp4MoEMethod(Mxfp4MoEMethod):
1099
+class XpuMxfp4MoEMethod(Mxfp4MoEMethod):
1100
def __init__(self, moe_config: FusedMoEConfig):
1101
super().__init__(moe_config)
1102
self.moe_config = moe_config
@@ -1121,21 +1121,7 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod):
1121
self.original_hidden_size = hidden_size
1122
1123
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
1124
- import intel_extension_for_pytorch as ipex
1125
-
1126
- layer.w13_weight.data = layer.w13_weight.data.view(torch.int32)
1127
- layer.w2_weight.data = layer.w2_weight.data.view(torch.int32)
1128
- ep_rank_start = self.moe_config.ep_rank * self.moe_config.num_local_experts
1129
- layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
1130
- layer.w13_weight,
1131
- layer.w2_weight,
1132
- w1_scale_inv=layer.w13_weight_scale,
1133
- w2_scale_inv=layer.w2_weight_scale,
1134
- w13_bias=layer.w13_bias,
1135
- w2_bias=layer.w2_bias,
1136
- is_mxfp4=True,
1137
- experts_start_id=ep_rank_start,
1138
- )
1124
+ pass
1125
1126
@property
1127
def is_monolithic(self) -> bool:
@@ -1148,19 +1134,55 @@ class IpexMxfp4MoEMethod(Mxfp4MoEMethod):
1134
router_logits: torch.Tensor,
1135
) -> torch.Tensor:
1136
assert layer.activation == "swigluoai", (
1151
- "Only swiglu_oai activation is supported for IPEX MXFP4 MoE"
1137
+ "Only swiglu_oai activation is supported for XPU MXFP4 MoE"
1138
)
1153
- hidden_size_pad = round_up(self.original_hidden_size, 128)
1154
- x_pad = torch.nn.functional.pad(x, (0, hidden_size_pad - x.size(-1)))
1155
- hidden_states = layer.ipex_fusion(
1156
- x_pad,
1157
- layer.use_grouped_topk,
1158
- layer.top_k,
1159
- router_logits,
1160
- layer.renormalize,
1161
- layer.topk_group,
1162
- layer.num_expert_group,
1163
- activation="swiglu_oai",
1139
+ from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe
1140
+
1141
+ M, _ = x.size()
1142
+ routing_weights = torch.empty(
1143
+ M, layer.top_k, dtype=torch.float32, device=x.device
1144
+ )
1145
+ selected_experts = torch.empty(
1146
+ M, layer.top_k, dtype=torch.int32, device=x.device
1147
+ )
1148
+ token_expert_indices = torch.empty(
1149
+ M, layer.top_k, dtype=torch.int32, device=x.device
1150
+ )
1151
+
1152
+ if layer.use_grouped_topk:
1153
+ routing_weights, selected_experts = torch.ops._moe_C.fused_grouped_topk(
1154
+ x,
1155
+ router_logits,
1156
+ layer.top_k,
1157
+ layer.renormalize,
1158
+ n_expert_group=layer.num_expert_group,
1159
+ n_topk_group=layer.topk_group,
1160
+ scoring_func=layer.scoring_func,
1161
+ routed_scaling_factor=layer.routed_scaling_factor,
1162
+ bias=layer.e_score_correction_bias,
1163
+ )
1164
+ else:
1165
+ torch.ops._moe_C.topk_softmax(
1166
+ routing_weights,
1167
+ selected_experts,
1168
+ token_expert_indices,
1169
+ router_logits,
1170
+ layer.renormalize,
1171
+ layer.e_score_correction_bias,
1172
+ )
1173
+
1174
+ return xpu_fused_moe(
1175
+ hidden_states=x,
1176
+ w13=layer.w13_weight,
1177
+ w13_bias=layer.w13_bias if self.moe.has_bias else None,
1178
+ w13_scales=layer.w13_weight_scale,
1179
+ w2=layer.w2_weight,
1180
+ w2_bias=layer.w2_bias if self.moe.has_bias else None,
1181
+ w2_scales=layer.w2_weight_scale,
1182
+ topk_weights=routing_weights,
1183
+ topk_ids=selected_experts,
1184
+ n_experts_per_token=layer.top_k,
1185
+ activation=layer.activation,
1186
+ num_experts=layer.local_num_experts,
1187
+ is_mxfp4=True,
1188
)
1165
- hidden_states = hidden_states[..., : self.original_hidden_size].contiguous()
1166
- return hidden_states