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 2, 23:15 UTC
[XPU] reudce host overhead of XPU MOE (#42915)
Signed-off-by: mayuyuace <qiming1.zhang@intel.com> Co-authored-by: Kunshang Ji <kunshang.ji@intel.com>
Qiming Zhang committed
May 22, 2026 at 22:09 UTC
54d153637b5e8ecd3c75f8d4427f4c7309de0e2c
1 file changed
+23
-18
vllm/model_executor/layers/fused_moe/experts/xpu_moe.py
+23
-18
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
23
from vllm.platforms import current_platform
24
25
if current_platform.is_xpu():
26
- from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe
26
+ from vllm_xpu_kernels.fused_moe_interface import XpuFusedMoe
27
28
29
def prepare_fp8_moe_layer_for_xpu(
@@ -50,6 +50,7 @@ class XPUExperts(mk.FusedMoEExpertsModular):
50
self.is_fp8 = False
51
self.is_mxfp4 = False
52
self.is_mxfp8 = False
53
+ self.fused_moe_impl: XpuFusedMoe | None = None
54
55
@property
56
def expects_unquantized_inputs(self) -> bool:
@@ -132,26 +133,30 @@ class XPUExperts(mk.FusedMoEExpertsModular):
133
expert_tokens_meta: mk.ExpertTokensMetadata | None,
134
apply_router_weight_on_input: bool,
135
):
135
- topk = topk_ids.size(-1)
136
- xpu_fused_moe(
136
+ if self.fused_moe_impl is None:
137
+ topk = topk_ids.size(-1)
138
+ self.fused_moe_impl = XpuFusedMoe(
139
+ w13=w1,
140
+ w13_scales=self.w1_scale,
141
+ w13_bias=self.w1_bias,
142
+ w2=w2,
143
+ w2_scales=self.w2_scale,
144
+ w2_bias=self.w2_bias,
145
+ n_experts_per_token=topk,
146
+ activation=activation.value,
147
+ num_experts=self.moe_config.num_local_experts,
148
+ ep_rank=self.moe_config.ep_rank,
149
+ ep_size=self.moe_config.ep_size,
150
+ is_fp8=self.is_fp8,
151
+ is_mxfp4=self.is_mxfp4,
152
+ is_mxfp8=self.is_mxfp8,
153
+ )
154
+ assert self.fused_moe_impl is not None
155
+ self.fused_moe_impl.apply(
156
+ output=output,
157
hidden_states=hidden_states,
138
- w13=w1,
139
- w13_scales=self.w1_scale,
140
- w13_bias=self.w1_bias,
141
- w2=w2,
142
- w2_scales=self.w2_scale,
143
- w2_bias=self.w2_bias,
158
topk_weights=topk_weights,
159
topk_ids=topk_ids,
146
- n_experts_per_token=topk,
147
- activation=activation.value,
148
- num_experts=self.moe_config.num_local_experts,
149
- ep_rank=self.moe_config.ep_rank,
150
- ep_size=self.moe_config.ep_size,
151
- output=output,
152
- is_fp8=self.is_fp8,
153
- is_mxfp4=self.is_mxfp4,
154
- is_mxfp8=self.is_mxfp8,
160
)
161
162