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, 22:45 UTC
[XPU][2/N] add support unquantized moe support for xpu (#33659)
Signed-off-by: Kunshang Ji <kunshang.ji@intel.com>
Kunshang Ji committed
Feb 4, 2026 at 18:12 UTC
f79f777803ae70bc3ee2a4cfd2822e26a378dcdb
6 files changed
+139
-34
.buildkite/scripts/hardware_ci/run-xpu-test.sh
+2
@@ -39,6 +39,8 @@ docker run \
39
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend ray
40
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager -tp 2 --distributed-executor-backend mp
41
python3 examples/offline_inference/basic/generate.py --model facebook/opt-125m --block-size 64 --enforce-eager --attention-backend=TRITON_ATTN
42
+ python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2
43
+ python3 examples/offline_inference/basic/generate.py --model ibm-research/PowerMoE-3b --block-size 64 --enforce-eager -tp 2 --enable-expert-parallel
44
cd tests
45
pytest -v -s v1/core --ignore=v1/core/test_reset_prefix_cache_e2e.py
46
pytest -v -s v1/engine
requirements/xpu.txt
+1
-1
@@ -15,4 +15,4 @@ torch==2.10.0+xpu
15
torchaudio
16
torchvision
17
18
-vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.0/vllm_xpu_kernels-0.1.0-cp312-cp312-linux_x86_64.whl
\ No newline at end of file
18
+vllm_xpu_kernels @ https://github.com/vllm-project/vllm-xpu-kernels/releases/download/v0.1.1/vllm_xpu_kernels-0.1.1-cp312-cp312-linux_x86_64.whl
\ No newline at end of file
vllm/model_executor/layers/fused_moe/__init__.py
+4
@@ -100,6 +100,9 @@ if HAS_TRITON:
100
from vllm.model_executor.layers.fused_moe.triton_deep_gemm_moe import (
101
TritonOrDeepGemmExperts,
102
)
103
+ from vllm.model_executor.layers.fused_moe.xpu_fused_moe import (
104
+ XPUExperts,
105
+ )
106
107
__all__ += [
108
"AiterExperts",
@@ -117,6 +120,7 @@ if HAS_TRITON:
120
"DeepGemmExperts",
121
"BatchedDeepGemmExperts",
122
"TritonOrDeepGemmExperts",
123
+ "XPUExperts",
124
]
125
else:
126
# Some model classes directly use the custom ops. Add placeholders
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
+10
-1
@@ -46,7 +46,6 @@ class UnquantizedMoeBackend(Enum):
46
UNSUPPORTED_BACKEND = [
47
UnquantizedMoeBackend.FLASHINFER_TRTLLM,
48
UnquantizedMoeBackend.CPU,
49
- UnquantizedMoeBackend.XPU,
49
UnquantizedMoeBackend.TPU,
50
UnquantizedMoeBackend.OOT,
51
]
@@ -196,4 +195,14 @@ def make_unquantized_moe_kernel(
195
quant_config=quant_config,
196
),
197
)
198
+ elif backend == UnquantizedMoeBackend.XPU:
199
+ from vllm.model_executor.layers.fused_moe import XPUExperts
200
+
201
+ kernel = mk.FusedMoEModularKernel(
202
+ MoEPrepareAndFinalizeNoEP(),
203
+ XPUExperts(
204
+ moe_config=moe_config,
205
+ quant_config=quant_config,
206
+ ),
207
+ )
208
return kernel, use_inplace
vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py
+2
-32
@@ -40,7 +40,7 @@ from vllm.model_executor.utils import replace_parameter, set_weight_attrs
40
from vllm.platforms import current_platform
41
from vllm.platforms.interface import CpuArchEnum
42
43
-if current_platform.is_cuda_alike():
43
+if current_platform.is_cuda_alike() or current_platform.is_xpu():
44
from .fused_batched_moe import BatchedTritonExperts
45
from .fused_moe import TritonExperts
46
else:
@@ -71,7 +71,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
71
self.kernel: mk.FusedMoEModularKernel | None = None
72
self._is_monolithic = (
73
current_platform.is_cpu()
74
- or current_platform.is_xpu()
74
or self.unquantized_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM
75
)
76
@@ -82,8 +81,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
81
"""Select the monolithic implementation based on platform."""
82
if current_platform.is_cpu():
83
return self.forward_monolithic_cpu
85
- elif current_platform.is_xpu():
86
- return self.forward_monolithic_xpu
84
else:
85
return self.forward_monolithic_cuda
86
@@ -256,16 +253,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
253
)
254
layer.w13_weight = Parameter(w13_weights_shuffled, requires_grad=False)
255
layer.w2_weight = Parameter(w2_weights_shuffled, requires_grad=False)
259
- elif self.unquantized_backend == UnquantizedMoeBackend.XPU:
260
- import intel_extension_for_pytorch as ipex
261
-
262
- ep_rank_start = self.moe.ep_rank * self.moe.num_local_experts
263
- self.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
264
- layer.w13_weight,
265
- layer.w2_weight,
266
- use_prepack=True,
267
- experts_start_id=ep_rank_start,
268
- )
256
elif self.unquantized_backend == UnquantizedMoeBackend.CPU:
257
from vllm.model_executor.layers.fused_moe import cpu_fused_moe
258
@@ -297,7 +284,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
284
self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
285
else:
286
self.cpu_fused_moe = cpu_fused_moe.CPUFusedMOE(layer)
300
- elif current_platform.is_cuda_alike():
287
+ elif current_platform.is_cuda_alike() or current_platform.is_xpu():
288
self._setup_kernel(
289
layer=layer,
290
w13=layer.w13_weight,
@@ -399,20 +386,3 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
386
layer.apply_router_weight_on_input,
387
layer.activation,
388
)
402
-
403
- def forward_monolithic_xpu(
404
- self,
405
- layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821
406
- x: torch.Tensor,
407
- router_logits: torch.Tensor,
408
- ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
409
- return self.ipex_fusion(
410
- x,
411
- layer.use_grouped_topk,
412
- layer.top_k,
413
- router_logits,
414
- layer.renormalize,
415
- layer.topk_group,
416
- layer.num_expert_group,
417
- custom_routing_function=layer.custom_routing_function,
418
- )
vllm/model_executor/layers/fused_moe/xpu_fused_moe.py
new
+120
@@ -0,0 +1,120 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import torch
4
+
5
+import vllm.model_executor.layers.fused_moe.modular_kernel as mk
6
+from vllm.model_executor.layers.fused_moe.config import (
7
+ FusedMoEParallelConfig,
8
+)
9
+from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
10
+ TopKWeightAndReduceNoOP,
11
+)
12
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
13
+ QuantKey,
14
+ kFp8StaticTensorSym,
15
+)
16
+from vllm.platforms import current_platform
17
+
18
+if current_platform.is_xpu():
19
+ from vllm_xpu_kernels.fused_moe_interface import xpu_fused_moe
20
+
21
+
22
+class XPUExperts(mk.FusedMoEPermuteExpertsUnpermute):
23
+ @property
24
+ def expects_unquantized_inputs(self) -> bool:
25
+ return True
26
+
27
+ @staticmethod
28
+ def activation_format() -> mk.FusedMoEActivationFormat:
29
+ return mk.FusedMoEActivationFormat.Standard
30
+
31
+ @staticmethod
32
+ def _supports_current_device() -> bool:
33
+ return current_platform.is_xpu()
34
+
35
+ @staticmethod
36
+ def _supports_no_act_and_mul() -> bool:
37
+ return False
38
+
39
+ @staticmethod
40
+ def _supports_activation(activation: str) -> bool:
41
+ return activation in ["silu", "gelu", "swigluoai"]
42
+
43
+ @staticmethod
44
+ def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool:
45
+ return True
46
+
47
+ @staticmethod
48
+ def _supports_quant_scheme(
49
+ weight_key: QuantKey | None,
50
+ activation_key: QuantKey | None,
51
+ ) -> bool:
52
+ # TODO: dispatch based on device.
53
+ SUPPORTED_W_A = [
54
+ (None, None),
55
+ (kFp8StaticTensorSym, None),
56
+ ]
57
+ return (weight_key, activation_key) in SUPPORTED_W_A
58
+
59
+ def supports_chunking(self) -> bool:
60
+ return False
61
+
62
+ def supports_expert_map(self) -> bool:
63
+ return True
64
+
65
+ def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
66
+ return TopKWeightAndReduceNoOP()
67
+
68
+ def workspace_shapes(
69
+ self,
70
+ M: int,
71
+ N: int,
72
+ K: int,
73
+ topk: int,
74
+ global_num_experts: int,
75
+ local_num_experts: int,
76
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
77
+ activation: str,
78
+ ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
79
+ workspace1 = (0,)
80
+ workspace2 = (0,)
81
+ output = (M, K)
82
+ return (workspace1, workspace2, output)
83
+
84
+ def apply(
85
+ self,
86
+ output: torch.Tensor,
87
+ hidden_states: torch.Tensor,
88
+ w1: torch.Tensor,
89
+ w2: torch.Tensor,
90
+ topk_weights: torch.Tensor,
91
+ topk_ids: torch.Tensor,
92
+ activation: str,
93
+ global_num_experts: int,
94
+ expert_map: torch.Tensor | None,
95
+ a1q_scale: torch.Tensor | None,
96
+ a2_scale: torch.Tensor | None,
97
+ workspace13: torch.Tensor,
98
+ workspace2: torch.Tensor,
99
+ expert_tokens_meta: mk.ExpertTokensMetadata | None,
100
+ apply_router_weight_on_input: bool,
101
+ ):
102
+ topk = topk_ids.size(-1)
103
+ xpu_fused_moe(
104
+ hidden_states=hidden_states,
105
+ w13=w1,
106
+ w13_scales=a1q_scale,
107
+ w13_bias=self.w1_bias,
108
+ w2=w2,
109
+ w2_scales=a2_scale,
110
+ w2_bias=self.w2_bias,
111
+ topk_weights=topk_weights,
112
+ topk_ids=topk_ids,
113
+ n_experts_per_token=topk,
114
+ activation=activation,
115
+ num_experts=self.moe_config.num_local_experts,
116
+ ep_rank=self.moe_config.ep_rank,
117
+ ep_size=self.moe_config.ep_size,
118
+ output=output,
119
+ )
120
+ return