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 2, 03:30 UTC
[Kernel][MoE] Add GELU_TANH to CPU, CUTLASS, and WNA16 MoE backends (#42027)
Signed-off-by: lesj0610 <lesj0610@users.noreply.github.com> Co-authored-by: lesj0610 <lesj0610@users.noreply.github.com>
SeongJun Lee committed
Jun 3, 2026 at 06:12 UTC
3099de3617b5b2e569fb5273ab13863bc7e9cf49
7 files changed
+119
-7
csrc/cpu/cpu_fused_moe.cpp
+49
-1
@@ -30,7 +30,12 @@
30
}()
31
32
namespace {
33
-enum class FusedMOEAct { SiluAndMul, SwigluOAIAndMul, GeluAndMul };
33
+enum class FusedMOEAct {
34
+ SiluAndMul,
35
+ SwigluOAIAndMul,
36
+ GeluAndMul,
37
+ GeluTanhAndMul,
38
+};
39
40
FusedMOEAct get_act_type(const std::string& act) {
41
if (act == "silu") {
@@ -39,6 +44,8 @@ FusedMOEAct get_act_type(const std::string& act) {
44
return FusedMOEAct::SwigluOAIAndMul;
45
} else if (act == "gelu") {
46
return FusedMOEAct::GeluAndMul;
47
+ } else if (act == "gelu_tanh") {
48
+ return FusedMOEAct::GeluTanhAndMul;
49
} else {
50
TORCH_CHECK(false, "Invalid act type: " + act);
51
}
@@ -143,6 +150,44 @@ void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
150
}
151
}
152
153
+template <typename scalar_t>
154
+void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
155
+ const int32_t m_size, const int32_t n_size,
156
+ const int32_t input_stride,
157
+ const int32_t output_stride) {
158
+ using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
159
+ const int32_t dim = n_size / 2;
160
+ float* __restrict__ gate = input;
161
+ float* __restrict__ up = input + dim;
162
+ vec_op::FP32Vec16 one_vec(1.0);
163
+ vec_op::FP32Vec16 w1_vec(0.7978845608028654);
164
+ vec_op::FP32Vec16 w2_vec(0.5);
165
+ vec_op::FP32Vec16 w3_vec(0.044715);
166
+ alignas(64) float temp[16];
167
+
168
+ for (int32_t m = 0; m < m_size; ++m) {
169
+ for (int32_t n = 0; n < dim; n += 16) {
170
+ vec_op::FP32Vec16 gate_vec(gate + n);
171
+ vec_op::FP32Vec16 up_vec(up + n);
172
+ auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
173
+ auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
174
+
175
+ inner_vec.save(temp);
176
+ for (int32_t i = 0; i < 16; ++i) {
177
+ temp[i] = std::tanh(temp[i]);
178
+ }
179
+ vec_op::FP32Vec16 tanh_vec(temp);
180
+ auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
181
+ auto gated_output_fp32 = up_vec * gelu_tanh;
182
+ scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
183
+ gated_output.save(output + n);
184
+ }
185
+ gate += input_stride;
186
+ up += input_stride;
187
+ output += output_stride;
188
+ }
189
+}
190
+
191
template <typename scalar_t>
192
FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
193
float* __restrict__ input,
@@ -160,6 +205,9 @@ FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
205
case FusedMOEAct::GeluAndMul:
206
gelu_and_mul(input, output, m, n, input_stride, output_stride);
207
return;
208
+ case FusedMOEAct::GeluTanhAndMul:
209
+ gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
210
+ return;
211
default:
212
TORCH_CHECK(false, "Unsupported act type.");
213
}
tests/kernels/moe/test_cpu_fused_moe.py
+6
-1
@@ -20,7 +20,12 @@ EXPERT_NUM = [
20
HIDDEN_DIM = [128, 2880]
21
INTERMEDIATE_DIM = [128, 2880]
22
BATCH_SIZE = [1, 64, 256]
23
-ACT = [MoEActivation.SILU, MoEActivation.SWIGLUOAI, MoEActivation.GELU]
23
+ACT = [
24
+ MoEActivation.SILU,
25
+ MoEActivation.SWIGLUOAI,
26
+ MoEActivation.GELU,
27
+ MoEActivation.GELU_TANH,
28
+]
29
USE_BIAS = [True, False]
30
ISA = ["amx", "vec"] if torch.cpu._is_amx_tile_supported() else ["vec"]
31
DTYPE = [torch.bfloat16]
tests/kernels/moe/test_cutlass_moe.py
+7
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.fused_moe.config import (
22
fp8_w8a8_moe_quant_config,
23
)
24
from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
25
+ CutlassExpertsFp4,
26
CutlassExpertsFp8,
27
run_cutlass_moe_fp8,
28
)
@@ -52,6 +53,12 @@ MNK_FACTORS = [
53
vllm_config = VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
54
55
56
+def test_cutlass_moe_supports_gelu_tanh_activation_metadata():
57
+ assert CutlassExpertsFp8._supports_activation(MoEActivation.GELU_TANH)
58
+ assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH)
59
+ assert CutlassExpertsFp4._supports_activation(MoEActivation.GELU_TANH_NO_MUL)
60
+
61
+
62
@dataclasses.dataclass
63
class MOETensors:
64
a: torch.Tensor
tests/quantization/test_moe_wna16.py
new
+49
@@ -0,0 +1,49 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from types import SimpleNamespace
5
+
6
+import pytest
7
+import torch
8
+
9
+from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10
+from vllm.model_executor.layers.quantization.moe_wna16 import MoeWNA16Method
11
+from vllm.platforms import current_platform
12
+
13
+
14
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="Only test on CUDA")
15
+def test_moe_wna16_apply_passes_layer_activation(monkeypatch):
16
+ captured_kwargs = {}
17
+
18
+ def fake_fused_experts(*args, **kwargs):
19
+ captured_kwargs.update(kwargs)
20
+ return torch.empty(1, 2)
21
+
22
+ monkeypatch.setattr(
23
+ "vllm.model_executor.layers.fused_moe.fused_experts",
24
+ fake_fused_experts,
25
+ )
26
+
27
+ method = object.__new__(MoeWNA16Method)
28
+ method.moe = SimpleNamespace(disable_inplace=False)
29
+ method.moe_quant_config = object()
30
+ layer = SimpleNamespace(
31
+ w13_qweight=torch.empty(1, 2),
32
+ w2_qweight=torch.empty(1, 2),
33
+ activation=MoEActivation.GELU_TANH,
34
+ apply_router_weight_on_input=False,
35
+ global_num_experts=1,
36
+ expert_map=None,
37
+ )
38
+
39
+ output = method.apply(
40
+ layer,
41
+ x=torch.empty(1, 2),
42
+ topk_weights=torch.empty(1, 1),
43
+ topk_ids=torch.empty(1, 1, dtype=torch.int32),
44
+ shared_experts=None,
45
+ shared_experts_input=None,
46
+ )
47
+
48
+ assert output.shape == (1, 2)
49
+ assert captured_kwargs["activation"] is MoEActivation.GELU_TANH
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+4
@@ -53,6 +53,10 @@ _CPU_MOE_ACT_FN: dict[MoEActivation, Callable[[torch.Tensor], torch.Tensor]] = {
53
MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
54
MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
55
MoEActivation.GELU: _gelu_and_mul,
56
+ MoEActivation.GELU_TANH: (
57
+ lambda x: F.gelu(x[..., : x.shape[-1] // 2], approximate="tanh")
58
+ * x[..., x.shape[-1] // 2 :]
59
+ ),
60
}
61
62
vllm/model_executor/layers/fused_moe/experts/cutlass_moe.py
+3
@@ -322,6 +322,7 @@ class CutlassExpertsFp8Base(mk.FusedMoEExpertsModular):
322
return activation in [
323
MoEActivation.SILU,
324
MoEActivation.GELU,
325
+ MoEActivation.GELU_TANH,
326
MoEActivation.SWIGLUOAI,
327
]
328
@@ -719,10 +720,12 @@ class CutlassExpertsFp4(mk.FusedMoEExpertsModular):
720
return activation in [
721
MoEActivation.SILU,
722
MoEActivation.GELU,
723
+ MoEActivation.GELU_TANH,
724
MoEActivation.SWIGLUOAI,
725
MoEActivation.SWIGLUSTEP,
726
MoEActivation.SILU_NO_MUL,
727
MoEActivation.GELU_NO_MUL,
728
+ MoEActivation.GELU_TANH_NO_MUL,
729
MoEActivation.RELU2_NO_MUL,
730
]
731
vllm/model_executor/layers/quantization/moe_wna16.py
+1
-5
@@ -13,7 +13,6 @@ from vllm.model_executor.layers.fused_moe import (
13
RoutedExperts,
14
SharedExperts,
15
)
16
-from vllm.model_executor.layers.fused_moe.activation import MoEActivation
16
from vllm.model_executor.layers.fused_moe.config import (
17
FusedMoEQuantConfig,
18
int4_w4a16_moe_quant_config,
@@ -367,16 +366,13 @@ class MoeWNA16Method(FusedMoEMethodBase):
366
) -> torch.Tensor:
367
from vllm.model_executor.layers.fused_moe import fused_experts
368
370
- assert layer.activation == MoEActivation.SILU, (
371
- f"Only SiLU activation is supported, not {layer.activation}."
372
- )
373
-
369
return fused_experts(
370
x,
371
layer.w13_qweight,
372
layer.w2_qweight,
373
topk_weights=topk_weights,
374
topk_ids=topk_ids,
375
+ activation=layer.activation,
376
apply_router_weight_on_input=layer.apply_router_weight_on_input,
377
global_num_experts=layer.global_num_experts,
378
expert_map=layer.expert_map,