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 8, 12:30 UTC
[Kernel] Apply 256bit LDG/STG To Activation Kernels (#33022)
Signed-off-by: Dzerzhinsky <256908701+AstroVoyager7@users.noreply.github.com> Signed-off-by: Дзержи́нский <256908701+AstroVoyager7@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com>
Дзержи́нский committed
Feb 11, 2026 at 11:31 UTC
1485396abb7c575d0196c2f52f4cdff7f9280a19
1 file changed
+401
-123
csrc/activation_kernels.cu
+401
-123
@@ -9,6 +9,111 @@
9
10
namespace vllm {
11
12
+struct alignas(32) u32x8_t {
13
+ uint32_t u0, u1, u2, u3, u4, u5, u6, u7;
14
+};
15
+
16
+__device__ __forceinline__ void ld256(u32x8_t& val, const u32x8_t* ptr) {
17
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
18
+ asm volatile("ld.global.nc.v8.u32 {%0,%1,%2,%3,%4,%5,%6,%7}, [%8];\n"
19
+ : "=r"(val.u0), "=r"(val.u1), "=r"(val.u2), "=r"(val.u3),
20
+ "=r"(val.u4), "=r"(val.u5), "=r"(val.u6), "=r"(val.u7)
21
+ : "l"(ptr));
22
+#else
23
+ const uint4* uint_ptr = reinterpret_cast<const uint4*>(ptr);
24
+ uint4 top_half = __ldg(&uint_ptr[0]);
25
+ uint4 bottom_half = __ldg(&uint_ptr[1]);
26
+ val.u0 = top_half.x;
27
+ val.u1 = top_half.y;
28
+ val.u2 = top_half.z;
29
+ val.u3 = top_half.w;
30
+ val.u4 = bottom_half.x;
31
+ val.u5 = bottom_half.y;
32
+ val.u6 = bottom_half.z;
33
+ val.u7 = bottom_half.w;
34
+#endif
35
+}
36
+
37
+__device__ __forceinline__ void st256(u32x8_t& val, u32x8_t* ptr) {
38
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000
39
+ asm volatile("st.global.v8.u32 [%0], {%1,%2,%3,%4,%5,%6,%7,%8};\n"
40
+ :
41
+ : "l"(ptr), "r"(val.u0), "r"(val.u1), "r"(val.u2), "r"(val.u3),
42
+ "r"(val.u4), "r"(val.u5), "r"(val.u6), "r"(val.u7)
43
+ : "memory");
44
+#else
45
+ uint4* uint_ptr = reinterpret_cast<uint4*>(ptr);
46
+ uint_ptr[0] = make_uint4(val.u0, val.u1, val.u2, val.u3);
47
+ uint_ptr[1] = make_uint4(val.u4, val.u5, val.u6, val.u7);
48
+#endif
49
+}
50
+
51
+template <bool support_256>
52
+struct VecTraits;
53
+
54
+template <>
55
+struct VecTraits<true> {
56
+ static constexpr int ARCH_MAX_VEC_SIZE = 32;
57
+ using vec_t = u32x8_t;
58
+};
59
+
60
+template <>
61
+struct VecTraits<false> {
62
+ static constexpr int ARCH_MAX_VEC_SIZE = 16;
63
+ using vec_t = int4;
64
+};
65
+
66
+template <typename T>
67
+struct PackedTraits;
68
+
69
+template <>
70
+struct PackedTraits<c10::BFloat16> {
71
+ using packed_t = __nv_bfloat162;
72
+};
73
+
74
+template <>
75
+struct PackedTraits<c10::Half> {
76
+ using packed_t = __half2;
77
+};
78
+
79
+template <>
80
+struct PackedTraits<float> {
81
+ using packed_t = float2;
82
+};
83
+
84
+template <typename packed_t>
85
+__device__ __forceinline__ float2 cast_to_float2(const packed_t& val) {
86
+ if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
87
+ return __bfloat1622float2(val);
88
+ } else if constexpr (std::is_same_v<packed_t, __half2>) {
89
+ return __half22float2(val);
90
+ } else if constexpr (std::is_same_v<packed_t, float2>) {
91
+ return float2(val);
92
+ }
93
+}
94
+
95
+template <typename packed_t>
96
+__device__ __forceinline__ packed_t cast_to_packed(const float2& val) {
97
+ if constexpr (std::is_same_v<packed_t, __nv_bfloat162>) {
98
+ return __float22bfloat162_rn(val);
99
+ } else if constexpr (std::is_same_v<packed_t, __half2>) {
100
+ return __float22half2_rn(val);
101
+ } else if constexpr (std::is_same_v<packed_t, float2>) {
102
+ return float2(val);
103
+ }
104
+}
105
+
106
+template <typename packed_t>
107
+__device__ __forceinline__ packed_t packed_mul(const packed_t& x,
108
+ const packed_t& y) {
109
+ if constexpr (std::is_same_v<packed_t, __nv_bfloat162> ||
110
+ std::is_same_v<packed_t, __half2>) {
111
+ return __hmul2(x, y);
112
+ } else if constexpr (std::is_same_v<packed_t, float2>) {
113
+ return make_float2(x.x * y.x, x.y * y.y);
114
+ }
115
+}
116
+
117
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
118
bool act_first>
119
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
@@ -16,52 +121,69 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
121
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
122
}
123
124
+template <typename packed_t, packed_t (*PACKED_ACT_FN)(const packed_t&),
125
+ bool act_first>
126
+__device__ __forceinline__ packed_t packed_compute(const packed_t& x,
127
+ const packed_t& y) {
128
+ return act_first ? packed_mul(PACKED_ACT_FN(x), y)
129
+ : packed_mul(x, PACKED_ACT_FN(y));
130
+}
131
+
132
// Check if all pointers are 16-byte aligned for int4 vectorized access
20
-__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
133
+__host__ __device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
134
return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
135
}
136
137
+// Check if all pointers are 16-byte aligned for longlong4_32a vectorized access
138
+__host__ __device__ __forceinline__ bool is_32byte_aligned(const void* ptr) {
139
+ return (reinterpret_cast<uintptr_t>(ptr) & 31) == 0;
140
+}
141
+
142
// Activation and gating kernel template.
25
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
26
- bool act_first>
143
+template <typename scalar_t, typename packed_t,
144
+ scalar_t (*ACT_FN)(const scalar_t&),
145
+ packed_t (*PACKED_ACT_FN)(const packed_t&), bool act_first,
146
+ bool use_vec, bool use_256b = false>
147
__global__ void act_and_mul_kernel(
148
scalar_t* __restrict__ out, // [..., d]
149
const scalar_t* __restrict__ input, // [..., 2, d]
150
const int d) {
31
- constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
32
- const int64_t token_idx = blockIdx.x;
33
- const scalar_t* x_ptr = input + token_idx * 2 * d;
151
+ const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
152
const scalar_t* y_ptr = x_ptr + d;
35
- scalar_t* out_ptr = out + token_idx * d;
153
+ scalar_t* out_ptr = out + blockIdx.x * d;
154
37
- // Check alignment for 128-bit vectorized access.
38
- // All three pointers must be 16-byte aligned for safe int4 operations.
39
- const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
40
- is_16byte_aligned(out_ptr);
155
+ if constexpr (use_vec) {
156
+ // Fast path: 128-bit/256-bit vectorized loop
157
+ using vec_t = typename VecTraits<use_256b>::vec_t;
158
+ constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
159
+ constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
160
42
- if (aligned && d >= VEC_SIZE) {
43
- // Fast path: 128-bit vectorized loop
44
- const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
45
- const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
46
- int4* out_vec = reinterpret_cast<int4*>(out_ptr);
47
- const int num_vecs = d / VEC_SIZE;
48
- const int vec_end = num_vecs * VEC_SIZE;
161
+ const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
162
+ const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
163
+ vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
164
+ const int num_vecs = d / 2 / VEC_SIZE;
165
166
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
51
- int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
52
- auto* xp = reinterpret_cast<scalar_t*>(&x);
53
- auto* yp = reinterpret_cast<scalar_t*>(&y);
54
- auto* rp = reinterpret_cast<scalar_t*>(&r);
167
+ vec_t x, y;
168
+ if constexpr (use_256b) {
169
+ ld256(x, &x_vec[i]);
170
+ ld256(y, &y_vec[i]);
171
+ } else {
172
+ x = VLLM_LDG(&x_vec[i]);
173
+ y = VLLM_LDG(&y_vec[i]);
174
+ }
175
+ auto* xp = reinterpret_cast<packed_t*>(&x);
176
+ auto* yp = reinterpret_cast<packed_t*>(&y);
177
#pragma unroll
178
for (int j = 0; j < VEC_SIZE; j++) {
57
- rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
179
+ xp[j] =
180
+ packed_compute<packed_t, PACKED_ACT_FN, act_first>(xp[j], yp[j]);
181
+ }
182
+ if constexpr (use_256b) {
183
+ st256(x, &out_vec[i]);
184
+ } else {
185
+ out_vec[i] = x;
186
}
59
- out_vec[i] = r;
60
- }
61
- // Scalar cleanup for remaining elements
62
- for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
63
- out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(VLLM_LDG(&x_ptr[i]),
64
- VLLM_LDG(&y_ptr[i]));
187
}
188
} else {
189
// Scalar fallback for unaligned data or small d
@@ -79,6 +201,15 @@ __device__ __forceinline__ T silu_kernel(const T& x) {
201
return (T)(((float)x) / (1.0f + expf((float)-x)));
202
}
203
204
+template <typename packed_t>
205
+__device__ __forceinline__ packed_t packed_silu_kernel(const packed_t& val) {
206
+ // x * sigmoid(x)
207
+ float2 fval = cast_to_float2(val);
208
+ fval.x = fval.x / (1.0f + expf(-fval.x));
209
+ fval.y = fval.y / (1.0f + expf(-fval.y));
210
+ return cast_to_packed<packed_t>(fval);
211
+}
212
+
213
template <typename T>
214
__device__ __forceinline__ T gelu_kernel(const T& x) {
215
// Equivalent to PyTorch GELU with 'none' approximation.
@@ -89,6 +220,18 @@ __device__ __forceinline__ T gelu_kernel(const T& x) {
220
return (T)(f * 0.5f * (1.0f + ::erf(f * ALPHA)));
221
}
222
223
+template <typename packed_t>
224
+__device__ __forceinline__ packed_t packed_gelu_kernel(const packed_t& val) {
225
+ // Equivalent to PyTorch GELU with 'none' approximation.
226
+ // Refer to:
227
+ // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
228
+ constexpr float ALPHA = M_SQRT1_2;
229
+ float2 fval = cast_to_float2(val);
230
+ fval.x = fval.x * 0.5f * (1.0f + ::erf(fval.x * ALPHA));
231
+ fval.y = fval.y * 0.5f * (1.0f + ::erf(fval.y * ALPHA));
232
+ return cast_to_packed<packed_t>(fval);
233
+}
234
+
235
template <typename T>
236
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
237
// Equivalent to PyTorch GELU with 'tanh' approximation.
@@ -102,32 +245,83 @@ __device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
245
return (T)(0.5f * f * (1.0f + ::tanhf(inner)));
246
}
247
248
+template <typename packed_t>
249
+__device__ __forceinline__ packed_t
250
+packed_gelu_tanh_kernel(const packed_t& val) {
251
+ // Equivalent to PyTorch GELU with 'tanh' approximation.
252
+ // Refer to:
253
+ // https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
254
+ float2 fval = cast_to_float2(val);
255
+ constexpr float BETA = M_SQRT2 * M_2_SQRTPI * 0.5f;
256
+ constexpr float KAPPA = 0.044715;
257
+
258
+ float x_cube = fval.x * fval.x * fval.x;
259
+ float inner = BETA * (fval.x + KAPPA * x_cube);
260
+ fval.x = 0.5f * fval.x * (1.0f + ::tanhf(inner));
261
+
262
+ x_cube = fval.y * fval.y * fval.y;
263
+ inner = BETA * (fval.y + KAPPA * x_cube);
264
+ fval.y = 0.5f * fval.y * (1.0f + ::tanhf(inner));
265
+ return cast_to_packed<packed_t>(fval);
266
+}
267
+
268
} // namespace vllm
269
270
// Launch activation and gating kernel.
271
// Use ACT_FIRST (bool) indicating whether to apply the activation function
272
// first.
110
-#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST) \
111
- int d = input.size(-1) / 2; \
112
- int64_t num_tokens = input.numel() / input.size(-1); \
113
- dim3 grid(num_tokens); \
114
- dim3 block(std::min(d, 1024)); \
115
- if (num_tokens == 0) { \
116
- return; \
117
- } \
118
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
119
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
120
- VLLM_DISPATCH_FLOATING_TYPES( \
121
- input.scalar_type(), "act_and_mul_kernel", [&] { \
122
- vllm::act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST> \
123
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
124
- input.data_ptr<scalar_t>(), d); \
125
- });
273
+#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, PACKED_KERNEL, ACT_FIRST) \
274
+ auto dtype = input.scalar_type(); \
275
+ int d = input.size(-1) / 2; \
276
+ int64_t num_tokens = input.numel() / input.size(-1); \
277
+ if (num_tokens == 0) { \
278
+ return; \
279
+ } \
280
+ dim3 grid(num_tokens); \
281
+ int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
282
+ int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
283
+ int vec_size = support_vec / at::elementSize(dtype); \
284
+ const bool use_vec = (d % vec_size == 0); \
285
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
286
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
287
+ if (use_vec) { \
288
+ dim3 block(std::min(d / vec_size, 1024)); \
289
+ if (cc_major >= 10 && num_tokens > 128) { \
290
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
291
+ vllm::act_and_mul_kernel< \
292
+ scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
293
+ KERNEL<scalar_t>, \
294
+ PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
295
+ ACT_FIRST, true, true><<<grid, block, 0, stream>>>( \
296
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
297
+ }); \
298
+ } else { \
299
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
300
+ vllm::act_and_mul_kernel< \
301
+ scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
302
+ KERNEL<scalar_t>, \
303
+ PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
304
+ ACT_FIRST, true, false><<<grid, block, 0, stream>>>( \
305
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
306
+ }); \
307
+ } \
308
+ } else { \
309
+ dim3 block(std::min(d, 1024)); \
310
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel", [&] { \
311
+ vllm::act_and_mul_kernel< \
312
+ scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
313
+ KERNEL<scalar_t>, \
314
+ PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
315
+ ACT_FIRST, false><<<grid, block, 0, stream>>>( \
316
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
317
+ }); \
318
+ }
319
320
void silu_and_mul(torch::Tensor& out, // [..., d]
321
torch::Tensor& input) // [..., 2 * d]
322
{
130
- LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, true);
323
+ LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
324
+ true);
325
}
326
327
void mul_and_silu(torch::Tensor& out, // [..., d]
@@ -135,19 +329,22 @@ void mul_and_silu(torch::Tensor& out, // [..., d]
329
{
330
// The difference between mul_and_silu and silu_and_mul is that mul_and_silu
331
// applies the silu to the latter half of the input.
138
- LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, false);
332
+ LAUNCH_ACTIVATION_GATE_KERNEL(vllm::silu_kernel, vllm::packed_silu_kernel,
333
+ false);
334
}
335
336
void gelu_and_mul(torch::Tensor& out, // [..., d]
337
torch::Tensor& input) // [..., 2 * d]
338
{
144
- LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, true);
339
+ LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_kernel, vllm::packed_gelu_kernel,
340
+ true);
341
}
342
343
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
344
torch::Tensor& input) // [..., 2 * d]
345
{
150
- LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel, true);
346
+ LAUNCH_ACTIVATION_GATE_KERNEL(vllm::gelu_tanh_kernel,
347
+ vllm::packed_gelu_tanh_kernel, true);
348
}
349
350
namespace vllm {
@@ -158,42 +355,57 @@ __device__ __forceinline__ T fatrelu_kernel(const T& x, const float threshold) {
355
return (T)(f > threshold ? f : 0.0f);
356
}
357
161
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
358
+template <typename packed_t>
359
+__device__ __forceinline__ packed_t
360
+packed_fatrelu_kernel(const packed_t& val, const float threshold) {
361
+ float2 fval = cast_to_float2(val);
362
+ fval.x = fval.x > threshold ? fval.x : 0.0f;
363
+ fval.y = fval.y > threshold ? fval.y : 0.0f;
364
+ return cast_to_packed<packed_t>(fval);
365
+}
366
+
367
+template <typename scalar_t, typename packed_t,
368
+ scalar_t (*ACT_FN)(const scalar_t&, const float),
369
+ packed_t (*PACKED_ACT_FN)(const packed_t&, const float), bool use_vec,
370
+ bool use_256b = false>
371
__global__ void act_and_mul_kernel_with_param(
372
scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
373
const float param) {
165
- constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
166
- const int64_t token_idx = blockIdx.x;
167
- const scalar_t* x_ptr = input + token_idx * 2 * d;
374
+ const scalar_t* x_ptr = input + blockIdx.x * 2 * d;
375
const scalar_t* y_ptr = x_ptr + d;
169
- scalar_t* out_ptr = out + token_idx * d;
376
+ scalar_t* out_ptr = out + blockIdx.x * d;
377
171
- // Check alignment for 128-bit vectorized access
172
- const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
173
- is_16byte_aligned(out_ptr);
378
+ if constexpr (use_vec) {
379
+ // Fast path: 128-bit/256-bit vectorized loop
380
+ using vec_t = typename VecTraits<use_256b>::vec_t;
381
+ constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
382
+ constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(packed_t);
383
175
- if (aligned && d >= VEC_SIZE) {
176
- // Fast path: 128-bit vectorized loop
177
- const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
178
- const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
179
- int4* out_vec = reinterpret_cast<int4*>(out_ptr);
180
- const int num_vecs = d / VEC_SIZE;
181
- const int vec_end = num_vecs * VEC_SIZE;
384
+ const vec_t* x_vec = reinterpret_cast<const vec_t*>(x_ptr);
385
+ const vec_t* y_vec = reinterpret_cast<const vec_t*>(y_ptr);
386
+ vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
387
+ const int num_vecs = d / 2 / VEC_SIZE;
388
389
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
184
- int4 x = VLLM_LDG(&x_vec[i]), y = VLLM_LDG(&y_vec[i]), r;
185
- auto* xp = reinterpret_cast<scalar_t*>(&x);
186
- auto* yp = reinterpret_cast<scalar_t*>(&y);
187
- auto* rp = reinterpret_cast<scalar_t*>(&r);
390
+ vec_t x, y;
391
+ if constexpr (use_256b) {
392
+ ld256(x, &x_vec[i]);
393
+ ld256(y, &y_vec[i]);
394
+ } else {
395
+ x = VLLM_LDG(&x_vec[i]);
396
+ y = VLLM_LDG(&y_vec[i]);
397
+ }
398
+ auto* xp = reinterpret_cast<packed_t*>(&x);
399
+ auto* yp = reinterpret_cast<packed_t*>(&y);
400
#pragma unroll
401
for (int j = 0; j < VEC_SIZE; j++) {
190
- rp[j] = ACT_FN(xp[j], param) * yp[j];
402
+ xp[j] = packed_mul(PACKED_ACT_FN(xp[j], param), yp[j]);
403
+ }
404
+ if constexpr (use_256b) {
405
+ st256(x, &out_vec[i]);
406
+ } else {
407
+ out_vec[i] = x;
408
}
192
- out_vec[i] = r;
193
- }
194
- // Scalar cleanup for remaining elements
195
- for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
196
- out_ptr[i] = ACT_FN(VLLM_LDG(&x_ptr[i]), param) * VLLM_LDG(&y_ptr[i]);
409
}
410
} else {
411
// Scalar fallback for unaligned data or small d
@@ -276,20 +488,58 @@ __global__ void swigluoai_and_mul_kernel(
488
489
} // namespace vllm
490
279
-#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PARAM) \
280
- int d = input.size(-1) / 2; \
281
- int64_t num_tokens = input.numel() / input.size(-1); \
282
- dim3 grid(num_tokens); \
283
- dim3 block(std::min(d, 1024)); \
284
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
285
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
286
- VLLM_DISPATCH_FLOATING_TYPES( \
287
- input.scalar_type(), "act_and_mul_kernel_with_param", [&] { \
288
- vllm::act_and_mul_kernel_with_param<scalar_t, KERNEL<scalar_t>> \
289
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
290
- input.data_ptr<scalar_t>(), d, \
291
- PARAM); \
292
- });
491
+#define LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(KERNEL, PACKED_KERNEL, PARAM) \
492
+ auto dtype = input.scalar_type(); \
493
+ int d = input.size(-1) / 2; \
494
+ int64_t num_tokens = input.numel() / input.size(-1); \
495
+ if (num_tokens == 0) { \
496
+ return; \
497
+ } \
498
+ dim3 grid(num_tokens); \
499
+ int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
500
+ int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
501
+ int vec_size = support_vec / at::elementSize(dtype); \
502
+ const bool use_vec = (d % vec_size == 0); \
503
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
504
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
505
+ if (use_vec) { \
506
+ dim3 block(std::min(d / vec_size, 1024)); \
507
+ if (cc_major >= 10 && num_tokens > 128) { \
508
+ VLLM_DISPATCH_FLOATING_TYPES( \
509
+ dtype, "act_and_mul_kernel_with_param", [&] { \
510
+ vllm::act_and_mul_kernel_with_param< \
511
+ scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
512
+ KERNEL<scalar_t>, \
513
+ PACKED_KERNEL< \
514
+ typename vllm::PackedTraits<scalar_t>::packed_t>, \
515
+ true, true><<<grid, block, 0, stream>>>( \
516
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
517
+ PARAM); \
518
+ }); \
519
+ } else { \
520
+ VLLM_DISPATCH_FLOATING_TYPES( \
521
+ dtype, "act_and_mul_kernel_with_param", [&] { \
522
+ vllm::act_and_mul_kernel_with_param< \
523
+ scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
524
+ KERNEL<scalar_t>, \
525
+ PACKED_KERNEL< \
526
+ typename vllm::PackedTraits<scalar_t>::packed_t>, \
527
+ true, false><<<grid, block, 0, stream>>>( \
528
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, \
529
+ PARAM); \
530
+ }); \
531
+ } \
532
+ } else { \
533
+ dim3 block(std::min(d, 1024)); \
534
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "act_and_mul_kernel_with_param", [&] { \
535
+ vllm::act_and_mul_kernel_with_param< \
536
+ scalar_t, typename vllm::PackedTraits<scalar_t>::packed_t, \
537
+ KERNEL<scalar_t>, \
538
+ PACKED_KERNEL<typename vllm::PackedTraits<scalar_t>::packed_t>, \
539
+ false><<<grid, block, 0, stream>>>( \
540
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d, PARAM); \
541
+ }); \
542
+ }
543
544
#define LAUNCH_SIGLUOAI_AND_MUL(KERNEL, ALPHA, LIMIT) \
545
int d = input.size(-1) / 2; \
@@ -309,7 +559,8 @@ __global__ void swigluoai_and_mul_kernel(
559
void fatrelu_and_mul(torch::Tensor& out, // [..., d],
560
torch::Tensor& input, // [..., 2 * d]
561
double threshold) {
312
- LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(vllm::fatrelu_kernel, threshold);
562
+ LAUNCH_ACTIVATION_GATE_KERNEL_WITH_PARAM(
563
+ vllm::fatrelu_kernel, vllm::packed_fatrelu_kernel, threshold);
564
}
565
void swigluoai_and_mul(torch::Tensor& out, // [..., d]
566
torch::Tensor& input, // [..., 2 * d]
@@ -319,39 +570,41 @@ void swigluoai_and_mul(torch::Tensor& out, // [..., d]
570
namespace vllm {
571
572
// Element-wise activation kernel template.
322
-template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&)>
573
+template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&), bool use_vec,
574
+ bool use_256b = false>
575
__global__ void activation_kernel(
576
scalar_t* __restrict__ out, // [..., d]
577
const scalar_t* __restrict__ input, // [..., d]
578
const int d) {
327
- constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
328
- const int64_t token_idx = blockIdx.x;
329
- const scalar_t* in_ptr = input + token_idx * d;
330
- scalar_t* out_ptr = out + token_idx * d;
331
-
332
- // Check alignment for 128-bit vectorized access
333
- const bool aligned = is_16byte_aligned(in_ptr) && is_16byte_aligned(out_ptr);
334
-
335
- if (aligned && d >= VEC_SIZE) {
336
- // Fast path: 128-bit vectorized loop
337
- const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
338
- int4* out_vec = reinterpret_cast<int4*>(out_ptr);
579
+ const scalar_t* in_ptr = input + blockIdx.x * d;
580
+ scalar_t* out_ptr = out + blockIdx.x * d;
581
+
582
+ if constexpr (use_vec) {
583
+ // Fast path: 128-bit/256-bit vectorized loop
584
+ using vec_t = typename VecTraits<use_256b>::vec_t;
585
+ constexpr int ARCH_MAX_VEC_SIZE = VecTraits<use_256b>::ARCH_MAX_VEC_SIZE;
586
+ constexpr int VEC_SIZE = ARCH_MAX_VEC_SIZE / sizeof(scalar_t);
587
+ const vec_t* in_vec = reinterpret_cast<const vec_t*>(in_ptr);
588
+ vec_t* out_vec = reinterpret_cast<vec_t*>(out_ptr);
589
const int num_vecs = d / VEC_SIZE;
340
- const int vec_end = num_vecs * VEC_SIZE;
590
591
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
343
- int4 v = VLLM_LDG(&in_vec[i]), r;
592
+ vec_t v;
593
+ if constexpr (use_256b) {
594
+ ld256(v, &in_vec[i]);
595
+ } else {
596
+ v = VLLM_LDG(&in_vec[i]);
597
+ }
598
auto* vp = reinterpret_cast<scalar_t*>(&v);
345
- auto* rp = reinterpret_cast<scalar_t*>(&r);
599
#pragma unroll
600
for (int j = 0; j < VEC_SIZE; j++) {
348
- rp[j] = ACT_FN(vp[j]);
601
+ vp[j] = ACT_FN(vp[j]);
602
+ }
603
+ if constexpr (use_256b) {
604
+ st256(v, &out_vec[i]);
605
+ } else {
606
+ out_vec[i] = v;
607
}
350
- out_vec[i] = r;
351
- }
352
- // Scalar cleanup for remaining elements
353
- for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
354
- out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[i]));
608
}
609
} else {
610
// Scalar fallback for unaligned data or small d
@@ -365,18 +618,43 @@ __global__ void activation_kernel(
618
} // namespace vllm
619
620
// Launch element-wise activation kernel.
368
-#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \
369
- int d = input.size(-1); \
370
- int64_t num_tokens = input.numel() / d; \
371
- dim3 grid(num_tokens); \
372
- dim3 block(std::min(d, 1024)); \
373
- const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
374
- const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
375
- VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "activation_kernel", [&] { \
376
- vllm::activation_kernel<scalar_t, KERNEL<scalar_t>> \
377
- <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
378
- input.data_ptr<scalar_t>(), d); \
379
- });
621
+#define LAUNCH_ACTIVATION_KERNEL(KERNEL) \
622
+ auto dtype = input.scalar_type(); \
623
+ int d = input.size(-1); \
624
+ int64_t num_tokens = input.numel() / input.size(-1); \
625
+ if (num_tokens == 0) { \
626
+ return; \
627
+ } \
628
+ dim3 grid(num_tokens); \
629
+ int cc_major = at::cuda::getCurrentDeviceProperties()->major; \
630
+ int support_vec = (cc_major >= 10 && num_tokens > 128) ? 32 : 16; \
631
+ int vec_size = support_vec / at::elementSize(dtype); \
632
+ const bool use_vec = (d % vec_size == 0); \
633
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
634
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
635
+ if (use_vec) { \
636
+ dim3 block(std::min(d / vec_size, 1024)); \
637
+ if (cc_major >= 10 && num_tokens > 128) { \
638
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
639
+ vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, true> \
640
+ <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
641
+ input.data_ptr<scalar_t>(), d); \
642
+ }); \
643
+ } else { \
644
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
645
+ vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, true, false> \
646
+ <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
647
+ input.data_ptr<scalar_t>(), d); \
648
+ }); \
649
+ } \
650
+ } else { \
651
+ dim3 block(std::min(d, 1024)); \
652
+ VLLM_DISPATCH_FLOATING_TYPES(dtype, "activation_kernel", [&] { \
653
+ vllm::activation_kernel<scalar_t, KERNEL<scalar_t>, false> \
654
+ <<<grid, block, 0, stream>>>(out.data_ptr<scalar_t>(), \
655
+ input.data_ptr<scalar_t>(), d); \
656
+ }); \
657
+ }
658
659
namespace vllm {
660