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 12, 14:45 UTC
[fix][cpu] Use a SwigluOAI impl which supports interleaved gate-up wei (#29273)
Signed-off-by: Fadi Arafeh <fadi.arafeh@arm.com>
Fadi Arafeh committed
Nov 25, 2025 at 07:11 UTC
98caeadd54599c8038fab5b19cc8ef5688b7b03a
1 file changed
+8
-21
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+8
-21
@@ -6,22 +6,7 @@ import torch
6
from torch.nn import functional as F
7
8
from vllm import _custom_ops as ops
9
-
10
-
11
-def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
12
- d = x.shape[-1] // 2
13
- return F.silu(x[..., :d]) * x[..., d:]
14
-
15
-
16
-def swigluoai_and_mul(
17
- x: torch.Tensor, alpha: float = 1.702, limit: float = 7.0
18
-) -> torch.Tensor:
19
- d = x.shape[-1] // 2
20
- gate, up = x[..., :d], x[..., d:]
21
- gate = gate.clamp(max=limit)
22
- up = up.clamp(min=-limit, max=limit)
23
- glu = gate * torch.sigmoid(alpha * gate)
24
- return (up + 1) * glu
9
+from vllm.model_executor.layers.activation import SiluAndMul, SwigluOAIAndMul
10
11
12
def grouped_topk(
@@ -227,6 +212,11 @@ class CPUFusedMOE:
212
layer.w13_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
213
layer.w2_weight = torch.nn.Parameter(torch.empty(0), requires_grad=False)
214
215
+ self.act_to_impl = {
216
+ "silu": SiluAndMul(),
217
+ "swigluoai": SwigluOAIAndMul(),
218
+ }
219
+
220
def __call__(
221
self,
222
layer: torch.nn.Module,
@@ -246,7 +236,7 @@ class CPUFusedMOE:
236
apply_router_weight_on_input: bool = False,
237
activation: str = "silu",
238
) -> torch.Tensor:
249
- assert activation in {"silu", "swigluoai"}, f"{activation} is not supported."
239
+ assert activation in self.act_to_impl, f"{activation} is not supported."
240
assert not apply_router_weight_on_input
241
topk_weights, topk_ids = select_experts(
242
hidden_states=x,
@@ -283,10 +273,7 @@ class CPUFusedMOE:
273
tokens_for_this_expert = sorted_tokens[start_idx:end_idx]
274
275
gate_up = layer.gate_up_linear[i](tokens_for_this_expert)
286
- if activation == "swigluoai":
287
- gate_up = swigluoai_and_mul(gate_up)
288
- else:
289
- gate_up = silu_and_mul(gate_up)
276
+ gate_up = self.act_to_impl[activation].forward_native(gate_up)
277
expert_out = layer.down_linear[i](gate_up)
278
outputs.append(expert_out)
279
start_idx = end_idx