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 4, 15:00 UTC
[DSV4] Add silu clamp limit to shared expert (#40950)
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Yongye Zhu committed
Apr 27, 2026 at 03:37 UTC
706a04d34ba64ea23d430d5e50038791aacfae96
7 files changed
+269
-29
csrc/activation_kernels.cu
+82
-25
@@ -11,29 +11,74 @@
11
namespace vllm {
12
13
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
14
- bool act_first>
14
+ bool act_first, bool HAS_CLAMP>
15
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
16
- const scalar_t& y) {
17
- return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
16
+ const scalar_t& y,
17
+ const float limit) {
18
+ if constexpr (act_first) {
19
+ scalar_t gate = x;
20
+ scalar_t up = y;
21
+ if constexpr (HAS_CLAMP) {
22
+ gate = (scalar_t)fminf((float)gate, limit);
23
+ up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
24
+ }
25
+ return ACT_FN(gate) * up;
26
+ } else {
27
+ scalar_t gate = x;
28
+ scalar_t up = y;
29
+ if constexpr (HAS_CLAMP) {
30
+ gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
31
+ up = (scalar_t)fminf((float)up, limit);
32
+ }
33
+ return gate * ACT_FN(up);
34
+ }
35
}
36
37
template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
21
- bool act_first>
38
+ bool act_first, bool HAS_CLAMP>
39
__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
23
- const packed_t& y) {
24
- return act_first ? packed_mul(PACKED_ACT_FN(x), y)
25
- : packed_mul(x, PACKED_ACT_FN(y));
40
+ const packed_t& y,
41
+ const float limit) {
42
+ if constexpr (act_first) {
43
+ packed_t gate = x;
44
+ packed_t up = y;
45
+ if constexpr (HAS_CLAMP) {
46
+ float2 g = cast_to_float2(gate);
47
+ float2 u = cast_to_float2(up);
48
+ g.x = fminf(g.x, limit);
49
+ g.y = fminf(g.y, limit);
50
+ u.x = fmaxf(fminf(u.x, limit), -limit);
51
+ u.y = fmaxf(fminf(u.y, limit), -limit);
52
+ gate = cast_to_packed<packed_t>(g);
53
+ up = cast_to_packed<packed_t>(u);
54
+ }
55
+ return packed_mul(PACKED_ACT_FN(gate), up);
56
+ } else {
57
+ packed_t gate = x;
58
+ packed_t up = y;
59
+ if constexpr (HAS_CLAMP) {
60
+ float2 g = cast_to_float2(gate);
61
+ float2 u = cast_to_float2(up);
62
+ g.x = fmaxf(fminf(g.x, limit), -limit);
63
+ g.y = fmaxf(fminf(g.y, limit), -limit);
64
+ u.x = fminf(u.x, limit);
65
+ u.y = fminf(u.y, limit);
66
+ gate = cast_to_packed<packed_t>(g);
67
+ up = cast_to_packed<packed_t>(u);
68
+ }
69
+ return packed_mul(gate, PACKED_ACT_FN(up));
70
+ }
71
}
72
73
// Activation and gating kernel template.
74
template <typename scalar_t, typename packed_t,
75
scalar_t (*ACT_FN)(const scalar_t&),
76
packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
32
- bool use_vec, bool use_256b = false>
77
+ bool use_vec, bool HAS_CLAMP, bool use_256b = false>
78
__global__ void act_and_mul_kernel(
79
scalar_t* __restrict__ out, // [..., d]
80
const scalar_t* __restrict__ input, // [..., 2, d]
36
- const int d) {
81
+ const int d, const float limit) {
82
const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
83
const scalar_t* y_ptr = x_ptr + d;
84
scalar_t* out_ptr = out + blockIdx.x * d;
@@ -58,8 +103,9 @@ __global__ void act_and_mul_kernel(
103
}
104
#pragma unroll
105
for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
61
- x.elts[j] = packed_compute<packed_t, PACKED_ACT_FN, act_first>(
62
- x.elts[j], y.elts[j]);
106
+ x.elts[j] =
107
+ packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
108
+ x.elts[j], y.elts[j], limit);
109
}
110
if constexpr (use_256b) {
111
st256(x, &out_vec[i]);
@@ -72,7 +118,8 @@ __global__ void act_and_mul_kernel(
118
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
119
const scalar_t x = VLLM_LDG(&x_ptr[idx]);
120
const scalar_t y = VLLM_LDG(&y_ptr[idx]);
75
- out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
121
+ out_ptr[idx] =
122
+ compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit);
123
}
124
}
125
}
@@ -151,8 +198,11 @@ packed_gelu_tanh_kernel(const packed_t& val) {
198
199
// Launch activation and gating kernel.
200
// Use ACT_FIRST (bool) indicating whether to apply the activation function
154
-// first.
155
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
201
+// first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
202
+// clamped (max only) and up input is clamped (both sides) before the
203
+// activation function is applied.
204
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
205
+ HAS_CLAMP, LIMIT) \
206
auto dtype = input.scalar_type(); \
207
int d = input.size(-1) / 2; \
208
int64_t num_tokens = input.numel() / input.size(-1); \
@@ -177,8 +227,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
227
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
228
KERNEL<scalar_t>, \
229
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
180
- ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
181
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
230
+ ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
231
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
232
}); \
233
} else { \
234
VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
@@ -186,8 +236,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
236
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
237
KERNEL<scalar_t>, \
238
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
189
- ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
190
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
239
+ ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
240
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
241
}); \
242
} \
243
} else { \
@@ -197,8 +247,8 @@ packed_gelu_tanh_kernel(const packed_t& val) {
247
scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
248
KERNEL<scalar_t>, \
249
PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
200
- ACT_FIRST, false><<<grid, block, 0, stream>>>( \
201
- out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
250
+ ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
251
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, LIMIT); \
252
}); \
253
}
254
@@ -206,7 +256,14 @@ void silu_and_mul(torch::Tensor& out, // [..., d]
256
torch::Tensor& input) // [..., 2 * d]
257
{
258
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
209
- true);
259
+ true, false, 0.0f);
260
+}
261
+
262
+void silu_and_mul_clamp(torch::Tensor& out, // [..., d]
263
+ torch::Tensor& input, // [..., 2 * d]
264
+ double limit) {
265
+ LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
266
+ true, true, (float)limit);
267
}
268
269
void mul_and_silu(torch::Tensor& out, // [..., d]
@@ -215,21 +272,21 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
272
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
273
// applies the silu to the latter half of the input.
274
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
218
- false);
275
+ false, false, 0.0f);
276
}
277
278
void gelu_and_mul(torch::Tensor& out, // [..., d]
279
torch::Tensor& input) // [..., 2 * d]
280
{
281
LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
225
- true);
282
+ true, false, 0.0f);
283
}
284
285
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
286
torch::Tensor& input) // [..., 2 * d]
287
{
231
- LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
232
- vllm::packed_gelu_tanh_kernel, true);
288
+ LAUNCH_ACTIVATION_GATE_KERNEL(
289
+ vllm::gelu_tanh_kernel, vllm::packed_gelu_tanh_kernel, true, false, 0.0f);
290
}
291
292
namespace vllm {
csrc/ops.h
+2
@@ -163,6 +163,8 @@ void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
163
164
void silu_and_mul(torch::Tensor& out, torch::Tensor& input);
165
166
+void silu_and_mul_clamp(torch::Tensor& out, torch::Tensor& input, double limit);
167
+
168
void silu_and_mul_quant(torch::Tensor& out, torch::Tensor& input,
169
torch::Tensor& scale);
170
csrc/torch_bindings.cpp
+6
@@ -106,6 +106,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
106
ops.def("silu_and_mul(Tensor! result, Tensor input) -> ()");
107
ops.impl("silu_and_mul", torch::kCUDA, &silu_and_mul);
108
109
+ // SwiGLU activation with input clamping.
110
+ ops.def(
111
+ "silu_and_mul_with_clamp(Tensor! result, Tensor input, float limit) "
112
+ "-> ()");
113
+ ops.impl("silu_and_mul_with_clamp", torch::kCUDA, &silu_and_mul_clamp);
114
+
115
ops.def(
116
"silu_and_mul_quant(Tensor! result, Tensor input, Tensor scale) -> ()");
117
ops.impl("silu_and_mul_quant", torch::kCUDA, &silu_and_mul_quant);
tests/kernels/core/test_activation.py
+80
@@ -16,6 +16,7 @@ from vllm.model_executor.layers.activation import (
16
NewGELU,
17
QuickGELU,
18
SiluAndMul,
19
+ SiluAndMulWithClamp,
20
SwigluOAIAndMul,
21
SwigluStepAndMul,
22
swiglustep_and_mul_triton,
@@ -116,6 +117,85 @@ def test_act_and_mul(
117
opcheck(fn, (out, x))
118
119
120
+SWIGLU_LIMITS = [3.0, 7.0, 15.0]
121
+
122
+
123
+@pytest.mark.parametrize("swiglu_limit", SWIGLU_LIMITS)
124
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
125
+@pytest.mark.parametrize("d", D)
126
+@pytest.mark.parametrize("dtype", DTYPES)
127
+@pytest.mark.parametrize("seed", SEEDS)
128
+@pytest.mark.parametrize("device", CUDA_DEVICES)
129
+@torch.inference_mode()
130
+def test_silu_and_mul_with_clamp(
131
+ default_vllm_config,
132
+ swiglu_limit: float,
133
+ num_tokens: int,
134
+ d: int,
135
+ dtype: torch.dtype,
136
+ seed: int,
137
+ device: str,
138
+) -> None:
139
+ """SiluAndMulWithClamp: cuda kernel must match native reference."""
140
+ set_random_seed(seed)
141
+ torch.set_default_device(device)
142
+ # Use large values to ensure clamping is exercised.
143
+ x = torch.randn(num_tokens, 2 * d, dtype=dtype) * swiglu_limit * 2
144
+
145
+ layer = SiluAndMulWithClamp(swiglu_limit, compile_native=False)
146
+ out = layer(x)
147
+ ref_out = layer.forward_native(x)
148
+
149
+ rtol = {
150
+ torch.float16: 2e-3,
151
+ torch.bfloat16: 2e-2,
152
+ torch.float: 1.3e-6,
153
+ }
154
+ torch.testing.assert_close(
155
+ out, ref_out, atol=get_default_atol(out), rtol=rtol[out.dtype]
156
+ )
157
+
158
+ # Verify clamping is actually being applied: the clamped output should
159
+ # differ from the unclamped SiluAndMul output when inputs are large.
160
+ unclamped_out = SiluAndMul.forward_native(x)
161
+ assert not torch.equal(ref_out.float(), unclamped_out.float()), (
162
+ "Input was not large enough to exercise the clamp; increase scale"
163
+ )
164
+
165
+ # Verify gate clamping semantics with a controlled scalar case.
166
+ # gate=large_val is clamped to limit first, then silu(limit) * 1.0.
167
+ x_gate = torch.tensor(
168
+ [[swiglu_limit * 20.0, 1.0]], dtype=torch.float32, device=device
169
+ )
170
+ out_gate = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_gate)
171
+ expected_gate = torch.nn.functional.silu(
172
+ torch.tensor(swiglu_limit, dtype=torch.float32)
173
+ ).item()
174
+ torch.testing.assert_close(
175
+ out_gate,
176
+ torch.tensor([[expected_gate]], dtype=torch.float32, device=device),
177
+ atol=1e-3,
178
+ rtol=1e-3,
179
+ )
180
+
181
+ # Verify up clamping semantics: up >> limit gets clamped to limit.
182
+ x_up = torch.tensor(
183
+ [[1.0, swiglu_limit * 20.0]], dtype=torch.float32, device=device
184
+ )
185
+ out_up = SiluAndMulWithClamp(swiglu_limit, compile_native=False)(x_up)
186
+ silu_1 = torch.nn.functional.silu(torch.tensor(1.0)).item()
187
+ torch.testing.assert_close(
188
+ out_up,
189
+ torch.tensor([[silu_1 * swiglu_limit]], dtype=torch.float32, device=device),
190
+ atol=1e-3,
191
+ rtol=1e-3,
192
+ )
193
+
194
+ # opcheck
195
+ out_buf = torch.empty(x.shape[:-1] + (d,), dtype=dtype, device=device)
196
+ opcheck(torch.ops._C.silu_and_mul_with_clamp, (out_buf, x, swiglu_limit))
197
+
198
+
199
@pytest.mark.parametrize(
200
"activation",
201
[
vllm/model_executor/layers/activation.py
+40
@@ -151,6 +151,46 @@ class SiluAndMul(CustomOp):
151
return self.forward_cuda(x)
152
153
154
+@CustomOp.register("silu_and_mul_with_clamp")
155
+class SiluAndMulWithClamp(CustomOp):
156
+ """SwiGLU activation with input clamping (used by some MoE shared experts).
157
+
158
+ Computes:
159
+ gate = clamp(x[..., :d], max=swiglu_limit)
160
+ up = clamp(x[..., d:], min=-swiglu_limit, max=swiglu_limit)
161
+ out = silu(gate) * up
162
+ where d = x.shape[-1] // 2.
163
+
164
+ Shapes:
165
+ x: (num_tokens, 2 * d) or (batch_size, seq_len, 2 * d)
166
+ return: (num_tokens, d) or (batch_size, seq_len, d)
167
+ """
168
+
169
+ def __init__(self, swiglu_limit: float, *, compile_native: bool = True):
170
+ super().__init__(compile_native=compile_native)
171
+ self.swiglu_limit = float(swiglu_limit)
172
+ if current_platform.is_cuda_alike() or current_platform.is_xpu():
173
+ self.op = torch.ops._C.silu_and_mul_with_clamp
174
+ elif current_platform.is_cpu():
175
+ self._forward_method = self.forward_native
176
+
177
+ def forward_native(self, x: torch.Tensor) -> torch.Tensor:
178
+ d = x.shape[-1] // 2
179
+ gate = torch.clamp(x[..., :d], max=self.swiglu_limit)
180
+ up = torch.clamp(x[..., d:], min=-self.swiglu_limit, max=self.swiglu_limit)
181
+ return F.silu(gate) * up
182
+
183
+ def forward_cuda(self, x: torch.Tensor) -> torch.Tensor:
184
+ d = x.shape[-1] // 2
185
+ output_shape = x.shape[:-1] + (d,)
186
+ out = torch.empty(output_shape, dtype=x.dtype, device=x.device)
187
+ self.op(out, x, self.swiglu_limit)
188
+ return out
189
+
190
+ def forward_xpu(self, x: torch.Tensor) -> torch.Tensor:
191
+ return self.forward_cuda(x)
192
+
193
+
194
# --8<-- [start:mul_and_silu]
195
@CustomOp.register("mul_and_silu")
196
class MulAndSilu(CustomOp):
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+1
-1
@@ -45,7 +45,7 @@ def _gelu_and_mul(
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,
48
+ MoEActivation.SILU: lambda x: SiluAndMul(compile_native=False).forward_native(x),
49
MoEActivation.SWIGLUOAI: _swigluoai_forward_native,
50
MoEActivation.GELU: _gelu_and_mul,
51
}
vllm/model_executor/models/deepseek_v4.py
+58
-3
@@ -17,6 +17,7 @@ from vllm.distributed import (
17
get_tensor_model_parallel_world_size,
18
)
19
from vllm.forward_context import get_forward_context
20
+from vllm.model_executor.layers.activation import SiluAndMul, SiluAndMulWithClamp
21
from vllm.model_executor.layers.deepseek_v4_attention import (
22
DeepseekV4Indexer,
23
DeepseekV4MLAModules,
@@ -34,7 +35,10 @@ from vllm.model_executor.layers.linear import (
35
RowParallelLinear,
36
)
37
from vllm.model_executor.layers.logits_processor import LogitsProcessor
37
-from vllm.model_executor.layers.quantization import QuantizationMethods
38
+from vllm.model_executor.layers.quantization import (
39
+ QuantizationConfig,
40
+ QuantizationMethods,
41
+)
42
from vllm.model_executor.layers.quantization.fp8 import Fp8Config
43
from vllm.model_executor.layers.quantization.mxfp4 import Mxfp4MoEMethod
44
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -46,7 +50,6 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
50
VocabParallelEmbedding,
51
)
52
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
49
-from vllm.model_executor.models.deepseek_v2 import DeepseekV2MLP
53
from vllm.model_executor.utils import set_weight_attrs
54
from vllm.platforms import current_platform
55
from vllm.sequence import IntermediateTensors
@@ -63,6 +66,57 @@ from .utils import (
66
)
67
68
69
+class DeepseekV4MLP(nn.Module):
70
+ def __init__(
71
+ self,
72
+ hidden_size: int,
73
+ intermediate_size: int,
74
+ hidden_act: str,
75
+ swiglu_limit: float | None = None,
76
+ quant_config: QuantizationConfig | None = None,
77
+ reduce_results: bool = True,
78
+ is_sequence_parallel: bool = False,
79
+ prefix: str = "",
80
+ ) -> None:
81
+ super().__init__()
82
+
83
+ # If is_sequence_parallel, the input and output tensors are sharded
84
+ # across the ranks within the tp_group. In this case the weights are
85
+ # replicated and no collective ops are needed.
86
+ # Otherwise we use standard TP with an allreduce at the end.
87
+ self.gate_up_proj = MergedColumnParallelLinear(
88
+ hidden_size,
89
+ [intermediate_size] * 2,
90
+ bias=False,
91
+ quant_config=quant_config,
92
+ disable_tp=is_sequence_parallel,
93
+ prefix=f"{prefix}.gate_up_proj",
94
+ )
95
+ self.down_proj = RowParallelLinear(
96
+ intermediate_size,
97
+ hidden_size,
98
+ bias=False,
99
+ quant_config=quant_config,
100
+ reduce_results=reduce_results,
101
+ disable_tp=is_sequence_parallel,
102
+ prefix=f"{prefix}.down_proj",
103
+ )
104
+ if hidden_act != "silu":
105
+ raise ValueError(
106
+ f"Unsupported activation: {hidden_act}. Only silu is supported for now."
107
+ )
108
+ if swiglu_limit is not None:
109
+ self.act_fn = SiluAndMulWithClamp(swiglu_limit)
110
+ else:
111
+ self.act_fn = SiluAndMul()
112
+
113
+ def forward(self, x):
114
+ gate_up, _ = self.gate_up_proj(x)
115
+ x = self.act_fn(gate_up)
116
+ x, _ = self.down_proj(x)
117
+ return x
118
+
119
+
120
class DeepseekV4FP8Config(Fp8Config):
121
"""FP8 config that routes MoE layers to MXFP4 quantization.
122
@@ -672,10 +726,11 @@ class DeepseekV4MoE(nn.Module):
726
else:
727
intermediate_size = config.moe_intermediate_size * config.n_shared_experts
728
675
- self.shared_experts = DeepseekV2MLP(
729
+ self.shared_experts = DeepseekV4MLP(
730
hidden_size=config.hidden_size,
731
intermediate_size=intermediate_size,
732
hidden_act=config.hidden_act,
733
+ swiglu_limit=self.swiglu_limit,
734
quant_config=quant_config,
735
reduce_results=self.use_mega_moe,
736
prefix=f"{prefix}.shared_experts",