@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cu 1,451 lines 64.9 KB
Raw
1 #include <cuda.h>
2 #include <cuda_runtime.h>
3 #include <torch/csrc/stable/tensor.h>
4 #include <torch/headeronly/util/Float8_e4m3fn.h>
5
6 #ifndef USE_ROCM
7 #include <cuda_fp8.h>
8 #endif
9
10 #include <cmath>
11 #include <type_traits>
12
13 #include "../cuda_compat.h"
14 #include "async_util.cuh"
15 #include "cuda_vec_utils.cuh"
16 #include "dispatch_utils.h"
17 #include "torch_utils.h"
18
19 namespace vllm {
20
21 // `alpha` and `beta` are applied to opposite operands:
22 // - alpha lives INSIDE the activation (the activated half): the gated
23 // activation computes act_half * sigmoid(alpha * act_half).
24 // - beta is added to the OTHER (non-activated) half before the multiply.
25 // So the result is always ACT(act_half, alpha) * (other_half + beta).
26 // Which half is which depends on `act_first` (see below). Defaults
27 // alpha=1.0, beta=0.0 reproduce the plain SwiGLU/GeGLU behavior.
28 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
29 bool act_first, bool HAS_CLAMP>
30 __device__ __forceinline__ scalar_t compute(const scalar_t& x,
31 const scalar_t& y,
32 const float limit,
33 const float alpha,
34 const float beta) {
35 if constexpr (act_first) {
36 scalar_t gate = x;
37 scalar_t up = y;
38 if constexpr (HAS_CLAMP) {
39 gate = (scalar_t)fminf((float)gate, limit);
40 up = (scalar_t)fmaxf(fminf((float)up, limit), -limit);
41 }
42 // act_first: gate is the activated half -> alpha applies to gate;
43 // beta is added to up (the non-activated half).
44 return (scalar_t)(ACT_FN(gate, alpha) * ((float)up + beta));
45 } else {
46 scalar_t gate = x;
47 scalar_t up = y;
48 if constexpr (HAS_CLAMP) {
49 gate = (scalar_t)fmaxf(fminf((float)gate, limit), -limit);
50 up = (scalar_t)fminf((float)up, limit);
51 }
52 // !act_first: up is the activated half -> alpha applies to up;
53 // beta is added to gate (the non-activated half).
54 return (scalar_t)(((float)gate + beta) * ACT_FN(up, alpha));
55 }
56 }
57
58 template <typename packed_t,
59 packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
60 bool act_first, bool HAS_CLAMP>
61 __device__ __forceinline__ packed_t packed_compute(const packed_t& x,
62 const packed_t& y,
63 const float limit,
64 const float alpha,
65 const float beta) {
66 if constexpr (act_first) {
67 packed_t gate = x;
68 packed_t up = y;
69 float2 u = cast_to_float2(up);
70 if constexpr (HAS_CLAMP) {
71 float2 g = cast_to_float2(gate);
72 g.x = fminf(g.x, limit);
73 g.y = fminf(g.y, limit);
74 u.x = fmaxf(fminf(u.x, limit), -limit);
75 u.y = fmaxf(fminf(u.y, limit), -limit);
76 gate = cast_to_packed<packed_t>(g);
77 }
78 // act_first: gate is the activated half -> alpha applies to gate;
79 // beta is added to up (the non-activated half).
80 float2 activated = cast_to_float2(PACKED_ACT_FN(gate, alpha));
81 activated.x *= u.x + beta;
82 activated.y *= u.y + beta;
83 return cast_to_packed<packed_t>(activated);
84 } else {
85 packed_t gate = x;
86 packed_t up = y;
87 float2 g = cast_to_float2(gate);
88 if constexpr (HAS_CLAMP) {
89 float2 u = cast_to_float2(up);
90 g.x = fmaxf(fminf(g.x, limit), -limit);
91 g.y = fmaxf(fminf(g.y, limit), -limit);
92 u.x = fminf(u.x, limit);
93 u.y = fminf(u.y, limit);
94 up = cast_to_packed<packed_t>(u);
95 }
96 // !act_first: up is the activated half -> alpha applies to up;
97 // beta is added to gate (the non-activated half).
98 float2 activated = cast_to_float2(PACKED_ACT_FN(up, alpha));
99 activated.x *= g.x + beta;
100 activated.y *= g.y + beta;
101 return cast_to_packed<packed_t>(activated);
102 }
103 }
104
105 // Activation and gating kernel template.
106 template <typename scalar_t, typename packed_t,
107 scalar_t (*ACT_FN)(const scalar_t&, const float),
108 packed_t (*PACKED_ACT_FN)(const packed_t&, const float),
109 bool act_first, bool use_vec, bool HAS_CLAMP, bool use_256b = false>
110 __global__ void act_and_mul_kernel(
111 scalar_t* __restrict__ out, // [..., d]
112 const scalar_t* __restrict__ input, // [..., 2, d]
113 const int d, const float limit, const float alpha, const float beta) {
114 const int64_t token_idx = blockIdx.x;
115 const scalar_t* x_ptr = input + token_idx * 2 * d;
116 const scalar_t* y_ptr = x_ptr + d;
117 scalar_t* out_ptr = out + token_idx * d;
118
119 if constexpr (use_vec) {
120 using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
121 using pvec_t = PackedVec<cuda_t, use_256b>;
122
123 const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
124 const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
125 pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
126 const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
127
128 for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
129 pvec_t x, y;
130 if constexpr (use_256b) {
131 ld256(x, &x_vec[i]);
132 ld256(y, &y_vec[i]);
133 } else {
134 ld128(x, &x_vec[i]);
135 ld128(y, &y_vec[i]);
136 }
137 #pragma unroll
138 for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
139 x.elts[j] =
140 packed_compute<packed_t, PACKED_ACT_FN, act_first, HAS_CLAMP>(
141 x.elts[j], y.elts[j], limit, alpha, beta);
142 }
143 if constexpr (use_256b) {
144 st256(x, &out_vec[i]);
145 } else {
146 st128(x, &out_vec[i]);
147 }
148 }
149 } else {
150 // Scalar fallback for unaligned data or small d
151 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
152 const scalar_t x = VLLM_LDG(&x_ptr[idx]);
153 const scalar_t y = VLLM_LDG(&y_ptr[idx]);
154 out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(
155 x, y, limit, alpha, beta);
156 }
157 }
158 }
159
160 // Gated activations take an `alpha` argument that scales the sigmoid input
161 // (`x * sigmoid(alpha * x)`). alpha defaults to 1.0 at all call sites, which
162 // is exactly SiLU; only the clamp path (silu_and_mul_with_clamp) passes a
163 // non-default alpha. Activations that do not use alpha simply ignore it.
164 template <typename T>
165 __device__ __forceinline__ T silu_kernel(const T& x, const float alpha) {
166 // x * sigmoid(alpha * x)
167 return (T)(((float)x) / (1.0f + expf((float)-x * alpha)));
168 }
169
170 template <typename packed_t>
171 __device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val,
172 const float alpha) {
173 // x * sigmoid(alpha * x)
174 float2 fval = cast_to_float2(val);
175 fval.x = fval.x / (1.0f + expf(-fval.x * alpha));
176 fval.y = fval.y / (1.0f + expf(-fval.y * alpha));
177 return cast_to_packed<packed_t>(fval);
178 }
179
180 template <typename T>
181 __device__ __forceinline__ T gelu_kernel(const T& x, const float /*alpha*/) {
182 // Equivalent to PyTorch GELU with 'none' approximation.
183 // Refer to:
184 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
185 const float f = (float)x;
186 constexpr float ALPHA = M_SQRT1_2;
187 return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
188 }
189
190 template <typename packed_t>
191 __device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val,
192 const float /*alpha*/) {
193 // Equivalent to PyTorch GELU with 'none' approximation.
194 // Refer to:
195 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
196 constexpr float ALPHA = M_SQRT1_2;
197 float2 fval = cast_to_float2(val);
198 fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA));
199 fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA));
200 return cast_to_packed<packed_t>(fval);
201 }
202
203 template <typename T>
204 __device__ __forceinline__ T gelu_tanh_kernel(const T& x,
205 const float /*alpha*/) {
206 // Equivalent to PyTorch GELU with 'tanh' approximation.
207 // Refer to:
208 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
209 const float f = (float)x;
210 constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
211 constexpr float KAPPA = 0.044715;
212 float x_cube = f * f * f;
213 float inner = BETA * (f + KAPPA * x_cube);
214 return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
215 }
216
217 template <typename packed_t>
218 __device__ __forceinline__ packed_t
219 packed_gelu_tanh_kernel(const packed_t& val, const float /*alpha*/) {
220 // Equivalent to PyTorch GELU with 'tanh' approximation.
221 // Refer to:
222 // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
223 float2 fval = cast_to_float2(val);
224 constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
225 constexpr float KAPPA = 0.044715;
226
227 float x_cube = fval.x * fval.x * fval.x;
228 float inner = BETA * (fval.x + KAPPA * x_cube);
229 fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner));
230
231 x_cube = fval.y * fval.y * fval.y;
232 inner = BETA * (fval.y + KAPPA * x_cube);
233 fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner));
234 return cast_to_packed<packed_t>(fval);
235 }
236
237 } // namespace vllm
238
239 // Launch activation and gating kernel.
240 // Use ACT_FIRST (bool) indicating whether to apply the activation function
241 // first. HAS_CLAMP (bool) enables pre-activation clamping: gate input is
242 // clamped (max only) and up input is clamped (both sides) before the
243 // activation function is applied.
244 #define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST, \
245 HAS_CLAMP, LIMIT, ALPHA, BETA) \
246 auto dtype = input.scalar_type(); \
247 int d = input.size(-1) / 2; \
248 int64_t num_tokens = input.numel() / input.size(-1); \
249 if (num_tokens == 0) { \
250 return; \
251 } \
252 dim3 grid(num_tokens); \
253 int cc_major = get_device_prop()->major; \
254 int support_vec = \
255 (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
256 ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
257 : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
258 int vec_size = support_vec / input.element_size(); \
259 const bool use_vec = (d % vec_size == 0); \
260 const torch::stable::accelerator::DeviceGuard device_guard( \
261 input.get_device_index()); \
262 const cudaStream_t stream = get_current_cuda_stream(); \
263 if (use_vec) { \
264 dim3 block(std::min(d / vec_size, 1024)); \
265 if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
266 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
267 vllm::act_and_mul_kernel< \
268 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
269 KERNEL<scalar_t>, \
270 PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
271 ACT_FIRST, true, HAS_CLAMP, true><<<grid, block, 0, stream>>>( \
272 out.mutable_data_ptr<scalar_t>(), \
273 input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
274 }); \
275 } else { \
276 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
277 vllm::act_and_mul_kernel< \
278 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
279 KERNEL<scalar_t>, \
280 PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
281 ACT_FIRST, true, HAS_CLAMP, false><<<grid, block, 0, stream>>>( \
282 out.mutable_data_ptr<scalar_t>(), \
283 input.const_data_ptr<scalar_t>(), d, LIMIT, ALPHA, BETA); \
284 }); \
285 } \
286 } else { \
287 dim3 block(std::min(d, 1024)); \
288 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
289 vllm::act_and_mul_kernel< \
290 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
291 KERNEL<scalar_t>, \
292 PACKED_KERNEL<typename vllm::PackedTypeConverter<scalar_t>::Type>, \
293 ACT_FIRST, false, HAS_CLAMP><<<grid, block, 0, stream>>>( \
294 out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
295 d, LIMIT, ALPHA, BETA); \
296 }); \
297 }
298
299 void silu_and_mul(torch::stable::Tensor& out, // [..., d]
300 torch::stable::Tensor& input) // [..., 2 * d]
301 {
302 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
303 true, false, 0.0f, 1.0f, 0.0f);
304 }
305
306 void silu_and_mul_clamp(torch::stable::Tensor& out, // [..., d]
307 torch::stable::Tensor& input, // [..., 2 * d]
308 double limit, double alpha, double beta) {
309 // out = (gate.clamp(max=limit) * sigmoid(alpha * gate.clamp(max=limit)))
310 // * (up.clamp(+-limit) + beta)
311 // alpha=1.0, beta=0.0 reduce this to silu(gate) * up.
312 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
313 true, true, (float)limit, (float)alpha,
314 (float)beta);
315 }
316
317 void mul_and_silu(torch::stable::Tensor& out, // [..., d]
318 torch::stable::Tensor& input) // [..., 2 * d]
319 {
320 // The difference between mul_and_silu and silu_and_mul is that mul_and_silu
321 // applies the silu to the latter half of the input.
322 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
323 false, false, 0.0f, 1.0f, 0.0f);
324 }
325
326 void gelu_and_mul(torch::stable::Tensor& out, // [..., d]
327 torch::stable::Tensor& input) // [..., 2 * d]
328 {
329 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
330 true, false, 0.0f, 1.0f, 0.0f);
331 }
332
333 void gelu_tanh_and_mul(torch::stable::Tensor& out, // [..., d]
334 torch::stable::Tensor& input) // [..., 2 * d]
335 {
336 LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
337 vllm::packed_gelu_tanh_kernel, true, false,
338 0.0f, 1.0f, 0.0f);
339 }
340
341 namespace vllm {
342
343 template <typename T>
344 __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
345 const float f = (float)x;
346 return (T)(f > threshold ? f : 0.0f);
347 }
348
349 template <typename packed_t>
350 __device__ __forceinline__ packed_t
351 packed_fatrelu_kernel(const packed_t& val, const float threshold) {
352 float2 fval = cast_to_float2(val);
353 fval.x = fval.x > threshold ? fval.x : 0.0f;
354 fval.y = fval.y > threshold ? fval.y : 0.0f;
355 return cast_to_packed<packed_t>(fval);
356 }
357
358 template <typename scalar_t, typename packed_t,
359 scalar_t (*ACT_FN)(const scalar_t&, const float),
360 packed_t (*PACKED_ACT_FN)(const packed_t&, const float), bool use_vec,
361 bool use_256b = false>
362 __global__ void act_and_mul_kernel_with_param(
363 scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
364 const float param) {
365 const int64_t token_idx = blockIdx.x;
366 const scalar_t* x_ptr = input + token_idx * 2 * d;
367 const scalar_t* y_ptr = x_ptr + d;
368 scalar_t* out_ptr = out + token_idx * d;
369
370 if constexpr (use_vec) {
371 using cuda_t = typename CUDATypeConverter<scalar_t>::Type;
372 using pvec_t = PackedVec<cuda_t, use_256b>;
373
374 const pvec_t* x_vec = reinterpret_cast<const pvec_t*>(x_ptr);
375 const pvec_t* y_vec = reinterpret_cast<const pvec_t*>(y_ptr);
376 pvec_t* out_vec = reinterpret_cast<pvec_t*>(out_ptr);
377 const int num_vecs = d / 2 / pvec_t::NUM_ELTS;
378
379 for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
380 pvec_t x, y;
381 if constexpr (use_256b) {
382 ld256(x, &x_vec[i]);
383 ld256(y, &y_vec[i]);
384 } else {
385 ld128(x, &x_vec[i]);
386 ld128(y, &y_vec[i]);
387 }
388 #pragma unroll
389 for (int j = 0; j < pvec_t::NUM_ELTS; j++) {
390 x.elts[j] = packed_mul(PACKED_ACT_FN(x.elts[j], param), y.elts[j]);
391 }
392 if constexpr (use_256b) {
393 st256(x, &out_vec[i]);
394 } else {
395 st128(x, &out_vec[i]);
396 }
397 }
398 } else {
399 // Scalar fallback for unaligned data or small d
400 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
401 const scalar_t x = VLLM_LDG(&x_ptr[idx]);
402 const scalar_t y = VLLM_LDG(&y_ptr[idx]);
403 out_ptr[idx] = ACT_FN(x, param) * y;
404 }
405 }
406 }
407
408 template <typename T>
409 __device__ __forceinline__ T swigluoai_and_mul(const T& gate, const T& up,
410 float alpha, float limit) {
411 // Clamp gate to (-inf, limit] and up to [-limit, limit]
412 const float g = fminf((float)gate, limit);
413 const float u = fmaxf(fminf((float)up, limit), -limit);
414 // glu = gate * sigmoid(gate * alpha), then return (up + 1) * glu
415 return (T)((u + 1.0f) * g / (1.0f + expf(-g * alpha)));
416 }
417
418 // Interleaved gate/up: input has [gate0, up0, gate1, up1, ...].
419 template <typename scalar_t,
420 scalar_t (*ACT_FN)(const scalar_t&, const scalar_t&, const float,
421 const float)>
422 __global__ void swigluoai_and_mul_kernel(
423 scalar_t* __restrict__ out, // [..., d]
424 const scalar_t* __restrict__ input, // [..., 2 * d] (interleaved)
425 const int d, const float alpha, const float limit) {
426 // For interleaved data: input has 2*d elements per token (gate/up pairs)
427 // output has d elements per token
428 constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
429 constexpr int PAIRS = VEC_SIZE / 2; // Number of gate/up pairs per int4 load
430 const int64_t token_idx = blockIdx.x;
431 const scalar_t* in_ptr = input + token_idx * 2 * d;
432 scalar_t* out_ptr = out + token_idx * d;
433
434 // Check alignment for 128-bit vectorized access on input.
435 // For output we use int2 (64-bit) which has 8-byte alignment requirement.
436 const bool in_aligned = is_16byte_aligned(in_ptr);
437 const bool out_aligned =
438 (reinterpret_cast<uintptr_t>(out_ptr) & 7) == 0; // 8-byte for int2
439
440 if (in_aligned && out_aligned && d >= PAIRS) {
441 // Fast path: vectorized loop
442 // Each int4 load gives VEC_SIZE elements = PAIRS gate/up pairs
443 // Each int2 store writes PAIRS output elements
444 const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
445 int2* out_vec = reinterpret_cast<int2*>(out_ptr);
446 const int num_vecs = d / PAIRS;
447 const int vec_end = num_vecs * PAIRS;
448
449 for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
450 int4 v = VLLM_LDG(&in_vec[i]);
451 int2 r;
452 auto* vp = reinterpret_cast<scalar_t*>(&v);
453 auto* rp = reinterpret_cast<scalar_t*>(&r);
454 #pragma unroll
455 for (int j = 0; j < PAIRS; j++) {
456 rp[j] = ACT_FN(vp[2 * j], vp[2 * j + 1], alpha, limit);
457 }
458 out_vec[i] = r;
459 }
460 // Scalar cleanup for remaining elements
461 for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
462 out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[2 * i]),
463 VLLM_LDG(&in_ptr[2 * i + 1]), alpha, limit);
464 }
465 } else {
466 // Scalar fallback for unaligned data or small d
467 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
468 // gate = x[..., ::2] (even indices)
469 const scalar_t gate = VLLM_LDG(&in_ptr[2 * idx]);
470 // up = x[..., 1::2] (odd indices)
471 const scalar_t up = VLLM_LDG(&in_ptr[2 * idx + 1]);
472 out_ptr[idx] = ACT_FN(gate, up, alpha, limit);
473 }
474 }
475 }
476
477 // SITU (Kimi SituGLU) gated activation. Non-interleaved layout:
478 // input = [gate(d), up(d)] per token.
479 // gate_out = beta * tanh(gate / beta) * sigmoid(gate)
480 // up_out = (linear_beta > 0) ? linear_beta * tanh(up / linear_beta) : up
481 // out = gate_out * up_out
482 __device__ __forceinline__ float situ_tanh(float x) { return tanhf(x); }
483
484 // Kimi-K3 SITU params; baked into the fused LDG kernel to fold at compile time.
485 static constexpr float SITU_BETA = 4.0f;
486 static constexpr float SITU_LINEAR_BETA = 25.0f;
487 __device__ __forceinline__ float situ_activation(float g, float u, float beta,
488 float linear_beta,
489 bool clamp_up, float inv_beta,
490 float inv_linear_beta) {
491 // sigmoid(g) == (1 + tanh(g/2)) / 2.
492 const float gate_out =
493 (0.5f * beta) * situ_tanh(g * inv_beta) * (1.0f + situ_tanh(g * 0.5f));
494 const float up_out =
495 clamp_up ? linear_beta * situ_tanh(u * inv_linear_beta) : u;
496 return gate_out * up_out;
497 }
498
499 template <typename scalar_t>
500 __global__ void situ_and_mul_kernel(
501 scalar_t* __restrict__ out, // [..., d]
502 const scalar_t* __restrict__ input, // [..., 2, d]
503 const int d, const float beta, const float linear_beta) {
504 const int64_t row = blockIdx.x;
505 const scalar_t* gate_ptr = input + row * 2 * d;
506 const scalar_t* up_ptr = gate_ptr + d;
507 scalar_t* out_ptr = out + row * d;
508 const bool clamp_up = linear_beta > 0.0f;
509 const float inv_beta = 1.0f / beta;
510 const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
511 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
512 const float g = (float)VLLM_LDG(&gate_ptr[idx]);
513 const float u = (float)VLLM_LDG(&up_ptr[idx]);
514 out_ptr[idx] = (scalar_t)situ_activation(g, u, beta, linear_beta, clamp_up,
515 inv_beta, inv_linear_beta);
516 }
517 }
518
519 // Match Humming's hardware FP8 conversion. c10::Float8_e4m3fn's software cast
520 // can round ties differently from `cvt.rn.satfinite.e4m3.f32`.
521 __device__ __forceinline__ c10::Float8_e4m3fn cvt_fp8_hw(float x) {
522 #if !defined(USE_ROCM) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
523 return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(x, __NV_SATFINITE, __NV_E4M3),
524 c10::Float8_e4m3fn::from_bits());
525 #else
526 return static_cast<c10::Float8_e4m3fn>(x);
527 #endif
528 }
529
530 template <typename fp8_type>
531 __device__ __forceinline__ fp8_type quant_to_fp8(float val, float inv_scale,
532 float fp8_max) {
533 float x = val * inv_scale;
534 x = fmaxf(-fp8_max, fminf(x, fp8_max));
535 if constexpr (std::is_same_v<fp8_type, c10::Float8_e4m3fn>) {
536 return cvt_fp8_hw(x);
537 } else {
538 return static_cast<fp8_type>(x);
539 }
540 }
541
542 __device__ __forceinline__ float warp_reduce_max(float v) {
543 #pragma unroll
544 for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
545 v = fmaxf(v, VLLM_SHFL_XOR_SYNC(v, offset));
546 }
547 return v;
548 }
549
550 __device__ __forceinline__ float warp_reduce_absmax(float v) {
551 return warp_reduce_max(v);
552 }
553
554 template <typename scalar_t, typename fp8_type, int THREADS, int NUM_STAGES,
555 int GROUP_SIZE, int GRID_DIM, int D>
556 __global__ void situ_and_mul_quant_group_pipelined_kernel(
557 fp8_type* __restrict__ out, float* __restrict__ scale_out,
558 const scalar_t* __restrict__ input, const int64_t num_rows,
559 const int32_t* __restrict__ num_valid_tokens_ptr, const int64_t topk) {
560 const int tid = threadIdx.x;
561 const int warp_id = tid / WARP_SIZE;
562 const int lane_id = tid % WARP_SIZE;
563 static constexpr int NUM_WARPS = THREADS / WARP_SIZE;
564 static constexpr int NUM_GROUPS = D / GROUP_SIZE;
565 const int64_t row_bound =
566 num_valid_tokens_ptr != nullptr
567 ? max((int64_t)0,
568 min((int64_t)(*num_valid_tokens_ptr) * topk, num_rows))
569 : num_rows;
570
571 static_assert(NUM_GROUPS % 4 == 0,
572 "float4 scale fill needs NUM_GROUPS % 4 == 0");
573 static constexpr int NG4 = NUM_GROUPS / 4;
574 const int64_t pad_start4 = row_bound * NG4;
575 const int64_t pad_end4 = num_rows * NG4;
576 float4* scale4 = reinterpret_cast<float4*>(scale_out);
577 constexpr float4 ones{1.0f, 1.0f, 1.0f, 1.0f};
578 for (int64_t i = pad_start4 + (int64_t)blockIdx.x * THREADS + tid;
579 i < pad_end4; i += (int64_t)GRID_DIM * THREADS) {
580 scale4[i] = ones;
581 }
582
583 if constexpr (sizeof(scalar_t) == 2) {
584 static constexpr int LD_ELTS = 16 / sizeof(scalar_t);
585 static constexpr int ELTS_PER_LANE = GROUP_SIZE / WARP_SIZE;
586 static constexpr int STAGE_ELTS = 2 * GROUP_SIZE;
587
588 extern __shared__ __align__(16) unsigned char smem_raw[];
589 scalar_t* warp_smem = reinterpret_cast<scalar_t*>(smem_raw) +
590 (size_t)warp_id * NUM_STAGES * STAGE_ELTS;
591
592 static constexpr float beta = SITU_BETA;
593 static constexpr float linear_beta = SITU_LINEAR_BETA;
594 static constexpr bool clamp_up = linear_beta > 0.0f;
595 static constexpr float inv_beta = 1.0f / beta;
596 static constexpr float inv_linear_beta =
597 clamp_up ? 1.0f / linear_beta : 0.0f;
598 static constexpr float FP8_MAX =
599 std::is_same_v<fp8_type, c10::Float8_e4m3fn> ? 448.0f : 224.0f;
600 static constexpr float inv_fp8_max = 1.0f / FP8_MAX;
601
602 static_assert(
603 NUM_GROUPS % NUM_WARPS == 0,
604 "constexpr num_iters requires groups evenly split across warps");
605 static constexpr int num_iters = NUM_GROUPS / NUM_WARPS;
606
607 const bool up_half = lane_id >= WARP_SIZE / 2;
608 const int lane_l = up_half ? lane_id - WARP_SIZE / 2 : lane_id;
609 const int lane_src_off = (up_half ? D : 0) + lane_l * LD_ELTS;
610 const int lane_dst_off = (up_half ? GROUP_SIZE : 0) + lane_l * LD_ELTS;
611 static constexpr int warp_stride = NUM_WARPS * GROUP_SIZE;
612
613 const scalar_t* src_ptr = input + warp_id * GROUP_SIZE + lane_src_off;
614 static constexpr int src_row_stride = GRID_DIM * 2 * D;
615 const scalar_t* row_src = src_ptr + blockIdx.x * 2 * D;
616
617 uint32_t* out_ptr =
618 reinterpret_cast<uint32_t*>(out + warp_id * GROUP_SIZE) + lane_id;
619 static constexpr int out_row_stride = GRID_DIM * D / 4;
620 uint32_t* row_out = out_ptr + blockIdx.x * (D / 4);
621 float* scale_ptr = scale_out + warp_id;
622 static constexpr int scale_row_stride = GRID_DIM * NUM_GROUPS;
623 float* row_scale = scale_ptr + blockIdx.x * NUM_GROUPS;
624
625 #pragma unroll 1
626 for (int64_t row = blockIdx.x; row < row_bound; row += GRID_DIM,
627 row_src += src_row_stride, row_out += out_row_stride,
628 row_scale += scale_row_stride) {
629 const scalar_t* src = row_src;
630
631 auto issue_load = [&](int it, int slot) {
632 if (it < num_iters) {
633 cuda_async::cp_async_shared_global_16_cg(
634 warp_smem + (size_t)slot * STAGE_ELTS + lane_dst_off, src);
635 src += warp_stride;
636 }
637 cuda_async::cp_async_commit_group();
638 };
639
640 int load_slot = 0;
641 auto bump = [](int s) { return s + 1 == NUM_STAGES ? 0 : s + 1; };
642 #pragma unroll
643 for (int s = 0; s < NUM_STAGES - 1; s++) {
644 issue_load(s, load_slot);
645 load_slot = bump(load_slot);
646 }
647
648 uint32_t* out_st = row_out;
649 float* scale_st = row_scale;
650
651 int comp_slot = 0;
652 for (int it = 0; it < num_iters; it++) {
653 issue_load(it + NUM_STAGES - 1, load_slot);
654 load_slot = bump(load_slot);
655 cuda_async::cp_async_wait_group<NUM_STAGES - 1>();
656
657 const scalar_t* stage = warp_smem + (size_t)comp_slot * STAGE_ELTS;
658 comp_slot = bump(comp_slot);
659 static_assert(ELTS_PER_LANE == 4,
660 "expects GROUP_SIZE == 4 * WARP_SIZE");
661 union V {
662 float2 f2;
663 scalar_t s[ELTS_PER_LANE];
664 };
665 const float2* gate2 = reinterpret_cast<const float2*>(stage);
666 const float2* up2 = reinterpret_cast<const float2*>(stage + GROUP_SIZE);
667 const V gv{gate2[lane_id]};
668 const V uv{up2[lane_id]};
669 const scalar_t* gs = gv.s;
670 const scalar_t* us = uv.s;
671
672 float acts[ELTS_PER_LANE];
673 float thread_max = 0.0f;
674 #pragma unroll
675 for (int e = 0; e < ELTS_PER_LANE; e++) {
676 acts[e] = (float)(scalar_t)situ_activation(
677 (float)gs[e], (float)us[e], beta, linear_beta, clamp_up, inv_beta,
678 inv_linear_beta);
679 thread_max = fmaxf(thread_max, fabsf(acts[e]));
680 }
681 const float absmax = fmaxf(warp_reduce_absmax(thread_max), 1e-30f);
682 const float scale = absmax * inv_fp8_max;
683 if (lane_id == 0) *scale_st = scale;
684 scale_st += NUM_WARPS;
685 const float inv_scale = __fdividef(1.0f, scale);
686
687 union O {
688 uint32_t u;
689 fp8_type f[ELTS_PER_LANE];
690 };
691 O o;
692 #pragma unroll
693 for (int e = 0; e < ELTS_PER_LANE; e++) {
694 o.f[e] = quant_to_fp8<fp8_type>(acts[e], inv_scale, FP8_MAX);
695 }
696 out_st[0] = o.u;
697 out_st += NUM_WARPS * GROUP_SIZE / 4;
698 }
699 cuda_async::cp_async_wait_group<0>();
700 }
701 }
702 }
703
704 // Scalar fallback for the group path (odd d or the fp32 dispatch branch).
705 template <typename scalar_t, typename fp8_type, int GROUP_SIZE>
706 __global__ void situ_and_mul_quant_group_scalar_kernel(
707 fp8_type* __restrict__ out, // [num_tokens, d]
708 float* __restrict__ scale_out, // [num_tokens, num_groups]
709 const scalar_t* __restrict__ input, // [num_tokens, 2, d]
710 const int d, const int num_groups, const float beta,
711 const float linear_beta, const int64_t num_rows,
712 const int32_t* __restrict__ num_valid_tokens_ptr, const int64_t topk) {
713 static constexpr int ELTS_PER_LANE = GROUP_SIZE / WARP_SIZE;
714 const int tid = threadIdx.x;
715 const int warp_id = tid / WARP_SIZE;
716 const int lane_id = tid % WARP_SIZE;
717 const int num_warps = blockDim.x / WARP_SIZE;
718 const int64_t row_bound =
719 num_valid_tokens_ptr != nullptr
720 ? max((int64_t)0,
721 min((int64_t)(*num_valid_tokens_ptr) * topk, num_rows))
722 : num_rows;
723
724 const bool clamp_up = linear_beta > 0.0f;
725 const float inv_beta = __fdividef(1.0f, beta);
726 const float inv_linear_beta = clamp_up ? __fdividef(1.0f, linear_beta) : 0.0f;
727 static constexpr float FP8_MAX =
728 std::is_same_v<fp8_type, c10::Float8_e4m3fn> ? 448.0f : 224.0f;
729 static constexpr float inv_fp8_max = 1.0f / FP8_MAX;
730
731 for (int64_t row = blockIdx.x; row < row_bound; row += gridDim.x) {
732 const scalar_t* gate_ptr = input + row * 2 * (int64_t)d;
733 const scalar_t* up_ptr = gate_ptr + d;
734 fp8_type* out_ptr = out + row * (int64_t)d;
735 float* scale_row = scale_out + row * (int64_t)num_groups;
736
737 for (int g = warp_id; g < num_groups; g += num_warps) {
738 const int base = g * GROUP_SIZE;
739 float acts[ELTS_PER_LANE];
740 float thread_max = 0.0f;
741 #pragma unroll
742 for (int e = 0; e < ELTS_PER_LANE; e++) {
743 const int idx = base + e * WARP_SIZE + lane_id;
744 const float gv = (float)VLLM_LDG(&gate_ptr[idx]);
745 const float uv = (float)VLLM_LDG(&up_ptr[idx]);
746 acts[e] = (float)(scalar_t)situ_activation(
747 gv, uv, beta, linear_beta, clamp_up, inv_beta, inv_linear_beta);
748 thread_max = fmaxf(thread_max, fabsf(acts[e]));
749 }
750 const float absmax = fmaxf(warp_reduce_max(thread_max), 1e-30f);
751 const float scale = absmax * inv_fp8_max;
752 if (lane_id == 0) scale_row[g] = scale;
753 const float inv_scale = __fdividef(1.0f, scale);
754 #pragma unroll
755 for (int e = 0; e < ELTS_PER_LANE; e++) {
756 const int idx = base + e * WARP_SIZE + lane_id;
757 out_ptr[idx] = quant_to_fp8<fp8_type>(acts[e], inv_scale, FP8_MAX);
758 }
759 }
760 }
761
762 // Fill skipped padding-row scales with 1 so they don't feed NaN into w2.
763 const int64_t pad_start = row_bound * (int64_t)num_groups;
764 const int64_t pad_end = num_rows * (int64_t)num_groups;
765 for (int64_t i = pad_start + (int64_t)blockIdx.x * blockDim.x + tid;
766 i < pad_end; i += (int64_t)gridDim.x * blockDim.x) {
767 scale_out[i] = 1.0f;
768 }
769 }
770
771 constexpr int kMaxMaskedTokenBlocks = 32;
772
773 template <bool BATCHED_EXPERTS>
774 __device__ __forceinline__ bool get_masked_row_range(
775 const int* __restrict__ valid_token_counts, const int max_num_tokens,
776 int64_t& first_row, int64_t& end_row, int64_t& row_stride) {
777 if constexpr (BATCHED_EXPERTS) {
778 // [E, T, *]: z lanes grid-stride over one expert's valid token prefix.
779 const int expert = blockIdx.y;
780 const int num_tokens =
781 max(0, min(valid_token_counts[expert], max_num_tokens));
782 const int token_block = blockIdx.z;
783 if (token_block >= num_tokens) {
784 return false;
785 }
786 const int64_t expert_first_row =
787 static_cast<int64_t>(expert) * max_num_tokens;
788 first_row = expert_first_row + token_block;
789 end_row = expert_first_row + num_tokens;
790 row_stride = gridDim.z;
791 } else {
792 // [T, *]: y lanes grid-stride over the single valid token prefix.
793 const int num_tokens = max(0, min(valid_token_counts[0], max_num_tokens));
794 const int token_block = blockIdx.y;
795 if (token_block >= num_tokens) {
796 return false;
797 }
798 first_row = token_block;
799 end_row = num_tokens;
800 row_stride = gridDim.y;
801 }
802 return first_row < end_row;
803 }
804
805 __device__ __forceinline__ int get_masked_feature_index() {
806 return blockIdx.x * blockDim.x + threadIdx.x;
807 }
808
809 template <typename scalar_t, bool BATCHED_EXPERTS>
810 __global__ void masked_situ_and_mul_kernel(
811 scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
812 const int* __restrict__ valid_token_counts, const int max_num_tokens,
813 const int d, const float beta, const float linear_beta) {
814 int64_t first_row, end_row, row_stride;
815 const int idx = get_masked_feature_index();
816 if (idx >= d ||
817 !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
818 first_row, end_row, row_stride)) {
819 return;
820 }
821
822 const bool clamp_up = linear_beta > 0.0f;
823 const float inv_beta = 1.0f / beta;
824 const float inv_linear_beta = clamp_up ? 1.0f / linear_beta : 0.0f;
825 for (int64_t row = first_row; row < end_row; row += row_stride) {
826 const scalar_t* gate_ptr = input + row * 2 * d;
827 const scalar_t* up_ptr = gate_ptr + d;
828 scalar_t* out_ptr = out + row * d;
829 const float g = (float)VLLM_LDG(&gate_ptr[idx]);
830 const float u = (float)VLLM_LDG(&up_ptr[idx]);
831 out_ptr[idx] = (scalar_t)situ_activation(g, u, beta, linear_beta, clamp_up,
832 inv_beta, inv_linear_beta);
833 }
834 }
835
836 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float),
837 bool act_first, bool HAS_CLAMP, bool BATCHED_EXPERTS>
838 __global__ void masked_act_and_mul_kernel(
839 scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
840 const int* __restrict__ valid_token_counts, const int max_num_tokens,
841 const int d, const float limit, const float alpha, const float beta) {
842 int64_t first_row, end_row, row_stride;
843 const int idx = get_masked_feature_index();
844 if (idx >= d ||
845 !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
846 first_row, end_row, row_stride)) {
847 return;
848 }
849
850 for (int64_t row = first_row; row < end_row; row += row_stride) {
851 const scalar_t* x_ptr = input + row * 2 * d;
852 const scalar_t* y_ptr = x_ptr + d;
853 scalar_t* out_ptr = out + row * d;
854 const scalar_t x = VLLM_LDG(&x_ptr[idx]);
855 const scalar_t y = VLLM_LDG(&y_ptr[idx]);
856 out_ptr[idx] = compute<scalar_t, ACT_FN, act_first, HAS_CLAMP>(x, y, limit,
857 alpha, beta);
858 }
859 }
860
861 template <typename scalar_t, bool BATCHED_EXPERTS>
862 __global__ void masked_swigluoai_and_mul_kernel(
863 scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
864 const int* __restrict__ valid_token_counts, const int max_num_tokens,
865 const int d, const float alpha, const float limit) {
866 int64_t first_row, end_row, row_stride;
867 const int idx = get_masked_feature_index();
868 if (idx >= d ||
869 !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
870 first_row, end_row, row_stride)) {
871 return;
872 }
873
874 for (int64_t row = first_row; row < end_row; row += row_stride) {
875 const scalar_t* in_ptr = input + row * 2 * d;
876 scalar_t* out_ptr = out + row * d;
877 out_ptr[idx] =
878 swigluoai_and_mul(in_ptr[2 * idx], in_ptr[2 * idx + 1], alpha, limit);
879 }
880 }
881
882 template <typename scalar_t, bool BATCHED_EXPERTS>
883 __global__ void masked_swiglustep_and_mul_kernel(
884 scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
885 const int* __restrict__ valid_token_counts, const int max_num_tokens,
886 const int d, const float limit) {
887 int64_t first_row, end_row, row_stride;
888 const int idx = get_masked_feature_index();
889 if (idx >= d ||
890 !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
891 first_row, end_row, row_stride)) {
892 return;
893 }
894
895 for (int64_t row = first_row; row < end_row; row += row_stride) {
896 const scalar_t* gate_ptr = input + row * 2 * d;
897 const scalar_t* up_ptr = gate_ptr + d;
898 scalar_t* out_ptr = out + row * d;
899 const float gate = (float)VLLM_LDG(&gate_ptr[idx]);
900 const float up = (float)VLLM_LDG(&up_ptr[idx]);
901 const float gate_silu = gate / (1.0f + expf(-gate));
902 const float gate_clamped = fminf(gate_silu, limit);
903 const float up_clamped = fmaxf(fminf(up, limit), -limit);
904 out_ptr[idx] = (scalar_t)(gate_clamped * up_clamped);
905 }
906 }
907
908 template <typename T>
909 __device__ __forceinline__ T silu_no_mul_kernel(const T& x) {
910 return silu_kernel(x, 1.0f);
911 }
912
913 template <typename T>
914 __device__ __forceinline__ T gelu_no_mul_kernel(const T& x) {
915 return gelu_kernel(x, 1.0f);
916 }
917
918 template <typename T>
919 __device__ __forceinline__ T gelu_tanh_no_mul_kernel(const T& x) {
920 return gelu_tanh_kernel(x, 1.0f);
921 }
922
923 template <typename T>
924 __device__ __forceinline__ T relu2_no_mul_kernel(const T& x) {
925 const float value = fmaxf((float)x, 0.0f);
926 return (T)(value * value);
927 }
928
929 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
930 bool BATCHED_EXPERTS>
931 __global__ void masked_activation_kernel(
932 scalar_t* __restrict__ out, const scalar_t* __restrict__ input,
933 const int* __restrict__ valid_token_counts, const int max_num_tokens,
934 const int d) {
935 int64_t first_row, end_row, row_stride;
936 const int idx = get_masked_feature_index();
937 if (idx >= d ||
938 !get_masked_row_range<BATCHED_EXPERTS>(valid_token_counts, max_num_tokens,
939 first_row, end_row, row_stride)) {
940 return;
941 }
942
943 for (int64_t row = first_row; row < end_row; row += row_stride) {
944 const int64_t offset = row * d + idx;
945 out[offset] = ACT_FN(VLLM_LDG(&input[offset]));
946 }
947 }
948
949 } // namespace vllm
950
951 #define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
952 auto dtype = input.scalar_type(); \
953 int d = input.size(-1) / 2; \
954 int64_t num_tokens = input.numel() / input.size(-1); \
955 if (num_tokens == 0) { \
956 return; \
957 } \
958 dim3 grid(num_tokens); \
959 int cc_major = get_device_prop()->major; \
960 int support_vec = \
961 (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
962 ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
963 : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
964 int vec_size = support_vec / input.element_size(); \
965 const bool use_vec = (d % vec_size == 0); \
966 const torch::stable::accelerator::DeviceGuard device_guard( \
967 input.get_device_index()); \
968 const cudaStream_t stream = get_current_cuda_stream(); \
969 if (use_vec) { \
970 dim3 block(std::min(d / vec_size, 1024)); \
971 if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
972 VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
973 dtype, "act_and_mul_kernel_with_param", [&] { \
974 vllm::act_and_mul_kernel_with_param< \
975 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
976 KERNEL<scalar_t>, \
977 PACKED_KERNEL< \
978 typename vllm::PackedTypeConverter<scalar_t>::Type>, \
979 true, true><<<grid, block, 0, stream>>>( \
980 out.mutable_data_ptr<scalar_t>(), \
981 input.const_data_ptr<scalar_t>(), d, PARAM); \
982 }); \
983 } else { \
984 VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
985 dtype, "act_and_mul_kernel_with_param", [&] { \
986 vllm::act_and_mul_kernel_with_param< \
987 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
988 KERNEL<scalar_t>, \
989 PACKED_KERNEL< \
990 typename vllm::PackedTypeConverter<scalar_t>::Type>, \
991 true, false><<<grid, block, 0, stream>>>( \
992 out.mutable_data_ptr<scalar_t>(), \
993 input.const_data_ptr<scalar_t>(), d, PARAM); \
994 }); \
995 } \
996 } else { \
997 dim3 block(std::min(d, 1024)); \
998 VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
999 dtype, "act_and_mul_kernel_with_param", [&] { \
1000 vllm::act_and_mul_kernel_with_param< \
1001 scalar_t, typename vllm::PackedTypeConverter<scalar_t>::Type, \
1002 KERNEL<scalar_t>, \
1003 PACKED_KERNEL< \
1004 typename vllm::PackedTypeConverter<scalar_t>::Type>, \
1005 false><<<grid, block, 0, stream>>>( \
1006 out.mutable_data_ptr<scalar_t>(), \
1007 input.const_data_ptr<scalar_t>(), d, PARAM); \
1008 }); \
1009 }
1010
1011 #define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \
1012 int d = input.size(-1) / 2; \
1013 int64_t num_tokens = input.numel() / input.size(-1); \
1014 dim3 grid(num_tokens); \
1015 dim3 block(std::min(d, 1024)); \
1016 const torch::stable::accelerator::DeviceGuard device_guard( \
1017 input.get_device_index()); \
1018 const cudaStream_t stream = get_current_cuda_stream(); \
1019 VLLM_STABLE_DISPATCH_FLOATING_TYPES( \
1020 input.scalar_type(), "clamp_swiglu_kernel_with_params", [&] { \
1021 vllm::swigluoai_and_mul_kernel<scalar_t, KERNEL<scalar_t>> \
1022 <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
1023 input.const_data_ptr<scalar_t>(), d, \
1024 ALPHA, LIMIT); \
1025 });
1026
1027 void fatrelu_and_mul(torch::stable::Tensor& out, // [..., d],
1028 torch::stable::Tensor& input, // [..., 2 * d]
1029 double threshold) {
1030 LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
1031 vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
1032 }
1033 void swigluoai_and_mul(torch::stable::Tensor& out, // [..., d]
1034 torch::stable::Tensor& input, // [..., 2 * d]
1035 double alpha, double limit) {
1036 LAUNCH_SIGLUOAI_AND_MUL(vllm::swigluoai_and_mul, alpha, limit);
1037 }
1038
1039 // Kimi SITU gated activation. `linear_beta <= 0` means "unset" (up passed
1040 // through), matching SituAndMul(linear_beta=None) on the Python side.
1041 void situ_and_mul(torch::stable::Tensor& out, // [..., d]
1042 torch::stable::Tensor& input, // [..., 2 * d]
1043 double beta, double linear_beta) {
1044 int d = input.size(-1) / 2;
1045 int64_t num_tokens = input.numel() / input.size(-1);
1046 if (num_tokens == 0) {
1047 return;
1048 }
1049 dim3 grid(num_tokens);
1050 dim3 block(std::min(d, 1024));
1051 const torch::stable::accelerator::DeviceGuard device_guard(
1052 input.get_device_index());
1053 const cudaStream_t stream = get_current_cuda_stream();
1054 VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1055 input.scalar_type(), "situ_and_mul_kernel", [&] {
1056 vllm::situ_and_mul_kernel<scalar_t><<<grid, block, 0, stream>>>(
1057 out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(),
1058 d, (float)beta, (float)linear_beta);
1059 });
1060 }
1061
1062 void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d]
1063 torch::stable::Tensor& input, // [E, T, 2 * d]
1064 const torch::stable::Tensor& expert_num_tokens,
1065 double beta, double linear_beta) {
1066 int num_experts = input.size(0);
1067 int max_num_tokens = input.size(1);
1068 int d = input.size(2) / 2;
1069 if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
1070 return;
1071 }
1072 constexpr int block_size = 256;
1073 dim3 grid((d + block_size - 1) / block_size, num_experts);
1074 dim3 block(block_size);
1075 const torch::stable::accelerator::DeviceGuard device_guard(
1076 input.get_device_index());
1077 const cudaStream_t stream = get_current_cuda_stream();
1078 VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1079 input.scalar_type(), "masked_situ_and_mul_kernel", [&] {
1080 vllm::masked_situ_and_mul_kernel<scalar_t, true>
1081 <<<grid, block, 0, stream>>>(
1082 out.mutable_data_ptr<scalar_t>(),
1083 input.const_data_ptr<scalar_t>(),
1084 expert_num_tokens.const_data_ptr<int>(), max_num_tokens, d,
1085 (float)beta, (float)linear_beta);
1086 });
1087 }
1088
1089 #define LAUNCH_MASKED_ACT_AND_MUL(KERNEL, ACT_FIRST, HAS_CLAMP, LIMIT, ALPHA, \
1090 BETA) \
1091 vllm::masked_act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST, \
1092 HAS_CLAMP, BATCHED_EXPERTS> \
1093 <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
1094 input.const_data_ptr<scalar_t>(), \
1095 valid_token_counts.const_data_ptr<int>(), \
1096 max_num_tokens, d, LIMIT, ALPHA, BETA)
1097
1098 #define LAUNCH_MASKED_ACTIVATION(KERNEL) \
1099 vllm::masked_activation_kernel<scalar_t, KERNEL<scalar_t>, BATCHED_EXPERTS> \
1100 <<<grid, block, 0, stream>>>( \
1101 out.mutable_data_ptr<scalar_t>(), input.const_data_ptr<scalar_t>(), \
1102 valid_token_counts.const_data_ptr<int>(), max_num_tokens, d)
1103
1104 void masked_moe_activation(
1105 torch::stable::Tensor& out, // [T, d] or [E, T, d]
1106 torch::stable::Tensor& input, // [T, d] or [E, T, d], gated has 2 * d
1107 const torch::stable::Tensor& valid_token_counts,
1108 const std::string& activation, double clamp_limit, double alpha,
1109 double beta, double situ_beta, double situ_linear_beta) {
1110 const bool batched_experts = input.dim() == 3;
1111 const int num_experts = batched_experts ? input.size(0) : 1;
1112 const int max_num_tokens = batched_experts ? input.size(1) : input.size(0);
1113 const int d = out.size(-1);
1114 if (num_experts == 0 || max_num_tokens == 0 || d == 0) {
1115 return;
1116 }
1117
1118 constexpr int block_size = 256;
1119 const int feature_blocks = (d + block_size - 1) / block_size;
1120 const int token_blocks =
1121 std::min(max_num_tokens, vllm::kMaxMaskedTokenBlocks);
1122 // Batched grid: (feature tile, expert, token lane); flat grid: (feature
1123 // tile, token lane). Token lanes grid-stride each valid prefix.
1124 dim3 grid = batched_experts ? dim3(feature_blocks, num_experts, token_blocks)
1125 : dim3(feature_blocks, token_blocks);
1126 dim3 block(block_size);
1127 const torch::stable::accelerator::DeviceGuard device_guard(
1128 input.get_device_index());
1129 const cudaStream_t stream = get_current_cuda_stream();
1130 auto launch = [&]<bool BATCHED_EXPERTS>() {
1131 VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1132 input.scalar_type(), "masked_moe_activation", [&] {
1133 if (activation == "silu") {
1134 LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, false, 0.0f,
1135 1.0f, 0.0f);
1136 } else if (activation == "silu_with_clamp") {
1137 LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, true,
1138 (float)clamp_limit, (float)alpha,
1139 (float)beta);
1140 } else if (activation == "gelu") {
1141 LAUNCH_MASKED_ACT_AND_MUL(vllm::gelu_kernel, true, false, 0.0f,
1142 1.0f, 0.0f);
1143 } else if (activation == "gelu_tanh") {
1144 LAUNCH_MASKED_ACT_AND_MUL(vllm::gelu_tanh_kernel, true, false, 0.0f,
1145 1.0f, 0.0f);
1146 } else if (activation == "situ") {
1147 vllm::masked_situ_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
1148 <<<grid, block, 0, stream>>>(
1149 out.mutable_data_ptr<scalar_t>(),
1150 input.const_data_ptr<scalar_t>(),
1151 valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
1152 (float)situ_beta, (float)situ_linear_beta);
1153 } else if (activation == "swigluoai") {
1154 vllm::masked_swigluoai_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
1155 <<<grid, block, 0, stream>>>(
1156 out.mutable_data_ptr<scalar_t>(),
1157 input.const_data_ptr<scalar_t>(),
1158 valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
1159 (float)alpha, (float)clamp_limit);
1160 } else if (activation == "swigluoai_uninterleave") {
1161 LAUNCH_MASKED_ACT_AND_MUL(vllm::silu_kernel, true, true,
1162 (float)clamp_limit, (float)alpha,
1163 (float)beta);
1164 } else if (activation == "swiglustep") {
1165 vllm::masked_swiglustep_and_mul_kernel<scalar_t, BATCHED_EXPERTS>
1166 <<<grid, block, 0, stream>>>(
1167 out.mutable_data_ptr<scalar_t>(),
1168 input.const_data_ptr<scalar_t>(),
1169 valid_token_counts.const_data_ptr<int>(), max_num_tokens, d,
1170 (float)clamp_limit);
1171 } else if (activation == "silu_no_mul") {
1172 LAUNCH_MASKED_ACTIVATION(vllm::silu_no_mul_kernel);
1173 } else if (activation == "gelu_no_mul") {
1174 LAUNCH_MASKED_ACTIVATION(vllm::gelu_no_mul_kernel);
1175 } else if (activation == "gelu_tanh_no_mul") {
1176 LAUNCH_MASKED_ACTIVATION(vllm::gelu_tanh_no_mul_kernel);
1177 } else if (activation == "relu2_no_mul") {
1178 LAUNCH_MASKED_ACTIVATION(vllm::relu2_no_mul_kernel);
1179 } else {
1180 STD_TORCH_CHECK(false,
1181 "Unsupported masked MoE activation: ", activation);
1182 }
1183 });
1184 };
1185 if (batched_experts) {
1186 launch.template operator()<true>();
1187 } else {
1188 launch.template operator()<false>();
1189 }
1190 }
1191
1192 #undef LAUNCH_MASKED_ACT_AND_MUL
1193 #undef LAUNCH_MASKED_ACTIVATION
1194 // Fused SITU activation and block-FP8 quantization for Humming w2.
1195 // `num_valid_tokens` is the int32 DeepEP token count (psum[-1]); rows = it *
1196 // `topk`, so padding is excluded on-device and skipped rows receive scale 1.
1197 void situ_and_mul_quant(torch::stable::Tensor& out, // [..., d] (fp8)
1198 torch::stable::Tensor& scale, // [..., 1 or d/128]
1199 torch::stable::Tensor& input, // [..., 2 * d]
1200 double beta, double linear_beta, int64_t group_size,
1201 std::optional<torch::stable::Tensor> num_valid_tokens,
1202 int64_t topk) {
1203 STD_TORCH_CHECK(
1204 out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fn ||
1205 out.scalar_type() == torch::headeronly::ScalarType::Float8_e4m3fnuz,
1206 "situ_and_mul_quant output must be FP8 e4m3");
1207 STD_TORCH_CHECK(
1208 input.scalar_type() == torch::headeronly::ScalarType::Half ||
1209 input.scalar_type() == torch::headeronly::ScalarType::BFloat16,
1210 "situ_and_mul_quant input must be FP16 or BF16");
1211 STD_TORCH_CHECK(scale.scalar_type() == torch::headeronly::ScalarType::Float,
1212 "situ_and_mul_quant scale must be float32");
1213 STD_TORCH_CHECK(input.is_cuda() && out.is_cuda() && scale.is_cuda() &&
1214 out.get_device_index() == input.get_device_index() &&
1215 scale.get_device_index() == input.get_device_index(),
1216 "situ_and_mul_quant: input, out and scale must be CUDA "
1217 "tensors on the same device");
1218 int d = input.size(-1) / 2;
1219 int64_t num_tokens = input.numel() / input.size(-1);
1220 if (num_tokens == 0) {
1221 return;
1222 }
1223 STD_TORCH_CHECK(out.size(-1) == d && out.numel() == num_tokens * (int64_t)d,
1224 "situ_and_mul_quant: out shape must be [num_tokens, d]");
1225 const int32_t* num_valid_tokens_ptr = nullptr;
1226 if (num_valid_tokens.has_value()) {
1227 STD_TORCH_CHECK(
1228 num_valid_tokens->is_cuda() &&
1229 num_valid_tokens->get_device_index() == input.get_device_index(),
1230 "situ_and_mul_quant: num_valid_tokens must be a CUDA tensor on the "
1231 "same device as input");
1232 STD_TORCH_CHECK(
1233 num_valid_tokens->scalar_type() == torch::headeronly::ScalarType::Int,
1234 "situ_and_mul_quant: num_valid_tokens must be int32 (psum count)");
1235 num_valid_tokens_ptr = num_valid_tokens->const_data_ptr<int32_t>();
1236 }
1237 const torch::stable::accelerator::DeviceGuard device_guard(
1238 input.get_device_index());
1239 const cudaStream_t stream = get_current_cuda_stream();
1240 static constexpr int BLOCKS_PER_SM = 8; // matches kernel __launch_bounds__
1241 static constexpr int SM_COUNT = 132; // H200 (GH100, 132 SMs)
1242 static constexpr int GRID_DIM = SM_COUNT * BLOCKS_PER_SM;
1243
1244 STD_TORCH_CHECK(group_size == 128,
1245 "situ_and_mul_quant: only group_size 128 (block-FP8) "
1246 "supported, got ",
1247 group_size);
1248 STD_TORCH_CHECK(d % group_size == 0, "situ_and_mul_quant: d (", d,
1249 ") must be divisible by group_size ", group_size);
1250 {
1251 const int num_groups = d / (int)group_size;
1252 STD_TORCH_CHECK(scale.size(-1) == num_groups &&
1253 scale.numel() == num_tokens * (int64_t)num_groups,
1254 "situ_and_mul_quant: scale shape must be "
1255 "[num_tokens, d/group_size]");
1256 dim3 grid(GRID_DIM);
1257 VLLM_STABLE_DISPATCH_FLOATING_TYPES(
1258 input.scalar_type(), "situ_and_mul_quant_group_kernel", [&] {
1259 VLLM_STABLE_DISPATCH_FP8_TYPES(
1260 out.scalar_type(), "situ_and_mul_quant_group_kernel_fp8", [&] {
1261 #ifndef USE_ROCM
1262 // The pipelined kernel's float2-per-lane geometry assumes a
1263 // 32-lane warp (GROUP_SIZE == 4 * WARP_SIZE); on HIP (64-lane)
1264 // fall back to the WARP_SIZE-generic scalar kernel.
1265 constexpr int THREADS = 256;
1266 constexpr int SITU_D = 3072; // Kimi-K3 fused w2 input dim
1267 if constexpr (sizeof(scalar_t) == 2) {
1268 if (d == SITU_D && (float)beta == vllm::SITU_BETA &&
1269 (float)linear_beta == vllm::SITU_LINEAR_BETA) {
1270 constexpr int D = SITU_D;
1271 constexpr int GROUP_STAGES = 4;
1272 constexpr int NUM_WARPS = THREADS / 32;
1273 constexpr int STAGE_ELTS = 2 * 128;
1274 dim3 block(THREADS);
1275 size_t smem_bytes = (size_t)NUM_WARPS * GROUP_STAGES *
1276 STAGE_ELTS * sizeof(scalar_t);
1277 vllm::situ_and_mul_quant_group_pipelined_kernel<
1278 scalar_t, fp8_t, THREADS, GROUP_STAGES, 128, GRID_DIM,
1279 D><<<grid, block, smem_bytes, stream>>>(
1280 out.mutable_data_ptr<fp8_t>(),
1281 scale.mutable_data_ptr<float>(),
1282 input.const_data_ptr<scalar_t>(), num_tokens,
1283 num_valid_tokens_ptr, topk);
1284 return;
1285 }
1286 }
1287 #endif
1288 const int num_warps = std::min(num_groups, 1024 / WARP_SIZE);
1289 dim3 block(num_warps * WARP_SIZE);
1290 vllm::situ_and_mul_quant_group_scalar_kernel<scalar_t, fp8_t,
1291 128>
1292 <<<grid, block, 0, stream>>>(
1293 out.mutable_data_ptr<fp8_t>(),
1294 scale.mutable_data_ptr<float>(),
1295 input.const_data_ptr<scalar_t>(), d, num_groups,
1296 (float)beta, (float)linear_beta, num_tokens,
1297 num_valid_tokens_ptr, topk);
1298 });
1299 });
1300 }
1301 }
1302 namespace vllm {
1303
1304 // Element-wise activation kernel template.
1305 template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), bool use_vec,
1306 bool use_256b = false>
1307 __global__ void activation_kernel(
1308 scalar_t* __restrict__ out, // [..., d]
1309 const scalar_t* __restrict__ input, // [..., d]
1310 const int d) {
1311 const int64_t token_idx = blockIdx.x;
1312 const scalar_t* in_ptr = input + token_idx * d;
1313 scalar_t* out_ptr = out + token_idx * d;
1314
1315 if constexpr (use_vec) {
1316 // Fast path: 128-bit/256-bit vectorized loop
1317 using vec_t = typename VecTraits<use_256b>::vec_t;
1318 constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
1319 constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t);
1320 const vec_t* in_vec = reinterpret_cast<const vec_t*>(in_ptr);
1321 vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
1322 const int num_vecs = d / VEC_SIZE;
1323
1324 for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
1325 vec_t v;
1326 if constexpr (use_256b) {
1327 ld256(v, &in_vec[i]);
1328 } else {
1329 v = VLLM_LDG(&in_vec[i]);
1330 }
1331 auto* vp = reinterpret_cast<scalar_t*>(&v);
1332 #pragma unroll
1333 for (int j = 0; j < VEC_SIZE; j++) {
1334 vp[j] = ACT_FN(vp[j]);
1335 }
1336 if constexpr (use_256b) {
1337 st256(v, &out_vec[i]);
1338 } else {
1339 out_vec[i] = v;
1340 }
1341 }
1342 } else {
1343 // Scalar fallback for unaligned data or small d
1344 for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
1345 const scalar_t x = VLLM_LDG(&in_ptr[idx]);
1346 out_ptr[idx] = ACT_FN(x);
1347 }
1348 }
1349 }
1350
1351 } // namespace vllm
1352
1353 // Launch element-wise activation kernel.
1354 #define LAUNCH_ACTIVATION_KERNEL(KERNEL) \
1355 auto dtype = input.scalar_type(); \
1356 int d = input.size(-1); \
1357 int64_t num_tokens = input.numel() / input.size(-1); \
1358 if (num_tokens == 0) { \
1359 return; \
1360 } \
1361 dim3 grid(num_tokens); \
1362 int cc_major = get_device_prop()->major; \
1363 int support_vec = \
1364 (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) \
1365 ? vllm::VecTraits<true>::ARCH_MAX_VEC_SIZE \
1366 : vllm::VecTraits<false>::ARCH_MAX_VEC_SIZE; \
1367 int vec_size = support_vec / input.element_size(); \
1368 const bool use_vec = (d % vec_size == 0); \
1369 const torch::stable::accelerator::DeviceGuard device_guard( \
1370 input.get_device_index()); \
1371 const cudaStream_t stream = get_current_cuda_stream(); \
1372 if (use_vec) { \
1373 dim3 block(std::min(d / vec_size, 1024)); \
1374 if (CUDA_VERSION >= 12090 && cc_major >= 10 && num_tokens > 128) { \
1375 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
1376 vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true> \
1377 <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
1378 input.const_data_ptr<scalar_t>(), d); \
1379 }); \
1380 } else { \
1381 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
1382 vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
1383 <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
1384 input.const_data_ptr<scalar_t>(), d); \
1385 }); \
1386 } \
1387 } else { \
1388 dim3 block(std::min(d, 1024)); \
1389 VLLM_STABLE_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
1390 vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false> \
1391 <<<grid, block, 0, stream>>>(out.mutable_data_ptr<scalar_t>(), \
1392 input.const_data_ptr<scalar_t>(), d); \
1393 }); \
1394 }
1395
1396 namespace vllm {
1397
1398 template <typename T>
1399 __device__ __forceinline__ T gelu_new_kernel(const T& x) {
1400 const float x3 = (float)(x * x * x);
1401 const T t = (T)tanhf((T)(0.79788456f * (float)(x + (T)(0.044715f * x3))));
1402 return ((T)0.5) * x * (((T)1.0) + t);
1403 }
1404
1405 template <typename T>
1406 __device__ __forceinline__ T gelu_fast_kernel(const T& x) {
1407 const float f = (float)x;
1408 const T t =
1409 (T)tanhf(((T)(f * 0.79788456f)) * (((T)1.0) + (T)(0.044715f * f) * x));
1410 return ((T)0.5) * x * (((T)1.0) + t);
1411 }
1412
1413 template <typename T>
1414 __device__ __forceinline__ T gelu_quick_kernel(const T& x) {
1415 // x * sigmoid(1.702 * x)
1416 return (T)(((float)x) / (1.0f + expf(-1.702f * (float)x)));
1417 }
1418
1419 template <typename T>
1420 __device__ __forceinline__ T relu_squared_kernel(const T& x) {
1421 // relu(x)^2 — introduced in https://arxiv.org/abs/2109.08668v2
1422 const float f = (float)x;
1423 const float val = f > 0.0f ? f : 0.0f;
1424 return (T)(val * val);
1425 }
1426
1427 } // namespace vllm
1428
1429 void gelu_new(torch::stable::Tensor& out, // [..., d]
1430 torch::stable::Tensor& input) // [..., d]
1431 {
1432 LAUNCH_ACTIVATION_KERNEL(vllm::gelu_new_kernel);
1433 }
1434
1435 void gelu_fast(torch::stable::Tensor& out, // [..., d]
1436 torch::stable::Tensor& input) // [..., d]
1437 {
1438 LAUNCH_ACTIVATION_KERNEL(vllm::gelu_fast_kernel);
1439 }
1440
1441 void gelu_quick(torch::stable::Tensor& out, // [..., d]
1442 torch::stable::Tensor& input) // [..., d]
1443 {
1444 LAUNCH_ACTIVATION_KERNEL(vllm::gelu_quick_kernel);
1445 }
1446
1447 void relu_squared(torch::stable::Tensor& out, // [..., d]
1448 torch::stable::Tensor& input) // [..., d]
1449 {
1450 LAUNCH_ACTIVATION_KERNEL(vllm::relu_squared_kernel);
1451 }