@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
cpp 413 lines 17.1 KB
Raw
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 // clang-format off
5
6 #include "common.h"
7 #include "gemm.h"
8 #include "moe.h"
9
10 template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
11 void fused_experts_fp_kernel_impl(
12 scalar_t* __restrict__ output,
13 scalar_t* __restrict__ ic0,
14 scalar_t* __restrict__ ic1,
15 scalar_t* __restrict__ ic2,
16 scalar_t* __restrict__ A_tmp,
17 scalar_t* __restrict__ B_tmp,
18 float* __restrict__ C_tmp,
19 const scalar_t* __restrict__ input,
20 const packed_t* __restrict__ packed_w1,
21 const packed_t* __restrict__ packed_w2,
22 const float* __restrict__ w1_bias,
23 const float* __restrict__ w2_bias,
24 const param_t* __restrict__ w1s,
25 const param_t* __restrict__ w2s,
26 int64_t block_size_N,
27 int64_t block_size_K,
28 const float* __restrict__ topk_weights,
29 const int32_t* __restrict__ sorted_ids,
30 const int32_t* __restrict__ expert_ids,
31 const int32_t* __restrict__ offsets,
32 int64_t M,
33 int64_t N,
34 int64_t K,
35 int64_t E,
36 int64_t topk,
37 int64_t num_tokens_post_pad,
38 float alpha,
39 float limit,
40 CPUActMethod act_func,
41 bool with_bias) {
42 constexpr int64_t BLOCK_M = block_size_m();
43 constexpr int64_t BLOCK_N = block_size_n();
44
45 // stage 1: intermediate_cache0 = hidden_states @ w1
46 const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
47 const int64_t NB = div_up(2 * N, BLOCK_N);
48 int64_t scale_size_N = div_up(2 * N, block_size_N);
49 int64_t scale_size_K = div_up(K, block_size_K);
50 int64_t blocks_n_per_group = block_size_N / BLOCK_N;
51 std::function<int64_t(int64_t)> scale_offset_per_block;
52 if constexpr (is_mxfp4) {
53 scale_offset_per_block = [&](int64_t a) { return a * BLOCK_N; };
54 } else {
55 scale_offset_per_block = [&](int64_t a) { return a / blocks_n_per_group; };
56 }
57
58 const int64_t packed_K = get_row_size<packed_t>(K);
59
60 const int64_t stride_e = 2 * N * packed_K;
61 const int64_t stride_n = packed_K;
62
63 int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
64
65 // here we only parallel on half of 2N to fuse silu_and_mul with gemm
66 parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
67 // get local pointers
68 int tid = get_thread_num();
69 scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
70 bool use_brgemm = false;
71
72 loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
73 int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
74
75 // B shape [K, n_size] in vnni format
76 int32_t expert_id = expert_ids[mb];
77 const packed_t* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
78 const param_t* __restrict__ Bs =
79 w1s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
80 const float* __restrict__ B_bias = with_bias ? w1_bias + expert_id * 2 * N + nb * BLOCK_N : nullptr;
81
82 // do unpacking for the first row or a new expert
83 int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
84 bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
85
86 int64_t m_size = offsets[mb + 1] - offsets[mb];
87 if (m_size == 0) {
88 return;
89 }
90 const bool brg = can_use_brgemm<packed_t>(m_size, n_size);
91 use_brgemm |= brg;
92
93 if (nb_offset == 0) {
94 // 1.a load A
95 const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
96 for (int64_t m = 0; m < m_size; ++m) {
97 int32_t index = A_ids[m] / topk;
98 copy_stub(A + m * K, input + index * K, K);
99 }
100 }
101
102 const int64_t offset = offsets[mb];
103 tinygemm_kernel<scalar_t>(
104 /* A */ A,
105 /* B */ B,
106 /* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
107 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
108 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
109 /* Bbias */ B_bias,
110 /* scale */ Bs,
111 /* M */ m_size,
112 /* N */ n_size,
113 /* K */ K,
114 /* lda */ K,
115 /* ldb */ n_size,
116 /* ldc */ 2 * N,
117 /* brg */ brg,
118 /* block_size_K */ block_size_K,
119 /* do_unpack */ do_unpack);
120 });
121
122 if (use_brgemm) {
123 at::native::cpublas::brgemm_release();
124 }
125 });
126
127 // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
128 if (act_func == CPUActMethod::silu_and_mul) {
129 at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
130 for (int64_t m = begin; m < end; ++m) {
131 silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
132 }
133 });
134 } else if (act_func == CPUActMethod::swiglu) {
135 at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
136 for (int64_t m = begin; m < end; ++m) {
137 clamp_sigmoid_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, N / 2, alpha, limit);
138 clamp_sigmoid_and_mul_stub(ic1 + m * N + N / 2, ic0 + m * 2 * N + N, N / 2, alpha, limit);
139 }
140 });
141 }
142 // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
143 // w2 : [E, K, N] as [E, OC, IC]
144 const int64_t OC = K; // rename K as OC
145 const int64_t IC = N; // rename N as IC
146 const int64_t MB2 = MB;
147 const int64_t NB2 = div_up(OC, BLOCK_N);
148 scale_size_N = div_up(K, block_size_N);
149 scale_size_K = div_up(N, block_size_K);
150 const int64_t packed_IC = get_row_size<packed_t>(IC);
151 const int64_t stride_e2 = OC * packed_IC;
152 const int64_t stride_oc = packed_IC;
153
154 // parallel on [MB2, NB2]
155 parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
156 int tid = get_thread_num();
157 alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
158 bool use_brgemm = false;
159
160 loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
161 int64_t m_size = offsets[mb + 1] - offsets[mb];
162 if (m_size == 0) {
163 return;
164 }
165 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
166 const bool brg = can_use_brgemm<packed_t>(m_size, n_size);
167 use_brgemm |= brg;
168
169 // A ptr from ic1 of [M * topk, N] in sorted order
170 // so as to avoid copy A to tmp buffer again
171 const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
172 const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
173
174 // B shape [IC, n_size] in vnni format
175 int32_t expert_id = expert_ids[mb];
176 const packed_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
177 const param_t* __restrict__ Bs =
178 w2s + expert_id * scale_size_N * scale_size_K + scale_offset_per_block(nb) * scale_size_K;
179 const float* __restrict__ B_bias = with_bias ? w2_bias + expert_id * OC + nb * BLOCK_N : nullptr;
180
181 // do unpacking for the first row or a new expert
182 int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
183 bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
184
185 tinygemm_kernel<scalar_t>(
186 /* A */ A,
187 /* B */ B,
188 /* C */ C,
189 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
190 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
191 /* Bbias */ B_bias,
192 /* scale */ Bs,
193 /* M */ m_size,
194 /* N */ n_size,
195 /* K */ IC,
196 /* lda */ IC,
197 /* ldb */ n_size,
198 /* ldc */ BLOCK_N,
199 /* brg */ brg,
200 /* block_size_K */ block_size_K,
201 /* do_unpack */ do_unpack);
202
203 // 2.b copy from C to ic2 in original order
204 // and also mul topk_weights in float32
205 for (int64_t m = 0; m < m_size; ++m) {
206 int32_t index = A_ids[m];
207 float weight = topk_weights[index];
208 copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
209 }
210 });
211
212 if (use_brgemm) {
213 at::native::cpublas::brgemm_release();
214 }
215 });
216 // stage 3: out = intermediate_cache2.sum(dim=1)
217 // from [M, topk, K] to [M, K]
218 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
219 for (int64_t m = begin; m < end; ++m) {
220 sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
221 }
222 });
223 }
224
225 #define INSTANTIATE_MOE_FP_TEMPLATE(TYPE1, TYPE2, TYPE3, IS_MXFP4) \
226 template void fused_experts_fp_kernel_impl<TYPE1, TYPE2, TYPE3, IS_MXFP4>( \
227 TYPE1* __restrict__ output, \
228 TYPE1* __restrict__ ic0, \
229 TYPE1* __restrict__ ic1, \
230 TYPE1* __restrict__ ic2, \
231 TYPE1* __restrict__ A_tmp, \
232 TYPE1* __restrict__ B_tmp, \
233 float* __restrict__ C_tmp, \
234 const TYPE1* __restrict__ input, \
235 const TYPE2* __restrict__ packed_w1, \
236 const TYPE2* __restrict__ packed_w2, \
237 const float* __restrict__ w1_bias, \
238 const float* __restrict__ w2_bias, \
239 const TYPE3* __restrict__ w1s, \
240 const TYPE3* __restrict__ w2s, \
241 int64_t block_size_N, \
242 int64_t block_size_K, \
243 const float* __restrict__ topk_weights, \
244 const int32_t* __restrict__ sorted_ids, \
245 const int32_t* __restrict__ expert_ids, \
246 const int32_t* __restrict__ offsets, \
247 int64_t M, \
248 int64_t N, \
249 int64_t K, \
250 int64_t E, \
251 int64_t topk, \
252 int64_t num_tokens_post_pad, \
253 float alpha, \
254 float limit, \
255 CPUActMethod act_func, \
256 bool with_bias)
257
258 INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float, false);
259 INSTANTIATE_MOE_FP_TEMPLATE(at::Half, at::Float8_e4m3fn, float, false);
260 INSTANTIATE_MOE_FP_TEMPLATE(at::BFloat16, uint8_t, uint8_t, true);
261 INSTANTIATE_MOE_FP_TEMPLATE(at::Half, uint8_t, uint8_t, true);
262
263 template <typename scalar_t>
264 void shared_expert_fp8_kernel_impl(
265 scalar_t* __restrict__ output,
266 scalar_t* __restrict__ ic0,
267 scalar_t* __restrict__ ic1,
268 scalar_t* __restrict__ B_tmp,
269 float* __restrict__ C_tmp,
270 const scalar_t* __restrict__ input,
271 const at::Float8_e4m3fn* __restrict__ packed_w1,
272 const at::Float8_e4m3fn* __restrict__ packed_w2,
273 const float* __restrict__ w1s,
274 const float* __restrict__ w2s,
275 int64_t block_size_N,
276 int64_t block_size_K,
277 const scalar_t* __restrict__ fused_experts_out,
278 float routed_scaling_factor,
279 int64_t M,
280 int64_t N,
281 int64_t K) {
282 constexpr int64_t BLOCK_M = block_size_m();
283 constexpr int64_t BLOCK_N = block_size_n();
284
285 // stage 1: intermediate_cache0 = hidden_states @ w1
286 const int64_t MB = div_up(M, BLOCK_M);
287 const int64_t NB = div_up(2 * N, BLOCK_N);
288 int64_t scale_size_K = div_up(K, block_size_K);
289 int64_t blocks_n_per_group = block_size_N / BLOCK_N;
290
291 const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
292 const bool apply_scaling_factor = fused_experts_out != nullptr;
293
294 int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
295
296 parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
297 int tid = get_thread_num();
298
299 loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
300 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
301 int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
302
303 // do unpacking for the first row
304 bool do_unpack = (mb == mb0);
305
306 tinygemm_kernel<scalar_t>(
307 /* A */ input + mb * BLOCK_M * K,
308 /* B */ packed_w1 + nb * BLOCK_N * K,
309 /* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
310 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
311 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
312 /* Bbias */ nullptr,
313 /* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
314 /* M */ m_size,
315 /* N */ n_size,
316 /* K */ K,
317 /* lda */ K,
318 /* ldb */ n_size,
319 /* ldc */ 2 * N,
320 /* brg */ use_brgemm,
321 /* block_size_K */ block_size_K,
322 /* do_unpack */ do_unpack);
323 });
324
325 if (use_brgemm) {
326 at::native::cpublas::brgemm_release();
327 }
328 });
329
330 // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
331 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
332 for (int64_t m = begin; m < end; ++m) {
333 silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
334 }
335 });
336
337 // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
338 // w2 : [K, N] as [OC, IC]
339 const int64_t OC = K; // rename K as OC
340 const int64_t IC = N; // rename N as IC
341 const int64_t MB2 = MB;
342 const int64_t NB2 = div_up(K, BLOCK_N);
343 scale_size_K = div_up(N, block_size_K);
344
345 // parallel on [MB2, NB2]
346 parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
347 int tid = get_thread_num();
348 alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
349
350 loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
351 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
352 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
353
354 // do unpacking for the first row
355 bool do_unpack = (mb == mb0);
356
357 // 2.a gemm: C = A @ B
358 tinygemm_kernel<scalar_t>(
359 /* A */ ic1 + mb * BLOCK_M * N,
360 /* B */ packed_w2 + nb * BLOCK_N * N,
361 /* C */ C,
362 /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
363 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
364 /* Bbias */ nullptr,
365 /* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
366 /* M */ m_size,
367 /* N */ n_size,
368 /* K */ IC,
369 /* lda */ IC,
370 /* ldb */ n_size,
371 /* ldc */ BLOCK_N,
372 /* brg */ use_brgemm,
373 /* block_size_K */ block_size_K,
374 /* do_unpack */ do_unpack);
375
376 // 2.b copy from C to output and add fused_experts_out
377 scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
378 const scalar_t* __restrict__ fused_out =
379 apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
380 for (int64_t m = 0; m < m_size; ++m) {
381 const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
382 add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
383 }
384 });
385 });
386
387 if (use_brgemm) {
388 at::native::cpublas::brgemm_release();
389 }
390 }
391
392 #define INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(TYPE) \
393 template void shared_expert_fp8_kernel_impl<TYPE>( \
394 TYPE* __restrict__ output, \
395 TYPE* __restrict__ ic0, \
396 TYPE* __restrict__ ic1, \
397 TYPE* __restrict__ B_tmp, \
398 float* __restrict__ C_tmp, \
399 const TYPE* __restrict__ input, \
400 const at::Float8_e4m3fn* __restrict__ packed_w1, \
401 const at::Float8_e4m3fn* __restrict__ packed_w2, \
402 const float* __restrict__ w1s, \
403 const float* __restrict__ w2s, \
404 int64_t block_size_N, \
405 int64_t block_size_K, \
406 const TYPE* __restrict__ fused_experts_out, \
407 float routed_scaling_factor, \
408 int64_t M, \
409 int64_t N, \
410 int64_t K)
411
412 INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::BFloat16);
413 INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::Half);