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 14, 05:30 UTC
Add SwigluOAI implementation for CPUFusedMOE (#26347)
Signed-off-by: Sharif Inamdar <sharif.inamdar@arm.com>
isharif168 committed
Oct 8, 2025 at 03:17 UTC
046118b93858fa70ef928c1c2501b15096f5e89e
1 file changed
+16
-2
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+16
-2
@@ -13,6 +13,17 @@ def silu_and_mul(x: torch.Tensor) -> torch.Tensor:
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
25
+
26
+
27
def grouped_topk(
28
hidden_states: torch.Tensor,
29
gating_output: torch.Tensor,
@@ -247,7 +258,7 @@ class CPUFusedMOE:
258
apply_router_weight_on_input: bool = False,
259
activation: str = "silu",
260
) -> torch.Tensor:
250
- assert activation == "silu", f"{activation} is not supported."
261
+ assert activation in {"silu", "swigluoai"}, f"{activation} is not supported."
262
assert not apply_router_weight_on_input
263
topk_weights, topk_ids = select_experts(
264
hidden_states=x,
@@ -293,7 +304,10 @@ class CPUFusedMOE:
304
gate_up = F.linear(
305
tokens_for_this_expert, layer_w13_weight, bias=layer_w13_bias
306
)
296
- gate_up = silu_and_mul(gate_up)
307
+ if activation == "swigluoai":
308
+ gate_up = swigluoai_and_mul(gate_up)
309
+ else:
310
+ gate_up = silu_and_mul(gate_up)
311
expert_out = F.linear(gate_up, layer_w2_weight, bias=layer_w2_bias)
312
outputs.append(expert_out)
313
start_idx = end_idx