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 6, 04:30 UTC
[XPU][MoE Refactor] Refactor xpu mxfp4 support into oracle (#37784)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Mar 23, 2026 at 19:10 UTC
debd6e768cee69e24a2b4de46674a71bed8ae4a1
3 files changed
+54
-101
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+23
-2
@@ -141,7 +141,10 @@ def backend_to_kernel_cls(
141
return [AiterExperts]
142
143
elif backend == Mxfp4MoeBackend.XPU:
144
- raise NotImplementedError("XPU backend uses XpuMxfp4MoEMethod directly.")
144
+ from vllm.model_executor.layers.fused_moe.xpu_fused_moe import XPUExpertsMXFp4
145
+
146
+ return [XPUExpertsMXFp4]
147
+
148
else:
149
raise ValueError(f"Unknown MXFP4 MoE backend: {backend.value}")
150
@@ -156,6 +159,7 @@ def map_mxfp4_backend(runner_backend: str) -> Mxfp4MoeBackend:
159
"triton": Mxfp4MoeBackend.TRITON,
160
"marlin": Mxfp4MoeBackend.MARLIN,
161
"ck": Mxfp4MoeBackend.CK,
162
+ "xpu": Mxfp4MoeBackend.XPU,
163
}
164
if backend := mapping.get(runner_backend):
165
return backend
@@ -178,6 +182,7 @@ def _get_priority_backends() -> list[Mxfp4MoeBackend]:
182
Mxfp4MoeBackend.TRITON_UNFUSED,
183
Mxfp4MoeBackend.MARLIN,
184
Mxfp4MoeBackend.BATCHED_MARLIN,
185
+ Mxfp4MoeBackend.XPU,
186
]
187
return _AVAILABLE_BACKENDS
188
@@ -351,7 +356,13 @@ def select_mxfp4_moe_backend(
356
if current_platform.is_xpu():
357
backend = Mxfp4MoeBackend.XPU
358
logger.info_once(_make_log_backend(backend))
354
- return backend, None
359
+ return _return_or_raise(
360
+ Mxfp4MoeBackend.XPU,
361
+ config,
362
+ kMxfp4Static,
363
+ None,
364
+ activation_format,
365
+ )
366
367
if current_platform.is_cuda() or current_platform.is_rocm():
368
raise NotImplementedError(
@@ -741,6 +752,16 @@ def convert_to_mxfp4_moe_kernel_format(
752
w13_bias,
753
w2_bias,
754
)
755
+ elif mxfp4_backend == Mxfp4MoeBackend.XPU:
756
+ # No additional transformation needed for XPU backend
757
+ return (
758
+ w13_weight,
759
+ w2_weight,
760
+ w13_weight_scale,
761
+ w2_weight_scale,
762
+ w13_bias,
763
+ w2_bias,
764
+ )
765
else:
766
raise ValueError(
767
f"Unsupported mxfp4_backend: {mxfp4_backend}: "
vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
+30
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
16
QuantKey,
17
kFp8DynamicTensorSym,
18
kFp8StaticTensorSym,
19
+ kMxfp4Static,
20
)
21
from vllm.platforms import current_platform
22
@@ -38,6 +39,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
39
num_dispatchers,
40
)
41
self.is_fp8 = False
42
+ self.is_mxfp4 = False
43
44
@property
45
def expects_unquantized_inputs(self) -> bool:
@@ -137,6 +139,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
139
ep_size=self.moe_config.ep_size,
140
output=output,
141
is_fp8=self.is_fp8,
142
+ is_mxfp4=self.is_mxfp4,
143
)
144
145
@@ -155,3 +158,30 @@ class XPUExpertsFp8(XPUExperts):
158
num_dispatchers,
159
)
160
self.is_fp8 = True
161
+
162
+
163
+class XPUExpertsMXFp4(XPUExperts):
164
+ def __init__(
165
+ self,
166
+ moe_config: FusedMoEConfig,
167
+ quant_config: FusedMoEQuantConfig,
168
+ max_num_tokens: int | None = None,
169
+ num_dispatchers: int | None = None,
170
+ ):
171
+ super().__init__(
172
+ moe_config,
173
+ quant_config,
174
+ max_num_tokens,
175
+ num_dispatchers,
176
+ )
177
+ self.is_mxfp4 = True
178
+
179
+ @staticmethod
180
+ def _supports_quant_scheme(
181
+ weight_key: QuantKey | None,
182
+ activation_key: QuantKey | None,
183
+ ) -> bool:
184
+ SUPPORTED_W_A = [
185
+ (kMxfp4Static, None),
186
+ ]
187
+ return (weight_key, activation_key) in SUPPORTED_W_A
vllm/model_executor/layers/quantization/mxfp4.py
+1
-99
@@ -10,7 +10,6 @@ from vllm.model_executor.layers.fused_moe import (
10
FusedMoE,
11
FusedMoEConfig,
12
FusedMoEMethodBase,
13
- MoEActivation,
13
)
14
from vllm.model_executor.layers.fused_moe import modular_kernel as mk
15
from vllm.model_executor.layers.fused_moe.config import (
@@ -33,7 +32,6 @@ from vllm.model_executor.layers.quantization.base_config import (
32
)
33
from vllm.model_executor.layers.quantization.utils.quant_utils import is_layer_skipped
34
from vllm.model_executor.utils import replace_parameter, set_weight_attrs
36
-from vllm.platforms import current_platform
35
36
logger = init_logger(__name__)
37
@@ -80,10 +78,7 @@ class Mxfp4Config(QuantizationConfig):
78
)
79
return UnquantizedLinearMethod()
80
elif isinstance(layer, FusedMoE):
83
- if current_platform.is_xpu():
84
- return XpuMxfp4MoEMethod(layer.moe_config)
85
- else:
86
- return Mxfp4MoEMethod(layer.moe_config)
81
+ return Mxfp4MoEMethod(layer.moe_config)
82
elif isinstance(layer, Attention):
83
logger.debug_once(
84
"MXFP4 attention layer is not implemented. "
@@ -420,96 +415,3 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
415
expert_map=layer.expert_map,
416
apply_router_weight_on_input=layer.apply_router_weight_on_input,
417
)
423
-
424
-
425
-class XpuMxfp4MoEMethod(Mxfp4MoEMethod):
426
- def __init__(self, moe_config: FusedMoEConfig):
427
- super().__init__(moe_config)
428
- self.moe_config = moe_config
429
-
430
- def create_weights(
431
- self,
432
- layer: torch.nn.Module,
433
- num_experts: int,
434
- hidden_size: int,
435
- intermediate_size_per_partition: int,
436
- params_dtype: torch.dtype,
437
- **extra_weight_attrs,
438
- ):
439
- super().create_weights(
440
- layer,
441
- num_experts,
442
- hidden_size,
443
- intermediate_size_per_partition,
444
- params_dtype,
445
- **extra_weight_attrs,
446
- )
447
- self.original_hidden_size = hidden_size
448
-
449
- def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
450
- pass
451
-
452
- @property
453
- def is_monolithic(self) -> bool:
454
- return True
455
-
456
- def apply_monolithic(
457
- self,
458
- layer: FusedMoE,
459
- x: torch.Tensor,
460
- router_logits: torch.Tensor,
461
- ) -> torch.Tensor:
462
- assert layer.activation == MoEActivation.SWIGLUOAI, (
463
- "Only swiglu_oai activation is supported for "
464
- f"XPU MXFP4 MoE, not {layer.activation}."
465
- )
466
- from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe
467
-
468
- M, _ = x.size()
469
- routing_weights = torch.empty(
470
- M, layer.top_k, dtype=torch.float32, device=x.device
471
- )
472
- selected_experts = torch.empty(
473
- M, layer.top_k, dtype=torch.int32, device=x.device
474
- )
475
- token_expert_indices = torch.empty(
476
- M, layer.top_k, dtype=torch.int32, device=x.device
477
- )
478
-
479
- if layer.use_grouped_topk:
480
- routing_weights, selected_experts = torch.ops._moe_C.fused_grouped_topk(
481
- x,
482
- router_logits,
483
- layer.top_k,
484
- layer.renormalize,
485
- n_expert_group=layer.num_expert_group,
486
- n_topk_group=layer.topk_group,
487
- scoring_func=layer.scoring_func,
488
- routed_scaling_factor=layer.routed_scaling_factor,
489
- bias=layer.e_score_correction_bias,
490
- )
491
- else:
492
- torch.ops._moe_C.topk_softmax(
493
- routing_weights,
494
- selected_experts,
495
- token_expert_indices,
496
- router_logits,
497
- layer.renormalize,
498
- layer.e_score_correction_bias,
499
- )
500
-
501
- return xpu_fused_moe(
502
- hidden_states=x,
503
- w13=layer.w13_weight,
504
- w13_bias=layer.w13_bias if self.moe.has_bias else None,
505
- w13_scales=layer.w13_weight_scale,
506
- w2=layer.w2_weight,
507
- w2_bias=layer.w2_bias if self.moe.has_bias else None,
508
- w2_scales=layer.w2_weight_scale,
509
- topk_weights=routing_weights,
510
- topk_ids=selected_experts,
511
- n_experts_per_token=layer.top_k,
512
- activation=layer.activation.value,
513
- num_experts=layer.local_num_experts,
514
- is_mxfp4=True,
515
- )