[Perf][DSv4] Use native CUDA SwiGLU clamp kernel for Humming MoE (throughput +1.40%) (#53685)
Signed-off-by: chaunceyjiang <chaunceyjiang@gmail.com>
Chauncey committed
Aug 28, 2026 at 04:52 UTC
7d5769b2d9cf1e999e3f2fc2c1da265cf0d5d794
2 files changed
+17
tests/kernels/core/test_activation.py
+15
@@ -22,6 +22,7 @@ from vllm.model_executor.layers.activation import (
22
SwigluStepAndMul,
23
swiglustep_and_mul_triton,
24
)
25
+from vllm.model_executor.layers.fused_moe.utils import swiglu_limit_func
26
from vllm.utils.torch_utils import set_random_seed
27
28
DTYPES = [torch.half, torch.bfloat16, torch.float]
@@ -121,6 +122,20 @@ def test_act_and_mul(
122
SWIGLU_LIMITS = [3.0, 7.0, 15.0]
123
124
125
+@torch.inference_mode()
126
+def test_swiglu_limit_func_without_routing_uses_output_buffer() -> None:
127
+ x = torch.randn(7, 1024, dtype=torch.bfloat16, device="cuda")
128
+ output = torch.empty(7, 512, dtype=x.dtype, device=x.device)
129
+
130
+ swiglu_limit_func(output, x, swiglu_limit=7.0)
131
+ gate, up = x.chunk(2, dim=-1)
132
+ expected = torch.nn.functional.silu(gate.clamp(max=7.0)) * up.clamp(
133
+ min=-7.0, max=7.0
134
+ )
135
+
136
+ torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
137
+
138
+
139
@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
140
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
141
@pytest.mark.parametrize("d", D)
vllm/model_executor/layers/fused_moe/utils.py
+2
@@ -647,6 +647,8 @@ def swiglu_limit_func(
647
# requires topk_ids. Fall back to the torch implementation otherwise.
648
if topk_ids is not None:
649
_swiglu_limit_pad_aware(output, input, topk_ids, swiglu_limit, expert_map)
650
+ elif current_platform.is_cuda():
651
+ torch.ops._C.silu_and_mul_with_clamp(output, input, swiglu_limit, 1.0, 0.0)
652
else:
653
_swiglu_limit_torch(output, input, swiglu_limit)
654