@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 215 lines 7.49 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 #ifndef CPU_FUSED_MOE_ACTIVATIONS_HPP
5 #define CPU_FUSED_MOE_ACTIVATIONS_HPP
6
7 #include <cmath>
8 #include <cstdint>
9 #include <string>
10
11 #include "cpu/cpu_arch_macros.h"
12 #include "cpu/utils.hpp"
13
14 #if defined(DEFINE_FAST_EXP)
15 #define DEFINE_CPU_FUSED_MOE_EXP DEFINE_FAST_EXP
16 #else
17 #define DEFINE_CPU_FUSED_MOE_EXP \
18 auto fast_exp = [](const vec_op::FP32Vec16& vec) { return vec.exp(); };
19 #endif
20
21 namespace cpu_fused_moe_utils {
22 enum class FusedMOEAct {
23 SiluAndMul,
24 SwigluOAIAndMul,
25 GeluAndMul,
26 GeluTanhAndMul,
27 };
28
29 inline FusedMOEAct get_act_type(const std::string& act) {
30 if (act == "silu") {
31 return FusedMOEAct::SiluAndMul;
32 } else if (act == "swigluoai") {
33 return FusedMOEAct::SwigluOAIAndMul;
34 } else if (act == "gelu") {
35 return FusedMOEAct::GeluAndMul;
36 } else if (act == "gelu_tanh") {
37 return FusedMOEAct::GeluTanhAndMul;
38 } else {
39 TORCH_CHECK(false, "Invalid act type: " + act);
40 }
41 }
42
43 template <typename scalar_t>
44 void swigluoai_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
45 const int32_t m_size, const int32_t n_size,
46 const int32_t input_stride,
47 const int32_t output_stride) {
48 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
49 #if defined(__AVX512F__)
50 // For GPT-OSS interleaved gate-up weights
51 alignas(64) static int32_t index[16] = {0, 2, 4, 6, 8, 10, 12, 14,
52 16, 18, 20, 22, 24, 26, 28, 30};
53 vec_op::INT32Vec16 index_vec(index);
54 #endif
55 vec_op::FP32Vec16 gate_up_max_vec(7.0);
56 vec_op::FP32Vec16 up_min_vec(-7.0);
57 vec_op::FP32Vec16 alpha_vec(1.702);
58 vec_op::FP32Vec16 one_vec(1.0);
59
60 DEFINE_CPU_FUSED_MOE_EXP
61
62 for (int32_t m = 0; m < m_size; ++m) {
63 for (int32_t n = 0; n < n_size; n += 32) {
64 // Note: AdvSIMD does not support gather loads
65 #if defined(__aarch64__)
66 vec_op::FP32Vec16 gate_vec(vec_op::uninit);
67 vec_op::FP32Vec16 up_vec(vec_op::uninit);
68 vec_op::FP32Vec16::load_even_odd(input + n, gate_vec, up_vec);
69 #elif defined(__AVX512F__)
70 vec_op::FP32Vec16 gate_vec(input + n, index_vec);
71 vec_op::FP32Vec16 up_vec(input + n + 1, index_vec);
72 #else
73 alignas(64) float gate_values[16];
74 alignas(64) float up_values[16];
75 const float* interleaved = input + n;
76 for (int32_t i = 0; i < 16; ++i) {
77 gate_values[i] = interleaved[2 * i];
78 up_values[i] = interleaved[2 * i + 1];
79 }
80 vec_op::FP32Vec16 gate_vec(gate_values);
81 vec_op::FP32Vec16 up_vec(up_values);
82 #endif
83 gate_vec = gate_vec.min(gate_up_max_vec);
84 up_vec = up_vec.clamp(up_min_vec, gate_up_max_vec);
85 auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec * alpha_vec));
86 auto glu = gate_vec * sigmoid_vec;
87 auto gated_output_fp32 = (one_vec + up_vec) * glu;
88 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
89 gated_output.save(output + n / 2);
90 }
91 input += input_stride;
92 output += output_stride;
93 }
94 }
95
96 template <typename scalar_t>
97 void silu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
98 const int32_t m_size, const int32_t n_size,
99 const int32_t input_stride, const int32_t output_stride) {
100 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
101 const int32_t dim = n_size / 2;
102 float* __restrict__ gate = input;
103 float* __restrict__ up = input + dim;
104 vec_op::FP32Vec16 one_vec(1.0);
105
106 DEFINE_CPU_FUSED_MOE_EXP
107
108 for (int32_t m = 0; m < m_size; ++m) {
109 for (int32_t n = 0; n < dim; n += 16) {
110 vec_op::FP32Vec16 gate_vec(gate + n);
111 vec_op::FP32Vec16 up_vec(up + n);
112 auto sigmoid_vec = one_vec / (one_vec + fast_exp(-gate_vec));
113 auto silu = gate_vec * sigmoid_vec;
114 auto gated_output_fp32 = up_vec * silu;
115 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
116 gated_output.save(output + n);
117 }
118 gate += input_stride;
119 up += input_stride;
120 output += output_stride;
121 }
122 }
123
124 template <typename scalar_t>
125 void gelu_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
126 const int32_t m_size, const int32_t n_size,
127 const int32_t input_stride, const int32_t output_stride) {
128 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
129 const int32_t dim = n_size / 2;
130 float* __restrict__ gate = input;
131 float* __restrict__ up = input + dim;
132 vec_op::FP32Vec16 one_vec(1.0);
133 vec_op::FP32Vec16 w1_vec(M_SQRT1_2);
134 vec_op::FP32Vec16 w2_vec(0.5);
135
136 for (int32_t m = 0; m < m_size; ++m) {
137 for (int32_t n = 0; n < dim; n += 16) {
138 vec_op::FP32Vec16 gate_vec(gate + n);
139 vec_op::FP32Vec16 up_vec(up + n);
140 auto er_input_vec = gate_vec * w1_vec;
141 auto er_vec = er_input_vec.er();
142 auto gelu = gate_vec * w2_vec * (one_vec + er_vec);
143 auto gated_output_fp32 = up_vec * gelu;
144 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
145 gated_output.save(output + n);
146 }
147 gate += input_stride;
148 up += input_stride;
149 output += output_stride;
150 }
151 }
152
153 template <typename scalar_t>
154 void gelu_tanh_and_mul(float* __restrict__ input, scalar_t* __restrict__ output,
155 const int32_t m_size, const int32_t n_size,
156 const int32_t input_stride,
157 const int32_t output_stride) {
158 using scalar_vec_t = typename cpu_utils::VecTypeTrait<scalar_t>::vec_t;
159 const int32_t dim = n_size / 2;
160 float* __restrict__ gate = input;
161 float* __restrict__ up = input + dim;
162 vec_op::FP32Vec16 one_vec(1.0);
163 vec_op::FP32Vec16 w1_vec(0.7978845608028654);
164 vec_op::FP32Vec16 w2_vec(0.5);
165 vec_op::FP32Vec16 w3_vec(0.044715);
166
167 for (int32_t m = 0; m < m_size; ++m) {
168 for (int32_t n = 0; n < dim; n += 16) {
169 vec_op::FP32Vec16 gate_vec(gate + n);
170 vec_op::FP32Vec16 up_vec(up + n);
171 auto gate_pow3_vec = gate_vec * gate_vec * gate_vec;
172 auto inner_vec = w1_vec * (gate_vec + w3_vec * gate_pow3_vec);
173 // Note: can't use fast_exp form because diffusiongemma will generate
174 // wrong results
175 auto tanh_vec = inner_vec.tanh();
176 auto gelu_tanh = gate_vec * w2_vec * (one_vec + tanh_vec);
177 auto gated_output_fp32 = up_vec * gelu_tanh;
178 scalar_vec_t gated_output = scalar_vec_t(gated_output_fp32);
179 gated_output.save(output + n);
180 }
181 gate += input_stride;
182 up += input_stride;
183 output += output_stride;
184 }
185 }
186
187 template <typename scalar_t>
188 FORCE_INLINE void apply_gated_act(const FusedMOEAct act,
189 float* __restrict__ input,
190 scalar_t* __restrict__ output,
191 const int32_t m, const int32_t n,
192 const int32_t input_stride,
193 const int32_t output_stride) {
194 switch (act) {
195 case FusedMOEAct::SwigluOAIAndMul:
196 swigluoai_and_mul(input, output, m, n, input_stride, output_stride);
197 return;
198 case FusedMOEAct::SiluAndMul:
199 silu_and_mul(input, output, m, n, input_stride, output_stride);
200 return;
201 case FusedMOEAct::GeluAndMul:
202 gelu_and_mul(input, output, m, n, input_stride, output_stride);
203 return;
204 case FusedMOEAct::GeluTanhAndMul:
205 gelu_tanh_and_mul(input, output, m, n, input_stride, output_stride);
206 return;
207 default:
208 TORCH_CHECK(false, "Unsupported act type.");
209 }
210 }
211 } // namespace cpu_fused_moe_utils
212
213 #undef DEFINE_CPU_FUSED_MOE_EXP
214
215 #endif