@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 449 lines 11.5 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 #pragma once
7 #include "common.h"
8 #include "blas_gemm.h"
9
10 #if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
11 #define CPU_CAPABILITY_AVX512
12 #endif
13
14 // amx-bf16
15 #define TILE_M 16
16 #define TILE_N 16
17 #define TILE_K 32
18
19 // block size for AMX gemm
20 constexpr int block_size_m() {
21 return 2 * TILE_M;
22 }
23 constexpr int block_size_n() {
24 return 2 * TILE_N;
25 }
26
27 constexpr bool brgemm_supported() {
28 #if defined(CPU_CAPABILITY_AVX512)
29 return true;
30 #else
31 return false;
32 #endif
33 }
34
35 // define threshold using brgemm (intel AMX)
36 template <typename T>
37 inline bool can_use_brgemm(int M);
38 template <>
39 inline bool can_use_brgemm<at::BFloat16>(int M) {
40 return brgemm_supported() && M > 4;
41 }
42 template <>
43 inline bool can_use_brgemm<at::Half>(int M) {
44 return brgemm_supported();
45 }
46 // this requires PyTorch 2.7 or above
47 template <>
48 inline bool can_use_brgemm<int8_t>(int M) {
49 return brgemm_supported() && M > 4;
50 }
51
52 template <>
53 inline bool can_use_brgemm<uint8_t>(int M) {
54 return brgemm_supported() && M > 4;
55 }
56
57 template <>
58 inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
59 return brgemm_supported() && M > 4;
60 }
61
62 template <typename T>
63 inline bool can_use_brgemm(int M, int N) {
64 return brgemm_supported() && (can_use_brgemm<T>(M) || N != block_size_n());
65 }
66
67 // work around compiler internal error
68 #define BLOCK_K 128 // 4 * TILE_K
69
70 // adjust leading dimension size for K
71 template <typename T>
72 inline int64_t get_row_size(int64_t K) {
73 return K;
74 }
75
76 template <>
77 inline int64_t get_row_size<int8_t>(int64_t K) {
78 return K + sizeof(int32_t);
79 }
80
81 // uint8: mxfp4 or int4
82 template <>
83 inline int64_t get_row_size<uint8_t>(int64_t K) {
84 return K >> 1;
85 }
86
87 inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
88 return use_int8_w8a8 ? K + sizeof(int32_t) : K;
89 }
90
91 enum class CPUActMethod : int {
92 silu_and_mul = 0,
93 swiglu = 1,
94 gelu_and_mul = 2,
95 };
96
97 enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3, MXFP4 = 4, FP8_W8A8 = 5 };
98
99 constexpr bool operator==(CPUQuantMethod a, int64_t b) {
100 return static_cast<int64_t>(a) == b;
101 }
102
103 constexpr bool operator==(int64_t a, CPUQuantMethod b) {
104 return a == static_cast<int64_t>(b);
105 }
106
107 enum class CPUQuantAlgo : int64_t { AWQ = 0, GPTQ = 1 };
108
109 constexpr bool operator==(CPUQuantAlgo a, int64_t b) {
110 return static_cast<int64_t>(a) == b;
111 }
112
113 constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
114 return a == static_cast<int64_t>(b);
115 }
116
117 inline int64_t get_row_size(CPUQuantMethod quant, int64_t K) {
118 switch (quant) {
119 case CPUQuantMethod::INT8_W8A8:
120 return K + sizeof(int32_t);
121 case CPUQuantMethod::MXFP4:
122 return K >> 1;
123 default:
124 return K;
125 }
126 }
127
128 inline int64_t get_4bit_block_k_size(int64_t group_size) {
129 return group_size > 128 ? 128 : group_size;
130 }
131
132 // pack weight to vnni format
133 at::Tensor convert_weight_packed(at::Tensor& weight);
134
135 // pack scale to blocked format for mxfp4
136 at::Tensor convert_scale_packed(at::Tensor& scale);
137
138 // pack weight to vnni format for int4
139 std::tuple<at::Tensor, at::Tensor, at::Tensor>
140 convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
141
142 // moe implementations for int8 w8a8
143 template <typename scalar_t>
144 void fused_experts_int8_kernel_impl(
145 scalar_t* __restrict__ output,
146 scalar_t* __restrict__ ic1,
147 scalar_t* __restrict__ ic2,
148 uint8_t* __restrict__ A_tmp,
149 float* __restrict__ C_tmp,
150 uint8_t* __restrict__ Aq_tmp,
151 float* __restrict__ As_tmp,
152 const scalar_t* __restrict__ input,
153 const int8_t* __restrict__ packed_w1,
154 const int8_t* __restrict__ packed_w2,
155 const float* __restrict__ w1s,
156 const float* __restrict__ w2s,
157 const float* __restrict__ topk_weights,
158 const int32_t* __restrict__ sorted_ids,
159 const int32_t* __restrict__ expert_ids,
160 const int32_t* __restrict__ offsets,
161 int64_t M,
162 int64_t N,
163 int64_t K,
164 int64_t E,
165 int64_t topk,
166 int64_t num_tokens_post_pad);
167
168 // moe implementations for fp8 w8a16 and mxfp4
169 template <typename scalar_t, typename packed_t, typename param_t, bool is_mxfp4>
170 void fused_experts_fp_kernel_impl(
171 scalar_t* __restrict__ output,
172 scalar_t* __restrict__ ic0,
173 scalar_t* __restrict__ ic1,
174 scalar_t* __restrict__ ic2,
175 scalar_t* __restrict__ A_tmp,
176 scalar_t* __restrict__ B_tmp,
177 float* __restrict__ C_tmp,
178 const scalar_t* __restrict__ input,
179 const packed_t* __restrict__ packed_w1,
180 const packed_t* __restrict__ packed_w2,
181 const float* __restrict__ w1_bias,
182 const float* __restrict__ w2_bias,
183 const param_t* __restrict__ w1s,
184 const param_t* __restrict__ w2s,
185 int64_t block_size_N,
186 int64_t block_size_K,
187 const float* __restrict__ topk_weights,
188 const int32_t* __restrict__ sorted_ids,
189 const int32_t* __restrict__ expert_ids,
190 const int32_t* __restrict__ offsets,
191 int64_t M,
192 int64_t N,
193 int64_t K,
194 int64_t E,
195 int64_t topk,
196 int64_t num_tokens_post_pad,
197 float alpha,
198 float limit,
199 CPUActMethod act_func,
200 bool with_bias);
201
202 // shared expert implementation for int8 w8a8
203 template <typename scalar_t>
204 void shared_expert_int8_kernel_impl(
205 scalar_t* __restrict__ output,
206 scalar_t* __restrict__ ic1,
207 float* __restrict__ C_tmp,
208 uint8_t* __restrict__ Aq_tmp,
209 float* __restrict__ As_tmp,
210 const scalar_t* __restrict__ input,
211 const int8_t* __restrict__ packed_w1,
212 const int8_t* __restrict__ packed_w2,
213 const float* __restrict__ w1s,
214 const float* __restrict__ w2s,
215 const scalar_t* __restrict__ fused_experts_out,
216 float routed_scaling_factor,
217 int64_t M,
218 int64_t N,
219 int64_t K);
220
221 template <typename scalar_t>
222 void fused_experts_int4_w4a8_kernel_impl(
223 scalar_t* __restrict__ output,
224 scalar_t* __restrict__ ic0,
225 scalar_t* __restrict__ ic1,
226 scalar_t* __restrict__ ic2,
227 uint8_t* __restrict__ A_tmp,
228 uint8_t* __restrict__ Aq_tmp,
229 float* __restrict__ As_tmp,
230 int32_t* __restrict__ Azp_tmp,
231 float* __restrict__ C_tmp,
232 int8_t* __restrict__ dqB_tmp,
233 const scalar_t* __restrict__ input,
234 const uint8_t* __restrict__ packed_w1,
235 const uint8_t* __restrict__ packed_w2,
236 const int8_t* __restrict__ w1z,
237 const int8_t* __restrict__ w2z,
238 const float* __restrict__ w1s,
239 const float* __restrict__ w2s,
240 int group_size,
241 const float* __restrict__ topk_weights,
242 const int32_t* __restrict__ sorted_ids,
243 const int32_t* __restrict__ expert_ids,
244 const int32_t* __restrict__ offsets,
245 int64_t M,
246 int64_t N,
247 int64_t K,
248 int64_t E,
249 int64_t topk,
250 int64_t num_tokens_post_pad);
251
252 template <typename scalar_t>
253 void shared_expert_fp8_kernel_impl(
254 scalar_t* __restrict__ output,
255 scalar_t* __restrict__ ic0,
256 scalar_t* __restrict__ ic1,
257 scalar_t* __restrict__ B_tmp,
258 float* __restrict__ C_tmp,
259 const scalar_t* __restrict__ input,
260 const at::Float8_e4m3fn* __restrict__ packed_w1,
261 const at::Float8_e4m3fn* __restrict__ packed_w2,
262 const float* __restrict__ w1s,
263 const float* __restrict__ w2s,
264 int64_t block_size_N,
265 int64_t block_size_K,
266 const scalar_t* __restrict__ fused_experts_out,
267 float routed_scaling_factor,
268 int64_t M,
269 int64_t N,
270 int64_t K);
271
272 // tinygemm interface
273 template <typename scalar_t>
274 void tinygemm_kernel(
275 const scalar_t* __restrict__ A,
276 const scalar_t* __restrict__ B,
277 scalar_t* __restrict__ C,
278 float* __restrict__ Ctmp,
279 int64_t M,
280 int64_t N,
281 int64_t K,
282 int64_t lda,
283 int64_t ldb,
284 int64_t ldc,
285 bool brg);
286
287 template <typename scalar_t>
288 void tinygemm_kernel(
289 const uint8_t* __restrict__ A,
290 const int8_t* __restrict__ B,
291 scalar_t* __restrict__ C,
292 int32_t* __restrict__ Ctmp,
293 const float* __restrict__ As,
294 const float* __restrict__ Bs,
295 int64_t M,
296 int64_t N,
297 int64_t K,
298 int64_t lda,
299 int64_t ldb,
300 int64_t ldc,
301 bool brg);
302
303 // block quantization
304 template <typename scalar_t>
305 void tinygemm_kernel(
306 const scalar_t* __restrict__ A,
307 const at::Float8_e4m3fn* __restrict__ B,
308 scalar_t* __restrict__ C,
309 scalar_t* __restrict__ Btmp,
310 float* __restrict__ Ctmp,
311 const float* __restrict__ Bbias,
312 const float* __restrict__ scale,
313 int64_t M,
314 int64_t N,
315 int64_t K,
316 int64_t lda,
317 int64_t ldb,
318 int64_t ldc,
319 bool brg,
320 int64_t block_size_K,
321 bool do_unpack = true);
322
323 // per tensor quantization
324 template <typename scalar_t>
325 void tinygemm_kernel(
326 const scalar_t* __restrict__ A,
327 const at::Float8_e4m3fn* __restrict__ B,
328 scalar_t* __restrict__ C,
329 scalar_t* __restrict__ Btmp,
330 float* __restrict__ Ctmp,
331 float scale,
332 int64_t M,
333 int64_t N,
334 int64_t K,
335 int64_t lda,
336 int64_t ldb,
337 int64_t ldc,
338 bool brg);
339
340 // mxfp4
341 template <typename scalar_t>
342 void tinygemm_kernel(
343 const scalar_t* __restrict__ A,
344 const uint8_t* __restrict__ B,
345 scalar_t* __restrict__ C,
346 scalar_t* __restrict__ Btmp,
347 float* __restrict__ Ctmp,
348 const float* __restrict__ Bbias,
349 const uint8_t* __restrict__ scale,
350 int64_t M,
351 int64_t N,
352 int64_t K,
353 int64_t lda,
354 int64_t ldb,
355 int64_t ldc,
356 bool brg,
357 int64_t block_size_K,
358 bool do_unpack = true);
359
360 template <typename scalar_t>
361 void tinygemm_kernel(
362 scalar_t* C,
363 float* C_temp,
364 const uint8_t* A,
365 const float* scales_a,
366 const int32_t* qzeros_a,
367 const uint8_t* B,
368 const float* scales_b,
369 const int8_t* qzeros_b,
370 const int32_t* compensation,
371 int8_t* dqB_tmp,
372 int64_t M,
373 int64_t K,
374 int64_t lda,
375 int64_t ldc_f,
376 int64_t ldc_s,
377 bool store_out,
378 bool use_brgemm);
379
380 // mxfp4
381 template <typename scalar_t>
382 void tinygemm_kernel(
383 const scalar_t* __restrict__ A,
384 const uint8_t* __restrict__ B,
385 scalar_t* __restrict__ C,
386 scalar_t* __restrict__ Btmp,
387 float* __restrict__ Ctmp,
388 const uint8_t* __restrict__ scale,
389 int64_t M,
390 int64_t N,
391 int64_t K,
392 int64_t lda,
393 int64_t ldb,
394 int64_t ldc,
395 bool brg,
396 int64_t block_size_K,
397 bool do_unpack = true);
398
399 // FP8 W8A8 tinygemm kernel (defined in gemm_fp8_w8a8.cpp)
400 void tinygemm_kernel(
401 float* C,
402 const at::Float8_e4m3fn* A,
403 const float* scales_a,
404 const at::Float8_e4m3fn* B,
405 const float* scales_b,
406 int64_t M,
407 int64_t K,
408 int64_t lda,
409 int64_t ldc,
410 int64_t ldsa,
411 float* ukernel_buf,
412 at::BFloat16* dqA_buf,
413 at::BFloat16* dqB_buf);
414
415 // FP8 W8A8 quantize activation (defined in gemm_fp8_w8a8.cpp)
416 std::tuple<at::Tensor, at::Tensor> quantize_fp8e4m3_vec(
417 const at::Tensor& t,
418 bool channelwise,
419 c10::optional<at::Tensor> scale_opt);
420
421 // FP8 W8A8 fused MoE kernel (defined in moe_fp8_w8a8.cpp)
422 template <typename scalar_t>
423 void fused_experts_fp8_a8_kernel_impl(
424 scalar_t* __restrict__ output,
425 scalar_t* __restrict__ ic0,
426 scalar_t* __restrict__ ic1,
427 scalar_t* __restrict__ ic2,
428 at::Float8_e4m3fn* __restrict__ A_tmp,
429 scalar_t* __restrict__ B_tmp,
430 float* __restrict__ C_tmp,
431 float* __restrict__ Ukernel_tmp,
432 const at::Float8_e4m3fn* __restrict__ input,
433 const at::Float8_e4m3fn* __restrict__ packed_w1,
434 const at::Float8_e4m3fn* __restrict__ packed_w2,
435 const float* __restrict__ As,
436 const float* __restrict__ w1s,
437 const float* __restrict__ w2s,
438 int64_t block_size_N,
439 int64_t block_size_K,
440 const float* __restrict__ topk_weights,
441 const int32_t* __restrict__ sorted_ids,
442 const int32_t* __restrict__ expert_ids,
443 const int32_t* __restrict__ offsets,
444 int64_t M,
445 int64_t N,
446 int64_t K,
447 int64_t E,
448 int64_t topk,
449 int64_t num_tokens_post_pad);