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 5, 16:00 UTC
[CPU] Support gelu act in cpu_fused_moe (#38770)
Signed-off-by: jiang1.li <jiang1.li@intel.com>
Li, Jiang committed
Apr 2, 2026 at 14:14 UTC
c6f722b93e8e795065751172812ee6a5540e5901
5 files changed
+59
-3
csrc/cpu/cpu_fused_moe.cpp
+43
-1
@@ -30,13 +30,15 @@
30
}()
31
32
namespace {
33
-enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul };
33
+enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul };
34
35
FusedMOEAct get_act_type(const std::string& act) {
36
if (act == "silu") {
37
return FusedMOEAct::SiluAndMul;
38
} else if (act == "swigluoai") {
39
return FusedMOEAct::SwigluOAIAndMul;
40
+ } else if (act == "gelu") {
41
+ return FusedMOEAct::GeluAndMul;
42
} else {
43
TORCH_CHECK(false, "Invalid act type: " + act);
44
}
@@ -104,6 +106,43 @@ void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
106
}
107
}
108
109
+template <typename scalar_t>
110
+void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
111
+ const int32_t m_size, const int32_t n_size,
112
+ const int32_t input_stride, const int32_t output_stride) {
113
+ using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
114
+ const int32_t dim = n_size / 2;
115
+ float* __restrict__ gate = input;
116
+ float* __restrict__ up = input + dim;
117
+ vec_op::FP32Vec16 one_vec(1.0);
118
+ vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
119
+ vec_op::FP32Vec16 w2_vec(0.5);
120
+ alignas(64) float temp[16];
121
+
122
+ DEFINE_FAST_EXP
123
+
124
+ for (int32_t m = 0; m < m_size; ++m) {
125
+ for (int32_t n = 0; n < dim; n += 16) {
126
+ vec_op::FP32Vec16 gate_vec(gate + n);
127
+ vec_op::FP32Vec16 up_vec(up + n);
128
+ auto er_input_vec = gate_vec * w1_vec;
129
+
130
+ er_input_vec.save(temp);
131
+ for (int32_t i = 0; i < 16; ++i) {
132
+ temp[i] = std::erf(temp[i]);
133
+ }
134
+ vec_op::FP32Vec16 er_vec(temp);
135
+ auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
136
+ auto gated_output_fp32 = up_vec * gelu;
137
+ scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
138
+ gated_output.save(output + n);
139
+ }
140
+ gate += input_stride;
141
+ up += input_stride;
142
+ output += output_stride;
143
+ }
144
+}
145
+
146
template <typename scalar_t>
147
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
148
float* __restrict__ input,
@@ -118,6 +157,9 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
157
case FusedMOEAct::SiluAndMul:
158
silu_and_mul(input, output, m, n, input_stride, output_stride);
159
return;
160
+ case FusedMOEAct::GeluAndMul:
161
+ gelu_and_mul(input, output, m, n, input_stride, output_stride);
162
+ return;
163
default:
164
TORCH_CHECK(false, "Unsupported act type.");
165
}
tests/kernels/moe/test_cpu_fused_moe.py
+1
-1
@@ -20,7 +20,7 @@ EXPERT_NUM = [
20
HIDDEN_DIM = [128, 2880]
21
INTERMEDIATE_DIM = [128, 2880]
22
BATCH_SIZE = [1, 64, 256]
23
-ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI]
23
+ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI, MoEActivation.GELU]
24
USE_BIAS = [True, False]
25
ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
26
DTYPE = [torch.bfloat16]
vllm/envs.py
+5
@@ -51,6 +51,7 @@ if TYPE_CHECKING:
51
VLLM_CPU_OMP_THREADS_BIND: str = "auto"
52
VLLM_CPU_NUM_OF_RESERVED_CPU: int | None = None
53
VLLM_CPU_SGL_KERNEL: bool = False
54
+ VLLM_CPU_ATTN_SPLIT_KV: bool = True
55
VLLM_ZENTORCH_WEIGHT_PREPACK: bool = True
56
VLLM_CPU_INT4_W4A8: bool = True
57
VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
@@ -726,6 +727,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
727
else None,
728
# (CPU backend only) whether to use SGL kernels, optimized for small batch.
729
"VLLM_CPU_SGL_KERNEL": lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
730
+ # (CPU backend only) whether to enable attention spilt KV.
731
+ "VLLM_CPU_ATTN_SPLIT_KV": lambda: bool(
732
+ int(os.getenv("VLLM_CPU_ATTN_SPLIT_KV", "1"))
733
+ ),
734
# (Zen CPU backend) eagerly prepack weights into ZenDNN blocked layout
735
# at model load time. Eliminates per-inference layout conversion overhead.
736
"VLLM_ZENTORCH_WEIGHT_PREPACK": lambda: bool(
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+8
@@ -34,12 +34,20 @@ def _swigluoai_forward_native(
34
return gated_output
35
36
37
+def _gelu_and_mul(
38
+ x: torch.Tensor,
39
+) -> torch.Tensor:
40
+ d = x.shape[-1] // 2
41
+ return F.gelu(x[..., :d], approximate="none") * x[..., d:]
42
+
43
+
44
# Map activation names to their native forward functions.
45
# Uses static methods or standalone functions to avoid instantiating CustomOp
46
# classes, which would call get_current_vllm_config() before config is set.
47
_CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
48
MoEActivation.SILU: SiluAndMul.forward_native,
49
MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
50
+ MoEActivation.GELU: _gelu_and_mul,
51
}
52
53
vllm/v1/attention/backends/cpu_attn.py
+2
-1
@@ -6,6 +6,7 @@ from typing import ClassVar
6
import torch
7
8
from vllm import _custom_ops as ops
9
+from vllm import envs
10
from vllm.config import VllmConfig
11
from vllm.logger import init_logger
12
from vllm.platforms import CpuArchEnum, current_platform
@@ -181,7 +182,7 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
182
causal=causal,
183
sliding_window_size=self.window_size,
184
isa=self.isa,
184
- enable_kv_split=True,
185
+ enable_kv_split=envs.VLLM_CPU_ATTN_SPLIT_KV,
186
)
187
188
attn_metadata = CPUAttentionMetadata(