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 11, 08:45 UTC
[Perf][Kernels] Vectorize `csrc/activations_kernels.cu` (#29512)
Signed-off-by: mgoin <mgoin64@gmail.com>
Michael Goin committed
Dec 16, 2025 at 17:56 UTC
0a1ab1e565fce5070bc1c1b1f3374537e437550c
2 files changed
+176
-38
benchmarks/kernels/benchmark_activation.py
+2
-2
@@ -13,8 +13,8 @@ from vllm.triton_utils import triton
13
from vllm.utils.argparse_utils import FlexibleArgumentParser
14
from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE
15
16
-batch_size_range = [1, 16, 32, 64, 128]
17
-seq_len_range = [1, 16, 64, 128, 256, 512, 1024, 2048, 4096]
16
+batch_size_range = [1, 16, 128]
17
+seq_len_range = [1, 16, 64, 1024, 4096]
18
intermediate_size = [3072, 9728, 12288]
19
configs = list(itertools.product(batch_size_range, seq_len_range, intermediate_size))
20
csrc/activation_kernels.cu
+174
-36
@@ -15,19 +15,61 @@ __device__ __forceinline__ scalar_t compute(const scalar_t& x,
15
const scalar_t& y) {
16
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
17
}
18
-// Activation and gating kernel template.
18
19
+// Check if all pointers are 16-byte aligned for int4 vectorized access
20
+__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
21
+ return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
22
+}
23
+
24
+// Activation and gating kernel template.
25
template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&),
26
bool act_first>
27
__global__ void act_and_mul_kernel(
28
scalar_t* __restrict__ out, // [..., d]
29
const scalar_t* __restrict__ input, // [..., 2, d]
30
const int d) {
31
+ constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
32
const int64_t token_idx = blockIdx.x;
27
- for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
28
- const scalar_t x = VLLM_LDG(&input[token_idx * 2 * d + idx]);
29
- const scalar_t y = VLLM_LDG(&input[token_idx * 2 * d + d + idx]);
30
- out[token_idx * d + idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
33
+ const scalar_t* x_ptr = input + token_idx * 2 * d;
34
+ const scalar_t* y_ptr = x_ptr + d;
35
+ scalar_t* out_ptr = out + token_idx * d;
36
+
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);
41
+
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;
49
+
50
+ 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);
55
+#pragma unroll
56
+ for (int j = 0; j < VEC_SIZE; j++) {
57
+ rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
58
+ }
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]));
65
+ }
66
+ } else {
67
+ // Scalar fallback for unaligned data or small d
68
+ for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
69
+ const scalar_t x = VLLM_LDG(&x_ptr[idx]);
70
+ const scalar_t y = VLLM_LDG(&y_ptr[idx]);
71
+ out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
72
+ }
73
}
74
}
75
@@ -120,50 +162,115 @@ template <typename scalar_t, scalar_t (*ACT_FN)(const scalar_t&, const float)>
162
__global__ void act_and_mul_kernel_with_param(
163
scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const int d,
164
const float param) {
165
+ constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
166
const int64_t token_idx = blockIdx.x;
124
- for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
125
- const scalar_t x = VLLM_LDG(&input[token_idx * 2 * d + idx]);
126
- const scalar_t y = VLLM_LDG(&input[token_idx * 2 * d + d + idx]);
127
- out[token_idx * d + idx] = ACT_FN(x, param) * y;
167
+ const scalar_t* x_ptr = input + token_idx * 2 * d;
168
+ const scalar_t* y_ptr = x_ptr + d;
169
+ scalar_t* out_ptr = out + token_idx * d;
170
+
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);
174
+
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;
182
+
183
+ 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);
188
+#pragma unroll
189
+ for (int j = 0; j < VEC_SIZE; j++) {
190
+ rp[j] = ACT_FN(xp[j], param) * yp[j];
191
+ }
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]);
197
+ }
198
+ } else {
199
+ // Scalar fallback for unaligned data or small d
200
+ for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
201
+ const scalar_t x = VLLM_LDG(&x_ptr[idx]);
202
+ const scalar_t y = VLLM_LDG(&y_ptr[idx]);
203
+ out_ptr[idx] = ACT_FN(x, param) * y;
204
+ }
205
}
206
}
207
208
template <typename T>
209
__device__ __forceinline__ T swigluoai_and_mul(const T& gate, const T& up,
210
float alpha, float limit) {
134
- // clamp gate: min=None, max=limit
135
- const float gate_f = (float)gate;
136
- const float clamped_gate = gate_f > limit ? limit : gate_f;
137
-
138
- // clamp up: min=-limit, max=limit
139
- const float up_f = (float)up;
140
- const float clamped_up =
141
- up_f > limit ? limit : (up_f < -limit ? -limit : up_f);
142
-
143
- // glu = gate * sigmoid(gate * alpha)
144
- const float sigmoid_val = 1.0f / (1.0f + expf(-clamped_gate * alpha));
145
- const float glu = clamped_gate * sigmoid_val;
146
-
147
- // (up + 1) * glu
148
- return (T)((clamped_up + 1.0f) * glu);
211
+ // Clamp gate to (-inf, limit] and up to [-limit, limit]
212
+ const float g = fminf((float)gate, limit);
213
+ const float u = fmaxf(fminf((float)up, limit), -limit);
214
+ // glu = gate * sigmoid(gate * alpha), then return (up + 1) * glu
215
+ return (T)((u + 1.0f) * g / (1.0f + expf(-g * alpha)));
216
}
217
218
+// Interleaved gate/up: input has [gate0, up0, gate1, up1, ...].
219
template <typename scalar_t,
220
scalar_t (*ACT_FN)(const scalar_t&, const scalar_t&, const float,
221
const float)>
222
__global__ void swigluoai_and_mul_kernel(
223
scalar_t* __restrict__ out, // [..., d]
156
- const scalar_t* __restrict__ input, // [..., 2, d]
224
+ const scalar_t* __restrict__ input, // [..., 2 * d] (interleaved)
225
const int d, const float alpha, const float limit) {
226
+ // For interleaved data: input has 2*d elements per token (gate/up pairs)
227
+ // output has d elements per token
228
+ constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
229
+ constexpr int PAIRS = VEC_SIZE / 2; // Number of gate/up pairs per int4 load
230
const int64_t token_idx = blockIdx.x;
159
- // TODO: Vectorize loads and stores.
160
- for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
161
- // gate = x[..., ::2] (even indices)
162
- const scalar_t gate = VLLM_LDG(&input[token_idx * 2 * d + 2 * idx]);
163
- // up = x[..., 1::2] (odd indices)
164
- const scalar_t up = VLLM_LDG(&input[token_idx * 2 * d + 2 * idx + 1]);
165
-
166
- out[token_idx * d + idx] = ACT_FN(gate, up, alpha, limit);
231
+ const scalar_t* in_ptr = input + token_idx * 2 * d;
232
+ scalar_t* out_ptr = out + token_idx * d;
233
+
234
+ // Check alignment for 128-bit vectorized access on input.
235
+ // For output we use int2 (64-bit) which has 8-byte alignment requirement.
236
+ const bool in_aligned = is_16byte_aligned(in_ptr);
237
+ const bool out_aligned =
238
+ (reinterpret_cast<uintptr_t>(out_ptr) & 7) == 0; // 8-byte for int2
239
+
240
+ if (in_aligned && out_aligned && d >= PAIRS) {
241
+ // Fast path: vectorized loop
242
+ // Each int4 load gives VEC_SIZE elements = PAIRS gate/up pairs
243
+ // Each int2 store writes PAIRS output elements
244
+ const int4* in_vec = reinterpret_cast<const int4*>(in_ptr);
245
+ int2* out_vec = reinterpret_cast<int2*>(out_ptr);
246
+ const int num_vecs = d / PAIRS;
247
+ const int vec_end = num_vecs * PAIRS;
248
+
249
+ for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
250
+ int4 v = VLLM_LDG(&in_vec[i]);
251
+ int2 r;
252
+ auto* vp = reinterpret_cast<scalar_t*>(&v);
253
+ auto* rp = reinterpret_cast<scalar_t*>(&r);
254
+#pragma unroll
255
+ for (int j = 0; j < PAIRS; j++) {
256
+ rp[j] = ACT_FN(vp[2 * j], vp[2 * j + 1], alpha, limit);
257
+ }
258
+ out_vec[i] = r;
259
+ }
260
+ // Scalar cleanup for remaining elements
261
+ for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
262
+ out_ptr[i] = ACT_FN(VLLM_LDG(&in_ptr[2 * i]),
263
+ VLLM_LDG(&in_ptr[2 * i + 1]), alpha, limit);
264
+ }
265
+ } else {
266
+ // Scalar fallback for unaligned data or small d
267
+ for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
268
+ // gate = x[..., ::2] (even indices)
269
+ const scalar_t gate = VLLM_LDG(&in_ptr[2 * idx]);
270
+ // up = x[..., 1::2] (odd indices)
271
+ const scalar_t up = VLLM_LDG(&in_ptr[2 * idx + 1]);
272
+ out_ptr[idx] = ACT_FN(gate, up, alpha, limit);
273
+ }
274
}
275
}
276
@@ -217,10 +324,41 @@ __global__ void activation_kernel(
324
scalar_t* __restrict__ out, // [..., d]
325
const scalar_t* __restrict__ input, // [..., d]
326
const int d) {
327
+ constexpr int VEC_SIZE = 16 / sizeof(scalar_t);
328
const int64_t token_idx = blockIdx.x;
221
- for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
222
- const scalar_t x = VLLM_LDG(&input[token_idx * d + idx]);
223
- out[token_idx * d + idx] = ACT_FN(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);
339
+ const int num_vecs = d / VEC_SIZE;
340
+ const int vec_end = num_vecs * VEC_SIZE;
341
+
342
+ for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
343
+ int4 v = VLLM_LDG(&in_vec[i]), r;
344
+ auto* vp = reinterpret_cast<scalar_t*>(&v);
345
+ auto* rp = reinterpret_cast<scalar_t*>(&r);
346
+#pragma unroll
347
+ for (int j = 0; j < VEC_SIZE; j++) {
348
+ rp[j] = ACT_FN(vp[j]);
349
+ }
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]));
355
+ }
356
+ } else {
357
+ // Scalar fallback for unaligned data or small d
358
+ for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
359
+ const scalar_t x = VLLM_LDG(&in_ptr[idx]);
360
+ out_ptr[idx] = ACT_FN(x);
361
+ }
362
}
363
}
364