@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 175 lines 5.85 KB
Raw
1 #include "cpu_types.hpp"
2
3 namespace {
4 template <typename scalar_t, vec_op::FP32Vec8 (*func)(const vec_op::FP32Vec8&),
5 bool is_gated>
6 void activation_kernel(int num_tokens, int d, scalar_t* __restrict__ input,
7 scalar_t* __restrict__ output) {
8 using scalar_vec_t = vec_op::vec_t<scalar_t>;
9 constexpr int VEC_ELEM_NUM = scalar_vec_t::get_elem_num();
10
11 TORCH_CHECK(d % VEC_ELEM_NUM == 0);
12
13 #pragma omp parallel for
14 for (int i = 0; i < num_tokens; ++i) {
15 for (int j = 0; j < d; j += VEC_ELEM_NUM) {
16 int start = i * d;
17 if constexpr (is_gated) {
18 start *= 2;
19 }
20
21 const scalar_vec_t x(input + start + j);
22 const vec_op::FP32Vec8 f32_x(x);
23 vec_op::FP32Vec8 f32_ans = func(f32_x);
24
25 if constexpr (is_gated) {
26 const scalar_vec_t y(input + start + d + j);
27 const vec_op::FP32Vec8 f32_y(y);
28 f32_ans = f32_y * f32_ans;
29 }
30
31 const scalar_vec_t result(f32_ans);
32 result.save(output + i * d + j);
33 }
34 }
35 }
36
37 FORCE_INLINE vec_op::FP32Vec8 silu_act(const vec_op::FP32Vec8& x) {
38 const vec_op::FP32Vec8 zeros(0.0);
39 const vec_op::FP32Vec8 ones(1.0);
40 return x / (ones + (zeros - x).exp());
41 }
42
43 FORCE_INLINE vec_op::FP32Vec8 gelu_new_act(const vec_op::FP32Vec8& x) {
44 const vec_op::FP32Vec8 ones(1.0);
45 const vec_op::FP32Vec8 w1(0.79788456f);
46 const vec_op::FP32Vec8 w2(0.044715f);
47 const vec_op::FP32Vec8 w3(0.5);
48 const vec_op::FP32Vec8 x3 = x * x * x;
49 const vec_op::FP32Vec8 t = (w1 * (x + w2 * x3)).tanh();
50 return w3 * x * (ones + t);
51 }
52
53 FORCE_INLINE vec_op::FP32Vec8 gelu_fast_act(const vec_op::FP32Vec8& x) {
54 const vec_op::FP32Vec8 ones(1.0);
55 const vec_op::FP32Vec8 w1(0.79788456f);
56 const vec_op::FP32Vec8 w2(0.044715f);
57 const vec_op::FP32Vec8 w3(0.5);
58 const vec_op::FP32Vec8 t = (x * w1 * (ones + x * w2 * x)).tanh();
59 return w3 * x * (ones + t);
60 }
61
62 FORCE_INLINE vec_op::FP32Vec8 gelu_quick_act(const vec_op::FP32Vec8& x) {
63 const vec_op::FP32Vec8 zeros(0.0);
64 const vec_op::FP32Vec8 ones(1.0);
65 const vec_op::FP32Vec8 w1(1.702f);
66 return x / (ones + (zeros - w1 * x).exp());
67 }
68
69 FORCE_INLINE vec_op::FP32Vec8 gelu_act(const vec_op::FP32Vec8& x) {
70 const vec_op::FP32Vec8 ones(1.0);
71 const vec_op::FP32Vec8 w1(M_SQRT1_2);
72 const vec_op::FP32Vec8 w2(0.5);
73 return x * w2 * (ones + (x * w1).er());
74 }
75
76 FORCE_INLINE vec_op::FP32Vec8 gelu_tanh_act(const vec_op::FP32Vec8& x) {
77 const vec_op::FP32Vec8 ones(1.0);
78 const vec_op::FP32Vec8 w1(M_SQRT2 * M_2_SQRTPI * 0.5);
79 const vec_op::FP32Vec8 w2(0.5);
80 const vec_op::FP32Vec8 w3(0.044715);
81 const vec_op::FP32Vec8 x_3 = x * x * x;
82 const vec_op::FP32Vec8 inner = w1 * (x + x_3 * w3);
83 return x * w2 * (ones + inner.tanh());
84 }
85 }; // namespace
86
87 void silu_and_mul(torch::Tensor& out, torch::Tensor& input) {
88 int num_tokens = input.numel() / input.size(-1);
89 int d = input.size(-1) / 2;
90
91 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "silu_and_mul_impl", [&] {
92 CPU_KERNEL_GUARD_IN(silu_and_mul_impl)
93 activation_kernel<scalar_t, silu_act, true>(
94 num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
95 CPU_KERNEL_GUARD_OUT(silu_and_mul_impl)
96 });
97 }
98
99 void gelu_and_mul(torch::Tensor& out, // [..., d]
100 torch::Tensor& input) // [..., 2 * d]
101 {
102 int num_tokens = input.numel() / input.size(-1);
103 int d = input.size(-1) / 2;
104
105 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_and_mul_impl", [&] {
106 CPU_KERNEL_GUARD_IN(gelu_and_mul_impl)
107 activation_kernel<scalar_t, gelu_act, true>(
108 num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
109 CPU_KERNEL_GUARD_OUT(gelu_and_mul_impl)
110 });
111 }
112
113 void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
114 torch::Tensor& input) // [..., 2 * d]
115 {
116 int num_tokens = input.numel() / input.size(-1);
117 int d = input.size(-1) / 2;
118
119 VLLM_DISPATCH_FLOATING_TYPES(
120 input.scalar_type(), "gelu_tanh_and_mul_impl", [&] {
121 CPU_KERNEL_GUARD_IN(gelu_tanh_and_mul_impl)
122 activation_kernel<scalar_t, gelu_tanh_act, true>(
123 num_tokens, d, input.data_ptr<scalar_t>(),
124 out.data_ptr<scalar_t>());
125 CPU_KERNEL_GUARD_OUT(gelu_tanh_and_mul_impl)
126 });
127 }
128
129 void gelu_tanh(torch::Tensor& out, torch::Tensor& input) {
130 int num_tokens = input.numel() / input.size(-1);
131 int d = input.size(-1);
132
133 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_tanh_impl", [&] {
134 CPU_KERNEL_GUARD_IN(gelu_tanh_impl)
135 activation_kernel<scalar_t, gelu_tanh_act, false>(
136 num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
137 CPU_KERNEL_GUARD_OUT(gelu_tanh_impl)
138 });
139 }
140
141 void gelu_new(torch::Tensor& out, torch::Tensor& input) {
142 int num_tokens = input.numel() / input.size(-1);
143 int d = input.size(-1);
144
145 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_new_impl", [&] {
146 CPU_KERNEL_GUARD_IN(gelu_new_impl)
147 activation_kernel<scalar_t, gelu_new_act, false>(
148 num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
149 CPU_KERNEL_GUARD_OUT(gelu_new_impl)
150 });
151 }
152
153 void gelu_fast(torch::Tensor& out, torch::Tensor& input) {
154 int num_tokens = input.numel() / input.size(-1);
155 int d = input.size(-1);
156
157 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_fast_impl", [&] {
158 CPU_KERNEL_GUARD_IN(gelu_fast_impl)
159 activation_kernel<scalar_t, gelu_fast_act, false>(
160 num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
161 CPU_KERNEL_GUARD_OUT(gelu_fast_impl)
162 });
163 }
164
165 void gelu_quick(torch::Tensor& out, torch::Tensor& input) {
166 int num_tokens = input.numel() / input.size(-1);
167 int d = input.size(-1);
168
169 VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "gelu_quick_impl", [&] {
170 CPU_KERNEL_GUARD_IN(gelu_quick_impl)
171 activation_kernel<scalar_t, gelu_quick_act, false>(
172 num_tokens, d, input.data_ptr<scalar_t>(), out.data_ptr<scalar_t>());
173 CPU_KERNEL_GUARD_OUT(gelu_quick_impl)
174 });
175 }