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 13, 17:00 UTC
[CPU]Improve cpu fused moe perf (#27244)
Signed-off-by: Zhang Xiangze <Xiangze.Zhang@arm.com>
xiangze-arm committed
Nov 6, 2025 at 19:04 UTC
c757a15f0f8ed54b7f85b849d5e075226fabbcd9
1 file changed
+40
-12
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+40
-12
@@ -5,6 +5,7 @@ from collections.abc import Callable
5
import torch
6
from torch.nn import functional as F
7
8
+from vllm import _custom_ops as ops
9
from vllm import envs
10
11
@@ -237,7 +238,43 @@ class SGLFusedMOE:
238
239
class CPUFusedMOE:
240
def __init__(self, layer: torch.nn.Module) -> None:
240
- pass
241
+ use_onednn_mm = ops._supports_onednn and ops.is_onednn_acl_supported()
242
+
243
+ num_experts = layer.w13_weight.size(0)
244
+ has_w13_bias = hasattr(layer, "w13_bias")
245
+ has_w2_bias = hasattr(layer, "w2_bias")
246
+
247
+ layer.gate_up_linear = []
248
+ layer.down_linear = []
249
+
250
+ for i in range(num_experts):
251
+ layer_w13_weight = layer.w13_weight[i]
252
+ layer_w13_bias = layer.w13_bias[i] if has_w13_bias else None
253
+ layer_w2_weight = layer.w2_weight[i]
254
+ layer_w2_bias = layer.w2_bias[i] if has_w2_bias else None
255
+ if use_onednn_mm:
256
+ gate_up_handle = ops.create_onednn_mm(layer_w13_weight.t(), 32)
257
+ layer.gate_up_linear.append(
258
+ lambda x, handle=gate_up_handle, bias=layer_w13_bias: ops.onednn_mm(
259
+ handle, x, bias
260
+ )
261
+ )
262
+ down_handle = ops.create_onednn_mm(layer_w2_weight.t(), 32)
263
+ layer.down_linear.append(
264
+ lambda x, handle=down_handle, bias=layer_w2_bias: ops.onednn_mm(
265
+ handle, x, bias
266
+ )
267
+ )
268
+ else:
269
+ layer.gate_up_linear.append(
270
+ lambda x, w=layer_w13_weight, b=layer_w13_bias: F.linear(x, w, b)
271
+ )
272
+ layer.down_linear.append(
273
+ lambda x, w=layer_w2_weight, b=layer_w2_bias: F.linear(x, w, b)
274
+ )
275
+ if use_onednn_mm: # remove weight
276
+ layer.w13_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
277
+ layer.w2_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
278
279
def __call__(
280
self,
@@ -287,8 +324,6 @@ class CPUFusedMOE:
324
325
outputs = []
326
start_idx = 0
290
- has_w13_bias = hasattr(layer, "w13_bias")
291
- has_w2_bias = hasattr(layer, "w2_bias")
327
328
for i, num_tokens in enumerate(tokens_per_expert):
329
end_idx = start_idx + num_tokens
@@ -296,19 +331,12 @@ class CPUFusedMOE:
331
continue
332
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
333
299
- layer_w13_weight = layer.w13_weight[i]
300
- layer_w13_bias = layer.w13_bias[i] if has_w13_bias else None
301
- layer_w2_weight = layer.w2_weight[i]
302
- layer_w2_bias = layer.w2_bias[i] if has_w2_bias else None
303
-
304
- gate_up = F.linear(
305
- tokens_for_this_expert, layer_w13_weight, bias=layer_w13_bias
306
- )
334
+ gate_up = layer.gate_up_linear[i](tokens_for_this_expert)
335
if activation == "swigluoai":
336
gate_up = swigluoai_and_mul(gate_up)
337
else:
338
gate_up = silu_and_mul(gate_up)
311
- expert_out = F.linear(gate_up, layer_w2_weight, bias=layer_w2_bias)
339
+ expert_out = layer.down_linear[i](gate_up)
340
outputs.append(expert_out)
341
start_idx = end_idx
342