@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / b3945cc316
Fork 0

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 4, 03:45 UTC

[CPU] Bump up to the latest CPU kernels (#41924)

Signed-off-by: jiang1.li <jiang1.li@intel.com>

Li, Jiang committed May 7, 2026 at 20:45 UTC b3945cc3160119111927816841950715c5410708
25 files changed +6225 -1642
.buildkite/hardware_tests/cpu.yaml
+1
@@ -65,6 +65,7 @@ steps:
65 - vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_int8.py
66 - vllm/model_executor/layers/quantization/kernels/scaled_mm/cpu.py
67 - vllm/model_executor/layers/quantization/kernels/mixed_precision/cpu.py
68 + - vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
69 - tests/quantization/test_compressed_tensors.py
70 - tests/quantization/test_cpu_wna16.py
71 commands:
cmake/cpu_extension.cmake
+3
@@ -409,12 +409,15 @@ endif()
409
410 if (ENABLE_X86_ISA)
411 set(VLLM_EXT_SRC_SGL
412 + "csrc/cpu/sgl-kernels/fla.cpp"
413 + "csrc/cpu/sgl-kernels/conv.cpp"
414 "csrc/cpu/sgl-kernels/gemm.cpp"
415 "csrc/cpu/sgl-kernels/gemm_int8.cpp"
416 "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
417 "csrc/cpu/sgl-kernels/gemm_int4.cpp"
418 "csrc/cpu/sgl-kernels/moe.cpp"
419 "csrc/cpu/sgl-kernels/moe_int8.cpp"
420 + "csrc/cpu/sgl-kernels/moe_int4.cpp"
421 "csrc/cpu/sgl-kernels/moe_fp8.cpp")
422
423 set(VLLM_EXT_SRC_AVX512
csrc/cpu/sgl-kernels/common.h
+247 -68
@@ -1,13 +1,12 @@
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
8 #include <ATen/ATen.h>
9 #include <ATen/Parallel.h>
8 -#include <ATen/record_function.h>
9 -
10 -// clang-format off
10
11 #if defined(_OPENMP)
12 #include <omp.h>
@@ -16,40 +15,157 @@
15 namespace {
16
17 // dispatch bool
19 -#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \
20 - [&] { \
21 - if (BOOL_V) { \
22 - constexpr bool BOOL_NAME = true; \
23 - return __VA_ARGS__(); \
24 - } else { \
25 - constexpr bool BOOL_NAME = false; \
26 - return __VA_ARGS__(); \
27 - } \
18 +#define AT_DISPATCH_BOOL(BOOL_V, BOOL_NAME, ...) \
19 + [&] { \
20 + if (BOOL_V) { \
21 + constexpr bool BOOL_NAME = true; \
22 + return __VA_ARGS__(); \
23 + } else { \
24 + constexpr bool BOOL_NAME = false; \
25 + return __VA_ARGS__(); \
26 + } \
27 + }()
28 +
29 +#define AT_DISPATCH_BOOL2(BOOL_V1, BOOL_NAME1, BOOL_V2, BOOL_NAME2, ...) \
30 + [&] { \
31 + if (BOOL_V1) { \
32 + constexpr bool BOOL_NAME1 = true; \
33 + if (BOOL_V2) { \
34 + constexpr bool BOOL_NAME2 = true; \
35 + return __VA_ARGS__(); \
36 + } else { \
37 + constexpr bool BOOL_NAME2 = false; \
38 + return __VA_ARGS__(); \
39 + } \
40 + } else { \
41 + constexpr bool BOOL_NAME1 = false; \
42 + if (BOOL_V2) { \
43 + constexpr bool BOOL_NAME2 = true; \
44 + return __VA_ARGS__(); \
45 + } else { \
46 + constexpr bool BOOL_NAME2 = false; \
47 + return __VA_ARGS__(); \
48 + } \
49 + } \
50 + }()
51 +
52 +// dispatch: bfloat16, float16, int8_t, fp8_e4m3, uint8_t(mxfp4/int4)
53 +#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...) \
54 + [&] { \
55 + switch (TYPE) { \
56 + case at::ScalarType::BFloat16: { \
57 + using packed_t = at::BFloat16; \
58 + return __VA_ARGS__(); \
59 + } \
60 + case at::ScalarType::Half: { \
61 + using packed_t = at::Half; \
62 + return __VA_ARGS__(); \
63 + } \
64 + case at::ScalarType::Char: { \
65 + using packed_t = int8_t; \
66 + return __VA_ARGS__(); \
67 + } \
68 + case at::ScalarType::Float8_e4m3fn: { \
69 + using packed_t = at::Float8_e4m3fn; \
70 + return __VA_ARGS__(); \
71 + } \
72 + case at::ScalarType::Byte: { \
73 + using packed_t = uint8_t; \
74 + return __VA_ARGS__(); \
75 + } \
76 + default: \
77 + TORCH_CHECK(false, "Unsupported floating data type.\n"); \
78 + } \
79 + }()
80 +
81 +// Helper MICRO for CPU_DISPATCH_FLOATING_TYPES_EXT:
82 +// TYPE1: the primary dtype (input, output, weight);
83 +// TYPE2: defined as PARAM_T input
84 +#define CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, PARAM_T, ...) \
85 + switch (TYPE1) { \
86 + case at::ScalarType::BFloat16: { \
87 + using scalar_t = at::BFloat16; \
88 + using param_t = PARAM_T; \
89 + return __VA_ARGS__(); \
90 + } \
91 + case at::ScalarType::Half: { \
92 + using scalar_t = at::Half; \
93 + using param_t = PARAM_T; \
94 + return __VA_ARGS__(); \
95 + } \
96 + case at::ScalarType::Float: { \
97 + using scalar_t = float; \
98 + using param_t = PARAM_T; \
99 + return __VA_ARGS__(); \
100 + } \
101 + default: \
102 + TORCH_CHECK(false, "Unsupported floating data type."); \
103 + }
104 +
105 +// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
106 +// TYPE1: the primary dtype (input, output, weight);
107 +// TYPE2: defined as PARAM_T input
108 +#define CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, PARAM_T, ...) \
109 + switch (TYPE1) { \
110 + case at::ScalarType::BFloat16: { \
111 + using scalar_t = at::BFloat16; \
112 + using param_t = PARAM_T; \
113 + return __VA_ARGS__(); \
114 + } \
115 + case at::ScalarType::Half: { \
116 + using scalar_t = at::Half; \
117 + using param_t = PARAM_T; \
118 + return __VA_ARGS__(); \
119 + } \
120 + default: \
121 + TORCH_CHECK(false, "Unsupported floating data type."); \
122 + }
123 +
124 +// Helper MICRO for CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT:
125 +// TYPE1: the dtype both for scalar_t and param_t
126 +#define CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, ...) \
127 + switch (TYPE1) { \
128 + case at::ScalarType::BFloat16: { \
129 + using scalar_t = at::BFloat16; \
130 + using param_t = at::BFloat16; \
131 + return __VA_ARGS__(); \
132 + } \
133 + case at::ScalarType::Half: { \
134 + using scalar_t = at::Half; \
135 + using param_t = at::Half; \
136 + return __VA_ARGS__(); \
137 + } \
138 + default: \
139 + TORCH_CHECK(false, "Unsupported reduced floating data type."); \
140 + }
141 +
142 +// dispatch with mixed dtypes (TYPE1, TYPE2):
143 +// TYPE1: the primary dtype (input, output, weight);
144 +// TYPE2: the secondary dtype (bias, etc.).
145 +#define CPU_DISPATCH_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
146 + [&] { \
147 + if (TYPE2 == at::kFloat) { \
148 + CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, float, __VA_ARGS__) \
149 + } else if (TYPE2 == at::ScalarType::BFloat16) { \
150 + CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::BFloat16, __VA_ARGS__) \
151 + } else if (TYPE2 == at::ScalarType::Half) { \
152 + CPU_DISPATCH_TYPE1_WITH_PARAM(TYPE1, at::Half, __VA_ARGS__) \
153 + } else { \
154 + TORCH_CHECK(false, "Unsupported floating data type."); \
155 + } \
156 }()
157
30 -// dispatch: bfloat16, float16, int8_t, fp8_e4m3
31 -#define CPU_DISPATCH_PACKED_TYPES(TYPE, ...) \
32 - [&] { \
33 - switch (TYPE) { \
34 - case at::ScalarType::BFloat16 : { \
35 - using packed_t = at::BFloat16; \
36 - return __VA_ARGS__(); \
37 - } \
38 - case at::ScalarType::Half: { \
39 - using packed_t = at::Half; \
40 - return __VA_ARGS__(); \
41 - } \
42 - case at::ScalarType::Char : { \
43 - using packed_t = int8_t; \
44 - return __VA_ARGS__(); \
45 - } \
46 - case at::ScalarType::Float8_e4m3fn : { \
47 - using packed_t = at::Float8_e4m3fn; \
48 - return __VA_ARGS__(); \
49 - } \
50 - default: \
51 - TORCH_CHECK(false, "Unsupported floating data type.\n"); \
52 - } \
158 +// dispatch with mixed dtypes (reduced one, no float for TYPE1) (TYPE1, TYPE2):
159 +// TYPE1: the primary dtype (input, output, weight);
160 +// TYPE2: the secondary dtype (bias, etc.).
161 +#define CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(TYPE1, TYPE2, ...) \
162 + [&] { \
163 + if (TYPE2 == at::kFloat) { \
164 + CPU_DISPATCH_TYPE1_WITH_PARAM_REDUCED(TYPE1, float, __VA_ARGS__) \
165 + } else { \
166 + TORCH_CHECK(TYPE1 == TYPE2); \
167 + CPU_DISPATCH_TYPE1_WITH_SAME_PARAM_REDUCED(TYPE1, __VA_ARGS__) \
168 + } \
169 }()
170
171 #define UNUSED(x) (void)(x)
@@ -70,13 +186,51 @@ namespace {
186 #define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor")
187
188 #define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
189 +#define CHECK_GT(a, b) TORCH_CHECK((a) > (b), "CHECK_GT(" #a ", " #b ") failed. ", a, " vs ", b)
190 +#define CHECK_GE(a, b) TORCH_CHECK((a) >= (b), "CHECK_GE(" #a ", " #b ") failed. ", a, " vs ", b)
191 +
192 +template <bool is_only_lastdim_contiguous>
193 +static inline void CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const at::IntArrayRef sizes, at::ScalarType st) {
194 + TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
195 + TORCH_CHECK(tensor.scalar_type() == st, "Input tensor dtype mismatch");
196 + if constexpr (is_only_lastdim_contiguous) {
197 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
198 + } else {
199 + CHECK_INPUT(tensor);
200 + }
201 +}
202
74 -// parallel routines
203 +// [NB] Parallel Routines
204 +//
205 +// * at::parallel_for - applies for most of generic use cases, this will be compiled
206 +// against openmp in default torch release.
207 +//
208 +// * parallel_for - same function as above, can choose payload partition scheme in
209 +// balance211.
210 +//
211 +// * parallel_2d - parallel for 2 dimensions, used in GEMM, etc.
212 +// this one will do payload balance across 2 dimensions.
213 +//
214 +
215 +// grain size for each thread
216 constexpr int GRAIN_SIZE = 1024;
217
218 template <typename T, typename std::enable_if<std::is_integral<T>::value, int>::type = 0>
78 -inline T div_up(T x, T y) { return (x + y - 1) / y; }
219 +inline T div_up(T x, T y) {
220 + return (x + y - 1) / y;
221 +}
222 +
223 +// you can only use at::get_thread_num() with at::parallel_for()
224 +// as it is lazy initialized, otherwise it will always return 0.
225 +inline int get_thread_num() {
226 +#if defined(_OPENMP)
227 + return omp_get_thread_num();
228 +#else
229 + return 0;
230 +#endif
231 +}
232
233 +// balance payload across each thread
234 template <typename T>
235 inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
236 #if 0
@@ -94,10 +248,10 @@ inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
248 }
249 n_end += n_start;
250 #else
97 - // pytorch aten partition pattern
98 - T n_my = div_up(n, nth);
99 - n_start = ith * n_my;
100 - n_end = std::min(n_start + n_my, n);
251 + // pytorch aten partition pattern
252 + T n_my = div_up(n, nth);
253 + n_start = ith * n_my;
254 + n_end = std::min(n_start + n_my, n);
255 #endif
256 }
257
@@ -105,23 +259,15 @@ template <typename func_t>
259 inline void parallel_for(int n, const func_t& f) {
260 #if defined(_OPENMP)
261 #pragma omp parallel
108 -{
262 + {
263 int nth = omp_get_num_threads();
264 int ith = omp_get_thread_num();
265 int tbegin, tend;
266 balance211(n, nth, ith, tbegin, tend);
267 f(tbegin, tend);
114 -}
115 -#else
116 - f(0, n);
117 -#endif
118 -}
119 -
120 -inline int get_thread_num() {
121 -#if defined(_OPENMP)
122 - return omp_get_thread_num();
268 + }
269 #else
124 - return 0;
270 + f(0, n);
271 #endif
272 }
273
@@ -137,7 +283,6 @@ int inline adjust_num_threads(int m) {
283
284 template <typename func_t>
285 inline void parallel_2d(int m, int n, const func_t& f) {
140 -
286 // make sure we have even num_threads
287 int nth = adjust_num_threads(m);
288
@@ -165,31 +310,59 @@ inline void parallel_2d(int m, int n, const func_t& f) {
310
311 #if defined(_OPENMP)
312 #pragma omp parallel num_threads(nth)
168 -{
169 - int ith = omp_get_thread_num();
170 - int ith_m = ith / nth_n;
171 - int ith_n = ith % nth_n;
313 + {
314 + int ith = omp_get_thread_num();
315 + int ith_m = ith / nth_n;
316 + int ith_n = ith % nth_n;
317
173 - int thread_block_m = div_up(m, nth_m);
174 - int thread_block_n = div_up(n, nth_n);
318 + int thread_block_m = div_up(m, nth_m);
319 + int thread_block_n = div_up(n, nth_n);
320
176 - int begin_m = ith_m * thread_block_m;
177 - int end_m = std::min(m, begin_m + thread_block_m);
178 - int begin_n = ith_n * thread_block_n;
179 - int end_n = std::min(n, begin_n + thread_block_n);
321 + int begin_m = ith_m * thread_block_m;
322 + int end_m = std::min(m, begin_m + thread_block_m);
323 + int begin_n = ith_n * thread_block_n;
324 + int end_n = std::min(n, begin_n + thread_block_n);
325
181 - f(begin_m, end_m, begin_n, end_n);
182 -}
326 + f(begin_m, end_m, begin_n, end_n);
327 + }
328 #else
329 f(0, m, 0, n);
330 #endif
331 }
332
333 +// limit max cache blocks
334 +// when we need to do pre-unpack for weights, e.g. fp8
335 +#define MAX_CACHE_BLOCK_SIZE 4
336 +
337 template <typename T>
189 -int get_cache_blocks(int BLOCK_SIZE, int K) {
338 +inline int get_cache_blocks(int chunk_size) {
339 // L2 2MB and ratio of 50%
340 const int L2_size = 2048 * 1024 >> 1;
192 - return std::max(1, int(L2_size / (BLOCK_SIZE * K * sizeof(T))));
341 + return std::max(1, int(L2_size / (chunk_size * sizeof(T))));
342 +}
343 +
344 +template <>
345 +inline int get_cache_blocks<at::Float8_e4m3fn>(int chunk_size) {
346 + // fp8 uses bf16 as accumulate type
347 + int cache_block_size = get_cache_blocks<at::BFloat16>(chunk_size);
348 + return std::min(MAX_CACHE_BLOCK_SIZE, cache_block_size);
349 +}
350 +
351 +// 2d sequential loop in range : [mb0, mb1), [nb0, nb1)
352 +template <typename T, typename func_t>
353 +inline void loop_2d(int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1, int64_t chunk_size, const func_t& f) {
354 + // get number of blocks for L2 in most inner loop
355 + int64_t cache_blocks_nb = get_cache_blocks<T>(chunk_size);
356 +
357 + // loop order: [NB / cache_blocks_nb, MB, cache_blocks_nb]
358 + // TODO: implement reverse order of [MB / cache_blocks_mb, NB, cache_blocks_mb]
359 + for (int64_t nbb = nb0; nbb < nb1; nbb += cache_blocks_nb) {
360 + for (int64_t mb = mb0; mb < mb1; ++mb) {
361 + for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, nb1); ++nb) {
362 + f(mb, nb, nb - nbb);
363 + }
364 + }
365 + }
366 }
367
368 // data indexing for dimension collapse
@@ -243,4 +416,10 @@ struct Unroll<1> {
416 }
417 };
418
246 -} // anonymous namespace
419 +// conditional data ptr for optional tensor
420 +template <typename T>
421 +inline T* conditional_data_ptr(const std::optional<at::Tensor>& opt) {
422 + return opt.has_value() ? opt.value().data_ptr<T>() : nullptr;
423 +}
424 +
425 +} // anonymous namespace
csrc/cpu/sgl-kernels/conv.cpp new
+709
@@ -0,0 +1,709 @@
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 "vec.h"
9 +
10 +namespace {
11 +
12 +template <typename scalar_t>
13 +inline void copy_stub(scalar_t* __restrict__ y, const scalar_t* __restrict__ x, int64_t size) {
14 + using Vec = at::vec::Vectorized<scalar_t>;
15 + const bool is_padding = (x == nullptr);
16 + for (int64_t d = 0; d < size; d += Vec::size()) {
17 + Vec data_vec = is_padding ? Vec(0.f) : Vec::loadu(x + d);
18 + data_vec.store(y + d);
19 + }
20 +}
21 +
22 +// no remainder
23 +template <typename scalar_t>
24 +void inline update_conv_state(
25 + scalar_t* __restrict__ conv_states,
26 + const scalar_t* __restrict__ input,
27 + int64_t width,
28 + int64_t dim,
29 + int64_t seqlen,
30 + bool has_initial_states) {
31 + // width for `conv_states`
32 + int64_t width1 = width - 1;
33 + int64_t w = 0;
34 + for (; w < width1 - seqlen; ++w) {
35 + scalar_t* y = conv_states + w * dim;
36 + const scalar_t* x = has_initial_states ? conv_states + (w + seqlen) * dim : nullptr;
37 + copy_stub(y, x, dim);
38 + }
39 + for (; w < width1; ++w) {
40 + scalar_t* y = conv_states + w * dim;
41 + const scalar_t* x = input + (w + seqlen - width1) * dim;
42 + copy_stub(y, x, dim);
43 + }
44 +}
45 +
46 +// A : [M, BLOCK_N]
47 +// B : [BLOCK_N, K], prepacked as [K/2, BLOCK_N, 2]
48 +// C : [M, BLOCK_N]
49 +// bias : [BLOCK_N]
50 +//
51 +// lda : leading dimension of `input` and `out`
52 +//
53 +template <typename scalar_t, int K, int BLOCK_N, bool has_bias, bool has_silu>
54 +struct tinygemm_kernel {
55 + static inline void apply(
56 + const scalar_t* __restrict__ A,
57 + const scalar_t* __restrict__ B,
58 + scalar_t* __restrict__ C,
59 + const scalar_t* __restrict__ bias,
60 + const scalar_t* __restrict__ conv_states,
61 + bool has_initial_state,
62 + int64_t M,
63 + int64_t lda,
64 + bool is_first_token) {
65 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
66 + }
67 +};
68 +
69 +#if defined(CPU_CAPABILITY_AVX512)
70 +template <int K, int BLOCK_N, bool has_bias, bool has_silu>
71 +struct tinygemm_kernel<at::BFloat16, K, BLOCK_N, has_bias, has_silu> {
72 + static inline void apply(
73 + const at::BFloat16* __restrict__ A,
74 + const at::BFloat16* __restrict__ B,
75 + at::BFloat16* __restrict__ C,
76 + const at::BFloat16* __restrict__ bias,
77 + const at::BFloat16* __restrict__ conv_states,
78 + bool has_initial_state,
79 + int64_t M,
80 + int64_t lda,
81 + bool is_first_token) {
82 + assert(K == 4);
83 + constexpr int ROWS = K;
84 + constexpr int COLS = BLOCK_N / block_size_n();
85 +
86 + // leading dimension size for b for next block [K/2, 32, 2]
87 + constexpr int ldb = block_size_n() * K;
88 +
89 + __m512bh va[ROWS * COLS];
90 + __m512bh vb[ROWS * COLS];
91 + __m512 vc[COLS * 2];
92 +
93 + // k: {-3, -2, -1} -> {0, 1, 2}
94 + auto set_conv_states = [&](int k, int col) -> __m512i {
95 + return has_initial_state ? _mm512_loadu_si512(conv_states + (k + K - 1) * lda + col * 32)
96 + : _mm512_setzero_si512();
97 + };
98 +
99 +#define MM512_LOAD_A(idx) \
100 + ((idx) < 0 && is_first_token) ? (__m512bh)(set_conv_states((idx), col)) \
101 + : (__m512bh)(_mm512_loadu_si512(A + (idx) * lda + col * 32))
102 +
103 +#define MM512_PACK_A(ap, bp, a, b) \
104 + do { \
105 + __m512i r0 = (__m512i)(a); \
106 + __m512i r1 = (__m512i)(b); \
107 + __m512i d0 = _mm512_unpacklo_epi16(r0, r1); \
108 + __m512i d1 = _mm512_unpackhi_epi16(r0, r1); \
109 + r0 = _mm512_shuffle_i32x4(d0, d1, 0x88); \
110 + r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd); \
111 + (ap) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0x88); \
112 + (bp) = (__m512bh)_mm512_shuffle_i32x4(r0, r1, 0xdd); \
113 + } while (0)
114 +
115 + // step 0 : preload a at time step [-3][-2][-1]
116 + auto preloada = [&](auto i) {
117 + constexpr int col = i;
118 + int64_t m = 0;
119 + va[1 * COLS + col] = MM512_LOAD_A(m - 3);
120 + va[2 * COLS + col] = MM512_LOAD_A(m - 2);
121 + va[3 * COLS + col] = MM512_LOAD_A(m - 1);
122 + };
123 + Unroll<COLS>{}(preloada);
124 +
125 + auto loada = [&](auto i, int64_t m) {
126 + constexpr int col = i;
127 + // update previous time step
128 + va[0 * COLS + col] = va[1 * COLS + col];
129 + va[1 * COLS + col] = va[2 * COLS + col];
130 + va[2 * COLS + col] = va[3 * COLS + col];
131 + // load current time step
132 + va[3 * COLS + col] = MM512_LOAD_A(m);
133 + };
134 +
135 + // step 1 : load weight for just once
136 + auto loadb = [&](auto i) {
137 + constexpr int row = i / COLS;
138 + constexpr int col = i % COLS;
139 + vb[row * COLS + col] = (__m512bh)(_mm512_loadu_si512(B + col * ldb + row * 32));
140 + };
141 + Unroll<ROWS * COLS>{}(loadb);
142 +
143 + // [NB] accumulates 4x32 bfloat16 blocks
144 + //
145 + // +------------+------------+
146 + // | col0 | col1 |
147 + // +------------+------------+
148 + // | va0 va1 | va0 va1 |
149 + // | va2 va3 | va2 va3 |
150 + // +------------+------------+
151 + // | vc0 vc1 | vc0 vc1 |
152 + // +------------+------------+
153 + //
154 + // * va and vb shares the same memory layout
155 + // * block_n 32 with 4 rows equals to 4 registers
156 + // * 37 uops with avx512bf16 v.s. 57 uops with avx512f
157 + //
158 + auto compute = [&](auto i) {
159 + constexpr int col = i;
160 +
161 + // init accumulators
162 + if constexpr (has_bias) {
163 + __m512i b16 = _mm512_loadu_si512(reinterpret_cast<const __m512i*>(bias + col * 32));
164 + vc[col * 2 + 0] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 0));
165 + vc[col * 2 + 1] = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(b16, 1));
166 + } else {
167 + vc[col * 2 + 0] = _mm512_set1_ps(0.f);
168 + vc[col * 2 + 1] = _mm512_set1_ps(0.f);
169 + }
170 +
171 + // convert to vnni2 format
172 + __m512bh va0, va1, va2, va3;
173 + MM512_PACK_A(va0, va1, va[0 * COLS + col], va[1 * COLS + col]);
174 + MM512_PACK_A(va2, va3, va[2 * COLS + col], va[3 * COLS + col]);
175 +
176 + // accumulate
177 + vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va0, vb[0 * COLS + col]);
178 + vc[col * 2 + 0] = _mm512_dpbf16_ps(vc[col * 2 + 0], va2, vb[2 * COLS + col]);
179 + vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va1, vb[1 * COLS + col]);
180 + vc[col * 2 + 1] = _mm512_dpbf16_ps(vc[col * 2 + 1], va3, vb[3 * COLS + col]);
181 + };
182 +
183 + using fVec = at::vec::Vectorized<float>;
184 + using bVec = at::vec::Vectorized<at::BFloat16>;
185 + const fVec one = fVec(1.f);
186 + auto storec = [&](auto i, int64_t m) {
187 + constexpr int col = i;
188 + fVec x0 = fVec(vc[col * 2 + 0]);
189 + fVec x1 = fVec(vc[col * 2 + 1]);
190 + if constexpr (has_silu) {
191 + x0 = x0 / (one + x0.neg().exp_u20());
192 + x1 = x1 / (one + x1.neg().exp_u20());
193 + }
194 + bVec out_vec = convert_from_float_ext<at::BFloat16>(x0, x1);
195 + out_vec.store(C + m * lda + col * 32);
196 + };
197 +
198 + for (int64_t m = 0; m < M; ++m) {
199 + // step 3.a : load a at current time step
200 + Unroll<COLS>{}(loada, m);
201 + // step 3.b : accumulate for window size (4)
202 + Unroll<COLS>{}(compute);
203 + // step 3.c : store c at current time step
204 + Unroll<COLS>{}(storec, m);
205 + }
206 + }
207 +};
208 +#endif
209 +
210 +#define LAUNCH_TINYGEMM_KERNEL(K, NB_SIZE) \
211 + tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
212 + input + bs * seqlen * dim + mb_start * dim + nb_start, \
213 + weight + nb_start * width, \
214 + out + bs * seqlen * dim + mb_start * dim + nb_start, \
215 + has_bias ? bias + nb_start : nullptr, \
216 + has_conv_states ? conv_states + conv_state_index * (K - 1) * dim + nb_start : nullptr, \
217 + has_initial_states_value, \
218 + mb_size, \
219 + dim, \
220 + mb_start == 0);
221 +
222 +template <typename scalar_t>
223 +void causal_conv1d_fwd_kernel_impl(
224 + scalar_t* __restrict__ out,
225 + const scalar_t* __restrict__ input,
226 + const scalar_t* __restrict__ weight,
227 + const scalar_t* __restrict__ bias,
228 + scalar_t* __restrict__ conv_states,
229 + const int32_t* __restrict__ conv_indices,
230 + const bool* __restrict__ has_initial_state,
231 + bool silu_activation,
232 + int64_t batch,
233 + int64_t dim,
234 + int64_t seqlen,
235 + int64_t width,
236 + int64_t num_seq_blocks) {
237 + // handle 32 x 64 per block
238 + constexpr int64_t BLOCK_M = block_size_m();
239 + constexpr int64_t BLOCK_N = block_size_n() * 2;
240 + const int64_t NB = div_up(dim, BLOCK_N);
241 +
242 + const int64_t num_blocks_per_seq = div_up(seqlen, BLOCK_M);
243 + const bool has_conv_states = conv_states != nullptr;
244 + const bool has_conv_indices = conv_indices != nullptr;
245 +
246 + // parallel on [batch, seq, NB]
247 + AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
248 + at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
249 + int64_t mb{0}, nb{0};
250 + data_index_init(begin, mb, num_seq_blocks, nb, NB);
251 +
252 + for (int64_t i = begin; i < end; ++i) {
253 + int64_t bs = mb / num_blocks_per_seq;
254 +
255 + int64_t mb_start = (mb % num_blocks_per_seq) * BLOCK_M;
256 + int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
257 + int64_t nb_start = nb * BLOCK_N;
258 + int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
259 +
260 + const bool has_initial_states_value = has_conv_states ? has_initial_state[bs] : false;
261 + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
262 +
263 + switch (width << 4 | nb_size >> 4) {
264 + case 0x42:
265 + LAUNCH_TINYGEMM_KERNEL(4, 32);
266 + break;
267 + case 0x44:
268 + LAUNCH_TINYGEMM_KERNEL(4, 64);
269 + break;
270 + default:
271 + TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
272 + }
273 +
274 + // move to the next index
275 + data_index_step(mb, num_seq_blocks, nb, NB);
276 + }
277 + });
278 + });
279 +
280 + // update conv_states if necessary
281 + if (has_conv_states) {
282 + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
283 + for (int64_t bs = begin; bs < end; ++bs) {
284 + update_conv_state(
285 + conv_states + bs * (width - 1) * dim, input + bs * seqlen * dim, width, dim, seqlen, has_initial_state[bs]);
286 + }
287 + });
288 + }
289 +}
290 +
291 +#define LAUNCH_TINYGEMM_VARLEN_KERNEL(K, NB_SIZE) \
292 + tinygemm_kernel<scalar_t, K, NB_SIZE, has_bias, has_silu>::apply( \
293 + input + batch_offset * dim + mb_start * dim + nb_start, \
294 + weight + nb_start * width, \
295 + out + batch_offset * dim + mb_start * dim + nb_start, \
296 + has_bias ? bias + nb_start : nullptr, \
297 + nullptr, \
298 + false, \
299 + mb_size, \
300 + dim, \
301 + mb_start == 0);
302 +
303 +// TODO: add `has_initial_state` support for varlen kernel
304 +template <typename scalar_t>
305 +void causal_conv1d_fwd_varlen_kernel_impl(
306 + scalar_t* __restrict__ out,
307 + const scalar_t* __restrict__ input,
308 + const scalar_t* __restrict__ weight,
309 + const scalar_t* __restrict__ bias,
310 + scalar_t* __restrict__ conv_states,
311 + const int32_t* __restrict__ query_start_loc,
312 + const int32_t* __restrict__ conv_indices,
313 + const bool* __restrict__ has_initial_state,
314 + const int32_t* __restrict__ block_indices,
315 + bool silu_activation,
316 + int64_t batch,
317 + int64_t dim,
318 + int64_t width,
319 + int64_t num_seq_blocks) {
320 + // handle 32 x 64 per block
321 + constexpr int64_t BLOCK_M = block_size_m();
322 + constexpr int64_t BLOCK_N = block_size_n() * 2;
323 + const int64_t NB = div_up(dim, BLOCK_N);
324 +
325 + const bool has_conv_states = conv_states != nullptr;
326 + const bool has_conv_indices = conv_indices != nullptr;
327 +
328 + // parallel on [batch, seq, NB]
329 + AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
330 + at::parallel_for(0, num_seq_blocks * NB, 0, [&](int64_t begin, int64_t end) {
331 + int64_t mb{0}, nb{0};
332 + data_index_init(begin, mb, num_seq_blocks, nb, NB);
333 +
334 + for (int64_t i = begin; i < end; ++i) {
335 + int32_t bs = block_indices[mb * 2 + 0];
336 + int32_t batch_offset = query_start_loc[bs];
337 + int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
338 +
339 + int64_t mb_start = block_indices[mb * 2 + 1] * BLOCK_M;
340 + int64_t mb_size = std::min(seqlen - mb_start, BLOCK_M);
341 + int64_t nb_start = nb * BLOCK_N;
342 + int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
343 +
344 + switch (width << 4 | nb_size >> 4) {
345 + case 0x42:
346 + LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 32);
347 + break;
348 + case 0x44:
349 + LAUNCH_TINYGEMM_VARLEN_KERNEL(4, 64);
350 + break;
351 + default:
352 + TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
353 + }
354 +
355 + // move to the next index
356 + data_index_step(mb, num_seq_blocks, nb, NB);
357 + }
358 + });
359 + });
360 +
361 + // update conv_states if necessary
362 + if (has_conv_states) {
363 + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
364 + for (int64_t bs = begin; bs < end; ++bs) {
365 + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
366 + int32_t seqlen = query_start_loc[bs + 1] - query_start_loc[bs];
367 + int32_t batch_offset = query_start_loc[bs];
368 + update_conv_state(
369 + conv_states + conv_state_index * (width - 1) * dim,
370 + input + batch_offset * dim,
371 + width,
372 + dim,
373 + seqlen,
374 + /* has_initial_state */ false);
375 + }
376 + });
377 + }
378 +}
379 +
380 +template <typename scalar_t>
381 +void causal_conv1d_update_kernel_impl(
382 + scalar_t* __restrict__ out,
383 + const scalar_t* __restrict__ input,
384 + scalar_t* __restrict__ conv_states,
385 + const scalar_t* __restrict__ weight,
386 + const scalar_t* __restrict__ bias,
387 + const int32_t* __restrict__ conv_indices,
388 + bool silu_activation,
389 + int64_t batch,
390 + int64_t dim,
391 + int64_t seqlen,
392 + int64_t width) {
393 + // handle 32 x 64 per block
394 + constexpr int64_t BLOCK_M = block_size_m();
395 + constexpr int64_t BLOCK_N = block_size_n() * 2;
396 + const int64_t NB = div_up(dim, BLOCK_N);
397 +
398 + const bool has_conv_states = conv_states != nullptr;
399 + const bool has_conv_indices = conv_indices != nullptr;
400 +
401 + // parallel on [batch, NB]
402 + AT_DISPATCH_BOOL2(bias != nullptr, has_bias, silu_activation, has_silu, [&] {
403 + at::parallel_for(0, batch * NB, 0, [&](int64_t begin, int64_t end) {
404 + int64_t bs{0}, nb{0};
405 + data_index_init(begin, bs, batch, nb, NB);
406 +
407 + for (int64_t i = begin; i < end; ++i) {
408 + int64_t mb_start = 0;
409 + int64_t mb_size = 1;
410 + int64_t nb_start = nb * BLOCK_N;
411 + int64_t nb_size = std::min(dim - nb_start, BLOCK_N);
412 +
413 + const bool has_initial_states_value = true;
414 + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
415 +
416 + switch (width << 4 | nb_size >> 4) {
417 + case 0x42:
418 + LAUNCH_TINYGEMM_KERNEL(4, 32);
419 + break;
420 + case 0x44:
421 + LAUNCH_TINYGEMM_KERNEL(4, 64);
422 + break;
423 + default:
424 + TORCH_CHECK(false, "Unexpected block size, ", width, " x ", nb_size);
425 + }
426 +
427 + // move to the next index
428 + data_index_step(bs, batch, nb, NB);
429 + }
430 + });
431 + });
432 +
433 +#define CONV_STATE_INDEXR(w) conv_states + conv_state_index*(width - 1) * dim + (w) * dim
434 +
435 + // update conv_states
436 + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
437 + for (int64_t bs = begin; bs < end; ++bs) {
438 + // update old states, range [1, width - 1)
439 + int32_t conv_state_index = has_conv_indices ? conv_indices[bs] : bs;
440 + for (int64_t w = 1; w < width - 1; ++w) {
441 + std::memcpy(CONV_STATE_INDEXR(w - 1), CONV_STATE_INDEXR(w), dim * sizeof(scalar_t));
442 + }
443 + // copy new states
444 + std::memcpy(CONV_STATE_INDEXR(width - 2), input + bs * dim, dim * sizeof(scalar_t));
445 + }
446 + });
447 +}
448 +
449 +} // anonymous namespace
450 +
451 +// from [dim, width] or [N, K]
452 +// to [N/BLOCK_N, K/2, BLOCK_N, 2]
453 +at::Tensor causal_conv1d_weight_pack(const at::Tensor& weight) {
454 + CHECK_INPUT(weight);
455 +
456 + int64_t dim = weight.size(0);
457 + int64_t width = weight.size(1);
458 + constexpr int64_t BLOCK_N = block_size_n();
459 + TORCH_CHECK(width == 4, "causal_conv1d_weight_pack: support only width of 4");
460 + TORCH_CHECK(dim % BLOCK_N == 0, "causal_conv1d_weight_pack: invalid dim size ", dim);
461 +
462 + const int64_t N = dim, K2 = width >> 1;
463 + const int64_t NB = div_up(N, BLOCK_N);
464 +
465 + auto packed_weight = at::empty_like(weight);
466 + AT_DISPATCH_REDUCED_FLOATING_TYPES(weight.scalar_type(), "causal_conv1d_fwd_kernel_impl", [&] {
467 + // cast to float32 as vnni size is 2
468 + const float* w_data = reinterpret_cast<float*>(weight.data_ptr<scalar_t>());
469 + float* packed_data = reinterpret_cast<float*>(packed_weight.data_ptr<scalar_t>());
470 +
471 + at::parallel_for(0, NB * K2 * BLOCK_N, 0, [&](int64_t begin, int64_t end) {
472 + int64_t nb{0}, k2{0}, n{0};
473 + data_index_init(begin, nb, NB, k2, K2, n, BLOCK_N);
474 +
475 + // TODO: optimize this if we need to online prepacking.
476 + for (int64_t i = begin; i < end; ++i) {
477 + packed_data[i] = w_data[nb * BLOCK_N * K2 + n * K2 + k2];
478 +
479 + // move to the next index
480 + data_index_step(nb, NB, k2, K2, n, BLOCK_N);
481 + }
482 + });
483 + });
484 + return packed_weight;
485 +}
486 +
487 +#define CHECK_OPTIONAL_SHAPE_DTYPE(OPT, SIZE, DTYPE) \
488 + if (OPT.has_value()) { \
489 + const auto tensor = OPT.value(); \
490 + CHECK_CONTIGUOUS(tensor); \
491 + CHECK_EQ(tensor.size(0), SIZE); \
492 + CHECK_EQ(tensor.scalar_type(), DTYPE); \
493 + }
494 +
495 +template <int BLOCK_M>
496 +int64_t get_block_count(const std::optional<at::Tensor>& offsets, int64_t batch, int64_t seqlen) {
497 + if (offsets.has_value()) {
498 + const int32_t* offsets_data = offsets.value().data_ptr<int32_t>();
499 + int32_t num_seq_blocks = 0;
500 + for (int64_t row = 0; row < batch; ++row) {
501 + num_seq_blocks += div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
502 + }
503 + return num_seq_blocks;
504 + }
505 + return batch * div_up(seqlen, int64_t(BLOCK_M));
506 +}
507 +
508 +template <int BLOCK_M>
509 +at::Tensor get_block_indices(const std::optional<at::Tensor>& offsets, int64_t num_seq_blocks) {
510 + if (!offsets.has_value()) {
511 + return at::Tensor();
512 + }
513 +
514 + const at::Tensor& offsets_ = offsets.value();
515 + at::Tensor indices = at::empty({num_seq_blocks, 2}, offsets_.options());
516 +
517 + int64_t batch = offsets_.size(0) - 1;
518 +
519 + const int32_t* offsets_data = offsets_.data_ptr<int32_t>();
520 + int32_t* indices_data = indices.data_ptr<int32_t>();
521 +
522 + int64_t idx = 0;
523 + for (int32_t row = 0; row < batch; ++row) {
524 + int32_t blocks = div_up(offsets_data[row + 1] - offsets_data[row], BLOCK_M);
525 +
526 + for (int32_t col = 0; col < blocks; ++col) {
527 + indices_data[idx * 2 + 0] = row;
528 + indices_data[idx * 2 + 1] = col;
529 + idx++;
530 + }
531 + }
532 + return indices;
533 +}
534 +
535 +// API aligned with GPUs
536 +//
537 +// x: (batch, dim, seqlen) or (dim, cu_seq_len) for varlen
538 +// weight: (dim, width)
539 +// bias: (dim,)
540 +// query_start_loc: (batch + 1) int32
541 +// cache_indices: (batch) int32
542 +// has_initial_state: (batch) bool
543 +// conv_states: (..., dim, width - 1) itype
544 +// activation: either None or "silu" or "swish"
545 +// pad_slot_id: int
546 +//
547 +at::Tensor causal_conv1d_fwd_cpu(
548 + const at::Tensor& x,
549 + const at::Tensor& weight,
550 + const std::optional<at::Tensor>& bias,
551 + const std::optional<at::Tensor>& conv_states,
552 + const std::optional<at::Tensor>& query_start_loc,
553 + const std::optional<at::Tensor>& conv_state_indices,
554 + const std::optional<at::Tensor>& has_initial_state,
555 + bool silu_activation,
556 + int64_t pad_slot_id,
557 + bool is_vnni) {
558 + CHECK_CONTIGUOUS(weight);
559 + auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
560 +
561 + const bool is_var_seqlen = query_start_loc.has_value();
562 + const int64_t input_ndim = is_var_seqlen ? 2 : 3;
563 + TORCH_CHECK(x.dim() == input_ndim, "causal_conv1d_fwd_cpu: expect x to be ", input_ndim, "D tensor.");
564 + TORCH_CHECK(x.stride(-2) == 1 && x.stride(-1) == x.size(-2), "causal_conv1d_fwd_cpu: expect x to be transposed.");
565 +
566 + const int64_t batch = is_var_seqlen ? query_start_loc.value().size(0) - 1 : x.size(0);
567 + const int64_t dim = x.size(-2);
568 + const int64_t seqlen = x.size(-1);
569 + const int64_t width = weight.size(-1);
570 +
571 + const auto scalar_type = x.scalar_type();
572 + CHECK_EQ(weight.scalar_type(), scalar_type);
573 + CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
574 + CHECK_OPTIONAL_SHAPE_DTYPE(query_start_loc, batch + 1, at::kInt);
575 + CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
576 + CHECK_OPTIONAL_SHAPE_DTYPE(has_initial_state, batch, at::kBool);
577 +
578 + if (conv_states.has_value()) {
579 + auto& conv_states_val = conv_states.value();
580 + int64_t padded_batch = conv_states_val.size(0);
581 + CHECK_EQ(conv_states_val.scalar_type(), scalar_type);
582 + CHECK_GE(padded_batch, batch);
583 + CHECK_EQ(conv_states_val.size(1), dim);
584 + CHECK_EQ(conv_states_val.size(2), width - 1);
585 +
586 + // adjust `conv_states` to be contiguous on `dim`
587 + // should happen only once
588 + if (conv_states_val.stride(-2) != 1) {
589 + auto conv_states_copy = conv_states_val.clone();
590 + conv_states_val.as_strided_({padded_batch, dim, width - 1}, {(width - 1) * dim, 1, dim});
591 + conv_states_val.copy_(conv_states_copy);
592 + }
593 + }
594 +
595 + // block size for sequence blocks, 32
596 + constexpr int64_t BLOCK_M = block_size_m();
597 +
598 + // total number of sequence blocks
599 + int64_t num_seq_blocks = get_block_count<BLOCK_M>(query_start_loc, batch, seqlen);
600 +
601 + at::Tensor out = at::empty_like(x);
602 + AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_fwd_kernel_impl", [&] {
603 + if (is_var_seqlen) {
604 + // record seq blocks in Coordinate format, aka [num_seq_blocks, 2]
605 + at::Tensor block_indices = get_block_indices<BLOCK_M>(query_start_loc, num_seq_blocks);
606 +
607 + causal_conv1d_fwd_varlen_kernel_impl(
608 + out.data_ptr<scalar_t>(),
609 + x.data_ptr<scalar_t>(),
610 + packed_w.data_ptr<scalar_t>(),
611 + conditional_data_ptr<scalar_t>(bias),
612 + conditional_data_ptr<scalar_t>(conv_states),
613 + conditional_data_ptr<int32_t>(query_start_loc),
614 + conditional_data_ptr<int32_t>(conv_state_indices),
615 + conditional_data_ptr<bool>(has_initial_state),
616 + block_indices.data_ptr<int32_t>(),
617 + silu_activation,
618 + batch,
619 + dim,
620 + width,
621 + num_seq_blocks);
622 + } else {
623 + causal_conv1d_fwd_kernel_impl<scalar_t>(
624 + out.data_ptr<scalar_t>(),
625 + x.data_ptr<scalar_t>(),
626 + packed_w.data_ptr<scalar_t>(),
627 + conditional_data_ptr<scalar_t>(bias),
628 + conditional_data_ptr<scalar_t>(conv_states),
629 + conditional_data_ptr<int32_t>(conv_state_indices),
630 + conditional_data_ptr<bool>(has_initial_state),
631 + silu_activation,
632 + batch,
633 + dim,
634 + seqlen,
635 + width,
636 + num_seq_blocks);
637 + }
638 + });
639 + return out;
640 +}
641 +
642 +// API aligned with GPUs
643 +//
644 +// x: (batch, dim) or (batch, dim, seqlen)
645 +// conv_state: (..., dim, state_len), where state_len >= width - 1
646 +// weight: (dim, width)
647 +// bias: (dim,)
648 +// cache_seqlens: (batch,), dtype int32.
649 +// conv_state_indices: (batch,), dtype int32
650 +// pad_slot_id: int
651 +// out: (batch, dim) or (batch, dim, seqlen)
652 +//
653 +at::Tensor causal_conv1d_update_cpu(
654 + const at::Tensor& x,
655 + const at::Tensor& conv_states,
656 + const at::Tensor& weight,
657 + const std::optional<at::Tensor>& bias,
658 + bool silu_activation,
659 + const std::optional<at::Tensor>& cache_seqlens,
660 + const std::optional<at::Tensor>& conv_state_indices,
661 + int64_t pad_slot_id,
662 + bool is_vnni) {
663 + CHECK_CONTIGUOUS(x);
664 + CHECK_CONTIGUOUS(weight);
665 + auto packed_w = is_vnni ? weight : causal_conv1d_weight_pack(weight);
666 +
667 + // TODO: add multi-token prediction support
668 + TORCH_CHECK(x.dim() == 2, "causal_conv1d_update_cpu: expect x to be 2D tensor.");
669 + TORCH_CHECK(!cache_seqlens.has_value(), "causal_conv1d_update_cpu: don't support cache_seqlens.");
670 +
671 + int64_t batch = x.size(0);
672 + int64_t dim = x.size(1);
673 + int64_t seqlen = 1;
674 + int64_t width = weight.size(-1);
675 +
676 + const auto scalar_type = x.scalar_type();
677 + CHECK_EQ(weight.scalar_type(), scalar_type);
678 + CHECK_OPTIONAL_SHAPE_DTYPE(bias, dim, scalar_type);
679 + CHECK_OPTIONAL_SHAPE_DTYPE(conv_state_indices, batch, at::kInt);
680 +
681 + CHECK_EQ(conv_states.scalar_type(), scalar_type);
682 + CHECK_EQ(conv_states.size(1), dim);
683 + CHECK_EQ(conv_states.size(2), width - 1);
684 +
685 + // adjust `conv_states` to be contiguous on `dim`
686 + if (conv_states.stride(-2) != 1) {
687 + int64_t num_cache_lines = conv_states.size(0);
688 + auto conv_states_copy = conv_states.clone();
689 + conv_states.as_strided_({num_cache_lines, dim, width - 1}, {(width - 1) * dim, 1, dim});
690 + conv_states.copy_(conv_states_copy);
691 + }
692 +
693 + at::Tensor out = at::empty_like(x);
694 + AT_DISPATCH_REDUCED_FLOATING_TYPES(scalar_type, "causal_conv1d_update_kernel_impl", [&] {
695 + causal_conv1d_update_kernel_impl<scalar_t>(
696 + out.data_ptr<scalar_t>(),
697 + x.data_ptr<scalar_t>(),
698 + conv_states.data_ptr<scalar_t>(),
699 + packed_w.data_ptr<scalar_t>(),
700 + conditional_data_ptr<scalar_t>(bias),
701 + conditional_data_ptr<int32_t>(conv_state_indices),
702 + silu_activation,
703 + batch,
704 + dim,
705 + seqlen,
706 + width);
707 + });
708 + return out;
709 +}
csrc/cpu/sgl-kernels/fla.cpp new
+1394
@@ -0,0 +1,1394 @@
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 "vec.h"
9 +#include "vec_pack.h"
10 +
11 +namespace {
12 +// For this cpu kernel, we have some innovations aside from the existing gpu kernels:
13 +// 1) Use less parallel loops, i.e. 4 including l2_norm.
14 +// 2) Fuse part of l2_norm with the rest of the computation.
15 +
16 +#define THREAD_BUFFER_ALLOC(dst, base_ptr, offset, type, size) \
17 + type* dst = reinterpret_cast<type*>((base_ptr) + (offset)); \
18 + offset += (size);
19 +
20 +template <typename scalar_t>
21 +inline void fill_stub(scalar_t* __restrict__ out, float val, int size) {
22 + using Vec = at::vec::Vectorized<scalar_t>;
23 + constexpr int kVecSize = Vec::size();
24 + const Vec data_vec = Vec(static_cast<scalar_t>(val));
25 + int d = 0;
26 +#pragma GCC unroll 4
27 + for (; d <= size - kVecSize; d += kVecSize) {
28 + data_vec.store(out + d);
29 + }
30 + if (size - d > 0) {
31 + data_vec.store(out + d, size - d);
32 + }
33 +}
34 +
35 +template <typename scalar_t, int64_t chunk_size = 64>
36 +void chunk_gated_delta_rule_kernel_impl(
37 + scalar_t* __restrict__ out, // [B, T, HV, EV]
38 + float* __restrict__ final_state_data, // [N, HV, EK, EV]
39 + const scalar_t* __restrict__ q_orig, // [B, T, HK, EK]
40 + const scalar_t* __restrict__ k_orig, // [B, T, HK, EK]
41 + const scalar_t* __restrict__ v_orig, // [B, T, HV, EV]
42 + const float* __restrict__ g_orig, // [B, T, HV] FP32
43 + const scalar_t* __restrict__ b_orig, // [B, T, HV]
44 + const int32_t* __restrict__ cu_seqlens_ptr, // [N + 1] INT32
45 + float* __restrict__ buff,
46 + scalar_t* __restrict__ reduced_buff,
47 + scalar_t* __restrict__ thread_buff,
48 + const int32_t* __restrict__ chunk_offsets_ptr,
49 + const int32_t* __restrict__ chunk_indices_ptr,
50 + bool use_qk_l2norm_in_kernel,
51 + const int64_t& batch_size,
52 + const int64_t& global_seq_len,
53 + const int64_t& qk_num_head,
54 + const int64_t& v_num_head,
55 + const int64_t& qk_head_size,
56 + const int64_t& v_head_size,
57 + const int64_t& qStrideH,
58 + const int64_t& qStrideT,
59 + const int64_t& kStrideH,
60 + const int64_t& kStrideT,
61 + const int64_t& vStrideH,
62 + const int64_t& vStrideT,
63 + const int64_t& oStrideH,
64 + const int64_t& oStrideT,
65 + const int64_t& global_total_seq_length,
66 + const int64_t& global_num_chunk,
67 + const int64_t& buff_size_16bit_per_thread,
68 + double eps = 1e-5) {
69 + int64_t gStrideH = 1;
70 + int64_t gStrideT = v_num_head;
71 + int64_t bStrideH = 1;
72 + int64_t bStrideT = v_num_head;
73 + int64_t final_state_StrideN = v_num_head * qk_head_size * v_head_size;
74 + int64_t final_state_StrideH = qk_head_size * v_head_size;
75 + int64_t final_state_StrideE = v_head_size;
76 + int64_t head_group = v_num_head / qk_num_head;
77 + float scale = 1.0 / std::sqrt(qk_head_size);
78 + using bVec = at::vec::Vectorized<scalar_t>;
79 + using fVec = at::vec::Vectorized<float>;
80 + constexpr int64_t VecSize = bVec::size();
81 + constexpr int64_t fVecSize = fVec::size();
82 +
83 + // Data pointers
84 + float* g_pad = buff;
85 + float* core_attn_out = g_pad + v_num_head * global_total_seq_length;
86 + float* decay_mask = core_attn_out + batch_size * v_num_head * global_total_seq_length * v_head_size;
87 + float* v_beta_attn = decay_mask + v_num_head * global_total_seq_length * chunk_size;
88 +
89 + scalar_t* q_pad = reduced_buff;
90 + scalar_t* k_pad = q_pad + qk_num_head * global_total_seq_length * qk_head_size;
91 + scalar_t* v_pad = k_pad + qk_num_head * global_total_seq_length * qk_head_size;
92 + scalar_t* k_beta = v_pad + v_num_head * global_total_seq_length * v_head_size;
93 + scalar_t* v_beta = k_beta + v_num_head * global_total_seq_length * qk_head_size;
94 + scalar_t* k_cumdecay_reduced = v_beta + v_num_head * global_total_seq_length * v_head_size;
95 + scalar_t* q_norm_sum = k_cumdecay_reduced + v_num_head * global_total_seq_length * qk_head_size;
96 + scalar_t* k_norm_sum = q_norm_sum + qk_num_head * global_seq_len;
97 +
98 + if (use_qk_l2norm_in_kernel) {
99 + at::parallel_for(0, qk_num_head * global_seq_len, 0, [&](int64_t begin, int64_t end) {
100 + int64_t h_qk = 0, l = 0;
101 + data_index_init(begin, h_qk, qk_num_head, l, global_seq_len);
102 + for (int64_t i = begin; i < end; ++i) {
103 + auto q_norm_sum_ptr = q_norm_sum + h_qk * global_seq_len + l;
104 + auto k_norm_sum_ptr = k_norm_sum + h_qk * global_seq_len + l;
105 + float sum_q = float(0);
106 + float sum_k = float(0);
107 + fVec sum_q_fvec = fVec(float(0));
108 + fVec sum_k_fvec = fVec(float(0));
109 + int64_t q_offset = l * qStrideT + h_qk * qStrideH;
110 + int64_t k_offset = l * qStrideT + h_qk * qStrideH;
111 + int64_t d;
112 + for (d = 0; d <= qk_head_size - VecSize; d += VecSize) {
113 + bVec q_bvec = bVec::loadu(q_orig + q_offset + d);
114 + fVec q_fvec0, q_fvec1;
115 + std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec);
116 + sum_q_fvec += q_fvec0 * q_fvec0;
117 + sum_q_fvec += q_fvec1 * q_fvec1;
118 + bVec k_bvec = bVec::loadu(k_orig + k_offset + d);
119 + fVec k_fvec0, k_fvec1;
120 + std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec);
121 + sum_k_fvec += k_fvec0 * k_fvec0;
122 + sum_k_fvec += k_fvec1 * k_fvec1;
123 + }
124 + sum_q += vec_reduce_sum(sum_q_fvec);
125 + sum_k += vec_reduce_sum(sum_k_fvec);
126 + q_norm_sum_ptr[0] = static_cast<scalar_t>(float(1) / std::sqrt(sum_q + eps));
127 + k_norm_sum_ptr[0] = static_cast<scalar_t>(float(1) / std::sqrt(sum_k + eps));
128 + data_index_step(h_qk, qk_num_head, l, global_seq_len);
129 + }
130 + });
131 + }
132 +
133 + // query = query * scale
134 + // k_beta = key * beta.unsqueeze(-1)
135 + // v_beta = value * beta.unsqueeze(-1)
136 + // Padding for q/k/v/beta
137 + at::parallel_for(0, qk_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) {
138 + int ompIdx = at::get_thread_num();
139 + int64_t h_qk = 0, c = 0;
140 + data_index_init(begin, h_qk, qk_num_head, c, global_num_chunk);
141 + for ([[maybe_unused]] auto z : c10::irange(begin, end)) {
142 + int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch
143 + int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk
144 + int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size;
145 + int64_t l = c * chunk_size;
146 + bool is_tail = (c + 1 == chunk_offsets_ptr[ib + 1]);
147 + int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib];
148 + int64_t real_chunk_size = is_tail ? seq_len - ic * chunk_size : chunk_size;
149 + auto q_orig_ptr = q_orig + h_qk * qStrideH + l_orig * qStrideT;
150 + auto k_orig_ptr = k_orig + h_qk * kStrideH + l_orig * kStrideT;
151 + auto v_orig_ptr = v_orig + l_orig * vStrideT;
152 + auto b_orig_ptr = b_orig + l_orig * bStrideT;
153 + auto q_pad_ptr = q_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size;
154 + auto k_pad_ptr = k_pad + h_qk * global_total_seq_length * qk_head_size + l * qk_head_size;
155 + auto v_pad_ptr = v_pad + l * v_head_size;
156 + auto k_beta_ptr = k_beta + l * qk_head_size;
157 + auto v_beta_ptr = v_beta + l * v_head_size;
158 +
159 + for (int64_t j = 0; j < real_chunk_size; j++) {
160 + auto curr_q_orig = q_orig_ptr + j * qStrideT;
161 + auto curr_k_orig = k_orig_ptr + j * kStrideT;
162 + auto curr_q_pad = q_pad_ptr + j * qk_head_size;
163 + auto curr_k_pad = k_pad_ptr + j * qk_head_size;
164 + auto q_scale =
165 + use_qk_l2norm_in_kernel ? *(q_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast<scalar_t>(1);
166 + auto k_scale =
167 + use_qk_l2norm_in_kernel ? *(k_norm_sum + h_qk * global_seq_len + l_orig + j) : static_cast<scalar_t>(1);
168 + auto q_scale_vec = bVec(q_scale);
169 + auto k_scale_vec = bVec(k_scale);
170 + int64_t i = 0;
171 + scalar_t scale_reduced = static_cast<scalar_t>(scale);
172 + auto vec_scale_reduced = bVec(scale_reduced);
173 + for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) {
174 + auto tmp0 = bVec::loadu(curr_q_orig + i, fVecSize);
175 + auto tmp1 = tmp0 * q_scale_vec * vec_scale_reduced;
176 + tmp1.store(curr_q_pad + i, fVecSize);
177 + auto tmp3 = bVec::loadu(curr_k_orig + i, fVecSize);
178 + auto tmp4 = tmp3 * k_scale_vec;
179 + tmp4.store(curr_k_pad + i, fVecSize);
180 + }
181 +
182 + for (auto hi = 0; hi < head_group; hi++) {
183 + int64_t h = h_qk * head_group + hi;
184 + auto curr_v_orig = v_orig_ptr + h * vStrideH + j * vStrideT;
185 + auto curr_b_orig = b_orig_ptr + h * bStrideH + j * bStrideT;
186 + scalar_t b_orig_val_reduced = *(curr_b_orig);
187 + auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
188 + auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size;
189 + auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
190 +
191 + // query = query * scale
192 + // k_beta = key * beta.unsqueeze(-1)
193 + int64_t i = 0;
194 + auto vec_b_reduced = bVec(b_orig_val_reduced);
195 + for (; i < fVecSize * (qk_head_size / fVecSize); i += fVecSize) {
196 + auto tmp0 = bVec::loadu(curr_k_orig + i, fVecSize);
197 + auto tmp2 = tmp0 * k_scale_vec * vec_b_reduced;
198 + tmp2.store(curr_k_beta + i, fVecSize);
199 + }
200 + // v_beta = value * beta.unsqueeze(-1)
201 + i = 0;
202 + for (; i < VecSize * (v_head_size / VecSize); i += VecSize) {
203 + auto tmp3 = bVec::loadu(curr_v_orig + i);
204 + tmp3.store(curr_v_pad + i);
205 + auto tmp5 = tmp3 * vec_b_reduced;
206 + tmp5.store(curr_v_beta + i);
207 + }
208 + }
209 + }
210 +
211 + for (int64_t j = real_chunk_size; j < chunk_size; j++) {
212 + auto curr_q_pad = q_pad_ptr + j * qk_head_size;
213 + auto curr_k_pad = k_pad_ptr + j * qk_head_size;
214 + int64_t i = 0;
215 + auto vec_zero = bVec(0.0);
216 + for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
217 + vec_zero.store(curr_q_pad + i);
218 + vec_zero.store(curr_k_pad + i);
219 + }
220 + for (auto hi = 0; hi < head_group; hi++) {
221 + int64_t h = h_qk * head_group + hi;
222 + auto curr_v_pad = v_pad_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
223 + auto curr_k_beta = k_beta_ptr + h * global_total_seq_length * qk_head_size + j * qk_head_size;
224 + auto curr_v_beta = v_beta_ptr + h * global_total_seq_length * v_head_size + j * v_head_size;
225 + int64_t i = 0;
226 + for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
227 + vec_zero.store(curr_k_beta + i);
228 + }
229 + i = 0;
230 + for (; i < VecSize * (v_head_size / VecSize); i += VecSize) {
231 + vec_zero.store(curr_v_pad + i);
232 + vec_zero.store(curr_v_beta + i);
233 + }
234 + }
235 + }
236 + // Move to the next query
237 + data_index_step(h_qk, qk_num_head, c, global_num_chunk);
238 + }
239 + });
240 +
241 + at::parallel_for(0, v_num_head * global_num_chunk, 1, [&](int64_t begin, int64_t end) {
242 + int64_t h = 0, c = 0;
243 + data_index_init(begin, h, v_num_head, c, global_num_chunk);
244 + int ompIdx = at::get_thread_num();
245 + int64_t offset = 0;
246 + scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread;
247 + THREAD_BUFFER_ALLOC(k_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size);
248 + THREAD_BUFFER_ALLOC(v_pack, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size);
249 + THREAD_BUFFER_ALLOC(k_beta_g, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
250 + THREAD_BUFFER_ALLOC(k_beta_g_pack, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
251 + THREAD_BUFFER_ALLOC(curr_attn, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2);
252 + THREAD_BUFFER_ALLOC(curr_attn_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size);
253 + THREAD_BUFFER_ALLOC(k_cumdecay, thread_buff_ptr, offset, float, chunk_size* qk_head_size * 2);
254 + THREAD_BUFFER_ALLOC(row, thread_buff_ptr, offset, float, chunk_size * 2);
255 + THREAD_BUFFER_ALLOC(updated, thread_buff_ptr, offset, float, chunk_size * 2);
256 + for ([[maybe_unused]] auto z : c10::irange(begin, end)) {
257 + int64_t ib = chunk_indices_ptr[c * 2]; // idx_batch
258 + int64_t ic = chunk_indices_ptr[c * 2 + 1]; // idx_chunk
259 + int64_t l_orig = cu_seqlens_ptr[ib] + ic * chunk_size;
260 + int64_t seq_len = cu_seqlens_ptr[ib + 1] - cu_seqlens_ptr[ib];
261 + int64_t h_qk = h / head_group;
262 + auto curr_g_orig = g_orig + h * gStrideH + l_orig * gStrideT;
263 + auto curr_g_pad = g_pad + h * global_total_seq_length + c * chunk_size;
264 + auto curr_decay_mask = decay_mask + h * global_total_seq_length * chunk_size + c * chunk_size * chunk_size;
265 + auto curr_k_pad = k_pad + h_qk * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size;
266 + auto curr_k_beta = k_beta + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size;
267 + auto curr_k_cumdecay_reduced =
268 + k_cumdecay_reduced + h * global_total_seq_length * qk_head_size + c * chunk_size * qk_head_size;
269 + auto curr_v_beta = v_beta + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size;
270 + auto curr_value = v_beta_attn + h * global_total_seq_length * v_head_size + c * chunk_size * v_head_size;
271 +
272 + float acc_val = 0;
273 + for (int64_t i = 0; i < chunk_size; i++) {
274 + // Padding for g
275 + // g = g.cumsum(dim=-1)
276 + // g: [B, HV, num_chunk, chunk_size]
277 + if (ic * chunk_size + i < seq_len) {
278 + acc_val += curr_g_orig[i * gStrideT];
279 + }
280 + curr_g_pad[i] = acc_val;
281 + // decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril()
282 + // decay_mask: [B, HV, num_chunk, chunk_size, chunk_size]
283 + float curr_g_pad_i = static_cast<float>(curr_g_pad[i]);
284 + auto vec_curr_g_pad_i = fVec(curr_g_pad_i);
285 + int64_t j = 0;
286 + int64_t len = i + 1;
287 + for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
288 + auto tmp0 = fVec::loadu(curr_g_pad + j);
289 + auto tmp1 = vec_curr_g_pad_i - tmp0;
290 + auto tmp2 = tmp1.exp_u20();
291 + tmp2.store(curr_decay_mask + i * chunk_size + j);
292 + }
293 + if (j < len) {
294 + auto tmp0 = fVec::loadu(curr_g_pad + j, len - j);
295 + auto tmp1 = vec_curr_g_pad_i - tmp0;
296 + auto tmp2 = tmp1.exp_u20();
297 + tmp2.store(curr_decay_mask + i * chunk_size + j, len - j);
298 + }
299 + }
300 +
301 + // attn = k_beta @ key.transpose(-1, -2)
302 + // attn: [B, HV, num_chunk, chunk_size, chunk_size]
303 + // transpose and pack for key
304 + pack_vnni<scalar_t>(
305 + /* dst */ k_transpose,
306 + /* src */ curr_k_pad,
307 + /* N */ chunk_size,
308 + /* K */ qk_head_size,
309 + /* ld_src */ qk_head_size,
310 + /* ld_dst */ chunk_size);
311 + // k_beta @ key.transpose(-1, -2)
312 + at::native::cpublas::brgemm(
313 + /* M */ chunk_size,
314 + /* N */ chunk_size,
315 + /* K */ qk_head_size,
316 + /* lda */ qk_head_size,
317 + /* ldb */ chunk_size,
318 + /* ldc */ chunk_size,
319 + /* add_C */ false,
320 + /* A */ curr_k_beta,
321 + /* B */ k_transpose,
322 + /* C */ curr_attn);
323 + // attn = attn * decay_mask
324 + for (int64_t m = 0; m < chunk_size; m++) {
325 + at::vec::map2<float>(
326 + [](fVec x, fVec y) { return fVec(0) - x * y; },
327 + curr_attn + m * chunk_size,
328 + curr_attn + m * chunk_size,
329 + curr_decay_mask + m * chunk_size,
330 + chunk_size);
331 + }
332 +
333 + // chunk decay
334 + // attn: [B, HV, num_chunk, chunk_size, chunk_size]
335 + // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0)
336 + // attn = -attn.masked_fill(mask, 0)
337 + // attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) [B, HV, num_chunk, i]
338 + // attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device)
339 + // attn = -attn.masked_fill(mask, 0)
340 + for (int i = 0; i < chunk_size; i++) {
341 + const auto vec_zero = fVec(0);
342 + int64_t len = chunk_size - i;
343 + int64_t front = len % fVecSize;
344 + int64_t j = i;
345 + // first masked vec for alignment
346 + if (front > 0) {
347 + vec_zero.store(curr_attn + i * chunk_size + j, front);
348 + j += front;
349 + }
350 + for (; j < fVecSize * (chunk_size / fVecSize); j += fVecSize) {
351 + vec_zero.store(curr_attn + i * chunk_size + j);
352 + }
353 + }
354 + for (int i = 1; i < chunk_size; i++) {
355 + // row = attn[..., i, :i] [B, HK, num_chunk, i]
356 + int64_t j = 0;
357 + int64_t len = i;
358 + for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
359 + auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j);
360 + tmp0.store(row + j);
361 + }
362 + if (j < len) {
363 + auto tmp0 = fVec::loadu(curr_attn + i * chunk_size + j, len - j);
364 + tmp0.store(row + j, len - j);
365 + }
366 + // (row.unsqueeze(-1) * sub).sum(-2)
367 + fill_stub(updated, 0, i);
368 + for (int k = 0; k < i; k++) {
369 + float row_k = row[k];
370 + auto vec_row_k = fVec(row_k);
371 + int64_t j = 0;
372 + int64_t len = i;
373 + for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
374 + auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j);
375 + auto tmp1 = vec_row_k * tmp0;
376 + auto tmp2 = fVec::loadu(updated + j);
377 + auto tmp3 = tmp1 + tmp2;
378 + tmp3.store(updated + j);
379 + }
380 + if (j < len) {
381 + auto tmp0 = fVec::loadu(curr_attn + k * chunk_size + j, len - j);
382 + auto tmp1 = vec_row_k * tmp0;
383 + auto tmp2 = fVec::loadu(updated + j);
384 + auto tmp3 = tmp1 + tmp2;
385 + tmp3.store(updated + j, len - j);
386 + }
387 + }
388 + // attn[..., i, :i] = row + sum(...)
389 + j = 0;
390 + len = i;
391 + for (; j < fVecSize * (len / fVecSize); j += fVecSize) {
392 + auto tmp0 = fVec::loadu(row + j);
393 + auto tmp1 = fVec::loadu(updated + j);
394 + auto tmp2 = tmp0 + tmp1;
395 + tmp2.store(curr_attn + i * chunk_size + j);
396 + }
397 + if (j < len) {
398 + auto tmp0 = fVec::loadu(row + j, len - j);
399 + auto tmp1 = fVec::loadu(updated + j, len - j);
400 + auto tmp2 = tmp0 + tmp1;
401 + tmp2.store(curr_attn + i * chunk_size + j, len - j);
402 + }
403 + }
404 + for (int i = 0; i < chunk_size; i++) {
405 + curr_attn[i * chunk_size + i] += 1.0f;
406 + at::vec::map<scalar_t>(
407 + [](fVec x) { return x; }, curr_attn_reduced + i * chunk_size, curr_attn + i * chunk_size, chunk_size);
408 + }
409 +
410 + // v_beta_attn = attn @ v_beta
411 + // k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1))
412 + // v_beta_attn: [B, HV, num_chunk, chunk_size, EV]
413 + // k_beta_g = k_beta * g: [B, HV, num_chunk, chunk_size, EK]
414 + // k_cumdecay: [B, HV, num_chunk, chunk_size, EK]
415 + // pack for value
416 + pack_vnni2<scalar_t>(
417 + /* dst */ v_pack,
418 + /* src */ curr_v_beta,
419 + /* N */ chunk_size,
420 + /* K */ v_head_size,
421 + /* ld_src */ v_head_size,
422 + /* ld_dst */ v_head_size);
423 + // value = attn @ v_beta
424 + at::native::cpublas::brgemm(
425 + /* M */ chunk_size,
426 + /* N */ v_head_size,
427 + /* K */ chunk_size,
428 + /* lda */ chunk_size,
429 + /* ldb */ v_head_size,
430 + /* ldc */ v_head_size,
431 + /* add_C */ false,
432 + /* A */ curr_attn_reduced,
433 + /* B */ v_pack,
434 + /* C */ curr_value);
435 + // k_beta_g = k_beta * g.exp().unsqueeze(-1)
436 + for (int64_t j = 0; j < chunk_size; j++) {
437 + int64_t i = 0;
438 + float g_exp = std::exp(curr_g_pad[j]);
439 + scalar_t g_exp_reduced = static_cast<scalar_t>(g_exp);
440 + auto vec_g_exp_reduced = bVec(g_exp_reduced);
441 + for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
442 + auto tmp0 = bVec::loadu(curr_k_beta + j * qk_head_size + i);
443 + auto tmp1 = tmp0 * vec_g_exp_reduced;
444 + tmp1.store(k_beta_g + j * qk_head_size + i);
445 + }
446 + }
447 + // pack for k_beta_g
448 + pack_vnni2<scalar_t>(
449 + /* dst */ k_beta_g_pack,
450 + /* src */ k_beta_g,
451 + /* N */ chunk_size,
452 + /* K */ qk_head_size,
453 + /* ld_src */ qk_head_size,
454 + /* ld_dst */ qk_head_size);
455 + // k_cumdecay = attn @ k_beta_g
456 + at::native::cpublas::brgemm(
457 + /* M */ chunk_size,
458 + /* N */ qk_head_size,
459 + /* K */ chunk_size,
460 + /* lda */ chunk_size,
461 + /* ldb */ qk_head_size,
462 + /* ldc */ qk_head_size,
463 + /* add_C */ false,
464 + /* A */ curr_attn_reduced,
465 + /* B */ k_beta_g_pack,
466 + /* C */ k_cumdecay);
467 + for (int i = 0; i < chunk_size; i++) {
468 + at::vec::map<scalar_t>(
469 + [](fVec x) { return x; },
470 + curr_k_cumdecay_reduced + i * qk_head_size,
471 + k_cumdecay + i * qk_head_size,
472 + qk_head_size);
473 + }
474 +
475 + // Move to the next query
476 + data_index_step(h, v_num_head, c, global_num_chunk);
477 + }
478 + });
479 +
480 + // for each chunk
481 + at::parallel_for(0, batch_size * v_num_head, 1, [&](int64_t begin, int64_t end) {
482 + int64_t b = 0, h = 0;
483 + data_index_init(begin, b, batch_size, h, v_num_head);
484 + int ompIdx = at::get_thread_num();
485 + int64_t offset =
486 + /* k_transpose */ qk_head_size * chunk_size +
487 + /* v_pack */ chunk_size * v_head_size +
488 + /* k_beta_g */ chunk_size * qk_head_size +
489 + /* k_beta_g_pack */ chunk_size * qk_head_size +
490 + /* attn */ chunk_size * chunk_size * 2 +
491 + /* attn_reduced */ chunk_size * chunk_size +
492 + /* k_cumdecay */ chunk_size * qk_head_size * 2 +
493 + /* row */ chunk_size * 2 +
494 + /* updated */ chunk_size * 2;
495 + scalar_t* thread_buff_ptr = thread_buff + ompIdx * buff_size_16bit_per_thread;
496 + THREAD_BUFFER_ALLOC(
497 + curr_last_recurrent_state_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size);
498 + THREAD_BUFFER_ALLOC(
499 + curr_last_recurrent_state_pack_reduced, thread_buff_ptr, offset, scalar_t, qk_head_size * v_head_size);
500 + THREAD_BUFFER_ALLOC(k_transpose_i, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size);
501 + THREAD_BUFFER_ALLOC(attn_i, thread_buff_ptr, offset, float, chunk_size* chunk_size * 2);
502 + THREAD_BUFFER_ALLOC(attn_i_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * chunk_size);
503 + THREAD_BUFFER_ALLOC(v_prime, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2);
504 + THREAD_BUFFER_ALLOC(v_prime_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size);
505 + THREAD_BUFFER_ALLOC(v_prime_pack_reduced, thread_buff_ptr, offset, scalar_t, chunk_size * v_head_size);
506 + THREAD_BUFFER_ALLOC(qg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
507 + THREAD_BUFFER_ALLOC(attn_inter, thread_buff_ptr, offset, float, chunk_size* v_head_size * 2);
508 + THREAD_BUFFER_ALLOC(kg, thread_buff_ptr, offset, scalar_t, chunk_size * qk_head_size);
509 + THREAD_BUFFER_ALLOC(kg_transpose, thread_buff_ptr, offset, scalar_t, qk_head_size * chunk_size);
510 + THREAD_BUFFER_ALLOC(kgv, thread_buff_ptr, offset, float, qk_head_size* v_head_size * 2);
511 +
512 + for ([[maybe_unused]] auto z : c10::irange(begin, end)) {
513 + int64_t start_q = cu_seqlens_ptr[b];
514 + int64_t seq_len = cu_seqlens_ptr[b + 1] - start_q;
515 + int64_t num_chunk = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b];
516 + int64_t chunk_offset = chunk_offsets_ptr[b];
517 + int64_t len_offset = chunk_offset * chunk_size;
518 +
519 + int64_t h_qk = h / head_group;
520 + auto out_ptr = out + start_q * oStrideT;
521 + auto curr_q = q_pad + len_offset * qk_head_size +
522 + h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK]
523 + auto curr_k = k_pad + len_offset * qk_head_size +
524 + h_qk * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK]
525 + auto curr_v = v_beta_attn + h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV]
526 + auto curr_decay_mask =
527 + decay_mask + h * global_total_seq_length * chunk_size; // [num_chunk, chunk_size, chunk_size]
528 + auto curr_k_cumdecay_reduced =
529 + k_cumdecay_reduced + h * global_total_seq_length * qk_head_size; // [num_chunk, chunk_size, EK]
530 + auto curr_last_recurrent_state =
531 + final_state_data + b * final_state_StrideN + h * final_state_StrideH; // [EK, EV]
532 + auto curr_g_pad = g_pad + len_offset + h * global_total_seq_length; // [num_chunk, chunk_size]
533 + auto curr_core_attn_out = core_attn_out + len_offset * v_head_size +
534 + h * global_total_seq_length * v_head_size; // [num_chunk, chunk_size, EV]
535 + for (int64_t c = 0; c < num_chunk; c++) {
536 + for (int i = 0; i < qk_head_size; i++) {
537 + at::vec::map<scalar_t>(
538 + [](fVec x) { return x; },
539 + curr_last_recurrent_state_reduced + i * v_head_size,
540 + curr_last_recurrent_state + i * v_head_size,
541 + v_head_size);
542 + }
543 + auto q_i = curr_q + c * chunk_size * qk_head_size; // [chunk_size, EK]
544 + auto k_i = curr_k + c * chunk_size * qk_head_size; // [chunk_size, EK]
545 + auto v_i = curr_v + (chunk_offset + c) * chunk_size * v_head_size; // [chunk_size, EV]
546 + auto decay_mask_i = curr_decay_mask + (chunk_offset + c) * chunk_size * chunk_size; // [chunk_size, chunk_size]
547 + auto k_cumdecay_i_reduced =
548 + curr_k_cumdecay_reduced + (chunk_offset + c) * chunk_size * qk_head_size; // [chunk_size, EK]
549 + auto g_pad_i = curr_g_pad + c * chunk_size; // [chunk_size]
550 + auto core_attn_out_i = curr_core_attn_out + c * chunk_size * v_head_size; // [chunk_size, EV]
551 +
552 + // attn_i = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0)
553 + // k_transpose_i = k_i.transpose(-1, -2)
554 + pack_vnni<scalar_t>(
555 + /* dst */ k_transpose_i,
556 + /* src */ k_i,
557 + /* N */ chunk_size,
558 + /* K */ qk_head_size,
559 + /* ld_src */ qk_head_size,
560 + /* ld_dst */ chunk_size);
561 + // attn_i = q_i @ k_transpose_i
562 + at::native::cpublas::brgemm(
563 + /* M */ chunk_size,
564 + /* N */ chunk_size,
565 + /* K */ qk_head_size,
566 + /* lda */ qk_head_size,
567 + /* ldb */ chunk_size,
568 + /* ldc */ chunk_size,
569 + /* add_C */ false,
570 + /* A */ q_i,
571 + /* B */ k_transpose_i,
572 + /* C */ attn_i);
573 + // attn_i = attn_i * decay_mask_i
574 + for (int64_t m = 0; m < chunk_size; m++) {
575 + auto attn_i_m = attn_i + m * chunk_size;
576 + auto attn_i_reduced_m = attn_i_reduced + m * chunk_size;
577 + auto decay_mask_i_m = decay_mask_i + m * chunk_size;
578 + int64_t n = 0;
579 + for (; n < fVecSize * (chunk_size / fVecSize); n += fVecSize) {
580 + auto tmp0 = fVec::loadu(attn_i_m + n);
581 + auto tmp1 = fVec::loadu(decay_mask_i_m + n);
582 + auto tmp2 = tmp0 * tmp1;
583 + auto tmp3 = at::vec::convert<scalar_t>(tmp2);
584 + tmp3.store(attn_i_reduced_m + n, fVecSize);
585 + }
586 + if (n < chunk_size) {
587 + auto tmp0 = fVec::loadu(attn_i_m + n, chunk_size - n);
588 + auto tmp1 = fVec::loadu(decay_mask_i_m + n, chunk_size - n);
589 + auto tmp2 = tmp0 * tmp1;
590 + auto tmp3 = at::vec::convert<scalar_t>(tmp2);
591 + tmp3.store(attn_i_reduced_m + n, chunk_size - n);
592 + }
593 + }
594 + // mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1)
595 + // attn_i = attn_i.masked_fill_(mask, 0)
596 + for (int i = 0; i < chunk_size - 1; i++) {
597 + const auto vec_zero = bVec(0);
598 + int64_t len = chunk_size - i - 1;
599 + int64_t front = len % VecSize;
600 + int64_t j = i + 1;
601 + // first masked vec for alignment
602 + if (front > 0) {
603 + vec_zero.store(attn_i_reduced + i * chunk_size + j, front);
604 + j += front;
605 + }
606 + for (; j < VecSize * (chunk_size / VecSize); j += VecSize) {
607 + vec_zero.store(attn_i_reduced + i * chunk_size + j);
608 + }
609 + }
610 +
611 + // pack for curr_last_recurrent_state
612 + pack_vnni2<scalar_t>(
613 + /* dst */ curr_last_recurrent_state_pack_reduced,
614 + /* src */ curr_last_recurrent_state_reduced,
615 + /* N */ qk_head_size,
616 + /* K */ v_head_size,
617 + /* ld_src */ v_head_size,
618 + /* ld_dst */ v_head_size);
619 +
620 + // v_prime = k_cumdecay_i @ curr_last_recurrent_state: [chunk_size, EV]
621 + // k_cumdecay_i: [chunk_size, EK]
622 + // curr_last_recurrent_state: [EK, EV]
623 + at::native::cpublas::brgemm(
624 + /* M */ chunk_size,
625 + /* N */ v_head_size,
626 + /* K */ qk_head_size,
627 + /* lda */ qk_head_size,
628 + /* ldb */ v_head_size,
629 + /* ldc */ v_head_size,
630 + /* add_C */ false,
631 + /* A */ k_cumdecay_i_reduced,
632 + /* B */ curr_last_recurrent_state_pack_reduced,
633 + /* C */ v_prime);
634 +
635 + // v_new = v_prime = v_i - v_prime
636 + // v_i: [chunk_size, EV]
637 + for (int64_t m = 0; m < chunk_size; m++) {
638 + int64_t i = 0;
639 + for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) {
640 + auto tmp0 = fVec::loadu(v_i + m * v_head_size + i);
641 + auto tmp1 = fVec::loadu(v_prime + m * v_head_size + i);
642 + auto tmp2 = tmp0 - tmp1;
643 + auto tmp3 = at::vec::convert<scalar_t>(tmp2);
644 + tmp3.store(v_prime_reduced + m * v_head_size + i, fVecSize);
645 + }
646 + }
647 +
648 + // attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state
649 + // qg = q_i * g[:, :, i, :, None].exp(): [chunk_size, EK]
650 + // q_i: [chunk_size, EK]
651 + // g[:, :, i, :, None]: [chunk_size, 1]
652 + for (int64_t m = 0; m < chunk_size; m++) {
653 + auto g_pad_i_m = g_pad_i + m;
654 + auto g_exp = std::exp(*g_pad_i_m);
655 + int64_t i = 0;
656 + scalar_t g_exp_reduced = static_cast<scalar_t>(g_exp);
657 + auto vec_g_exp_reduced = bVec(g_exp_reduced);
658 + for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
659 + auto tmp0 = bVec::loadu(q_i + m * qk_head_size + i);
660 + auto tmp2 = tmp0 * vec_g_exp_reduced;
661 + tmp2.store(qg + m * qk_head_size + i);
662 + }
663 + }
664 + // attn_inter = qg @ curr_last_recurrent_state: [chunk_size, EV]
665 + // curr_last_recurrent_state: [EK, EV]
666 + at::native::cpublas::brgemm(
667 + /* M */ chunk_size,
668 + /* N */ v_head_size,
669 + /* K */ qk_head_size,
670 + /* lda */ qk_head_size,
671 + /* ldb */ v_head_size,
672 + /* ldc */ v_head_size,
673 + /* add_C */ false,
674 + /* A */ qg,
675 + /* B */ curr_last_recurrent_state_pack_reduced,
676 + /* C */ attn_inter);
677 +
678 + // core_attn_out[:, :, i] = attn_inter + attn_i @ v_new
679 + // pack for v_prime
680 + pack_vnni2<scalar_t>(
681 + /* dst */ v_prime_pack_reduced,
682 + /* src */ v_prime_reduced,
683 + /* N */ chunk_size,
684 + /* K */ v_head_size,
685 + /* ld_src */ v_head_size,
686 + /* ld_dst */ v_head_size);
687 + // attn_inter = attn_inter + attn_i @ v_new: [chunk_size, EV]
688 + // attn_i: [chunk_size, chunk_size]
689 + // v_new: [chunk_size, EV]
690 + at::native::cpublas::brgemm(
691 + /* M */ chunk_size,
692 + /* N */ v_head_size,
693 + /* K */ chunk_size,
694 + /* lda */ chunk_size,
695 + /* ldb */ v_head_size,
696 + /* ldc */ v_head_size,
697 + /* add_C */ true,
698 + /* A */ attn_i_reduced,
699 + /* B */ v_prime_pack_reduced,
700 + /* C */ attn_inter);
701 +
702 + // core_attn_out[:, :, i] = attn_inter
703 + for (int64_t m = 0; m < chunk_size; m++) {
704 + at::vec::map<float>(
705 + [](fVec x) { return x; }, core_attn_out_i + m * v_head_size, attn_inter + m * v_head_size, v_head_size);
706 + }
707 +
708 + // last_recurrent_state = (
709 + // last_recurrent_state * g[:, :, i, -1, None, None].exp()
710 + // + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
711 + // )
712 + // 1) last_recurrent_state * g[:, :, i, -1, None, None].exp()
713 + // curr_last_recurrent_state: [EK, EV]
714 + // g[:, :, i, -1, None, None]: [1, 1]
715 + // last_recurrent_state * g[:, :, i, -1, None, None].exp(): [EK, EV]
716 + auto g_pad_i_last = g_pad_i + chunk_size - 1;
717 + auto g_exp_last = std::exp(g_pad_i_last[0]);
718 + for (int64_t m = 0; m < qk_head_size; m++) {
719 + int64_t i = 0;
720 + auto vec_g_exp_last = fVec(g_exp_last);
721 + for (; i < fVecSize * (v_head_size / fVecSize); i += fVecSize) {
722 + auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i);
723 + auto tmp1 = at::vec::convert<float>(tmp0);
724 + auto tmp2 = tmp1 * vec_g_exp_last;
725 + tmp2.store(curr_last_recurrent_state + m * v_head_size + i);
726 + }
727 + if (i < v_head_size) {
728 + auto tmp0 = bVec::loadu(curr_last_recurrent_state_reduced + m * v_head_size + i, v_head_size - i);
729 + auto tmp1 = at::vec::convert<float>(tmp0);
730 + auto tmp2 = tmp1 * vec_g_exp_last;
731 + tmp2.store(curr_last_recurrent_state + m * v_head_size + i, v_head_size - i);
732 + }
733 + }
734 + // 2) (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new
735 + // k_i: [chunk_size, EK]
736 + // g[:, :, i, -1, None]: [1]
737 + // g[:, :, i]: [chunk_size]
738 + // (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, 1]
739 + // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]: [chunk_size, EK]
740 + // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2): [EK, chunk_size]
741 + // v_new: [chunk_size, EV]
742 + // (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new: [EK, EV]
743 + // kg = k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]
744 + for (int64_t m = 0; m < chunk_size; m++) {
745 + auto g_exp = std::exp((g_pad_i_last[0] - g_pad_i[m]));
746 + int64_t i = 0;
747 + scalar_t g_exp_reduced = static_cast<scalar_t>(g_exp);
748 + auto vec_g_exp_reduced = bVec(g_exp_reduced);
749 + for (; i < VecSize * (qk_head_size / VecSize); i += VecSize) {
750 + auto tmp0 = bVec::loadu(k_i + m * qk_head_size + i);
751 + auto tmp2 = tmp0 * vec_g_exp_reduced;
752 + tmp2.store(kg + m * qk_head_size + i);
753 + }
754 + }
755 + // kg.transpose(-1, -2): [EK, chunk_size]
756 + at::native::utils::transpose<scalar_t>(
757 + /* M */ chunk_size,
758 + /* N */ qk_head_size,
759 + /* src */ kg,
760 + /* ld_src */ qk_head_size,
761 + /* dst */ kg_transpose,
762 + /* ld_dst */ chunk_size);
763 + // kgv = kg.transpose(-1, -2) @ v_new
764 + // v_new: [chunk_size, EV]
765 + at::native::cpublas::brgemm(
766 + /* M */ qk_head_size,
767 + /* N */ v_head_size,
768 + /* K */ chunk_size,
769 + /* lda */ chunk_size,
770 + /* ldb */ v_head_size,
771 + /* ldc */ v_head_size,
772 + /* add_C */ false,
773 + /* A */ kg_transpose,
774 + /* B */ v_prime_pack_reduced,
775 + /* C */ kgv);
776 + // last_recurrent_state = 1) + 2)
777 + for (int64_t m = 0; m < qk_head_size; m++) {
778 + at::vec::map2<float>(
779 + [](fVec x, fVec y) { return x + y; },
780 + curr_last_recurrent_state + m * v_head_size,
781 + curr_last_recurrent_state + m * v_head_size,
782 + kgv + m * v_head_size,
783 + v_head_size);
784 + }
785 + }
786 +
787 + // core_attn_out -> output
788 + // output: [B, T, HV, EV]
789 + // core_attn_out: [B, HV, padded_T, EV]
790 + auto curr_out = out_ptr + h * oStrideH;
791 + for (int64_t m = 0; m < seq_len; m++) {
792 + at::vec::map<scalar_t>(
793 + [](fVec x) { return x; }, curr_out + m * oStrideT, curr_core_attn_out + m * v_head_size, v_head_size);
794 + }
795 +
796 + // Move to the next query
797 + data_index_step(b, batch_size, h, v_num_head);
798 + }
799 + });
800 +}
801 +
802 +inline float softplus(float x, double threshold = 20.0) {
803 + if (x > threshold)
804 + return x;
805 + else if (x < -threshold)
806 + return std::exp(x);
807 + else
808 + return std::log1p(std::exp(x));
809 +}
810 +
811 +inline at::vec::Vectorized<float> softplus(const at::vec::Vectorized<float>& x, double threshold = 20.0) {
812 + using Vec = at::vec::Vectorized<float>;
813 + Vec mask_hi = x > Vec(threshold);
814 + Vec mask_lo = x < Vec(-threshold);
815 +
816 + Vec expx = x.exp_u20();
817 + Vec log1pex = (expx + Vec(1.0f)).log();
818 +
819 + return Vec::blendv(Vec::blendv(log1pex, expx, mask_lo), x, mask_hi);
820 +}
821 +
822 +template <typename scalar_t, typename param_t>
823 +void fused_sigmoid_gating_delta_rule_update_kernel_impl(
824 + const scalar_t* __restrict__ q_ptr,
825 + const scalar_t* __restrict__ k_ptr,
826 + const scalar_t* __restrict__ v_ptr,
827 + const param_t* __restrict__ A_log_ptr,
828 + const scalar_t* __restrict__ a_ptr,
829 + const scalar_t* __restrict__ dt_bias_ptr,
830 + const scalar_t* __restrict__ b_ptr,
831 + const int32_t* __restrict__ indices_ptr,
832 + float* __restrict__ state_ptr,
833 + scalar_t* __restrict__ o_ptr,
834 + float* __restrict__ qk_scale_buf,
835 + int64_t seq_len,
836 + int64_t batch_size,
837 + int64_t num_heads,
838 + int64_t head_dim,
839 + int64_t v_num_heads,
840 + int64_t v_head_dim,
841 + int64_t q_strideB,
842 + int64_t q_strideS,
843 + int64_t q_strideH,
844 + int64_t k_strideB,
845 + int64_t k_strideS,
846 + int64_t k_strideH,
847 + int64_t v_strideB,
848 + int64_t v_strideS,
849 + int64_t v_strideH,
850 + bool use_qk_l2norm_in_kernel,
851 + double softplus_threshold) {
852 + using bVec = at::vec::Vectorized<scalar_t>;
853 + using fVec = at::vec::Vectorized<float>;
854 +
855 + constexpr int64_t VecSize = bVec::size();
856 + constexpr int64_t fVecSize = fVec::size();
857 + int64_t group_size = v_num_heads / num_heads;
858 + double scale = 1 / std::sqrt(head_dim);
859 + fVec scale_vec = fVec(scale);
860 + if (use_qk_l2norm_in_kernel) {
861 + float eps = 1e-5;
862 + at::parallel_for(0, batch_size * seq_len * num_heads, 0, [&](int64_t begin, int64_t end) {
863 + int64_t bi{0}, si{0}, ni{0};
864 + data_index_init(begin, bi, batch_size, si, seq_len, ni, num_heads);
865 + for (int64_t i = begin; i < end; ++i) {
866 + float sum_q = float(0);
867 + float sum_k = float(0);
868 + fVec sum_q_fvec = fVec(float(0));
869 + fVec sum_k_fvec = fVec(float(0));
870 + int64_t q_offset = bi * q_strideB + si * q_strideS + ni * q_strideH;
871 + int64_t k_offset = bi * k_strideB + si * k_strideS + ni * k_strideH;
872 + int64_t q_scale_offset = bi * seq_len * num_heads + si * num_heads + ni;
873 + int64_t k_scale_offset = q_scale_offset + batch_size * seq_len * num_heads;
874 + int64_t d;
875 +#pragma GCC unroll 4
876 + for (d = 0; d <= head_dim - VecSize; d += VecSize) {
877 + bVec q_bvec = bVec::loadu(q_ptr + q_offset + d);
878 + fVec q_fvec0, q_fvec1;
879 + std::tie(q_fvec0, q_fvec1) = at::vec::convert_to_float(q_bvec);
880 + sum_q_fvec += q_fvec0 * q_fvec0;
881 + sum_q_fvec += q_fvec1 * q_fvec1;
882 + bVec k_bvec = bVec::loadu(k_ptr + k_offset + d);
883 + fVec k_fvec0, k_fvec1;
884 + std::tie(k_fvec0, k_fvec1) = at::vec::convert_to_float(k_bvec);
885 + sum_k_fvec += k_fvec0 * k_fvec0;
886 + sum_k_fvec += k_fvec1 * k_fvec1;
887 + }
888 +#pragma GCC unroll 4
889 + for (; d < head_dim; ++d) {
890 + float q_val = static_cast<float>(q_ptr[q_offset + d]);
891 + sum_q += q_val * q_val;
892 + float k_val = static_cast<float>(k_ptr[k_offset + d]);
893 + sum_k += k_val * k_val;
894 + }
895 +
896 + sum_q += vec_reduce_sum(sum_q_fvec);
897 + sum_k += vec_reduce_sum(sum_k_fvec);
898 + qk_scale_buf[q_scale_offset] = float(1) / std::sqrt(sum_q + eps);
899 + qk_scale_buf[k_scale_offset] = float(1) / std::sqrt(sum_k + eps);
900 +
901 + data_index_step(bi, batch_size, si, seq_len, ni, num_heads);
902 + }
903 + });
904 + }
905 + at::parallel_for(0, batch_size * seq_len * v_num_heads, 0, [&](int64_t begin, int64_t end) {
906 + int64_t bi{0}, si{0}, ni{0};
907 + data_index_init(begin, bi, batch_size, si, seq_len, ni, v_num_heads);
908 + for (int64_t i = begin; i < end; ++i) {
909 + int64_t cache_index = indices_ptr[bi];
910 + int64_t state_offset = (cache_index * v_num_heads + ni) * head_dim * v_head_dim;
911 + float g_val = -std::exp(float(A_log_ptr[ni])) *
912 + softplus(float(a_ptr[bi * v_num_heads + ni]) + float(dt_bias_ptr[ni]), softplus_threshold);
913 + float g_val_exp = std::exp(g_val);
914 + fVec g_val_exp_vec = fVec(g_val_exp);
915 + int64_t q_offset = si * q_strideS + bi * q_strideB + (ni / group_size) * q_strideH;
916 + int64_t k_offset = si * k_strideS + bi * k_strideB + (ni / group_size) * k_strideH;
917 + int64_t q_scale_offset = bi * seq_len * num_heads + si * num_heads + (ni / group_size);
918 + int64_t k_scale_offset = q_scale_offset + batch_size * seq_len * num_heads;
919 + float q_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[q_scale_offset] : 1.0f;
920 + float k_scale = use_qk_l2norm_in_kernel ? qk_scale_buf[k_scale_offset] : 1.0f;
921 + int64_t v_offset = si * v_strideS + bi * v_strideB + ni * v_strideH;
922 + int64_t o_offset = ((bi * seq_len + si) * v_num_heads + ni) * v_head_dim;
923 + float beta_val = 1 / (1 + std::exp(-b_ptr[ni]));
924 + fVec beta_vec = fVec(beta_val);
925 + int64_t dvi = 0;
926 + for (; dvi <= v_head_dim - VecSize; dvi += VecSize) {
927 + fVec kv_mem_vec0 = fVec(float(0));
928 + fVec kv_mem_vec1 = fVec(float(0));
929 + for (int di = 0; di < head_dim; ++di) {
930 + fVec k_val_vec = fVec(k_ptr[k_offset + di] * k_scale);
931 + fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi);
932 + fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
933 + kv_mem_vec0 = kv_mem_vec0 + state_vec0 * g_val_exp_vec * k_val_vec;
934 + kv_mem_vec1 = kv_mem_vec1 + state_vec1 * g_val_exp_vec * k_val_vec;
935 + }
936 + bVec v_bvec = bVec::loadu(v_ptr + v_offset + dvi);
937 + fVec v_vec0, v_vec1;
938 + std::tie(v_vec0, v_vec1) = at::vec::convert_to_float(v_bvec);
939 + fVec dt_vec0 = (v_vec0 - kv_mem_vec0) * beta_vec;
940 + fVec dt_vec1 = (v_vec1 - kv_mem_vec1) * beta_vec;
941 + fVec o_vec0 = fVec(float(0));
942 + fVec o_vec1 = fVec(float(0));
943 + for (int di = 0; di < head_dim; ++di) {
944 + fVec q_vec = fVec(q_ptr[q_offset + di] * q_scale);
945 + fVec k_vec = fVec(k_ptr[k_offset + di] * k_scale);
946 + fVec state_vec0 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi);
947 + fVec state_vec1 = fVec::loadu(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
948 + state_vec0 = state_vec0 * g_val_exp_vec + k_vec * dt_vec0;
949 + state_vec1 = state_vec1 * g_val_exp_vec + k_vec * dt_vec1;
950 + o_vec0 = o_vec0 + state_vec0 * q_vec * scale_vec;
951 + o_vec1 = o_vec1 + state_vec1 * q_vec * scale_vec;
952 + state_vec0.store(state_ptr + state_offset + di * v_head_dim + dvi);
953 + state_vec1.store(state_ptr + state_offset + di * v_head_dim + dvi + fVecSize);
954 + }
955 + bVec o_vec = at::vec::convert_from_float<scalar_t>(o_vec0, o_vec1);
956 + o_vec.store(o_ptr + o_offset + dvi);
957 + }
958 + for (; dvi < v_head_dim; ++dvi) {
959 + float kv_mem_val = 0;
960 + for (int di = 0; di < head_dim; ++di) {
961 + float k_val = k_ptr[k_offset + di] * k_scale;
962 + state_ptr[state_offset + di * v_head_dim + dvi] *= g_val_exp;
963 + kv_mem_val += state_ptr[state_offset + di * v_head_dim + dvi] * k_val;
964 + }
965 + float v_val = v_ptr[v_offset + dvi];
966 + float dt_val = (v_val - kv_mem_val) * beta_val;
967 + float o_val = 0;
968 + for (int di = 0; di < head_dim; ++di) {
969 + float q_val = q_ptr[q_offset + di] * q_scale;
970 + float k_val = k_ptr[k_offset + di] * k_scale;
971 + state_ptr[state_offset + di * v_head_dim + dvi] += k_val * dt_val;
972 + o_val += state_ptr[state_offset + di * v_head_dim + dvi] * q_val * scale;
973 + }
974 + o_ptr[o_offset + dvi] = o_val;
975 + }
976 + data_index_step(bi, batch_size, si, seq_len, ni, v_num_heads);
977 + }
978 + });
979 +}
980 +
981 +template <typename scalar_t>
982 +void fused_gdn_gating_kernel_impl(
983 + float* __restrict__ A_log,
984 + const scalar_t* __restrict__ a,
985 + const scalar_t* __restrict__ b,
986 + const scalar_t* __restrict__ dt_bias,
987 + float* __restrict__ out,
988 + scalar_t* __restrict__ beta,
989 + int64_t batch,
990 + int64_t num_heads) {
991 + using bVec = at::vec::Vectorized<scalar_t>;
992 + using fVec = at::vec::Vectorized<float>;
993 + constexpr int vec_size = bVec::size();
994 + constexpr int fvec_size = fVec::size();
995 + const fVec neg_one(-1.0f);
996 + const fVec one(1.0f);
997 + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
998 + for (int64_t i = begin; i < end; ++i) {
999 + int64_t j = 0;
1000 + for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
1001 + fVec A_log_vec0 = fVec::loadu(A_log + j);
1002 + fVec A_log_vec1 = fVec::loadu(A_log + j + fvec_size);
1003 + bVec dt_bias_vec = bVec::loadu(dt_bias + j);
1004 + bVec a_bvec = bVec::loadu(a + i * num_heads + j);
1005 + bVec b_bvec = bVec::loadu(b + i * num_heads + j);
1006 + fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
1007 + std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
1008 + std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
1009 + std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
1010 +
1011 + fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
1012 + fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
1013 + fVec beta0 = one / (one + (neg_one * b0).exp_u20());
1014 + fVec beta1 = one / (one + (neg_one * b1).exp_u20());
1015 +
1016 + g0.store(out + i * num_heads + j);
1017 + g1.store(out + i * num_heads + j + fvec_size);
1018 + bVec beta_vec = at::vec::convert_from_float<scalar_t>(beta0, beta1);
1019 + beta_vec.store(beta + i * num_heads + j);
1020 + }
1021 + for (; j < num_heads; ++j) {
1022 + out[i * num_heads + j] = -std::exp(A_log[j]) * softplus(float(a[i * num_heads + j]) + float(dt_bias[j]));
1023 + beta[i * num_heads + j] = 1 / (1 + std::exp(-b[i * num_heads + j]));
1024 + }
1025 + }
1026 + });
1027 +}
1028 +
1029 +template <typename scalar_t>
1030 +void fused_gdn_gating_kernel_impl(
1031 + scalar_t* __restrict__ A_log,
1032 + const scalar_t* __restrict__ a,
1033 + const scalar_t* __restrict__ b,
1034 + const scalar_t* __restrict__ dt_bias,
1035 + float* __restrict__ out,
1036 + scalar_t* __restrict__ beta,
1037 + int64_t batch,
1038 + int64_t num_heads) {
1039 + using bVec = at::vec::Vectorized<scalar_t>;
1040 + using fVec = at::vec::Vectorized<float>;
1041 + constexpr int vec_size = bVec::size();
1042 + constexpr int fvec_size = fVec::size();
1043 + const fVec neg_one(-1.0f);
1044 + const fVec one(1.0f);
1045 + at::parallel_for(0, batch, 0, [&](int64_t begin, int64_t end) {
1046 + for (int64_t i = begin; i < end; ++i) {
1047 + int64_t j = 0;
1048 + for (; j < num_heads - (num_heads % vec_size); j += vec_size) {
1049 + bVec A_log_bvec = bVec::loadu(A_log + j);
1050 + fVec A_log_vec0, A_log_vec1;
1051 + std::tie(A_log_vec0, A_log_vec1) = at::vec::convert_to_float(A_log_bvec);
1052 + bVec dt_bias_vec = bVec::loadu(dt_bias + j);
1053 + bVec a_bvec = bVec::loadu(a + i * num_heads + j);
1054 + bVec b_bvec = bVec::loadu(b + i * num_heads + j);
1055 + fVec a0, a1, dt_bias_vec0, dt_bias_vec1, b0, b1;
1056 + std::tie(a0, a1) = at::vec::convert_to_float(a_bvec);
1057 + std::tie(b0, b1) = at::vec::convert_to_float(b_bvec);
1058 + std::tie(dt_bias_vec0, dt_bias_vec1) = at::vec::convert_to_float(dt_bias_vec);
1059 +
1060 + fVec g0 = neg_one * A_log_vec0.exp_u20() * softplus(a0 + dt_bias_vec0);
1061 + fVec g1 = neg_one * A_log_vec1.exp_u20() * softplus(a1 + dt_bias_vec1);
1062 + fVec beta0 = one / (one + (neg_one * b0).exp_u20());
1063 + fVec beta1 = one / (one + (neg_one * b1).exp_u20());
1064 +
1065 + g0.store(out + i * num_heads + j);
1066 + g1.store(out + i * num_heads + j + fvec_size);
1067 + bVec beta_vec = at::vec::convert_from_float<scalar_t>(beta0, beta1);
1068 + beta_vec.store(beta + i * num_heads + j);
1069 + }
1070 + for (; j < num_heads; ++j) {
1071 + out[i * num_heads + j] = -std::exp(float(A_log[j])) * softplus(float(a[i * num_heads + j]) + float(dt_bias[j]));
1072 + beta[i * num_heads + j] = 1 / (1 + std::exp(-b[i * num_heads + j]));
1073 + }
1074 + }
1075 + });
1076 +}
1077 +
1078 +} // anonymous namespace
1079 +
1080 +template <bool is_last_dim_contiguous>
1081 +inline void
1082 +CHECK_INPUT_SHAPE_DTYPE(const at::Tensor& tensor, const int64_t& dim, const at::IntArrayRef& sizes, at::ScalarType st) {
1083 + TORCH_CHECK(tensor.sizes() == sizes, "Input tensor shape mismatch: expected ", sizes, ", got ", tensor.sizes());
1084 + TORCH_CHECK(tensor.dtype() == st, "Input tensor dtype mismatch");
1085 + CHECK_DIM(dim, tensor);
1086 + if (is_last_dim_contiguous) {
1087 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(tensor);
1088 + } else {
1089 + CHECK_CONTIGUOUS(tensor);
1090 + }
1091 +}
1092 +
1093 +// query: [B, T, HK, EK]
1094 +// key: [B, T, HK, EK]
1095 +// value: [B, T, HV, EV]
1096 +// g: [B, T, HV] FP32
1097 +// beta: [B, T, HV]
1098 +// initial_state: [N, HV, EK, EV] FP32
1099 +// output_final_state: bool
1100 +// cu_seqlens: [N + 1] INT32
1101 +// head_first: bool
1102 +// use_qk_l2norm_in_kernel: bool
1103 +std::tuple<at::Tensor, at::Tensor> chunk_gated_delta_rule_cpu(
1104 + const at::Tensor& query,
1105 + const at::Tensor& key,
1106 + const at::Tensor& value,
1107 + const at::Tensor& g,
1108 + const at::Tensor& beta,
1109 + const at::Tensor& initial_state,
1110 + bool output_final_state,
1111 + const at::Tensor& cu_seqlens,
1112 + bool head_first,
1113 + bool use_qk_l2norm_in_kernel,
1114 + double eps = 1e-5) {
1115 + TORCH_CHECK(head_first == false, "chunk_gated_delta_rule_cpu does not support head first");
1116 + int64_t B = query.size(0);
1117 + int64_t global_seq_len = query.size(1);
1118 + int64_t qk_num_head = query.size(2);
1119 + int64_t qk_head_size = query.size(3);
1120 + int64_t v_num_head = value.size(2);
1121 + int64_t v_head_size = value.size(3);
1122 + int64_t batch_size = initial_state.size(0);
1123 + CHECK_EQ(B, 1);
1124 + TORCH_CHECK(v_num_head % qk_num_head == 0, "expect v_num_head multiple of qk_num_head.");
1125 + TORCH_CHECK(qk_head_size % 32 == 0, "expect qk_head_size to be multiples of 32.");
1126 + TORCH_CHECK(v_head_size % 32 == 0, "expect v_head_size to be multiples of 32.");
1127 + CHECK_INPUT_SHAPE_DTYPE<true>(query, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16);
1128 + CHECK_INPUT_SHAPE_DTYPE<true>(key, 4, {B, global_seq_len, qk_num_head, qk_head_size}, at::kBFloat16);
1129 + CHECK_INPUT_SHAPE_DTYPE<true>(value, 4, {B, global_seq_len, v_num_head, v_head_size}, at::kBFloat16);
1130 + CHECK_INPUT_SHAPE_DTYPE<false>(g, 3, {B, global_seq_len, v_num_head}, at::kFloat);
1131 + CHECK_INPUT_SHAPE_DTYPE<false>(beta, 3, {B, global_seq_len, v_num_head}, at::kBFloat16);
1132 + CHECK_INPUT_SHAPE_DTYPE<false>(cu_seqlens, 1, {batch_size + 1}, at::kInt);
1133 + CHECK_INPUT_SHAPE_DTYPE<false>(initial_state, 4, {batch_size, v_num_head, qk_head_size, v_head_size}, at::kFloat);
1134 +
1135 + at::Tensor output = at::empty_like(value, value.options()); // [B, T, HV, EV]
1136 + at::Tensor final_state = initial_state.to(at::kFloat); // [N, HV, EK, EV]
1137 +
1138 + // Strides
1139 + int64_t qStrideH = query.stride(2);
1140 + int64_t qStrideT = query.stride(1);
1141 + int64_t kStrideH = key.stride(2);
1142 + int64_t kStrideT = key.stride(1);
1143 + int64_t vStrideH = value.stride(2);
1144 + int64_t vStrideT = value.stride(1);
1145 + int64_t oStrideH = output.stride(2);
1146 + int64_t oStrideT = output.stride(1);
1147 +
1148 + constexpr int64_t chunk_size = 64;
1149 + // Deduce the global chunks
1150 + // e.g. cu_seqlens: [0, 5, 13, 16], chunk_size = 4
1151 + // chunk_offsets: [0, 2, 4, 5]
1152 + // chunk_indices (batch_id, local_chunk_id): [[0, 0], [0, 1], [1, 0], [1, 1], [2, 0]]
1153 + at::Tensor chunk_offsets = at::empty(batch_size + 1, cu_seqlens.options());
1154 + auto chunk_offsets_ptr = chunk_offsets.data_ptr<int32_t>();
1155 + chunk_offsets_ptr[0] = 0;
1156 + int32_t* cu_seqlens_ptr = cu_seqlens.data_ptr<int32_t>();
1157 + int64_t s = 0;
1158 + int64_t e = 0;
1159 + int64_t s_pad = 0;
1160 + int64_t e_pad = 0;
1161 + for (int64_t b = 0; b < batch_size; b++) {
1162 + e = cu_seqlens_ptr[b + 1];
1163 + int64_t seq_len = e - s;
1164 + int64_t pad_size = (chunk_size - seq_len % chunk_size) % chunk_size;
1165 + int64_t total_seq_length = seq_len + pad_size;
1166 + e_pad = s_pad + total_seq_length;
1167 + chunk_offsets[b + 1] = e_pad / chunk_size;
1168 + s = e;
1169 + s_pad = e_pad;
1170 + }
1171 + int64_t global_total_seq_length = e_pad;
1172 + int64_t global_num_chunk = chunk_offsets_ptr[batch_size];
1173 + at::Tensor chunk_indices = at::empty(global_num_chunk * 2, cu_seqlens.options());
1174 + auto chunk_indices_ptr = chunk_indices.data_ptr<int32_t>();
1175 + int64_t curr_c = 0;
1176 + for (int64_t b = 0; b < batch_size; b++) {
1177 + int64_t batch_chunk_num = chunk_offsets_ptr[b + 1] - chunk_offsets_ptr[b];
1178 + for (int64_t c = 0; c < batch_chunk_num; c++) {
1179 + chunk_indices_ptr[curr_c * 2] = b;
1180 + chunk_indices_ptr[curr_c * 2 + 1] = c;
1181 + curr_c += 1;
1182 + }
1183 + }
1184 +
1185 + // Allocate buffer
1186 + int64_t buff_size = v_num_head * global_total_seq_length // g_pad_data
1187 + + batch_size * v_num_head * global_total_seq_length * v_head_size // core_attn
1188 + + v_num_head * global_total_seq_length * chunk_size // decay_mask
1189 + + v_num_head * global_total_seq_length * v_head_size; // v_beta_attn
1190 + at::Tensor buff_data = at::empty({buff_size}, query.options().dtype(at::kFloat));
1191 + int64_t reduced_buff_size = qk_num_head * global_total_seq_length * qk_head_size // q_pad_data
1192 + + qk_num_head * global_total_seq_length * qk_head_size // k_pad_data
1193 + + v_num_head * global_total_seq_length * v_head_size // v_pad_data
1194 + + v_num_head * global_total_seq_length * qk_head_size // k_beta_data
1195 + + v_num_head * global_total_seq_length * v_head_size // v_beta_data
1196 + + v_num_head * global_total_seq_length * qk_head_size // k_cumdecay_reduced
1197 + + qk_num_head * global_seq_len // q_norm_sum
1198 + + qk_num_head * global_seq_len; // k_norm_sum
1199 + at::Tensor reduced_buff_data = at::empty({reduced_buff_size}, query.options());
1200 + int64_t num_thread = at::get_num_threads();
1201 + int64_t buff_size_16bit_per_thread =
1202 + /* k_transpose */ qk_head_size * chunk_size +
1203 + /* v_pack */ chunk_size * v_head_size +
1204 + /* k_beta_g */ chunk_size * qk_head_size +
1205 + /* k_beta_g_pack */ chunk_size * qk_head_size +
1206 + /* attn */ chunk_size * chunk_size * 2 +
1207 + /* attn_reduced */ chunk_size * chunk_size +
1208 + /* k_cumdecay */ chunk_size * qk_head_size * 2 +
1209 + /* row */ chunk_size * 2 +
1210 + /* updated */ chunk_size * 2 +
1211 + /* curr_last_recurrent_state_reduced */ qk_head_size * v_head_size +
1212 + /* curr_last_recurrent_state_pack_reduced */ qk_head_size * v_head_size +
1213 + /* k_transpose_i */ qk_head_size * chunk_size +
1214 + /* attn_i */ chunk_size * chunk_size * 2 +
1215 + /* attn_i_reduced */ chunk_size * chunk_size +
1216 + /* v_prime */ chunk_size * v_head_size * 2 +
1217 + /* v_prime_reduced */ chunk_size * v_head_size +
1218 + /* v_prime_pack_reduced */ chunk_size * v_head_size +
1219 + /* qg */ chunk_size * qk_head_size +
1220 + /* attn_inter */ chunk_size * v_head_size * 2 +
1221 + /* kg */ chunk_size * qk_head_size +
1222 + /* kg_transpose */ qk_head_size * chunk_size +
1223 + /* kgv */ qk_head_size * v_head_size * 2;
1224 + at::Tensor thread_buff_data = at::empty({num_thread, buff_size_16bit_per_thread}, query.options());
1225 +
1226 + AT_DISPATCH_REDUCED_FLOATING_TYPES(query.scalar_type(), "chunk_gated_delta_rule_kernel", [&] {
1227 + chunk_gated_delta_rule_kernel_impl<scalar_t, chunk_size>(
1228 + output.data_ptr<scalar_t>(),
1229 + final_state.data_ptr<float>(),
1230 + query.data_ptr<scalar_t>(),
1231 + key.data_ptr<scalar_t>(),
1232 + value.data_ptr<scalar_t>(),
1233 + g.data_ptr<float>(),
1234 + beta.data_ptr<scalar_t>(),
1235 + cu_seqlens_ptr,
1236 + buff_data.data_ptr<float>(),
1237 + reduced_buff_data.data_ptr<scalar_t>(),
1238 + thread_buff_data.data_ptr<scalar_t>(),
1239 + chunk_offsets_ptr,
1240 + chunk_indices_ptr,
1241 + use_qk_l2norm_in_kernel,
1242 + batch_size,
1243 + global_seq_len,
1244 + qk_num_head,
1245 + v_num_head,
1246 + qk_head_size,
1247 + v_head_size,
1248 + qStrideH,
1249 + qStrideT,
1250 + kStrideH,
1251 + kStrideT,
1252 + vStrideH,
1253 + vStrideT,
1254 + oStrideH,
1255 + oStrideT,
1256 + global_total_seq_length,
1257 + global_num_chunk,
1258 + buff_size_16bit_per_thread,
1259 + eps);
1260 + });
1261 + return std::make_tuple(std::move(output), std::move(final_state));
1262 +}
1263 +
1264 +// A_log: [v_num_heads]
1265 +// dt_bias: [v_num_heads]
1266 +// query: [seq_len, batch_size, num_heads, head_dim]
1267 +// key: [seq_len, batch_size, num_heads, head_dim]
1268 +// value: [seq_len, batch_size, v_num_heads, v_head_dim]
1269 +// a: [batch_size, v_num_heads]
1270 +// b: [batch_size, v_num_heads]
1271 +// initial_state_source:[num_tokens, v_num_heads, head_dim, v_head_dim]
1272 +// initial_state_indices: [batch_size]
1273 +// cu_seqlens: [batch_size + 1]
1274 +at::Tensor fused_sigmoid_gating_delta_rule_update_cpu(
1275 + const at::Tensor& A_log,
1276 + const at::Tensor& dt_bias,
1277 + const at::Tensor& q,
1278 + const at::Tensor& k,
1279 + const at::Tensor& v,
1280 + const at::Tensor& a,
1281 + const at::Tensor& b,
1282 + at::Tensor& initial_state_source,
1283 + const at::Tensor& initial_state_indices,
1284 + const at::Tensor& cu_seqlens,
1285 + bool use_qk_l2norm_in_kernel,
1286 + double softplus_beta = 1.0,
1287 + double softplus_threshold = 20.0) {
1288 + CHECK_DIM(4, q);
1289 + CHECK_DIM(4, v);
1290 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(q);
1291 + int64_t seq_len = q.size(0);
1292 + int64_t batch_size = q.size(1);
1293 + int64_t num_heads = q.size(2);
1294 + int64_t head_dim = q.size(3);
1295 + int64_t v_num_heads = v.size(2);
1296 + int64_t v_head_dim = v.size(3);
1297 + CHECK_INPUT_SHAPE_DTYPE<true>(k, {seq_len, batch_size, num_heads, head_dim}, q.scalar_type());
1298 + CHECK_INPUT_SHAPE_DTYPE<true>(v, {seq_len, batch_size, v_num_heads, v_head_dim}, q.scalar_type());
1299 + CHECK_INPUT_SHAPE_DTYPE<true>(a, {batch_size, v_num_heads}, q.scalar_type());
1300 + CHECK_INPUT_SHAPE_DTYPE<true>(dt_bias, {v_num_heads}, q.scalar_type());
1301 + CHECK_INPUT_SHAPE_DTYPE<true>(b, {batch_size, v_num_heads}, q.scalar_type());
1302 + CHECK_INPUT_SHAPE_DTYPE<true>(initial_state_indices, {batch_size}, at::kInt);
1303 + CHECK_INPUT_SHAPE_DTYPE<true>(cu_seqlens, {batch_size + 1}, at::kInt);
1304 + CHECK_INPUT_SHAPE_DTYPE<true>(
1305 + initial_state_source, {initial_state_source.size(0), v_num_heads, head_dim, v_head_dim}, at::kFloat);
1306 + CHECK(initial_state_source.size(0) >= batch_size);
1307 + CHECK_EQ(v_num_heads % num_heads, 0);
1308 + TORCH_CHECK(
1309 + A_log.sizes() == at::IntArrayRef({v_num_heads}),
1310 + "Input tensor shape mismatch: expected ",
1311 + at::IntArrayRef({v_num_heads}),
1312 + ", got ",
1313 + A_log.sizes());
1314 +
1315 + int64_t q_strideB = q.stride(1);
1316 + int64_t q_strideS = q.stride(0);
1317 + int64_t q_strideH = q.stride(2);
1318 + int64_t k_strideB = k.stride(1);
1319 + int64_t k_strideS = k.stride(0);
1320 + int64_t k_strideH = k.stride(2);
1321 + int64_t v_strideB = v.stride(1);
1322 + int64_t v_strideS = v.stride(0);
1323 + int64_t v_strideH = v.stride(2);
1324 + at::Tensor core_attn_out = at::empty({batch_size, seq_len, v_num_heads, v_head_dim}, q.options());
1325 + at::Tensor qk_scale_buf = at::empty({2 * batch_size, seq_len, num_heads}, at::kFloat);
1326 +
1327 + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(
1328 + q.scalar_type(), A_log.scalar_type(), "fused_sigmoid_gating_delta_rule_update_kernel_impl", [&] {
1329 + fused_sigmoid_gating_delta_rule_update_kernel_impl<scalar_t, param_t>(
1330 + q.data_ptr<scalar_t>(),
1331 + k.data_ptr<scalar_t>(),
1332 + v.data_ptr<scalar_t>(),
1333 + A_log.data_ptr<param_t>(),
1334 + a.data_ptr<scalar_t>(),
1335 + dt_bias.data_ptr<scalar_t>(),
1336 + b.data_ptr<scalar_t>(),
1337 + initial_state_indices.data_ptr<int32_t>(),
1338 + initial_state_source.data_ptr<float>(),
1339 + core_attn_out.data_ptr<scalar_t>(),
1340 + qk_scale_buf.data_ptr<float>(),
1341 + seq_len,
1342 + batch_size,
1343 + num_heads,
1344 + head_dim,
1345 + v_num_heads,
1346 + v_head_dim,
1347 + q_strideB,
1348 + q_strideS,
1349 + q_strideH,
1350 + k_strideB,
1351 + k_strideS,
1352 + k_strideH,
1353 + v_strideB,
1354 + v_strideS,
1355 + v_strideH,
1356 + use_qk_l2norm_in_kernel,
1357 + softplus_threshold);
1358 + });
1359 + return core_attn_out;
1360 +}
1361 +
1362 +// A_log: [num_v_heads]
1363 +// a: [batch, num_v_heads]
1364 +// b: [batch, num_v_heads]
1365 +// dt_bias: [num_v_heads]
1366 +// -A_log.float().exp() * F.softplus(a.float() + dt_bias)
1367 +std::tuple<at::Tensor, at::Tensor>
1368 +fused_gdn_gating_cpu(const at::Tensor& A_log, const at::Tensor& a, const at::Tensor& b, const at::Tensor& dt_bias) {
1369 + CHECK_DIM(1, A_log);
1370 + CHECK_DIM(2, a);
1371 + CHECK_DIM(2, b);
1372 + CHECK_DIM(1, dt_bias);
1373 + CHECK_CONTIGUOUS(a);
1374 + CHECK_EQ(A_log.size(0), a.size(1));
1375 + CHECK_EQ(A_log.size(0), dt_bias.size(0));
1376 + int batch = a.size(0);
1377 + int num_heads = a.size(1);
1378 + CHECK_EQ(b.size(0), batch);
1379 + CHECK_EQ(b.size(1), num_heads);
1380 + at::Tensor out = at::empty({1, batch, num_heads}, a.options().dtype(at::kFloat));
1381 + at::Tensor beta = at::empty({1, batch, num_heads}, b.options());
1382 + CPU_DISPATCH_REDUCED_FLOATING_TYPES_EXT(a.scalar_type(), A_log.scalar_type(), "fused_gdn_gating_kernel", [&] {
1383 + fused_gdn_gating_kernel_impl<scalar_t>(
1384 + A_log.data_ptr<param_t>(),
1385 + a.data_ptr<scalar_t>(),
1386 + b.data_ptr<scalar_t>(),
1387 + dt_bias.data_ptr<scalar_t>(),
1388 + out.data_ptr<float>(),
1389 + beta.data_ptr<scalar_t>(),
1390 + batch,
1391 + num_heads);
1392 + });
1393 + return std::make_tuple(out, beta);
1394 +}
csrc/cpu/sgl-kernels/gemm.cpp
+476 -87
@@ -1,11 +1,12 @@
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 -#include "common.h"
5 -#include "vec.h"
4 +// clang-format off
5 +
6 #include "gemm.h"
7
8 -// clang-format off
8 +#include "common.h"
9 +#include "vec.h"
10
11 namespace {
12
@@ -26,13 +27,13 @@ inline void s8s8_compensation(int8_t* __restrict__ packed, int K) {
27 const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
28 for (int k = 0; k < K / 4; ++k) {
29 for (int col = 0; col < COLS; ++col) {
29 - __m512i vb = _mm512_loadu_si512((const __m512i *)(packed + k * BLOCK_N * 4 + col * 64));
30 + __m512i vb = _mm512_loadu_si512((const __m512i*)(packed + k * BLOCK_N * 4 + col * 64));
31 vcomp[col] = _mm512_dpbusd_epi32(vcomp[col], off, vb);
32 }
33 }
34
35 for (int col = 0; col < COLS; ++col) {
35 - _mm512_storeu_si512((__m512i *)(packed + offset + col * 64), vcomp[col]);
36 + _mm512_storeu_si512((__m512i*)(packed + offset + col * 64), vcomp[col]);
37 }
38 #else
39 TORCH_CHECK(false, "s8s8_compensation not implemented!");
@@ -69,6 +70,43 @@ inline void pack_vnni<int8_t>(int8_t* __restrict__ packed, const int8_t* __restr
70 s8s8_compensation<BLOCK_N>(packed, K);
71 }
72
73 +// uint8_t: mxfp4 or int4
74 +// pack to vnni2 format as they are computed with bfloat16
75 +//
76 +// from [N, K'/2, 2] to [K'/2, N, 2], view 2x int4 as unit8:
77 +// from [N, K ] to [K, N ] where K = K'/2
78 +//
79 +template <>
80 +inline void pack_vnni<uint8_t>(uint8_t* __restrict__ packed, const uint8_t* __restrict__ weight, int N, int K) {
81 + constexpr int BLOCK_N = block_size_n();
82 +
83 + uint8_t unpacked[2 * BLOCK_N];
84 +
85 + // 32-way pack (align with BLOCK_N), faster for avx512 unpacking
86 + //
87 + // for a range of (64):
88 + // {0, 1, 2, ..., 63}
89 + //
90 + // original format:
91 + // { 1|0, 3|2, ..., 63|62}
92 + //
93 + // packed format:
94 + // {32|0, 31|1, ..., 63|31}
95 + //
96 + for (int k = 0; k < K; ++k) {
97 + // unpack first
98 + for (int n = 0; n < N; ++n) {
99 + uint8_t value = weight[n * K + k];
100 + unpacked[n * 2 + 0] = value & 0xF; // lower 4 bits
101 + unpacked[n * 2 + 1] = value >> 4; // higher 4 bits
102 + }
103 + // re-pack to 32-way
104 + for (int n = 0; n < N; ++n) {
105 + packed[k * N + n] = (unpacked[n + BLOCK_N] << 4) | unpacked[n];
106 + }
107 + }
108 +}
109 +
110 template <typename scalar_t>
111 inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
112 using bVec = at::vec::Vectorized<scalar_t>;
@@ -76,7 +114,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
114 constexpr int kVecSize = bVec::size();
115
116 int64_t d;
79 - #pragma GCC unroll 4
117 +#pragma GCC unroll 4
118 for (d = 0; d <= size - kVecSize; d += kVecSize) {
119 fVec data0 = fVec::loadu(input + d);
120 fVec data1 = fVec::loadu(input + d + fVec::size());
@@ -89,13 +127,34 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
127 }
128
129 template <typename scalar_t>
92 -inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
130 +inline void copy_stub(float* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
131 using bVec = at::vec::Vectorized<scalar_t>;
132 using fVec = at::vec::Vectorized<float>;
133 constexpr int kVecSize = bVec::size();
134
135 int64_t d;
98 - #pragma GCC unroll 4
136 +#pragma GCC unroll 4
137 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
138 + fVec data0, data1;
139 + bVec b_vec = bVec::loadu(input + d);
140 + std::tie(data0, data1) = at::vec::convert_to_float(b_vec);
141 + data0.store(out + d);
142 + data1.store(out + d + fVec::size());
143 + }
144 + for (; d < size; ++d) {
145 + out[d] = static_cast<float>(input[d]);
146 + }
147 +}
148 +
149 +template <typename scalar_t>
150 +inline void copy_add_stub(
151 + scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
152 + using bVec = at::vec::Vectorized<scalar_t>;
153 + using fVec = at::vec::Vectorized<float>;
154 + constexpr int kVecSize = bVec::size();
155 +
156 + int64_t d;
157 +#pragma GCC unroll 4
158 for (d = 0; d <= size - kVecSize; d += kVecSize) {
159 fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
160 fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
@@ -107,11 +166,51 @@ inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__
166 }
167 }
168
169 +template <typename scalar_t, bool has_bias>
170 +inline void scalar_sigmoid_and_mul(
171 + scalar_t* __restrict__ out,
172 + const float* __restrict__ input,
173 + const float* __restrict__ bias,
174 + const scalar_t* __restrict__ mul,
175 + int SIZE) {
176 + using bVec = at::vec::Vectorized<scalar_t>;
177 + using fVec = at::vec::Vectorized<float>;
178 + // scalar sigmoid
179 + const fVec one = fVec(1.f);
180 + fVec X;
181 + if constexpr (has_bias) {
182 + assert(bias != nullptr);
183 + X = fVec(input[0] + bias[0]);
184 + } else {
185 + X = fVec(input[0]);
186 + }
187 + X = one / (one + X.neg().exp_u20());
188 +
189 + // vec mul
190 + constexpr int kVecSize = bVec::size();
191 + for (int d = 0; d < SIZE; d += kVecSize) {
192 + bVec m_bvec = bVec::loadu(mul + d);
193 + fVec m_fvec0, m_fvec1;
194 + std::tie(m_fvec0, m_fvec1) = at::vec::convert_to_float(m_bvec);
195 + m_fvec0 = m_fvec0 * X;
196 + m_fvec1 = m_fvec1 * X;
197 +
198 + bVec out_vec = convert_from_float_ext<scalar_t>(m_fvec0, m_fvec1);
199 + out_vec.store(out + d);
200 + }
201 +}
202 +
203 template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
204 struct tinygemm_kernel_nn {
205 static inline void apply(
113 - const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
114 - const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
206 + const scalar_t* __restrict__ A,
207 + const scalar_t* __restrict__ B,
208 + scalar_t* __restrict__ C,
209 + const float* __restrict__ bias,
210 + int64_t K,
211 + int64_t lda,
212 + int64_t ldb,
213 + int64_t ldc) {
214 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
215 }
216 };
@@ -120,9 +219,14 @@ struct tinygemm_kernel_nn {
219 template <bool has_bias, int BLOCK_M, int BLOCK_N>
220 struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
221 static inline void apply(
123 - const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B, at::BFloat16* __restrict__ C,
124 - const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
125 -
222 + const at::BFloat16* __restrict__ A,
223 + const at::BFloat16* __restrict__ B,
224 + at::BFloat16* __restrict__ C,
225 + const float* __restrict__ bias,
226 + int64_t K,
227 + int64_t lda,
228 + int64_t ldb,
229 + int64_t ldc) {
230 constexpr int ROWS = BLOCK_M;
231 constexpr int COLS = BLOCK_N / 16;
232
@@ -145,7 +249,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
249
250 const int64_t K2 = K >> 1;
251 const int64_t lda2 = lda >> 1;
148 - const int64_t ldb2 = ldb; // ldb * 2 >> 1;
252 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
253 const float* a_ptr = reinterpret_cast<const float*>(A);
254 const float* b_ptr = reinterpret_cast<const float*>(B);
255
@@ -180,9 +284,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
284 (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
285 }
286 } else {
183 - _mm256_storeu_si256(
184 - reinterpret_cast<__m256i*>(C + row * ldc + col * 16),
185 - (__m256i)(_mm512_cvtneps_pbh(vc[i])));
287 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(C + row * ldc + col * 16), (__m256i)(_mm512_cvtneps_pbh(vc[i])));
288 }
289 };
290 Unroll<ROWS * COLS>{}(storec);
@@ -190,22 +292,33 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
292 };
293 #endif
294
193 -#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
194 - tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
195 - A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
196 - has_bias ? bias + nb_start : nullptr, K, lda, ldb, ldc);
295 +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
296 + tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
297 + A + mb_start * lda, \
298 + B + nb_start * 2, \
299 + C + mb_start * ldc + nb_start, \
300 + has_bias ? bias + nb_start : nullptr, \
301 + K, \
302 + lda, \
303 + ldb, \
304 + ldc);
305
306 template <typename scalar_t, bool has_bias>
307 struct brgemm {
308 static inline void apply(
201 - const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
202 - float* __restrict__ Ctmp, const float* __restrict__ bias,
203 - int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
204 -
309 + const scalar_t* __restrict__ A,
310 + const scalar_t* __restrict__ B,
311 + scalar_t* __restrict__ C,
312 + float* __restrict__ Ctmp,
313 + const float* __restrict__ bias,
314 + int64_t M,
315 + int64_t N,
316 + int64_t K,
317 + int64_t lda,
318 + int64_t ldb,
319 + int64_t ldc) {
320 constexpr int BLOCK_N = block_size_n();
206 - at::native::cpublas::brgemm(
207 - M, N, K, lda, ldb, BLOCK_N, /* add_C */false,
208 - A, B, Ctmp);
321 + at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
322
323 // copy from Ctmp to C
324 for (int64_t m = 0; m < M; ++m) {
@@ -216,6 +329,21 @@ struct brgemm {
329 }
330 }
331 }
332 + static inline void apply(
333 + const float* __restrict__ A,
334 + const float* __restrict__ B,
335 + scalar_t* __restrict__ C,
336 + float* __restrict__ Ctmp,
337 + const float* __restrict__ bias,
338 + int64_t M,
339 + int64_t N,
340 + int64_t K,
341 + int64_t lda,
342 + int64_t ldb,
343 + int64_t ldc) {
344 + constexpr int BLOCK_N = block_size_n();
345 + at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
346 + }
347 };
348
349 template <typename scalar_t, bool has_bias>
@@ -232,15 +360,12 @@ void tinygemm_kernel(
360 int64_t ldb,
361 int64_t ldc,
362 bool brg) {
235 -
363 if (brg) {
237 - brgemm<scalar_t, has_bias>::apply(
238 - A, B, C, Ctmp, bias,
239 - M, N, K, lda, ldb, ldc);
364 + brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
365 return;
366 }
367
243 - // pattern: 1-4-16
368 + // pattern: 1-4-16, N = 16, 32, 48, 64
369 constexpr int64_t BLOCK_M = 4;
370 constexpr int64_t BLOCK_N = 64;
371 const int64_t MB = div_up(M, BLOCK_M);
@@ -252,25 +377,88 @@ void tinygemm_kernel(
377 int64_t nb_start = nb * BLOCK_N;
378 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
379
255 - switch(mb_size << 4 | nb_size >> 4) {
380 + switch (mb_size << 4 | nb_size >> 4) {
381 // mb_size = 1
257 - case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
258 - case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
382 + case 0x11:
383 + LAUNCH_TINYGEMM_KERNEL_NN(1, 16);
384 + break;
385 + case 0x12:
386 + LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
387 + break;
388 + case 0x13:
389 + LAUNCH_TINYGEMM_KERNEL_NN(1, 48);
390 + break;
391 + case 0x14:
392 + LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
393 + break;
394 // mb_size = 2
260 - case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
261 - case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
395 + case 0x21:
396 + LAUNCH_TINYGEMM_KERNEL_NN(2, 16);
397 + break;
398 + case 0x22:
399 + LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
400 + break;
401 + case 0x23:
402 + LAUNCH_TINYGEMM_KERNEL_NN(2, 48);
403 + break;
404 + case 0x24:
405 + LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
406 + break;
407 // mb_size = 3
263 - case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
264 - case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
408 + case 0x31:
409 + LAUNCH_TINYGEMM_KERNEL_NN(3, 16);
410 + break;
411 + case 0x32:
412 + LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
413 + break;
414 + case 0x33:
415 + LAUNCH_TINYGEMM_KERNEL_NN(3, 48);
416 + break;
417 + case 0x34:
418 + LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
419 + break;
420 // mb_size = 4
266 - case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
267 - case 0x44: LAUNCH_TINYGEMM_KERNEL_NN(4, 64); break;
268 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
421 + case 0x41:
422 + LAUNCH_TINYGEMM_KERNEL_NN(4, 16);
423 + break;
424 + case 0x42:
425 + LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
426 + break;
427 + case 0x43:
428 + LAUNCH_TINYGEMM_KERNEL_NN(4, 48);
429 + break;
430 + case 0x44:
431 + LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
432 + break;
433 + default:
434 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, " x ", nb_size);
435 }
436 }
437 }
438 }
439
440 +template <typename scalar_t, bool has_bias>
441 +void tinygemm_kernel(
442 + const float* __restrict__ A,
443 + const float* __restrict__ B,
444 + scalar_t* __restrict__ C,
445 + float* __restrict__ Ctmp,
446 + const float* __restrict__ bias,
447 + int64_t M,
448 + int64_t N,
449 + int64_t K,
450 + int64_t lda,
451 + int64_t ldb,
452 + int64_t ldc,
453 + bool brg) {
454 + TORCH_CHECK(brg, "Expected to use fp32 brgemm for small N GEMM");
455 + if (brg) {
456 + brgemm<scalar_t, has_bias>::apply(A, B, C, Ctmp, bias, M, N, K, lda, ldb, ldc);
457 + return;
458 + }
459 + // TODO : add intrinsic path
460 +}
461 +
462 template <typename scalar_t>
463 void weight_packed_linear_kernel_impl(
464 scalar_t* __restrict__ out,
@@ -282,29 +470,20 @@ void weight_packed_linear_kernel_impl(
470 int64_t K,
471 int64_t mat1_strideM,
472 int64_t out_strideM) {
285 -
473 constexpr int64_t BLOCK_M = block_size_m();
474 constexpr int64_t BLOCK_N = block_size_n();
475 const int64_t MB = div_up(M, BLOCK_M);
476 const int64_t NB = div_up(N, BLOCK_N);
477
291 - // use avx512-bf16 when a) M is small; b) dtype is bfloat16, otherwise use amx
292 - const bool use_brgemm = (M > 4) || (!std::is_same_v<scalar_t, at::BFloat16>);
293 -
294 - // l2 cache block for n
295 - int64_t cache_blocks_nb = get_cache_blocks<scalar_t>(BLOCK_N, K);
478 + const bool use_brgemm = can_use_brgemm<scalar_t>(M);
479
480 // parallel on [MB, NB]
481 AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
299 - parallel_2d(MB, NB, [&](int64_t begin_mb, int64_t end_mb, int64_t begin_nb, int64_t end_nb) {
300 -
482 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
483 // for brgemm, use float32 for accumulate
484 alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
485
304 - for (int64_t nbb = begin_nb; nbb < end_nb; nbb += cache_blocks_nb) {
305 - for (int64_t mb = begin_mb; mb < end_mb; ++mb) {
306 - for (int64_t nb = nbb; nb < std::min(nbb + cache_blocks_nb, end_nb); ++nb) {
307 -
486 + loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
487 int64_t mb_start = mb * BLOCK_M;
488 int64_t mb_size = std::min(M - mb_start, BLOCK_M);
489 int64_t nb_start = nb * BLOCK_N;
@@ -323,7 +502,82 @@ void weight_packed_linear_kernel_impl(
502 /* ldb */ nb_size,
503 /* ldc */ out_strideM,
504 /* brg */ use_brgemm);
326 - }}}
505 + });
506 +
507 + if (use_brgemm) {
508 + at::native::cpublas::brgemm_release();
509 + }
510 + });
511 + });
512 +}
513 +
514 +template <typename scalar_t>
515 +void weight_packed_linear_kernel_impl(
516 + scalar_t* __restrict__ out,
517 + const scalar_t* __restrict__ mat1,
518 + const float* __restrict__ mat2,
519 + const float* __restrict__ bias,
520 + const scalar_t* __restrict__ post_mul_mat,
521 + int64_t M,
522 + int64_t N,
523 + int64_t K,
524 + int64_t mat1_strideM,
525 + int64_t out_strideM) {
526 + constexpr int64_t BLOCK_M = block_size_m();
527 + constexpr int64_t BLOCK_N = block_size_n();
528 + const int64_t MB = div_up(M, BLOCK_M);
529 + const int64_t NB = div_up(N, BLOCK_N);
530 +
531 + const bool use_brgemm = true; // TODO: add intrinsic path
532 + // parallel on [MB, NB]
533 + AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
534 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
535 + // for brgemm, use float32 for accumulate
536 + alignas(64) float Atmp[BLOCK_M * K];
537 + alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
538 +
539 + loop_2d<float>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
540 + int64_t mb_start = mb * BLOCK_M;
541 + int64_t mb_size = std::min(M - mb_start, BLOCK_M);
542 + int64_t nb_start = nb * BLOCK_N;
543 + int64_t nb_size = std::min(N - nb_start, BLOCK_N);
544 + for (int64_t m = 0; m < mb_size; ++m) {
545 + copy_stub<scalar_t>(Atmp + m * K, mat1 + mb_start * mat1_strideM + m * K, K);
546 + }
547 + tinygemm_kernel<scalar_t, has_bias>(
548 + /* A */ Atmp,
549 + /* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
550 + /* C */ out + mb_start * out_strideM + nb_start,
551 + /* Ctmp*/ Ctmp,
552 + /* bias*/ bias + nb_start,
553 + /* M */ mb_size,
554 + /* N */ nb_size,
555 + /* K */ K,
556 + /* lda */ mat1_strideM,
557 + /* ldb */ nb_size,
558 + /* ldc */ out_strideM,
559 + /* brg */ use_brgemm);
560 +
561 + if (post_mul_mat != nullptr) {
562 + for (int64_t m = 0; m < mb_size; ++m) {
563 + scalar_sigmoid_and_mul<scalar_t, has_bias>(
564 + out + mb_start * out_strideM + nb_start + m * out_strideM,
565 + Ctmp + m * BLOCK_N,
566 + bias + nb_start,
567 + post_mul_mat + mb_start * out_strideM + m * out_strideM,
568 + out_strideM);
569 + }
570 + } else {
571 + for (int64_t m = 0; m < mb_size; ++m) {
572 + if constexpr (has_bias) {
573 + copy_add_stub(
574 + out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, bias + nb_start, N);
575 + } else {
576 + copy_stub(out + mb_start * out_strideM + nb_start + m * out_strideM, Ctmp + m * BLOCK_N, N);
577 + }
578 + }
579 + }
580 + });
581
582 if (use_brgemm) {
583 at::native::cpublas::brgemm_release();
@@ -332,20 +586,38 @@ void weight_packed_linear_kernel_impl(
586 });
587 }
588
335 -} // anonymous namespace
589 +} // anonymous namespace
590
591 // tinygemm interface
592 template <typename scalar_t>
339 -void tinygemm_kernel(const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, scalar_t* __restrict__ C,
340 - float* __restrict__ Ctmp, int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg) {
593 +void tinygemm_kernel(
594 + const scalar_t* __restrict__ A,
595 + const scalar_t* __restrict__ B,
596 + scalar_t* __restrict__ C,
597 + float* __restrict__ Ctmp,
598 + int64_t M,
599 + int64_t N,
600 + int64_t K,
601 + int64_t lda,
602 + int64_t ldb,
603 + int64_t ldc,
604 + bool brg) {
605 tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, nullptr, M, N, K, lda, ldb, ldc, brg);
606 }
607
344 -#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
345 - template void tinygemm_kernel<TYPE>( \
346 - const TYPE* __restrict__ A, const TYPE* __restrict__ B, TYPE* __restrict__ C, \
347 - float* __restrict__ Ctmp, int64_t M, int64_t N, int64_t K, int64_t lda, \
348 - int64_t ldb, int64_t ldc, bool brg)
608 +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
609 + template void tinygemm_kernel<TYPE>( \
610 + const TYPE* __restrict__ A, \
611 + const TYPE* __restrict__ B, \
612 + TYPE* __restrict__ C, \
613 + float* __restrict__ Ctmp, \
614 + int64_t M, \
615 + int64_t N, \
616 + int64_t K, \
617 + int64_t lda, \
618 + int64_t ldb, \
619 + int64_t ldc, \
620 + bool brg)
621
622 INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
623 INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
@@ -359,14 +631,23 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
631
632 const int64_t ndim = weight.ndimension();
633 TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor.");
634 +
635 + if (ndim == 2 && weight.size(0) < TILE_N) {
636 + // for 2D weight and small OC shape, we use fma linear path, which needs transpose not pack
637 + return weight.to(at::kFloat).t().contiguous();
638 + }
639 +
640 const auto st = weight.scalar_type();
641 const int64_t E = ndim == 3 ? weight.size(0) : 1;
642 const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0);
643 const int64_t IC = ndim == 3 ? weight.size(2) : weight.size(1);
644
645 + // mxfp4 or int4 are packed with uint8
646 + const int64_t actual_IC = st == at::kByte ? IC * 2 : IC;
647 +
648 // we handle 2 TILE_N at a time.
649 TORCH_CHECK(OC % TILE_N == 0, "invalid weight out features ", OC);
369 - TORCH_CHECK(IC % TILE_K == 0, "invalid weight input features ", IC);
650 + TORCH_CHECK(actual_IC % TILE_K == 0, "invalid weight input features ", actual_IC);
651
652 constexpr int64_t BLOCK_N = block_size_n();
653 const int64_t NB = div_up(OC, BLOCK_N);
@@ -375,12 +656,14 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
656 auto packed_weight = at::empty({}, weight.options());
657 const int64_t stride = OC * IC;
658
378 - TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn,
379 - "expect weight to be bfloat16, float16, int8 or fp8_e4m3.");
659 + // Note: for `kByte` (uint8), it represents either `mxfp4` or `int4`.
660 + TORCH_CHECK(
661 + st == at::kBFloat16 || st == at::kHalf || st == at::kChar || st == at::kFloat8_e4m3fn || st == at::kByte,
662 + "expect weight to be bfloat16, float16, int8, fp8_e4m3 or uint8(mxfp4 or int4).");
663
664 CPU_DISPATCH_PACKED_TYPES(st, [&] {
665 // adjust most inner dimension size
383 - const int packed_row_size = get_row_size<packed_t>(IC);
666 + const int packed_row_size = get_row_size<packed_t>(actual_IC);
667 auto sizes = weight.sizes().vec();
668 sizes[ndim - 1] = packed_row_size;
669 packed_weight.resize_(sizes);
@@ -399,10 +682,7 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
682 int64_t n = nb * BLOCK_N;
683 int64_t n_size = std::min(BLOCK_N, OC - n);
684 pack_vnni<packed_t>(
402 - packed_data + e * OC * packed_row_size + n * packed_row_size,
403 - w_data + e * stride + n * IC,
404 - n_size,
405 - IC);
685 + packed_data + e * OC * packed_row_size + n * packed_row_size, w_data + e * stride + n * IC, n_size, IC);
686
687 // move to the next index
688 data_index_step(e, E, nb, NB);
@@ -412,33 +692,141 @@ at::Tensor convert_weight_packed(at::Tensor& weight) {
692 return packed_weight;
693 }
694
695 +at::Tensor convert_scale_packed(at::Tensor& scale) {
696 + CHECK_INPUT(scale);
697 +
698 + const int64_t ndim = scale.ndimension();
699 + TORCH_CHECK(ndim == 2 || ndim == 3, "expect scale to be 2d or 3d, got ", ndim, "d tensor.");
700 + const auto st = scale.scalar_type();
701 + const int64_t E = ndim == 3 ? scale.size(0) : 1;
702 + const int64_t N = ndim == 3 ? scale.size(1) : scale.size(0);
703 + // number of groups, e.g. K/32
704 + const int64_t G = ndim == 3 ? scale.size(2) : scale.size(1);
705 +
706 + constexpr int64_t BLOCK_N = block_size_n();
707 + TORCH_CHECK(N % BLOCK_N == 0, "invalid weight out features ", N);
708 + const int64_t NB = N / BLOCK_N;
709 +
710 + auto packed_scale = at::empty_like(scale);
711 + TORCH_CHECK(st == at::kByte, "expect scale to be uint8.");
712 +
713 + const uint8_t* s_data = scale.data_ptr<uint8_t>();
714 + uint8_t* packed_data = packed_scale.data_ptr<uint8_t>();
715 +
716 + // parallel on src {E, NB, BLOCK_N, G}, dst {E, NB, G, BLOCK_N}
717 + at::parallel_for(0, E * NB * BLOCK_N * G, 0, [&](int64_t begin, int64_t end) {
718 + int64_t e{0}, nb{0}, n{0}, g{0};
719 + data_index_init(begin, e, E, nb, NB, n, BLOCK_N, g, G);
720 +
721 + for (int64_t i = begin; i < end; ++i) {
722 + packed_data[e * N * G + nb * G * BLOCK_N + g * BLOCK_N + n] = s_data[i];
723 + // move to the next index
724 + data_index_step(e, E, nb, NB, n, BLOCK_N, g, G);
725 + }
726 + });
727 + return packed_scale;
728 +}
729 +
730 // mat1 : [M, K]
416 -// mat2 : [N, K]
731 +// mat2 : [N, K] ([K, N] if use_fma_gemm)
732 // bias : [N]
733 // out : [M, N]
734 //
420 -at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
421 - const std::optional<at::Tensor>& bias, bool is_vnni) {
422 - RECORD_FUNCTION(
423 - "sgl-kernel::weight_packed_linear", std::vector<c10::IValue>({mat1, mat2, bias}));
424 -
735 +at::Tensor
736 +weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2, const std::optional<at::Tensor>& bias, bool is_vnni) {
737 auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
738 + bool use_fma_gemm = false;
739 + if (packed_w.scalar_type() == at::kFloat) {
740 + use_fma_gemm = true;
741 + }
742 +
743 + int64_t M = mat1.size(0);
744 + int64_t K = mat1.size(1);
745 + int64_t N = use_fma_gemm ? mat2.size(1) : mat2.size(0);
746
747 CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
748 CHECK_INPUT(mat2);
429 -
430 - int64_t M = mat1.size(0);
431 - int64_t N = mat2.size(0);
432 - int64_t K = mat2.size(1);
433 - CHECK_EQ(mat1.size(1), K);
749 CHECK_DIM(2, mat1);
750 CHECK_DIM(2, mat2);
751 + if (!use_fma_gemm) {
752 + CHECK_EQ(mat1.size(1), K);
753 + }
754
755 + auto dispatch_type = mat1.scalar_type();
756 auto out = at::empty({M, N}, mat1.options());
438 -
757 // strides
440 - int64_t mat1_strideM = mat1.stride(0);
758 int64_t out_strideM = out.stride(0);
759 + int64_t mat1_strideM = mat1.stride(0);
760 +
761 + const bool has_bias = bias.has_value();
762 + const float* bias_data = nullptr;
763 + if (has_bias) {
764 + CHECK_EQ(bias.value().size(0), N);
765 + bias_data = bias.value().data_ptr<float>();
766 + }
767 +
768 + AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "weight_packed_linear_kernel_impl", [&] {
769 + if (use_fma_gemm) {
770 + weight_packed_linear_kernel_impl<scalar_t>(
771 + out.data_ptr<scalar_t>(),
772 + mat1.data_ptr<scalar_t>(),
773 + packed_w.data_ptr<float>(),
774 + bias_data,
775 + nullptr,
776 + M,
777 + N,
778 + K,
779 + mat1_strideM,
780 + out_strideM);
781 + } else {
782 + weight_packed_linear_kernel_impl<scalar_t>(
783 + out.data_ptr<scalar_t>(),
784 + mat1.data_ptr<scalar_t>(),
785 + packed_w.data_ptr<scalar_t>(),
786 + bias_data,
787 + M,
788 + N,
789 + K,
790 + mat1_strideM,
791 + out_strideM);
792 + }
793 + });
794 +
795 + return out;
796 +}
797 +
798 +// mat1 : [M, K]
799 +// mat2 : [K, 1]
800 +// post_mul_mat : [M, K]
801 +// bias : [N]
802 +// out : [M, N]
803 +//
804 +at::Tensor fused_linear_sigmoid_mul(
805 + at::Tensor& mat1,
806 + at::Tensor& mat2,
807 + const std::optional<at::Tensor>& bias,
808 + bool is_vnni,
809 + const at::Tensor& post_mul_mat) {
810 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
811 + TORCH_CHECK(packed_w.scalar_type() == at::kFloat, "fused_linear_sigmoid_mul requires packed float weight")
812 +
813 + int64_t M = mat1.size(0);
814 + int64_t K = mat1.size(1);
815 + int64_t N = mat2.size(1);
816 +
817 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
818 + CHECK_INPUT(mat2);
819 + CHECK_DIM(2, mat1);
820 + CHECK_DIM(2, mat2);
821 +
822 + int64_t out_strideM = post_mul_mat.size(1);
823 + int64_t mat1_strideM = mat1.stride(0);
824 + auto dispatch_type = mat1.scalar_type();
825 + auto out = at::empty({M, out_strideM}, mat1.options());
826 +
827 + TORCH_CHECK(
828 + N == 1 && out_strideM % 32 == 0,
829 + "post_mul_mat tensor size(1) should be 32 dividable, and the mat2 OC=1 (Mx1 as linear output shape)")
830
831 const bool has_bias = bias.has_value();
832 const float* bias_data = nullptr;
@@ -447,12 +835,13 @@ at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
835 bias_data = bias.value().data_ptr<float>();
836 }
837
450 - AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "weight_packed_linear_kernel_impl", [&] {
838 + AT_DISPATCH_REDUCED_FLOATING_TYPES(dispatch_type, "fused_linear_sigmoid_mul", [&] {
839 weight_packed_linear_kernel_impl<scalar_t>(
840 out.data_ptr<scalar_t>(),
841 mat1.data_ptr<scalar_t>(),
454 - packed_w.data_ptr<scalar_t>(),
842 + packed_w.data_ptr<float>(),
843 bias_data,
844 + post_mul_mat.data_ptr<scalar_t>(),
845 M,
846 N,
847 K,
csrc/cpu/sgl-kernels/gemm.h
+121 -86
@@ -1,8 +1,12 @@
1 -#pragma once
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 <ATen/native/CPUBlas.h>
8
5 -// clang-format off
9 +#include "common.h"
10
11 // amx-bf16
12 #define TILE_M 16
@@ -10,20 +14,42 @@
14 #define TILE_K 32
15
16 // block size for AMX gemm
13 -constexpr int block_size_m() { return 2 * TILE_M; }
14 -constexpr int block_size_n() { return 2 * TILE_N; }
17 +constexpr int block_size_m() {
18 + return 2 * TILE_M;
19 +}
20 +constexpr int block_size_n() {
21 + return 2 * TILE_N;
22 +}
23
24 // define threshold using brgemm (intel AMX)
17 -template <typename T> inline bool can_use_brgemm(int M);
18 -template <> inline bool can_use_brgemm<at::BFloat16>(int M) { return M > 4; }
19 -template <> inline bool can_use_brgemm<at::Half>(int M) { return true; }
20 -template <> inline bool can_use_brgemm<int8_t>(int M) { return M > 4; }
21 -template <> inline bool can_use_brgemm<uint8_t>(int M) { return M > 4; }
22 -template <> inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) { return M > 4; }
23 -template <> inline bool can_use_brgemm<at::quint4x2>(int M) { return M > 4; }
25 +template <typename T>
26 +inline bool can_use_brgemm(int M);
27 +template <>
28 +inline bool can_use_brgemm<at::BFloat16>(int M) {
29 + return M > 4;
30 +}
31 +template <>
32 +inline bool can_use_brgemm<at::Half>(int M) {
33 + return true;
34 +}
35 +// this requires PyTorch 2.7 or above
36 +template <>
37 +inline bool can_use_brgemm<int8_t>(int M) {
38 + return M > 4;
39 +}
40 +
41 +template <>
42 +inline bool can_use_brgemm<uint8_t>(int M) {
43 + return M > 4;
44 +}
45 +
46 +template <>
47 +inline bool can_use_brgemm<at::Float8_e4m3fn>(int M) {
48 + return M > 4;
49 +}
50
51 // work around compiler internal error
26 -#define BLOCK_K 128 // 4 * TILE_K
52 +#define BLOCK_K 128 // 4 * TILE_K
53
54 // adjust leading dimension size for K
55 template <typename T>
@@ -36,18 +62,44 @@ inline int64_t get_row_size<int8_t>(int64_t K) {
62 return K + sizeof(int32_t);
63 }
64
65 +// uint8: mxfp4 or int4
66 +template <>
67 +inline int64_t get_row_size<uint8_t>(int64_t K) {
68 + return K >> 1;
69 +}
70 +
71 inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
72 return use_int8_w8a8 ? K + sizeof(int32_t) : K;
73 }
74
75 +enum class CPUQuantMethod : int64_t { BF16 = 0, INT8_W8A8 = 1, FP8_W8A16 = 2, INT4_W4A8 = 3 };
76 +
77 +constexpr bool operator==(CPUQuantMethod a, int64_t b) {
78 + return static_cast<int64_t>(a) == b;
79 +}
80 +
81 +constexpr bool operator==(int64_t a, CPUQuantMethod b) {
82 + return a == static_cast<int64_t>(b);
83 +}
84 +
85 +enum class CPUQuantAlgo : int64_t { AWQ = 0, GPTQ = 1 };
86 +
87 +constexpr bool operator==(CPUQuantAlgo a, int64_t b) {
88 + return static_cast<int64_t>(a) == b;
89 +}
90 +
91 +constexpr bool operator==(int64_t a, CPUQuantAlgo b) {
92 + return a == static_cast<int64_t>(b);
93 +}
94 +
95 inline int64_t get_4bit_block_k_size(int64_t group_size) {
96 return group_size > 128 ? 128 : group_size;
97 }
98
47 -// pack weight into vnni format
99 +// pack weight to vnni format
100 at::Tensor convert_weight_packed(at::Tensor& weight);
101
50 -// pack weight to vnni format for int4 (adapted from sglang)
102 +// pack weight to vnni format for int4
103 std::tuple<at::Tensor, at::Tensor, at::Tensor>
104 convert_weight_packed_scale_zp(at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
105
@@ -105,35 +157,6 @@ void fused_experts_fp8_kernel_impl(
157 int64_t topk,
158 int64_t num_tokens_post_pad);
159
108 -// moe implementations for int4 w4a16
109 -template <typename scalar_t>
110 -void fused_experts_int4_w4a16_kernel_impl(
111 - scalar_t* __restrict__ output,
112 - scalar_t* __restrict__ ic0,
113 - scalar_t* __restrict__ ic1,
114 - scalar_t* __restrict__ ic2,
115 - scalar_t* __restrict__ A_tmp,
116 - scalar_t* __restrict__ B_tmp,
117 - float* __restrict__ C_tmp,
118 - const scalar_t* __restrict__ input,
119 - const at::quint4x2* __restrict__ packed_w1,
120 - const at::quint4x2* __restrict__ packed_w2,
121 - const uint8_t* __restrict__ w1z,
122 - const uint8_t* __restrict__ w2z,
123 - const scalar_t* __restrict__ w1s,
124 - const scalar_t* __restrict__ w2s,
125 - int group_size,
126 - const float* __restrict__ topk_weights,
127 - const int32_t* __restrict__ sorted_ids,
128 - const int32_t* __restrict__ expert_ids,
129 - const int32_t* __restrict__ offsets,
130 - int64_t M,
131 - int64_t N,
132 - int64_t K,
133 - int64_t E,
134 - int64_t topk,
135 - int64_t num_tokens_post_pad);
136 -
160 // shared expert implementation for int8 w8a8
161 template <typename scalar_t>
162 void shared_expert_int8_kernel_impl(
@@ -153,6 +176,37 @@ void shared_expert_int8_kernel_impl(
176 int64_t N,
177 int64_t K);
178
179 +template <typename scalar_t>
180 +void fused_experts_int4_w4a8_kernel_impl(
181 + scalar_t* __restrict__ output,
182 + scalar_t* __restrict__ ic0,
183 + scalar_t* __restrict__ ic1,
184 + scalar_t* __restrict__ ic2,
185 + uint8_t* __restrict__ A_tmp,
186 + uint8_t* __restrict__ Aq_tmp,
187 + float* __restrict__ As_tmp,
188 + int32_t* __restrict__ Azp_tmp,
189 + float* __restrict__ C_tmp,
190 + int8_t* __restrict__ dqB_tmp,
191 + const scalar_t* __restrict__ input,
192 + const uint8_t* __restrict__ packed_w1,
193 + const uint8_t* __restrict__ packed_w2,
194 + const int8_t* __restrict__ w1z,
195 + const int8_t* __restrict__ w2z,
196 + const float* __restrict__ w1s,
197 + const float* __restrict__ w2s,
198 + int group_size,
199 + const float* __restrict__ topk_weights,
200 + const int32_t* __restrict__ sorted_ids,
201 + const int32_t* __restrict__ expert_ids,
202 + const int32_t* __restrict__ offsets,
203 + int64_t M,
204 + int64_t N,
205 + int64_t K,
206 + int64_t E,
207 + int64_t topk,
208 + int64_t num_tokens_post_pad);
209 +
210 template <typename scalar_t>
211 void shared_expert_fp8_kernel_impl(
212 scalar_t* __restrict__ output,
@@ -204,6 +258,7 @@ void tinygemm_kernel(
258 int64_t ldc,
259 bool brg);
260
261 +// block quantization
262 template <typename scalar_t>
263 void tinygemm_kernel(
264 const scalar_t* __restrict__ A,
@@ -219,33 +274,26 @@ void tinygemm_kernel(
274 int64_t ldb,
275 int64_t ldc,
276 bool brg,
222 - int64_t block_size_K);
277 + int64_t block_size_K,
278 + bool do_unpack = true);
279
280 +// per tensor quantization
281 template <typename scalar_t>
282 void tinygemm_kernel(
283 const scalar_t* __restrict__ A,
227 - const at::quint4x2* __restrict__ B,
284 + const at::Float8_e4m3fn* __restrict__ B,
285 scalar_t* __restrict__ C,
229 - const uint8_t* __restrict__ Bz,
230 - const scalar_t* __restrict__ Bs,
286 scalar_t* __restrict__ Btmp,
287 float* __restrict__ Ctmp,
288 + float scale,
289 int64_t M,
290 int64_t N,
291 int64_t K,
236 - int group_size,
292 int64_t lda,
293 int64_t ldb,
294 int64_t ldc,
240 - int64_t strideBz,
241 - int64_t strideBs,
295 bool brg);
296
244 -// int4 scaled GEMM (adapted from sglang)
245 -at::Tensor int4_scaled_mm_cpu(
246 - at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias);
247 -
248 -// int4 tinygemm kernel interface(adapted from sglang)
297 template <typename scalar_t>
298 void tinygemm_kernel(
299 scalar_t* C,
@@ -266,34 +314,21 @@ void tinygemm_kernel(
314 bool store_out,
315 bool use_brgemm);
316
269 -// TODO: debug print, remove me later
270 -inline void print_16x32i(const __m512i x) {
271 - int32_t a[16];
272 - _mm512_storeu_si512((__m512i *)a, x);
273 -
274 - for (int i = 0; i < 16; i++){
275 - std::cout << a[i] << " ";
276 - }
277 - std::cout << std::endl;
278 -}
279 -
280 -inline void print_16x32(const __m512 x) {
281 - float a[16];
282 - _mm512_storeu_ps((__m512 *)a, x);
283 -
284 - for (int i = 0; i < 16; i++){
285 - std::cout << a[i] << " ";
286 - }
287 - std::cout << std::endl;
288 -}
289 -
290 -
291 -inline void print_32x8u(const __m256i x) {
292 - uint8_t a[32];
293 - _mm256_storeu_si256((__m256i *)a, x);
294 -
295 - for (int i = 0; i < 32; ++i) {
296 - std::cout << int32_t(a[i]) << " ";
297 - }
298 - std::cout << std::endl;
299 -}
317 +// mxfp4
318 +template <typename scalar_t>
319 +void tinygemm_kernel(
320 + const scalar_t* __restrict__ A,
321 + const uint8_t* __restrict__ B,
322 + scalar_t* __restrict__ C,
323 + scalar_t* __restrict__ Btmp,
324 + float* __restrict__ Ctmp,
325 + const uint8_t* __restrict__ scale,
326 + int64_t M,
327 + int64_t N,
328 + int64_t K,
329 + int64_t lda,
330 + int64_t ldb,
331 + int64_t ldc,
332 + bool brg,
333 + int64_t block_size_K,
334 + bool do_unpack = true);
csrc/cpu/sgl-kernels/gemm_fp8.cpp
+773 -143
@@ -1,14 +1,11 @@
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 -#include "common.h"
5 -#include "vec.h"
6 -#include "gemm.h"
7 -
4 // clang-format off
5
10 -// we use 4x32 for BLOCK_M
11 -#define BLOCK_SIZE_M_SCALE 4
6 +#include "common.h"
7 +#include "gemm.h"
8 +#include "vec.h"
9
10 namespace {
11
@@ -19,7 +16,7 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
16 constexpr int kVecSize = bVec::size();
17
18 int64_t d;
22 - #pragma GCC unroll 4
19 +#pragma GCC unroll 4
20 for (d = 0; d <= size - kVecSize; d += kVecSize) {
21 fVec data0 = fVec::loadu(input + d);
22 fVec data1 = fVec::loadu(input + d + fVec::size());
@@ -32,13 +29,14 @@ inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ inpu
29 }
30
31 template <typename scalar_t>
35 -inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
32 +inline void copy_add_stub(
33 + scalar_t* __restrict__ out, const float* __restrict__ input, const float* __restrict__ bias, int64_t size) {
34 using bVec = at::vec::Vectorized<scalar_t>;
35 using fVec = at::vec::Vectorized<float>;
36 constexpr int kVecSize = bVec::size();
37
38 int64_t d;
41 - #pragma GCC unroll 4
39 +#pragma GCC unroll 4
40 for (d = 0; d <= size - kVecSize; d += kVecSize) {
41 fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
42 fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
@@ -49,21 +47,41 @@ inline void copy_add_stub(scalar_t* __restrict__ out, const float* __restrict__
47 out[d] = static_cast<scalar_t>(input[d] + bias[d]);
48 }
49 }
50 +template <typename scalar_t>
51 +inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int size, float scale) {
52 + using bVec = at::vec::Vectorized<scalar_t>;
53 + using fVec = at::vec::Vectorized<float>;
54 + constexpr int kVecSize = bVec::size();
55 + const fVec vscale = fVec(scale);
56 +
57 + int d;
58 +#pragma GCC unroll 4
59 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
60 + fVec data0 = fVec::loadu(input + d) * vscale;
61 + fVec data1 = fVec::loadu(input + d + fVec::size()) * vscale;
62 + bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
63 + out_vec.store(out + d);
64 + }
65 + for (; d < size; ++d) {
66 + out[d] = static_cast<scalar_t>(input[d] * scale);
67 + }
68 +}
69
70 inline void unpack_B(
71 at::BFloat16* __restrict__ Btmp,
72 const at::Float8_e4m3fn* __restrict__ packed_B,
56 - int N,
57 - int K,
58 - int ldb,
59 - int ldb_tmp,
73 + int64_t N,
74 + int64_t K,
75 + int64_t ldb,
76 + int64_t ldb_tmp,
77 float scale) {
78 #if defined(CPU_CAPABILITY_AVX512)
79 // [K/2, N, 2]
63 - const int K2 = K >> 1;
64 - const int ldb2 = ldb; // ldb * 2 >> 1;
80 + const int64_t K2 = K >> 1;
81 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
82 const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(packed_B);
66 - const __m512 vd = _mm512_set1_ps(scale);
83 + const __m512 vexp = _mm512_castsi512_ps(_mm512_set1_epi32(kFP8_BIAS));
84 + const __m512 vd = _mm512_mul_ps(_mm512_set1_ps(scale), vexp);
85
86 constexpr int BLOCK_N = block_size_n();
87 static_assert(BLOCK_N == 32);
@@ -72,7 +90,7 @@ inline void unpack_B(
90 constexpr int PREFETCH_SIZE_K = 64;
91
92 #pragma GCC unroll 4
75 - for (int k = 0; k < K2; ++k) {
93 + for (int64_t k = 0; k < K2; ++k) {
94 __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2);
95 if constexpr (PREFETCH_SIZE_K > 0) {
96 _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0);
@@ -81,8 +99,8 @@ inline void unpack_B(
99 __m256i b8_0 = _mm512_extracti32x8_epi32(b8, 0);
100 __m256i b8_1 = _mm512_extracti32x8_epi32(b8, 1);
101
84 - __m512bh bf16_0 = CVT_FP8_TO_BF16(b8_0);
85 - __m512bh bf16_1 = CVT_FP8_TO_BF16(b8_1);
102 + __m512bh bf16_0 = CVT_FP8_TO_BF16_EXT(b8_0);
103 + __m512bh bf16_1 = CVT_FP8_TO_BF16_EXT(b8_1);
104
105 // Apply scale
106 __m512 f0_lo = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_0, 0));
@@ -106,26 +124,139 @@ inline void unpack_B(
124 #endif
125 }
126
109 -template <typename scalar_t, typename packed_t, bool has_bias, int BLOCK_M, int BLOCK_N>
127 +inline void unpack_B(
128 + at::BFloat16* __restrict__ Btmp,
129 + const at::Float8_e4m3fn* __restrict__ packed_B,
130 + int N,
131 + int K,
132 + int ldb,
133 + int ldb_tmp) {
134 +#if defined(CPU_CAPABILITY_AVX512)
135 + // [K/2, N, 2]
136 + const int K2 = K >> 1;
137 + const int ldb2 = ldb; // ldb * 2 >> 1;
138 + const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(packed_B);
139 +
140 + // prefetch distance
141 + constexpr int PREFETCH_SIZE_K = 64;
142 +#pragma GCC unroll 4
143 + for (int k = 0; k < K2; ++k) {
144 + __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2);
145 + if constexpr (PREFETCH_SIZE_K > 0) {
146 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0);
147 + }
148 +
149 + __m256i b8_0 = _mm512_extracti32x8_epi32(b8, 0);
150 + __m256i b8_1 = _mm512_extracti32x8_epi32(b8, 1);
151 +
152 + __m512bh bf16_0 = CVT_FP8_TO_BF16(b8_0);
153 + __m512bh bf16_1 = CVT_FP8_TO_BF16(b8_1);
154 + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 0, (__m512i)bf16_0);
155 + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 32, (__m512i)bf16_1);
156 + }
157 +#else
158 + TORCH_CHECK(false, "unpack_B: scalar path not implemented!");
159 +#endif
160 +}
161 +
162 +// mxfp4
163 +inline void unpack_B(
164 + at::BFloat16* __restrict__ Btmp,
165 + const uint8_t* __restrict__ packed_B,
166 + int64_t N,
167 + int64_t K,
168 + int64_t ldb,
169 + int64_t ldb_tmp,
170 + const uint8_t* __restrict__ scale) {
171 +#if defined(CPU_CAPABILITY_AVX512)
172 + // [K/2, N, 2]
173 + const int64_t K2 = K >> 1;
174 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
175 + const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(packed_B); // 2 * 4 bit = 8 bit
176 +
177 + constexpr int BLOCK_N = block_size_n();
178 + static_assert(BLOCK_N == 32);
179 +
180 + // prefetch distance
181 + constexpr int PREFETCH_SIZE_K = 64;
182 +
183 + // exponent bias 127
184 + const __m512i off = _mm512_set1_epi16(0x7F);
185 +
186 + // load 32 bytes only once for each block
187 + __m256i s8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(scale));
188 + __m512i s16 = _mm512_slli_epi16(_mm512_sub_epi16(_mm512_cvtepu8_epi16(s8), off), 0x7);
189 +
190 + // holds Nx2(64) scales, interleaved as 2 belongs to K dimension
191 + // e.g. vs0: { s0, s0, s1, s1, ..., s15, s15}
192 + // vs1: {s16, s16, s17, s17, ..., s31, s31}
193 + auto [vscale0, vscale1] = transpose_2x32_16bit(s16, s16);
194 +
195 +#pragma GCC unroll 4
196 + for (int64_t k = 0; k < K2; ++k) {
197 + __m256i b4 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(b_ptr + k * ldb2));
198 + if constexpr (PREFETCH_SIZE_K > 0) {
199 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0);
200 + }
201 + auto [vb0, vb1] = CVT_MXFP4_TO_BF16(b4, vscale0, vscale1);
202 +
203 + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 0, (__m512i)vb0);
204 + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 32, (__m512i)vb1);
205 + }
206 +#else
207 + TORCH_CHECK(false, "unpack_B: scalar path not implemented!");
208 +#endif
209 +}
210 +
211 +template <typename scalar_t, typename packed_t, typename param_t, bool has_bias, int BLOCK_M, int BLOCK_N>
212 struct tinygemm_kernel_nn {
213 static inline void apply(
112 - const scalar_t* __restrict__ A, const packed_t* __restrict__ B, scalar_t* __restrict__ C,
113 - const float* __restrict__ bias, const float* __restrict__ scale, int K, int lda, int ldb, int ldc, int64_t block_size_K) {
214 + const scalar_t* __restrict__ A,
215 + const packed_t* __restrict__ B,
216 + scalar_t* __restrict__ C,
217 + const float* __restrict__ bias,
218 + const param_t* __restrict__ scale,
219 + int64_t K,
220 + int64_t lda,
221 + int64_t ldb,
222 + int64_t ldc,
223 + int64_t block_size_K) {
224 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
225 }
226 };
227
228 +template <typename scalar_t, int BLOCK_M, int BLOCK_N>
229 +struct tinygemm_kernel_nn2 {
230 + static inline void apply(
231 + const scalar_t* __restrict__ A,
232 + const at::Float8_e4m3fn* __restrict__ B,
233 + scalar_t* __restrict__ C,
234 + float scale,
235 + int K,
236 + int lda,
237 + int ldb,
238 + int ldc) {
239 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
240 + }
241 +};
242 #if defined(CPU_CAPABILITY_AVX512)
243 template <bool has_bias, int BLOCK_M, int BLOCK_N>
120 -struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BLOCK_N> {
244 +struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, float, has_bias, BLOCK_M, BLOCK_N> {
245 static inline void apply(
122 - const at::BFloat16* __restrict__ A, const at::Float8_e4m3fn* __restrict__ B, at::BFloat16* __restrict__ C,
123 - const float* __restrict__ bias, const float* __restrict__ scale, int K, int lda, int ldb, int ldc, int64_t block_size_K) {
124 -
246 + const at::BFloat16* __restrict__ A,
247 + const at::Float8_e4m3fn* __restrict__ B,
248 + at::BFloat16* __restrict__ C,
249 + const float* __restrict__ bias,
250 + const float* __restrict__ scale,
251 + int64_t K,
252 + int64_t lda,
253 + int64_t ldb,
254 + int64_t ldc,
255 + int64_t block_size_K) {
256 constexpr int ROWS = BLOCK_M;
257 constexpr int COLS = BLOCK_N / 16;
258
128 - const int KB = div_up(K, BLOCK_K);
259 + const int64_t KB = div_up(K, (int64_t)BLOCK_K);
260
261 // prefetch distance
262 constexpr int PREFETCH_SIZE_K = 64;
@@ -139,6 +270,8 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BL
270 // block quant scale
271 __m512 vscale;
272
273 + const __m512 vexp = _mm512_castsi512_ps(_mm512_set1_epi32(kFP8_BIAS));
274 +
275 auto loadc = [&](auto i) {
276 constexpr int col = i % COLS;
277 if constexpr (has_bias) {
@@ -149,8 +282,8 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BL
282 };
283 Unroll<ROWS * COLS>{}(loadc);
284
152 - const int lda2 = lda >> 1;
153 - const int ldb2 = ldb; // ldb * 2 >> 1;
285 + const int64_t lda2 = lda >> 1;
286 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
287 const float* a_ptr = reinterpret_cast<const float*>(A);
288 const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(B);
289
@@ -170,34 +303,31 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BL
303 if constexpr (PREFETCH_SIZE_K > 0) {
304 _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
305 }
173 - vb[col + 0] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 0));
174 - vb[col + 1] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 1));
306 + vb[col + 0] = CVT_FP8_TO_BF16_EXT(_mm512_extracti32x8_epi32(b8, 0));
307 + vb[col + 1] = CVT_FP8_TO_BF16_EXT(_mm512_extracti32x8_epi32(b8, 1));
308 }
309 }
310 vsum[i] = _mm512_dpbf16_ps(vsum[i], va, vb[col]);
311 };
312
180 - constexpr int BLOCK_K2 = BLOCK_K >> 1;
181 - for (int kb = 0; kb < KB; ++kb) {
182 - int kb_start = kb * BLOCK_K2;
183 - int kb_end = std::min(K, kb_start + BLOCK_K2);
313 + constexpr int64_t BLOCK_K2 = BLOCK_K >> 1;
314 + for (int64_t kb = 0; kb < KB; ++kb) {
315 + int64_t kb_start = kb * BLOCK_K2;
316 + int64_t kb_end = std::min(K >> 1, kb_start + BLOCK_K2);
317 // 1. load scale vector
318 vscale = _mm512_set1_ps(scale[kb]);
319 + vscale = _mm512_mul_ps(vscale, vexp);
320 if constexpr (PREFETCH_SIZE_KB > 0) {
321 _mm_prefetch(scale + kb + PREFETCH_SIZE_KB, _MM_HINT_T0);
322 }
323 // 2. zero vsum for each block
190 - Unroll<ROWS * COLS>{}([&](auto i) {
191 - vsum[i] = _mm512_setzero_ps();
192 - });
324 + Unroll<ROWS * COLS>{}([&](auto i) { vsum[i] = _mm512_setzero_ps(); });
325 // 3. accumulate across each block
326 for (int k = kb_start; k < kb_end; ++k) {
327 Unroll<ROWS * COLS>{}(compute, k);
328 }
329 // 4. apply scale
198 - Unroll<ROWS * COLS>{}([&](auto i) {
199 - vc[i] = _mm512_fmadd_ps(vsum[i], vscale, vc[i]);
200 - });
330 + Unroll<ROWS * COLS>{}([&](auto i) { vc[i] = _mm512_fmadd_ps(vsum[i], vscale, vc[i]); });
331 }
332
333 auto storec = [&](auto i) {
@@ -213,14 +343,196 @@ struct tinygemm_kernel_nn<at::BFloat16, at::Float8_e4m3fn, has_bias, BLOCK_M, BL
343 Unroll<ROWS * COLS>{}(storec);
344 }
345 };
216 -#endif
346
218 -#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
219 - tinygemm_kernel_nn<scalar_t, at::Float8_e4m3fn, has_bias, MB_SIZE, NB_SIZE>::apply( \
220 - A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
221 - has_bias ? bias + nb_start : nullptr, scale, K, lda, ldb, ldc, block_size_K);
347 +template <int BLOCK_M, int BLOCK_N>
348 +struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
349 + static inline void apply(
350 + const at::BFloat16* __restrict__ A,
351 + const at::Float8_e4m3fn* __restrict__ B,
352 + at::BFloat16* __restrict__ C,
353 + float scale,
354 + int K,
355 + int lda,
356 + int ldb,
357 + int ldc) {
358 + constexpr int ROWS = BLOCK_M;
359 + constexpr int COLS = BLOCK_N / 16;
360 +
361 + // prefetch distance
362 + constexpr int PREFETCH_SIZE_K = 64;
363 +
364 + __m512bh va;
365 + __m512bh vb[COLS];
366 + __m512 vc[ROWS * COLS];
367 +
368 + const __m512 vscale = _mm512_set1_ps(scale);
369 +
370 + auto loadc = [&](auto i) { vc[i] = _mm512_setzero_ps(); };
371 + Unroll<ROWS * COLS>{}(loadc);
372 +
373 + const int K2 = K >> 1;
374 + const int lda2 = lda >> 1;
375 + const int ldb2 = ldb; // ldb * 2 >> 1;
376 + const float* a_ptr = reinterpret_cast<const float*>(A);
377 + const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(B);
378 +
379 + auto compute = [&](auto i, int k) {
380 + constexpr int row = i / COLS;
381 + constexpr int col = i % COLS;
382 +
383 + if constexpr (col == 0) {
384 + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
385 + }
386 + if constexpr (row == 0) {
387 + if constexpr (col % 2 == 0) {
388 + __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2 + col * 16);
389 + if constexpr (PREFETCH_SIZE_K > 0) {
390 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
391 + }
392 + vb[col + 0] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 0));
393 + vb[col + 1] = CVT_FP8_TO_BF16(_mm512_extracti32x8_epi32(b8, 1));
394 + }
395 + }
396 + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
397 + };
398 + for (int k = 0; k < K2; ++k) {
399 + Unroll<ROWS * COLS>{}(compute, k);
400 + }
401 +
402 + auto storec = [&](auto i) {
403 + constexpr int row = i / COLS;
404 + constexpr int col = i % COLS;
405 + // for COLS = 2, 4 use 512bit store
406 + if constexpr (col % 2 == 0) {
407 + __m512 vc0 = _mm512_mul_ps(vc[row * COLS + col + 0], vscale);
408 + __m512 vc1 = _mm512_mul_ps(vc[row * COLS + col + 1], vscale);
409 + _mm512_storeu_si512(
410 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
411 + }
412 + };
413 + Unroll<ROWS * COLS>{}(storec);
414 + }
415 +};
416 +
417 +template <bool has_bias, int BLOCK_M, int BLOCK_N>
418 +struct tinygemm_kernel_nn<at::BFloat16, uint8_t, uint8_t, has_bias, BLOCK_M, BLOCK_N> {
419 + static inline void apply(
420 + const at::BFloat16* __restrict__ A,
421 + const uint8_t* __restrict__ B,
422 + at::BFloat16* __restrict__ C,
423 + const float* __restrict__ bias,
424 + const uint8_t* __restrict__ scale,
425 + int K,
426 + int lda,
427 + int ldb,
428 + int ldc,
429 + int64_t block_size_K) {
430 + // mxfp4 supports only group size of 32
431 + // expect weight packed in 32-way, vnni2 format Nx2(64)
432 + assert(block_size_K == 32);
433 + assert(BLOCK_N == 32);
434 +
435 + constexpr int ROWS = BLOCK_M;
436 + constexpr int COLS = BLOCK_N / 16;
437 +
438 + // prefetch distance
439 + constexpr int PREFETCH_SIZE_K = 64;
440 + constexpr int PREFETCH_SIZE_KB = 1;
441 +
442 + __m512bh va;
443 + __m512bh vb[COLS];
444 + __m512 vc[ROWS * COLS];
445 +
446 + // holds Nx2(64) scales, interleaved as 2 belongs to K dimension
447 + // e.g. vs0: { s0, s0, s1, s1, ..., s15, s15}
448 + // vs1: {s16, s16, s17, s17, ..., s31, s31}
449 + __m512i vscale[COLS];
450 +
451 + // exponent bias 127
452 + const __m512i off = _mm512_set1_epi16(0x7F);
453 +
454 + auto loadc = [&](auto i) {
455 + constexpr int col = i % COLS;
456 + if constexpr (has_bias) {
457 + vc[i] = _mm512_loadu_ps(bias + col * 16);
458 + } else {
459 + vc[i] = _mm512_setzero_ps();
460 + }
461 + };
462 + Unroll<ROWS * COLS>{}(loadc);
463 +
464 + const int64_t K2 = K >> 1;
465 + const int64_t lda2 = lda >> 1;
466 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
467 + const float* a_ptr = reinterpret_cast<const float*>(A);
468 + const uint8_t* b_ptr = reinterpret_cast<const uint8_t*>(B);
469 +
470 + auto compute = [&](auto i, int k) {
471 + constexpr int row = i / COLS;
472 + constexpr int col = i % COLS;
473 +
474 + if constexpr (col == 0) {
475 + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
476 + if constexpr (PREFETCH_SIZE_K > 0) {
477 + _mm_prefetch(a_ptr + row * lda2 + k + PREFETCH_SIZE_K, _MM_HINT_T0);
478 + }
479 + }
480 + if constexpr (row == 0) {
481 + // load 32 * 2 (64) int4 at a time
482 + if constexpr (col % 2 == 0) {
483 + __m256i b4 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(b_ptr + k * ldb2 + col * 16));
484 + if constexpr (PREFETCH_SIZE_K > 0) {
485 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
486 + }
487 + std::tie(vb[col + 0], vb[col + 1]) = CVT_MXFP4_TO_BF16(b4, vscale[col + 0], vscale[col + 1]);
488 + }
489 + }
490 + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
491 + };
492 +
493 + for (int64_t k = 0; k < K2; ++k) {
494 + // update scales every 16x2 K
495 + if ((k & 15) == 0) {
496 + __m256i s8 = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(scale + (k >> 4) * 32));
497 + __m512i s16 = _mm512_slli_epi16(_mm512_sub_epi16(_mm512_cvtepu8_epi16(s8), off), 0x7);
498 + std::tie(vscale[0], vscale[1]) = transpose_2x32_16bit(s16, s16);
499 + }
500 + Unroll<ROWS * COLS>{}(compute, k);
501 + }
502 +
503 + auto storec = [&](auto i) {
504 + constexpr int row = i / COLS;
505 + constexpr int col = i % COLS;
506 + // for COLS = 2,4 use 512bit store
507 + if constexpr (col % 2 == 0) {
508 + _mm512_storeu_si512(
509 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
510 + (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
511 + }
512 + };
513 + Unroll<ROWS * COLS>{}(storec);
514 + }
515 +};
516 +#endif
517
223 -template <typename scalar_t, typename packed_t, bool has_bias>
518 +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
519 + tinygemm_kernel_nn<scalar_t, packed_t, param_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
520 + A + mb_start * lda, \
521 + B + nb_start * 2, \
522 + C + mb_start * ldc + nb_start, \
523 + has_bias ? bias + nb_start : nullptr, \
524 + scale, \
525 + K, \
526 + lda, \
527 + ldb, \
528 + ldc, \
529 + block_size_K);
530 +
531 +#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE) \
532 + tinygemm_kernel_nn2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
533 + A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, scale, K, lda, ldb, ldc);
534 +
535 +template <typename scalar_t, typename packed_t, typename param_t, bool has_bias>
536 struct brgemm {
537 static inline void apply(
538 const scalar_t* __restrict__ A,
@@ -229,19 +541,22 @@ struct brgemm {
541 scalar_t* __restrict__ Btmp,
542 float* __restrict__ Ctmp,
543 const float* __restrict__ bias,
232 - const float* __restrict__ scale,
544 + const param_t* __restrict__ scale,
545 int M,
546 int N,
547 int K,
548 int lda,
549 int ldb,
238 - int ldc) {
550 + int ldc,
551 + bool do_unpack = true) {
552 TORCH_CHECK(false, "struct brgemm: primary template not implemented!");
553 }
554 };
555 +template <typename scalar_t>
556 +struct brgemm2 {};
557
558 template <bool has_bias>
244 -struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
559 +struct brgemm<at::BFloat16, at::Float8_e4m3fn, float, has_bias> {
560 static inline void apply(
561 const at::BFloat16* __restrict__ A,
562 const at::Float8_e4m3fn* __restrict__ B,
@@ -255,22 +570,101 @@ struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
570 int K,
571 int lda,
572 int ldb,
258 - int ldc) {
259 -
573 + int ldc,
574 + bool do_unpack = true) {
575 constexpr int BLOCK_N = block_size_n();
576
577 // [K, BLOCK_N] -> [K / 2, BLOCK_N * 2]
578 const int ldb_tmp = BLOCK_N;
579
580 + if (do_unpack) {
581 + for (int k = 0; k < K; k += BLOCK_K) {
582 + int kb_size = std::min(BLOCK_K, K - k);
583 +
584 + int idx = k >> 7; // k / BLOCK_K where BLOCK_K = 128
585 + unpack_B(Btmp + k * ldb_tmp, B + k * ldb, N, kb_size, ldb, ldb_tmp, scale[idx]);
586 + }
587 + }
588 +
589 + at::native::cpublas::brgemm(M, N, K, lda, ldb_tmp, BLOCK_N, /* add_C */ false, A, Btmp, Ctmp);
590 +
591 + // copy from Ctmp to C
592 + for (int m = 0; m < M; ++m) {
593 + if constexpr (has_bias) {
594 + copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
595 + } else {
596 + copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
597 + }
598 + }
599 + }
600 +};
601 +
602 +template <>
603 +struct brgemm2<at::BFloat16> {
604 + static inline void apply(
605 + const at::BFloat16* __restrict__ A,
606 + const at::Float8_e4m3fn* __restrict__ B,
607 + at::BFloat16* __restrict__ C,
608 + at::BFloat16* __restrict__ Btmp,
609 + float* __restrict__ Ctmp,
610 + float scale,
611 + int M,
612 + int N,
613 + int K,
614 + int lda,
615 + int ldb,
616 + int ldc) {
617 + constexpr int BLOCK_N = block_size_n();
618 +
619 + // [BLOCK_K, BLOCK_N] -> [BLOCK_K / 2, BLOCK_N * 2]
620 + const int ldb_tmp = block_size_n();
621 +
622 + // accumulate across K per BLOCK_K
623 for (int k = 0; k < K; k += BLOCK_K) {
624 int kb_size = std::min(BLOCK_K, K - k);
625 + unpack_B(Btmp, B + k * ldb, N, kb_size, ldb, ldb_tmp);
626 +
627 + const bool add_C = (k != 0);
628 + at::native::cpublas::brgemm(M, N, kb_size, lda, ldb_tmp, BLOCK_N, add_C, A + k, Btmp, Ctmp);
629 + }
630 +
631 + // copy from Ctmp to C and mul scale
632 + for (int m = 0; m < M; ++m) {
633 + copy_mul_stub(C + m * ldc, Ctmp + m * BLOCK_N, N, scale);
634 + }
635 + }
636 +};
637 +
638 +template <bool has_bias>
639 +struct brgemm<at::BFloat16, uint8_t, uint8_t, has_bias> {
640 + static inline void apply(
641 + const at::BFloat16* __restrict__ A,
642 + const uint8_t* __restrict__ B,
643 + at::BFloat16* __restrict__ C,
644 + at::BFloat16* __restrict__ Btmp,
645 + float* __restrict__ Ctmp,
646 + const float* __restrict__ bias,
647 + const uint8_t* __restrict__ scale,
648 + int M,
649 + int N,
650 + int K,
651 + int lda,
652 + int ldb,
653 + int ldc,
654 + bool do_unpack = true) {
655 + constexpr int BLOCK_N = block_size_n();
656 +
657 + // [K, BLOCK_N] -> [K / 2, BLOCK_N * 2]
658 + const int ldb_tmp = BLOCK_N;
659
268 - int idx = k >> 7; // k / BLOCK_K where BLOCK_K = 128
269 - unpack_B(Btmp + k * ldb_tmp, B + k * ldb, N, kb_size, ldb, ldb_tmp, scale[idx]);
660 + if (do_unpack) {
661 + // group size 32 for mxfp4
662 + for (int k = 0; k < K; k += 32) {
663 + unpack_B(Btmp + k * ldb_tmp, B + k * (ldb >> 1), N, 32, ldb, ldb_tmp, scale + (k >> 5) * BLOCK_N);
664 + }
665 }
666
272 - at::native::cpublas::brgemm(
273 - M, N, K, lda, ldb_tmp, BLOCK_N, /* add_C */ false, A, Btmp, Ctmp);
667 + at::native::cpublas::brgemm(M, N, K, lda, ldb_tmp, BLOCK_N, /* add_C */ false, A, Btmp, Ctmp);
668
669 // copy from Ctmp to C
670 for (int m = 0; m < M; ++m) {
@@ -283,14 +677,14 @@ struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
677 }
678 };
679
286 -template <typename scalar_t, bool has_bias>
680 +template <typename scalar_t, typename packed_t, typename param_t, bool has_bias>
681 void tinygemm_kernel(
682 const scalar_t* __restrict__ A,
289 - const at::Float8_e4m3fn* __restrict__ B,
683 + const packed_t* __restrict__ B,
684 scalar_t* __restrict__ C,
685 scalar_t* __restrict__ Btmp,
686 float* __restrict__ Ctmp,
293 - const float* __restrict__ scale,
687 + const param_t* __restrict__ scale,
688 const float* __restrict__ bias,
689 int64_t M,
690 int64_t N,
@@ -299,11 +693,11 @@ void tinygemm_kernel(
693 int64_t ldb,
694 int64_t ldc,
695 bool brg,
302 - int64_t block_size_K) {
303 -
696 + int64_t block_size_K,
697 + bool do_unpack = true) {
698 if (brg) {
305 - brgemm<scalar_t, at::Float8_e4m3fn, has_bias>::apply(
306 - A, B, C, Btmp, Ctmp, bias, scale, M, N, K, lda, ldb, ldc);
699 + brgemm<scalar_t, packed_t, param_t, has_bias>::apply(
700 + A, B, C, Btmp, Ctmp, bias, scale, M, N, K, lda, ldb, ldc, do_unpack);
701 return;
702 }
703
@@ -319,23 +713,136 @@ void tinygemm_kernel(
713 int64_t nb_start = nb * BLOCK_N;
714 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
715
322 - switch(mb_size << 4 | nb_size >> 4) {
323 - case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
324 - case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
325 - case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
326 - case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
327 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
716 + switch (mb_size << 4 | nb_size >> 4) {
717 + case 0x12:
718 + LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
719 + break;
720 + case 0x22:
721 + LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
722 + break;
723 + case 0x32:
724 + LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
725 + break;
726 + case 0x42:
727 + LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
728 + break;
729 + default:
730 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
731 }
732 }
733 }
734 }
735
736 template <typename scalar_t>
334 -void fp8_scaled_mm_kernel_impl(
737 +void tinygemm_kernel2(
738 + const scalar_t* __restrict__ A,
739 + const at::Float8_e4m3fn* __restrict__ B,
740 + scalar_t* __restrict__ C,
741 + scalar_t* __restrict__ Btmp,
742 + float* __restrict__ Ctmp,
743 + float scale,
744 + int64_t M,
745 + int64_t N,
746 + int64_t K,
747 + int64_t lda,
748 + int64_t ldb,
749 + int64_t ldc,
750 + bool brg) {
751 + if (brg) {
752 + brgemm2<scalar_t>::apply(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc);
753 + return;
754 + }
755 +
756 + // pattern: 1-8-8
757 + if (M == 1) {
758 + constexpr int64_t BLOCK_N = 128;
759 + const int64_t NB = div_up(N, BLOCK_N);
760 + int64_t mb_start = 0;
761 +
762 + for (int64_t nb = 0; nb < NB; ++nb) {
763 + int64_t nb_start = nb * BLOCK_N;
764 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
765 +
766 + switch (nb_size >> 4) {
767 + case 2:
768 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 32);
769 + break;
770 + case 4:
771 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 64);
772 + break;
773 + case 6:
774 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 96);
775 + break;
776 + case 8:
777 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 128);
778 + break;
779 + default:
780 + TORCH_CHECK(false, "Unexpected block size, 1x", "nb_size");
781 + }
782 + }
783 + return;
784 + }
785 +
786 + // pattern: 1-4-16
787 + constexpr int64_t BLOCK_M = 4;
788 + constexpr int64_t BLOCK_N = 64;
789 + const int64_t MB = div_up(M, BLOCK_M);
790 + const int64_t NB = div_up(N, BLOCK_N);
791 + for (int64_t mb = 0; mb < MB; ++mb) {
792 + int64_t mb_start = mb * BLOCK_M;
793 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
794 + for (int64_t nb = 0; nb < NB; ++nb) {
795 + int64_t nb_start = nb * BLOCK_N;
796 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
797 +
798 + switch (mb_size << 4 | nb_size >> 4) {
799 + // mb_size = 1
800 + case 0x12:
801 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 32);
802 + break;
803 + case 0x14:
804 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 64);
805 + break;
806 + // mb_size = 2
807 + case 0x22:
808 + LAUNCH_TINYGEMM_KERNEL_NN2(2, 32);
809 + break;
810 + case 0x24:
811 + LAUNCH_TINYGEMM_KERNEL_NN2(2, 64);
812 + break;
813 + // mb_size = 3
814 + case 0x32:
815 + LAUNCH_TINYGEMM_KERNEL_NN2(3, 32);
816 + break;
817 + case 0x34:
818 + LAUNCH_TINYGEMM_KERNEL_NN2(3, 64);
819 + break;
820 + // mb_size = 4
821 + case 0x42:
822 + LAUNCH_TINYGEMM_KERNEL_NN2(4, 32);
823 + break;
824 + case 0x44:
825 + LAUNCH_TINYGEMM_KERNEL_NN2(4, 64);
826 + break;
827 + default:
828 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
829 + }
830 + }
831 + }
832 +}
833 +
834 +// NB: fp8/fp4 scaled mm kernel implementation
835 +//
836 +// scalar_t packed_t param_t
837 +// FP8 BF16 FP8 FP32
838 +// MXFP4 BF16 U8 U8
839 +//
840 +template <typename scalar_t, typename packed_t, typename param_t, typename func_t>
841 +void fp_scaled_mm_kernel_impl(
842 scalar_t* __restrict__ out,
843 const scalar_t* __restrict__ mat1,
337 - const at::Float8_e4m3fn* __restrict__ mat2,
338 - const float* __restrict__ scales2,
844 + const packed_t* __restrict__ mat2,
845 + const param_t* __restrict__ scales2,
846 const float* __restrict__ bias,
847 scalar_t* __restrict__ buffer,
848 int64_t M,
@@ -345,42 +852,41 @@ void fp8_scaled_mm_kernel_impl(
852 int64_t out_strideM,
853 int64_t block_size_N,
854 int64_t block_size_K,
348 - int64_t buffer_size_per_thread) {
349 -
350 - constexpr int64_t BLOCK_M = block_size_m() * BLOCK_SIZE_M_SCALE;
855 + int64_t buffer_size_per_thread,
856 + const func_t& scale_offset_per_block) {
857 + constexpr int64_t BLOCK_M = block_size_m();
858 constexpr int64_t BLOCK_N = block_size_n();
859 const int64_t MB = div_up(M, BLOCK_M);
860 const int64_t NB = div_up(N, BLOCK_N);
861
355 - const int64_t scale_size_K = div_up(K, block_size_K);
356 - const int64_t blocks_n_per_group = block_size_N / BLOCK_N;
862 + const bool use_brgemm = can_use_brgemm<packed_t>(M);
863
358 - const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
864 + // use K/2 for mxfp4 and K for fp8
865 + const int64_t packed_K = get_row_size<packed_t>(K);
866
867 // parallel on [MB, NB]
868 AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
362 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
363 - int64_t mb{0}, nb{0};
364 - data_index_init(begin, mb, MB, nb, NB);
365 -
366 - int tid = at::get_thread_num();
869 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
870 + int tid = get_thread_num();
871 scalar_t* __restrict__ Btmp = buffer + tid * buffer_size_per_thread;
368 - float* __restrict__ Ctmp = (float*)((void*)(Btmp + BLOCK_N * K));
872 + float* __restrict__ Ctmp = (float*)((void*)(Btmp + MAX_CACHE_BLOCK_SIZE * BLOCK_N * K));
873
370 - for (int64_t i = begin; i < end; ++i) {
371 - UNUSED(i);
372 - const float* scale_ptr = scales2 + (nb / blocks_n_per_group) * scale_size_K;
874 + loop_2d<packed_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
875 + const param_t* scale_ptr = scales2 + scale_offset_per_block(nb);
876
877 int64_t mb_start = mb * BLOCK_M;
878 int64_t mb_size = std::min(M - mb_start, BLOCK_M);
879 int64_t nb_start = nb * BLOCK_N;
880 int64_t nb_size = std::min(N - nb_start, BLOCK_N);
881
379 - tinygemm_kernel<scalar_t, has_bias>(
882 + // only do unpacking for the first row
883 + bool do_unpack = (mb == mb0);
884 +
885 + tinygemm_kernel<scalar_t, packed_t, param_t, has_bias>(
886 /* A */ mat1 + mb_start * mat1_strideM,
381 - /* B */ mat2 + nb_start * K, // nb * BLOCK_N * K
887 + /* B */ mat2 + nb_start * packed_K, // nb * BLOCK_N * K
888 /* C */ out + mb_start * out_strideM + nb_start,
383 - /* Btmp */ Btmp,
889 + /* Btmp */ Btmp + nb_offset * BLOCK_N * K,
890 /* Ctmp */ Ctmp,
891 /* scale */ scale_ptr,
892 /* bias */ bias + nb_start,
@@ -391,11 +897,9 @@ void fp8_scaled_mm_kernel_impl(
897 /* ldb */ nb_size,
898 /* ldc */ out_strideM,
899 /* brg */ use_brgemm,
394 - /* block_size_K */ block_size_K);
395 -
396 - // move to the next index
397 - data_index_step(mb, MB, nb, NB);
398 - }
900 + /* block_size_K */ block_size_K,
901 + /* do_unpack */ do_unpack);
902 + });
903
904 if (use_brgemm) {
905 at::native::cpublas::brgemm_release();
@@ -404,7 +908,7 @@ void fp8_scaled_mm_kernel_impl(
908 });
909 }
910
407 -} // anonymous namespace
911 +} // anonymous namespace
912
913 // tinygemm interface
914 template <typename scalar_t>
@@ -422,42 +926,126 @@ void tinygemm_kernel(
926 int64_t ldb,
927 int64_t ldc,
928 bool brg,
425 - int64_t block_size_K) {
426 - tinygemm_kernel<scalar_t, false>(A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K);
929 + int64_t block_size_K,
930 + bool do_unpack) {
931 + tinygemm_kernel<scalar_t, at::Float8_e4m3fn, float, false>(
932 + A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
933 }
934
429 -#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
935 +template <typename scalar_t>
936 +void tinygemm_kernel(
937 + const scalar_t* __restrict__ A,
938 + const at::Float8_e4m3fn* __restrict__ B,
939 + scalar_t* __restrict__ C,
940 + scalar_t* __restrict__ Btmp,
941 + float* __restrict__ Ctmp,
942 + float scale,
943 + int64_t M,
944 + int64_t N,
945 + int64_t K,
946 + int64_t lda,
947 + int64_t ldb,
948 + int64_t ldc,
949 + bool brg) {
950 + tinygemm_kernel2<scalar_t>(A, B, C, Btmp, Ctmp, scale, M, N, K, lda, ldb, ldc, brg);
951 +}
952 +
953 +template <typename scalar_t>
954 +void tinygemm_kernel(
955 + const scalar_t* __restrict__ A,
956 + const uint8_t* __restrict__ B,
957 + scalar_t* __restrict__ C,
958 + scalar_t* __restrict__ Btmp,
959 + float* __restrict__ Ctmp,
960 + const uint8_t* __restrict__ scale,
961 + int64_t M,
962 + int64_t N,
963 + int64_t K,
964 + int64_t lda,
965 + int64_t ldb,
966 + int64_t ldc,
967 + bool brg,
968 + int64_t block_size_K,
969 + bool do_unpack) {
970 + tinygemm_kernel<scalar_t, uint8_t, uint8_t, false>(
971 + A, B, C, Btmp, Ctmp, scale, nullptr, M, N, K, lda, ldb, ldc, brg, block_size_K, do_unpack);
972 +}
973 +
974 +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE_A, TYPE_B, TYPE_S) \
975 + template void tinygemm_kernel<TYPE_A>( \
976 + const TYPE_A* __restrict__ A, \
977 + const TYPE_B* __restrict__ B, \
978 + TYPE_A* __restrict__ C, \
979 + TYPE_A* __restrict__ Btmp, \
980 + float* __restrict__ Ctmp, \
981 + const TYPE_S* __restrict__ scale, \
982 + int64_t M, \
983 + int64_t N, \
984 + int64_t K, \
985 + int64_t lda, \
986 + int64_t ldb, \
987 + int64_t ldc, \
988 + bool brg, \
989 + int64_t block_size_K, \
990 + bool do_unpack)
991 +
992 +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16, at::Float8_e4m3fn, float);
993 +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half, at::Float8_e4m3fn, float);
994 +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16, uint8_t, uint8_t);
995 +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half, uint8_t, uint8_t);
996 +
997 +#define INSTANTIATE_TINYGEMM_TEMPLATE2(TYPE) \
998 template void tinygemm_kernel<TYPE>( \
999 const TYPE* __restrict__ A, \
1000 const at::Float8_e4m3fn* __restrict__ B, \
1001 TYPE* __restrict__ C, \
1002 TYPE* __restrict__ Btmp, \
1003 float* __restrict__ Ctmp, \
436 - const float* __restrict__ scale, \
1004 + float scale, \
1005 int64_t M, \
1006 int64_t N, \
1007 int64_t K, \
1008 int64_t lda, \
1009 int64_t ldb, \
1010 int64_t ldc, \
443 - bool brg, \
444 - int64_t block_size_K)
1011 + bool brg)
1012
446 -INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
447 -INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
1013 +INSTANTIATE_TINYGEMM_TEMPLATE2(at::BFloat16);
1014
449 -at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
450 - std::vector<int64_t> block_size, const std::optional<at::Tensor>& bias,
451 - at::ScalarType out_dtype, bool is_vnni) {
452 - RECORD_FUNCTION("sgl-kernel::fp8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, block_size, bias}));
1015 +inline const float* get_bias_data(const std::optional<at::Tensor>& bias, int64_t N) {
1016 + if (bias.has_value()) {
1017 + const auto& bias_ref = bias.value();
1018 + CHECK_EQ(bias_ref.size(0), N);
1019 + return bias_ref.data_ptr<float>();
1020 + }
1021 + return nullptr;
1022 +}
1023
1024 +// FP8 and MXFP4 WoQ uses the same pattern:
1025 +// Btmp : [T, BLOCK_N * K]
1026 +// Ctmp : [T, BLOCK_M * BLOCK_N]
1027 +inline at::Tensor alloc_thread_buffer(const at::TensorOptions& options, int64_t K) {
1028 + constexpr int64_t BLOCK_M = block_size_m();
1029 + constexpr int64_t BLOCK_N = block_size_n();
1030 + int num_threads = at::get_num_threads();
1031 + int64_t size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * K + BLOCK_M * BLOCK_N * 2;
1032 + return at::empty({num_threads, size_per_thread}, options);
1033 +}
1034 +
1035 +at::Tensor fp8_scaled_mm_cpu(
1036 + at::Tensor& mat1,
1037 + at::Tensor& mat2,
1038 + at::Tensor& scales2,
1039 + std::vector<int64_t> block_size,
1040 + const std::optional<at::Tensor>& bias,
1041 + at::ScalarType out_dtype,
1042 + bool is_vnni) {
1043 auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
1044
1045 CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
1046 CHECK_INPUT(mat2);
1047 CHECK_INPUT(scales2);
459 - TORCH_CHECK(scales2.scalar_type() == at::kFloat,
460 - "fp8_scaled_mm_cpu: expect scales2 to be float32.");
1048 + TORCH_CHECK(scales2.scalar_type() == at::kFloat, "fp8_scaled_mm_cpu: expect scales2 to be float32.");
1049
1050 int64_t M = mat1.size(0);
1051 int64_t N = mat2.size(0);
@@ -467,13 +1055,10 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& sca
1055 CHECK_DIM(2, mat1);
1056 CHECK_DIM(2, mat2);
1057
470 - TORCH_CHECK(block_size.size() == 2,
471 - "fp8_scaled_mm_cpu: expect block_size.size() to be 2.");
472 -
1058 + TORCH_CHECK(block_size.size() == 2, "fp8_scaled_mm_cpu: expect block_size.size() to be 2.");
1059 int64_t block_size_N = block_size[0];
1060 int64_t block_size_K = block_size[1];
1061
476 - constexpr int64_t BLOCK_M = block_size_m() * BLOCK_SIZE_M_SCALE;
1062 constexpr int64_t BLOCK_N = block_size_n();
1063 TORCH_CHECK(block_size_N % BLOCK_N == 0, "fp8_scaled_mm_cpu: expect block_size_N to be multiples of BLOCK_N");
1064 TORCH_CHECK(block_size_K == BLOCK_K, "fp8_scaled_mm_cpu: expect block_size_K equals to BLOCK_K");
@@ -481,49 +1066,94 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& sca
1066 CHECK_EQ(scales2.size(1), div_up(K, block_size_K));
1067
1068 const auto st = mat1.scalar_type();
484 - TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
485 - "fp8_scaled_mm_cpu: expect A to be bfloat16 or half.");
486 - TORCH_CHECK(st == out_dtype,
487 - "fp8_scaled_mm_cpu: expect A has same dtype with out_dtype.");
488 - TORCH_CHECK(mat2.scalar_type() == at::kFloat8_e4m3fn,
489 - "fp8_scaled_mm_cpu: expect mat2 to be fp8_e4m3.");
490 - TORCH_CHECK(scales2.scalar_type() == at::kFloat,
491 - "fp8_scaled_mm_cpu: expect scales to be float32.");
1069 + TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "fp8_scaled_mm_cpu: expect A to be bfloat16 or half.");
1070 + TORCH_CHECK(st == out_dtype, "fp8_scaled_mm_cpu: expect A has same dtype with out_dtype.");
1071 + TORCH_CHECK(mat2.scalar_type() == at::kFloat8_e4m3fn, "fp8_scaled_mm_cpu: expect mat2 to be fp8_e4m3.");
1072 + TORCH_CHECK(scales2.scalar_type() == at::kFloat, "fp8_scaled_mm_cpu: expect scales to be float32.");
1073 auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
1074
494 - // strides
495 - int64_t mat1_strideM = mat1.stride(0);
496 - int64_t out_strideM = out.stride(0);
497 -
498 - const bool has_bias = bias.has_value();
499 - const float* bias_data = nullptr;
500 - if (has_bias) {
501 - CHECK_EQ(bias.value().size(0), N);
502 - bias_data = bias.value().data_ptr<float>();
503 - }
504 -
505 - // Btmp : [T, BLOCK_N * K]
506 - // Ctmp : [T, BLOCK_M * BLOCK_N]
507 - int num_threads = at::get_num_threads();
508 - int64_t size_per_thread = BLOCK_N * K + BLOCK_M * BLOCK_N * 2;
509 - auto buffer = at::empty({num_threads, size_per_thread}, mat1.options());
1075 + auto buffer = alloc_thread_buffer(mat1.options(), K);
1076
1077 AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "fp8_scaled_mm_kernel_impl", [&] {
512 - fp8_scaled_mm_kernel_impl<scalar_t>(
1078 + // used for lambda computing scale offset for each block
1079 + // fp8 block gemm sale shape: [N/128, K/128]
1080 + // for each block: [1, K/128]
1081 + const int64_t scale_size_K = div_up(K, block_size_K);
1082 + const int64_t blocks_n_per_group = block_size_N / BLOCK_N;
1083 +
1084 + fp_scaled_mm_kernel_impl<scalar_t, at::Float8_e4m3fn, float>(
1085 out.data_ptr<scalar_t>(),
1086 mat1.data_ptr<scalar_t>(),
1087 packed_w.data_ptr<at::Float8_e4m3fn>(),
1088 scales2.data_ptr<float>(),
517 - bias_data,
1089 + get_bias_data(bias, N),
1090 buffer.data_ptr<scalar_t>(),
1091 M,
1092 N,
1093 K,
522 - mat1_strideM,
523 - out_strideM,
1094 + mat1.stride(0),
1095 + out.stride(0),
1096 block_size_N,
1097 block_size_K,
526 - size_per_thread);
1098 + buffer.size(-1),
1099 + [&](int64_t nb) { return (nb / blocks_n_per_group) * scale_size_K; });
1100 + });
1101 +
1102 + return out;
1103 +}
1104 +
1105 +// mat1 : [M, K] bfloat16
1106 +// mat2 : [N, K / 2] uint8, actual layout: [N / BLOCK_N, K / 2, BLOCK_N, 2]
1107 +// scales2: [N, K / G], actual layout: [N / BLOCK_N, K / G, BLOCK_N]
1108 +at::Tensor mxfp4_scaled_mm_cpu(
1109 + at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2, const std::optional<at::Tensor>& bias, bool is_vnni) {
1110 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
1111 +
1112 + CHECK_INPUT(mat1);
1113 + CHECK_INPUT(mat2);
1114 + CHECK_INPUT(scales2);
1115 +
1116 + int64_t M = mat1.size(0);
1117 + int64_t N = mat2.size(0);
1118 + int64_t K = mat2.size(1) * 2;
1119 +
1120 + // mxfp4 supports only group size of 32 (2^5)
1121 + constexpr int64_t group_size = 32;
1122 + constexpr int64_t BLOCK_N = block_size_n();
1123 +
1124 + CHECK_EQ(mat1.size(1), K);
1125 + CHECK_EQ(scales2.numel(), N * K >> 5);
1126 +
1127 + const auto st = mat1.scalar_type();
1128 + TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "mxfp4_scaled_mm_cpu: expect A to be bfloat16 or half.");
1129 + TORCH_CHECK(mat2.scalar_type() == at::kByte, "mxfp4_scaled_mm_cpu: expect mat2 to be uint8.");
1130 + TORCH_CHECK(scales2.scalar_type() == at::kByte, "mxfp4_scaled_mm_cpu: expect scales to be uint8.");
1131 + auto out = at::empty({M, N}, mat1.options());
1132 +
1133 + auto buffer = alloc_thread_buffer(mat1.options(), K);
1134 +
1135 + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "mxfp4_scaled_mm_kernel_impl", [&] {
1136 + // used for lambda computing scale offset for each block
1137 + // mxfp4 block gemm sale shape: [N/BLOCK_N, K/32, BLOCK_N]
1138 + // for each block: [K/32, BLOCK_N]
1139 + const int64_t s_strideN = (K >> 5) * BLOCK_N;
1140 +
1141 + fp_scaled_mm_kernel_impl<scalar_t, uint8_t, uint8_t>(
1142 + out.data_ptr<scalar_t>(),
1143 + mat1.data_ptr<scalar_t>(),
1144 + packed_w.data_ptr<uint8_t>(),
1145 + scales2.data_ptr<uint8_t>(),
1146 + get_bias_data(bias, N),
1147 + buffer.data_ptr<scalar_t>(),
1148 + M,
1149 + N,
1150 + K,
1151 + mat1.stride(0),
1152 + out.stride(0),
1153 + /* block_size_N */ 1,
1154 + /* block_size_K */ group_size,
1155 + buffer.size(-1),
1156 + [&](int64_t nb) { return nb * s_strideN; });
1157 });
1158
1159 return out;
csrc/cpu/sgl-kernels/gemm_int4.cpp
+369 -230
@@ -1,10 +1,10 @@
1 -// SPDX-License-Identifier: Apache-2.0
2 -// Adapted from sgl-project/sglang
3 -// https://github.com/sgl-project/sglang/pull/8226
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
5 -#include <ATen/ATen.h>
4 +// clang-format off
5 +
6 +#include <torch/all.h>
7
7 -#include "common.h"
8 #include "gemm.h"
9 #include "vec.h"
10
@@ -24,19 +24,22 @@ struct ActDtype<false> {
24 using type = uint8_t;
25 };
26
27 +#if defined(CPU_CAPABILITY_AVX512)
28 struct alignas(32) m256i_wrapper {
29 __m256i data;
30 };
31
31 -#if defined(CPU_CAPABILITY_AVX512)
32 -inline std::array<m256i_wrapper, 2> load_zps_4vnni(
33 - const int8_t* __restrict__ zps) {
34 - __m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast<const int64_t*>(zps));
35 - __m256i vzps_high =
36 - _mm256_set1_epi64x(*reinterpret_cast<const int64_t*>(zps + 8));
32 +inline std::array<m256i_wrapper, 2> load_zps_4vnni(const int8_t* __restrict__ zps) {
33 + // broadcast 01234567 to
34 + // 01234567012345670123456701234567
35 + __m256i vzps_low = _mm256_set1_epi64x(*reinterpret_cast<const long*>(zps));
36 + __m256i vzps_high = _mm256_set1_epi64x(*reinterpret_cast<const long*>(zps + 8));
37 + // shuffle from
38 + // 01234567012345670123456701234567
39 + // to
40 + // 00001111222233334444555566667777
41 __m256i shuffle_mask =
38 - _mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3,
39 - 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
42 + _mm256_set_epi8(7, 7, 7, 7, 6, 6, 6, 6, 5, 5, 5, 5, 4, 4, 4, 4, 3, 3, 3, 3, 2, 2, 2, 2, 1, 1, 1, 1, 0, 0, 0, 0);
43 vzps_low = _mm256_shuffle_epi8(vzps_low, shuffle_mask);
44 vzps_high = _mm256_shuffle_epi8(vzps_high, shuffle_mask);
45 m256i_wrapper vzps_low_wp, vzps_high_wp;
@@ -45,8 +48,7 @@ inline std::array<m256i_wrapper, 2> load_zps_4vnni(
48 return {vzps_low_wp, vzps_high_wp};
49 }
50
48 -inline std::array<m256i_wrapper, 2> load_uint4_as_int8(
49 - const uint8_t* __restrict__ qB) {
51 +inline std::array<m256i_wrapper, 2> load_uint4_as_int8(const uint8_t* __restrict__ qB) {
52 __m256i packed = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(qB));
53 const __m256i low_mask = _mm256_set1_epi8(0x0f);
54 __m256i high = _mm256_srli_epi16(packed, 4);
@@ -58,37 +60,42 @@ inline std::array<m256i_wrapper, 2> load_uint4_as_int8(
60 return {low_wp, high_wp};
61 }
62
61 -template <int N, int ldb>
62 -void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB,
63 - const int8_t* __restrict__ qzeros, int64_t K) {
64 - #pragma GCC unroll 2
63 +template <int64_t N, int64_t ldb>
64 +void _dequant_weight_zp_only(const uint8_t* __restrict__ B, int8_t* dqB, const int8_t* __restrict__ qzeros, int64_t K) {
65 + // unpack weight int8 -> two int4
66 + // subtract zero point
67 + // B shape = [K, ldb] = [K, N / 2], actual shape = [K / 4, N / 2, 4]
68 + // dqB shape = [K, N], actual shape = [K / 4, N, 4]
69 +#pragma GCC unroll 2
70 for (int n = 0; n < N; n += 16) {
71 auto [zps_low_wp, zps_high_wp] = load_zps_4vnni(&qzeros[n]);
72 auto zps_low = zps_low_wp.data;
73 auto zps_high = zps_high_wp.data;
74 for (int k = 0; k < K; k += 4) {
70 - auto [vb_low_wp, vb_high_wp] =
71 - load_uint4_as_int8(B + ldb * k + n / 2 * 4);
75 + auto [vb_low_wp, vb_high_wp] = load_uint4_as_int8(B + ldb * k + n / 2 * 4);
76 auto vb_low = vb_low_wp.data;
77 auto vb_high = vb_high_wp.data;
78 vb_high = _mm256_sub_epi8(vb_high, zps_high);
79 vb_low = _mm256_sub_epi8(vb_low, zps_low);
76 - _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4),
77 - vb_low);
78 - _mm256_storeu_si256(
79 - reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high);
80 + // store vb to B
81 + _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + n * 4), vb_low);
82 + _mm256_storeu_si256(reinterpret_cast<__m256i_u*>(dqB + N * k + (n + 8) * 4), vb_high);
83 }
84 }
85 }
86
84 -template <bool sym_quant_act, int N, bool accum>
85 -void _dequant_and_store(float* __restrict__ output,
86 - const int32_t* __restrict__ input,
87 - const float* __restrict__ scale_a,
88 - const int32_t* __restrict__ zp_a,
89 - const float* __restrict__ scale_b,
90 - const int32_t* __restrict__ comp_b, int M, int ldi,
91 - int ldo, int ldsa = 1) {
87 +template <bool accum, int64_t N, bool sym_quant_act>
88 +void _dequant_and_store(
89 + float* __restrict__ output,
90 + const int32_t* __restrict__ input,
91 + const float* __restrict__ scale_a,
92 + const int32_t* __restrict__ zp_a,
93 + const float* __restrict__ scale_b,
94 + const int32_t* __restrict__ comp_b,
95 + int M,
96 + int ldi,
97 + int ldo,
98 + int ldsa = 1) {
99 for (int m = 0; m < M; ++m) {
100 float a_scale = *(scale_a + m * ldsa);
101 __m512 va_scale = _mm512_set1_ps(a_scale);
@@ -99,7 +106,7 @@ void _dequant_and_store(float* __restrict__ output,
106 va_zp = _mm512_set1_epi32(a_zp);
107 }
108 int n = 0;
102 - #pragma GCC unroll 2
109 +#pragma GCC unroll 2
110 for (; n < N; n += 16) {
111 __m512i vc = _mm512_loadu_si512(input + m * ldi + n);
112 if constexpr (!sym_quant_act) {
@@ -122,8 +129,7 @@ void _dequant_and_store(float* __restrict__ output,
129 if constexpr (sym_quant_act) {
130 dq_val = (float)input[m * ldi + n] * a_scale * scale_b[n];
131 } else {
125 - dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale *
126 - scale_b[n];
132 + dq_val = (float)(input[m * ldi + n] - a_zp * comp_b[n]) * a_scale * scale_b[n];
133 }
134 if constexpr (accum) {
135 output[m * ldo + n] += dq_val;
@@ -135,9 +141,10 @@ void _dequant_and_store(float* __restrict__ output,
141 }
142
143 #else
138 -template <int N, int ldb>
139 -void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB,
140 - const int8_t* qzeros, int64_t K) {
144 +template <int64_t N, int64_t ldb>
145 +void _dequant_weight_zp_only(const uint8_t* B, int8_t* dqB, const int8_t* qzeros, int64_t K) {
146 + // B shape = [K, N / 2]
147 + // dqB shape = [K, N]
148 for (int k = 0; k < K; ++k) {
149 for (int n = 0; n < N / 2; ++n) {
150 int32_t b = (int32_t)B[k * ldb + n];
@@ -158,20 +165,31 @@ inline __m512i combine_m256i(std::array<m256i_wrapper, 2> two_256) {
165 return combine_m256i(two_256[0].data, two_256[1].data);
166 }
167
168 +// negate elements in a according to b's sign
169 static inline __m512i _mm512_sign_epi8(__m512i a, __m512i b) {
170 __m512i zero = _mm512_setzero_si512();
171 __mmask64 blt0 = _mm512_movepi8_mask(b);
172 return _mm512_mask_sub_epi8(a, blt0, zero, a);
173 }
174
167 -template <bool sym_quant_act, int M, int N, int ldb>
168 -void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
169 - const float* scales_a, const int32_t* qzeros_a,
170 - const uint8_t* B, const float* scales_b,
171 - const int8_t* qzeros_b, int64_t K, int64_t lda,
172 - int64_t ldc) {
175 +template <int64_t M, int64_t N, int64_t ldb, bool sym_quant_act>
176 +void _dequant_gemm_accum_small_M(
177 + float* __restrict__ C,
178 + const uint8_t* A,
179 + const float* scales_a,
180 + const int32_t* qzeros_a,
181 + const uint8_t* B,
182 + const float* scales_b,
183 + const int8_t* qzeros_b,
184 + int64_t K,
185 + int64_t lda,
186 + int64_t ldc) {
187 + // if sym_quant_act is true, A pointer type is passed in as uint8_t* but actually int8_t*.
188 +
189 constexpr int COLS = N / 16;
174 - __m512i ones = _mm512_set1_epi8(1);
190 + // Computing compensation is faster than loading it for small M
191 + // because it's memory bound.
192 + __m512i ones = _mm512_set1_epi8(1); // used for computing compensation
193 __m512i va;
194 __m512i vb[COLS];
195 __m512i vc[M * COLS];
@@ -179,6 +197,7 @@ void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
197 __m512i vzps[COLS];
198 __m512i vcompensate[COLS];
199
200 + // Load scales and zps
201 Unroll<COLS>{}([&](auto i) {
202 vscales[i] = _mm512_loadu_ps(scales_b + i * 16);
203 vzps[i] = combine_m256i(load_zps_4vnni(qzeros_b + i * 16));
@@ -214,25 +233,25 @@ void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
233 }
234 };
235
236 + // Accumulate along k
237 constexpr const int unroll = 4;
238 int k = 0;
239 for (; k < K / 4 / unroll; k++) {
220 - Unroll<unroll>{}(
221 - [&](auto i) { Unroll<M * COLS>{}(compute, 4 * (k * unroll + i)); });
240 + Unroll<unroll>{}([&](auto i) { Unroll<M * COLS>{}(compute, 4 * (k * unroll + i)); });
241 }
242 k *= 4 * unroll;
243 for (; k < K; k += 4) {
244 Unroll<M * COLS>{}(compute, k);
245 }
246
247 + // Store to C
248 auto store = [&](auto i) {
249 constexpr const int row = i / COLS;
250 constexpr const int col = i % COLS;
251 + // compute (qC - compensate * zp_a) * scale_a * scale_b
252 __m512 vc_float;
253 if constexpr (!sym_quant_act) {
233 - vc[i] = _mm512_sub_epi32(
234 - vc[i], _mm512_mullo_epi32(vcompensate[col],
235 - _mm512_set1_epi32(*(qzeros_a + row))));
254 + vc[i] = _mm512_sub_epi32(vc[i], _mm512_mullo_epi32(vcompensate[col], _mm512_set1_epi32(*(qzeros_a + row))));
255 }
256 vc_float = _mm512_cvtepi32_ps(vc[i]);
257 vc_float = _mm512_mul_ps(vc_float, _mm512_set1_ps(*(scales_a + row)));
@@ -245,17 +264,28 @@ void _dequant_gemm_accum_small_M(float* __restrict__ C, const uint8_t* A,
264 Unroll<M * COLS>{}(store);
265 }
266
248 - #define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M) \
249 - _dequant_gemm_accum_small_M<sym_quant_act, M, N, ldb>( \
250 - C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
267 +#define CALL_DEQUANT_GEMM_ACCUM_SMALL_M(M) \
268 + _dequant_gemm_accum_small_M<M, N, ldb, sym_quant_act>(C, A, scales_a, qzeros_a, B, scales_b, qzeros_b, K, lda, ldc);
269 #endif
270
253 -template <bool sym_quant_act, int N, int ldb>
254 -void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
255 - const int32_t* qzeros_a, const uint8_t* B,
256 - const float* scales_b, const int8_t* qzeros_b,
257 - const int32_t* compensation, int8_t* dqB, int64_t M,
258 - int64_t K, int64_t lda, int64_t ldc, bool use_brgemm) {
271 +template <int64_t N, int64_t ldb, bool sym_quant_act>
272 +void _dequant_gemm_accum(
273 + float* C,
274 + const uint8_t* A,
275 + const float* scales_a,
276 + const int32_t* qzeros_a,
277 + const uint8_t* B,
278 + const float* scales_b,
279 + const int8_t* qzeros_b,
280 + const int32_t* compensation,
281 + int8_t* dqB,
282 + int64_t M,
283 + int64_t K,
284 + int64_t lda,
285 + int64_t ldc,
286 + bool use_brgemm) {
287 + // Compute GEMM int8 * int8 -> int32
288 + // dequant result to float by applying scales/qzeros
289 #if defined(CPU_CAPABILITY_AVX512)
290 if (!use_brgemm) {
291 switch (M) {
@@ -282,14 +312,12 @@ void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
312 Tin* A_ptr = (Tin*)A;
313 if (use_brgemm) {
314 int32_t C_i32[M * N];
285 - at::native::cpublas::brgemm(M, N, K, lda, N /*ldb*/, N /*ldc*/,
286 - false /* add_C */, A_ptr, dqB, C_i32,
287 - true /* is_vnni */);
315 + at::native::cpublas::brgemm(
316 + M, N, K, lda, N /*ldb*/, N /*ldc*/, false /* add_C */, A_ptr, dqB, C_i32, true /* is_vnni */);
317 _mm_prefetch(B + N * K / 2, _MM_HINT_T0);
318 _mm_prefetch(A + K, _MM_HINT_T0);
290 - _dequant_and_store<sym_quant_act, N, true>(C, C_i32, scales_a, qzeros_a,
291 - scales_b, compensation, M,
292 - N /*ldi*/, ldc, 1 /*ldsa*/);
319 + _dequant_and_store<true, N, sym_quant_act>(
320 + C, C_i32, scales_a, qzeros_a, scales_b, compensation, M, N /*ldi*/, ldc, 1 /*ldsa*/);
321 } else
322 #endif
323 {
@@ -297,13 +325,13 @@ void _dequant_gemm_accum(float* C, const uint8_t* A, const float* scales_a,
325 }
326 }
327
300 -template <int N>
328 +template <int64_t N>
329 inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
330 if (bias_ptr) {
331 for (int i = 0; i < m; ++i) {
332 int j = 0;
333 #if defined(CPU_CAPABILITY_AVX512)
306 - #pragma GCC unroll 2
334 +#pragma GCC unroll 2
335 for (; j < N; j += 16) {
336 __m512 bias_vec = _mm512_loadu_ps(bias_ptr + j);
337 _mm512_storeu_ps(y_buf + i * N + j, bias_vec);
@@ -313,11 +341,11 @@ inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
341 y_buf[i * N + j] = bias_ptr[j];
342 }
343 }
316 - } else {
344 + } else { // initialize to zero
345 for (int i = 0; i < m; ++i) {
346 int j = 0;
347 #if defined(CPU_CAPABILITY_AVX512)
320 - #pragma GCC unroll 2
348 +#pragma GCC unroll 2
349 for (; j < N; j += 16) {
350 __m512 zero_vec = _mm512_setzero_ps();
351 _mm512_storeu_ps(y_buf + i * N + j, zero_vec);
@@ -330,14 +358,13 @@ inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m) {
358 }
359 }
360
333 -template <int N, typename out_dtype>
334 -inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
335 - int64_t lda) {
361 +template <typename out_dtype, int64_t N>
362 +inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m, /* int64_t n, */ int64_t lda) {
363 for (int i = 0; i < m; ++i) {
364 int j = 0;
365 if constexpr (std::is_same<out_dtype, float>::value) {
366 #if defined(CPU_CAPABILITY_AVX512)
340 - #pragma GCC unroll 2
367 +#pragma GCC unroll 2
368 for (; j < N; j += 16) {
369 __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
370 _mm512_storeu_ps(c_ptr + i * lda + j, y_vec);
@@ -348,12 +375,11 @@ inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
375 }
376 } else if constexpr (std::is_same<out_dtype, at::BFloat16>::value) {
377 #if defined(CPU_CAPABILITY_AVX512)
351 - #pragma GCC unroll 2
378 +#pragma GCC unroll 2
379 for (; j < N; j += 16) {
380 __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
381 __m256i y_bf16_vec = at::vec::cvtfp32_bf16(y_vec);
355 - _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j),
356 - y_bf16_vec);
382 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_bf16_vec);
383 }
384 #endif
385 for (; j < N; ++j) {
@@ -361,12 +387,11 @@ inline void store_out(const float* y_buf, out_dtype* c_ptr, int64_t m,
387 }
388 } else if constexpr (std::is_same<out_dtype, at::Half>::value) {
389 #if defined(CPU_CAPABILITY_AVX512)
364 - #pragma GCC unroll 2
390 +#pragma GCC unroll 2
391 for (; j < N; j += 16) {
392 __m512 y_vec = _mm512_loadu_ps(y_buf + i * N + j);
393 __m256i y_fp16_vec = at::vec::cvtfp32_fp16(y_vec);
368 - _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j),
369 - y_fp16_vec);
394 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(c_ptr + i * lda + j), y_fp16_vec);
395 }
396 #endif
397 for (; j < N; ++j) {
@@ -392,16 +417,25 @@ void fill_val_stub(int32_t* __restrict__ output, int32_t value, int64_t size) {
417 }
418 }
419
395 -template <bool sym_quant_act, typename act_dtype, typename out_dtype>
420 +template <typename act_dtype, typename out_dtype, bool sym_quant_act>
421 void _da8w4_linear_impl(
397 - act_dtype* __restrict__ input, const float* __restrict__ input_scales,
422 + act_dtype* __restrict__ input,
423 + const float* __restrict__ input_scales,
424 const int32_t* __restrict__ input_qzeros,
399 - const uint8_t* __restrict__ weight, const float* __restrict__ weight_scales,
400 - const int8_t* __restrict__ weight_qzeros, const float* __restrict__ bias,
401 - out_dtype* __restrict__ output, float* __restrict__ output_temp,
402 - int8_t* __restrict__ dequant_weight_temp, int64_t M, int64_t N, int64_t K,
425 + const uint8_t* __restrict__ weight,
426 + const float* __restrict__ weight_scales,
427 + const int8_t* __restrict__ weight_qzeros,
428 + const float* __restrict__ bias,
429 + out_dtype* __restrict__ output,
430 + float* __restrict__ output_temp,
431 + int8_t* __restrict__ dequant_weight_temp,
432 + int64_t M,
433 + int64_t N,
434 + int64_t K,
435 int64_t num_groups) {
404 - const bool use_brgemm = can_use_brgemm<act_dtype>(M);
436 + // weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)]
437 + // scales/qzeros shape = [Nc, G, BLOCK_N]
438 + const bool use_brgemm = can_use_brgemm<int8_t>(M);
439 int64_t block_m = [&]() -> long {
440 if (M <= 48) {
441 return M;
@@ -432,30 +466,24 @@ void _da8w4_linear_impl(
466 int64_t mc_end = parallel_on_M ? mc + 1 : Mc;
467
468 for (int mci = mc; mci < mc_end; ++mci) {
435 - int64_t m_size =
436 - mci * block_m + block_m > M ? M - mci * block_m : block_m;
469 + int64_t m_size = mci * block_m + block_m > M ? M - mci * block_m : block_m;
470 + // copy bias to y_buf if bias is not None
471 auto bias_data = bias ? bias + nc * BLOCK_N : nullptr;
472 copy_bias<BLOCK_N>(bias_data, C_tmp, m_size);
473 for (int kci = 0; kci < Kc; ++kci) {
474 int32_t* compensation_ptr =
475 sym_quant_act
476 ? nullptr
443 - : (int32_t*)(void*)(weight +
444 - (nc * Kc + kci) *
445 - (BLOCK_N *
446 - (_block_k / 2 + sizeof(int32_t))) +
447 - _block_k * BLOCK_N / 2);
448 - _dequant_gemm_accum<sym_quant_act, BLOCK_N, BLOCK_N / 2>(
477 + : (int32_t*)(void*)(weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
478 + _block_k * BLOCK_N / 2) /*Bcomp*/;
479 + _dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, sym_quant_act>(
480 /*C*/ C_tmp,
481 /*A*/ (uint8_t*)input + mci * block_m * K + kci * _block_k,
482 /*scales_a*/ input_scales + mci * block_m,
483 /*qzeros_a*/ input_qzeros + mci * block_m,
453 - /*B*/ weight + (nc * Kc + kci) *
454 - (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
455 - /*scales_b*/ weight_scales + nc * BLOCK_N * num_groups +
456 - kci / block_per_group * BLOCK_N,
457 - /*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups +
458 - kci / block_per_group * BLOCK_N,
484 + /*B*/ weight + (nc * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
485 + /*scales_b*/ weight_scales + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N,
486 + /*qzeros_b*/ weight_qzeros + nc * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N,
487 /*Bcomp*/ compensation_ptr,
488 /*dqB_tmp*/ dqB_tmp,
489 /*M*/ m_size,
@@ -464,8 +492,8 @@ void _da8w4_linear_impl(
492 /*ldc*/ BLOCK_N,
493 /*use_brgemm*/ use_brgemm);
494 }
467 - store_out<BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N,
468 - m_size, N /*lda*/);
495 + // store y_buf to output with dtype conversion
496 + store_out<out_dtype, BLOCK_N>(C_tmp, output + mci * block_m * N + nc * BLOCK_N, m_size, N /*lda*/);
497 }
498 }
499 if (use_brgemm) {
@@ -476,15 +504,16 @@ void _da8w4_linear_impl(
504
505 } // anonymous namespace
506
479 -std::tuple<at::Tensor, at::Tensor, at::Tensor>
480 -convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
481 - const at::Tensor& scales,
482 - const at::Tensor& qzeros) {
483 - TORCH_CHECK(weight.dim() == 2,
484 - "DA8W4 CPU: Weight should be a 2D tensor for packing");
485 - TORCH_CHECK(
486 - weight.size(1) % 2 == 0,
487 - "DA8W4 CPU: Weight should have even number of columns for packing");
507 +/*
508 +return: packed_weight, packed_scales, packed_qzeros
509 +*/
510 +std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_int4_weight_packed_with_compensation(
511 + const at::Tensor& weight, const at::Tensor& scales, const at::Tensor& qzeros) {
512 + // weight shape = [N, K]
513 + // scales shape = [N, G]
514 + // qzeros shape = [N, G]
515 + TORCH_CHECK(weight.dim() == 2, "DA8W4 CPU: Weight should be a 2D tensor for packing");
516 + TORCH_CHECK(weight.size(1) % 2 == 0, "DA8W4 CPU: Weight should have even number of columns for packing");
517
518 auto new_scales = scales;
519 auto new_qzeros = qzeros;
@@ -505,20 +534,21 @@ convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
534 int64_t Nc = N / block_n;
535 int64_t Kc = K / _block_k;
536
537 + // Reorder weight to [N/block_n, K/_block_k, _block_k, block_n]
538 + // Reorder scales/qzeros to [N/block_n, G, block_n]
539 + // weight + compensation shape = [Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
540 + // scales/qzeros shape = [Nc, G, block_n]
541 auto weight_view = weight.view({Nc, block_n, Kc, _block_k});
542 at::Tensor weight_reordered = weight_view.permute({0, 2, 3, 1}).contiguous();
543 at::Tensor blocked_weight;
511 - at::Tensor blocked_scales =
512 - new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
513 - at::Tensor blocked_qzeros =
514 - new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
515 - auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) -
516 - new_qzeros.view({Nc, block_n, G, -1});
544 + at::Tensor blocked_scales = new_scales.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
545 + at::Tensor blocked_qzeros = new_qzeros.view({Nc, block_n, G}).permute({0, 2, 1}).contiguous();
546 + // Compensation = Σ(k)(W[k][n] - ZP[n]) for each block.
547 + auto weight_sub_qzero = weight.view({Nc, block_n, G, -1}).to(at::kInt) - new_qzeros.view({Nc, block_n, G, -1});
548 weight_sub_qzero = weight_sub_qzero.view({Nc, block_n, Kc, _block_k});
549 at::Tensor compensation = weight_sub_qzero.sum(-1);
550 compensation = compensation.permute({0, 2, 1}).contiguous().to(at::kInt);
520 - int64_t buffer_size_nbytes =
521 - _block_k * block_n / 2 + block_n * sizeof(int32_t);
551 + int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t);
552 blocked_weight = at::empty({Nc, Kc, buffer_size_nbytes}, weight.options());
553
554 auto weight_ptr = weight_reordered.data_ptr<uint8_t>();
@@ -528,24 +558,25 @@ convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
558 at::parallel_for(0, num_blocks, 1, [&](int64_t begin, int64_t end) {
559 for (const auto i : c10::irange(begin, end)) {
560 auto in_ptr = weight_ptr + i * _block_k * block_n;
531 - auto out_ptr =
532 - blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t));
561 + auto out_ptr = blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t));
562 int32_t* comp_in_prt = compensation_ptr + i * block_n;
534 - int32_t* comp_out_prt =
535 - (int32_t*)(void*)(blocked_weight_ptr +
536 - i * block_n * (_block_k / 2 + sizeof(int32_t)) +
537 - _block_k * block_n / 2);
563 + int32_t* comp_out_prt = (int32_t*)(void*)(blocked_weight_ptr + i * block_n * (_block_k / 2 + sizeof(int32_t)) +
564 + _block_k * block_n / 2);
565 + // Reorder weight block to VNNI4 and pack two lanes along N
566 + // N=16 viewed as two lanes: a0, ...a7, b0, ...b7
567 + // pack two lanes: [a0, b0], ..., [a7, b7]
568 + // plain shape = [_block_k, block_n]
569 + // packed shape = [_block_k / 4, block_n / 2, 4] viewed as [_block_k, block_n / 2]
570 constexpr int n_group_size = 8;
571 constexpr int vnni_size = 4;
540 - constexpr int n_group = block_n / n_group_size;
572 + constexpr int n_group = block_n / n_group_size; // 4
573 for (int nb = 0; nb < n_group; nb += 2) {
574 for (int k = 0; k < _block_k; k += vnni_size) {
575 for (int ni = 0; ni < n_group_size; ++ni) {
576 for (int ki = 0; ki < vnni_size; ++ki) {
577 int src_idx_1 = nb * n_group_size + ni + (k + ki) * block_n;
578 int src_idx_2 = (nb + 1) * n_group_size + ni + (k + ki) * block_n;
547 - int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size +
548 - k * block_n / 2 + ki;
579 + int dst_idx = (nb / 2 * n_group_size + ni) * vnni_size + k * block_n / 2 + ki;
580 uint8_t src_1 = *(in_ptr + src_idx_1);
581 uint8_t src_2 = *(in_ptr + src_idx_2);
582 uint8_t dst = (src_1 & 0x0f) | ((src_2 & 0x0f) << 4);
@@ -554,39 +585,128 @@ convert_int4_weight_packed_with_compensation(const at::Tensor& weight,
585 }
586 }
587 }
588 + // compensation [block_n]
589 for (int nb = 0; nb < block_n; nb++) {
590 *(comp_out_prt + nb) = *(comp_in_prt + nb);
591 }
592 }
593 });
594
563 - return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales),
564 - std::move(blocked_qzeros));
595 + return std::make_tuple(std::move(blocked_weight), std::move(blocked_scales), std::move(blocked_qzeros));
596 }
597
567 -std::tuple<at::Tensor, at::Tensor> autoawq_to_int4pack(at::Tensor qweight,
568 - at::Tensor qzeros) {
569 - auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4;
570 - auto qweight_unsq = qweight.unsqueeze(-1);
571 - auto unpacked = at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF;
572 - auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte);
598 +std::tuple<at::Tensor, at::Tensor> unpack_4bit_to_32bit_signed(const at::Tensor& qweight, const at::Tensor& qzeros) {
599 + TORCH_CHECK(qweight.scalar_type() == at::kInt, "qweight must be int32");
600 + TORCH_CHECK(qzeros.scalar_type() == at::kInt, "qzeros must be int32");
601 + const auto W0 = qweight.size(0);
602 + const auto W1 = qweight.size(1);
603 + const auto Z0 = qzeros.size(0);
604 + const auto Z1 = qzeros.size(1);
605 +
606 + // unpacked_weights: (W0 * 8, W1), int8
607 + auto unpacked_weights = at::zeros({W0 * 8, W1}, at::TensorOptions().dtype(at::kChar));
608 + // unpacked_zeros: (Z0, Z1 * 8), int8
609 + auto unpacked_zeros = at::zeros({Z0, Z1 * 8}, at::TensorOptions().dtype(at::kChar));
610 +
611 + const int32_t* qw_ptr = qweight.data_ptr<int32_t>();
612 + const int32_t* qz_ptr = qzeros.data_ptr<int32_t>();
613 + int8_t* uw_ptr = unpacked_weights.data_ptr<int8_t>();
614 + int8_t* uz_ptr = unpacked_zeros.data_ptr<int8_t>();
615 +
616 + // ---- unpack qweight ----
617 + for (int64_t row = 0; row < W0 * 8; ++row) {
618 + const int i = row & 7; // row % 8
619 + const int src_row = row >> 3; // row // 8
620 + const int shift = 4 * i;
621 + for (int64_t col = 0; col < W1; ++col) {
622 + int32_t v = qw_ptr[src_row * W1 + col];
623 + uw_ptr[row * W1 + col] = static_cast<int8_t>((v >> shift) & 0xF);
624 + }
625 + }
626 + // ---- unpack qzeros ----
627 + for (int64_t col = 0; col < Z1 * 8; ++col) {
628 + const int i = col & 7;
629 + const int src_col = col >> 3;
630 + const int shift = 4 * i;
631 +
632 + for (int64_t row = 0; row < Z0; ++row) {
633 + int32_t v = qz_ptr[row * Z1 + src_col];
634 + uz_ptr[row * (Z1 * 8) + col] = static_cast<int8_t>((v >> shift) & 0xF);
635 + }
636 + }
637
574 - auto qzeros_unsq = qzeros.unsqueeze(-1);
575 - auto qzeros_unpacked = at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF;
576 - auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte);
638 + return std::make_tuple(unpacked_weights, unpacked_zeros + 1);
639 +}
640
578 - return std::make_tuple(qweight_final, qzeros_final);
641 +std::tuple<at::Tensor, at::Tensor>
642 +autogptq_to_int4pack(const at::Tensor& qweight_tensor, const at::Tensor& qzeros_tensor) {
643 + TORCH_CHECK(qweight_tensor.scalar_type() == at::kInt, "qweight_tensor must be int32");
644 + TORCH_CHECK(qzeros_tensor.scalar_type() == at::kInt, "qzeros_tensor must be int32");
645 + TORCH_CHECK(qweight_tensor.is_cpu(), "CPU only implementation");
646 + if (qweight_tensor.dim() == 3) {
647 + const int64_t B = qweight_tensor.size(0);
648 + std::vector<at::Tensor> qweight_list;
649 + std::vector<at::Tensor> qzeros_list;
650 + qweight_list.reserve(B);
651 + qzeros_list.reserve(B);
652 + for (int64_t i = 0; i < B; ++i) {
653 + auto outputs = unpack_4bit_to_32bit_signed(qweight_tensor[i], qzeros_tensor[i]);
654 + at::Tensor unpacked_qweight = std::get<0>(outputs);
655 + at::Tensor unpacked_qzeros = std::get<1>(outputs);
656 + qweight_list.push_back(unpacked_qweight.transpose(0, 1).contiguous().to(at::kByte));
657 + qzeros_list.push_back(unpacked_qzeros.contiguous().to(at::kByte));
658 + }
659 + return std::make_tuple(at::stack(qweight_list).detach(), at::stack(qzeros_list).detach());
660 + }
661 + auto outputs = unpack_4bit_to_32bit_signed(qweight_tensor, qzeros_tensor);
662 + at::Tensor unpacked_qweight = std::get<0>(outputs);
663 + at::Tensor unpacked_qzeros = std::get<1>(outputs);
664 + at::Tensor return_qweight = unpacked_qweight.transpose(0, 1).contiguous().to(at::kByte);
665 + at::Tensor return_qzeros = unpacked_qzeros.contiguous().to(at::kByte);
666 + return std::make_tuple(return_qweight, return_qzeros);
667 +}
668 +
669 +std::tuple<at::Tensor, at::Tensor> int4pack(at::Tensor qweight, at::Tensor qzeros, int64_t quant_method_4bit) {
670 + if (quant_method_4bit == CPUQuantAlgo::AWQ) {
671 + // autoawq unpacking
672 + qweight = qweight.contiguous();
673 + qzeros = qzeros.contiguous();
674 + // bitshifts: [0, 4, 1, 5, 2, 6, 3, 7] * 4
675 + auto bitshifts = at::tensor({0, 4, 1, 5, 2, 6, 3, 7}, at::kInt) * 4;
676 + auto qweight_unsq = qweight.unsqueeze(-1); // [..., K, N/8, 1]
677 + auto unpacked = (at::bitwise_right_shift(qweight_unsq, bitshifts) & 0xF).contiguous();
678 + auto qweight_final = unpacked.flatten(-2).transpose(-1, -2).to(at::kByte).clone();
679 + auto qzeros_unsq = qzeros.unsqueeze(-1);
680 + auto qzeros_unpacked = (at::bitwise_right_shift(qzeros_unsq, bitshifts) & 0xF).contiguous();
681 + auto qzeros_final = qzeros_unpacked.flatten(-2).to(at::kByte).clone();
682 + return std::make_tuple(qweight_final, qzeros_final);
683 + } else if (quant_method_4bit == CPUQuantAlgo::GPTQ) {
684 + // autogptq unpacking
685 + auto outputs = autogptq_to_int4pack(qweight, qzeros);
686 + at::Tensor unpacked_qweight = std::get<0>(outputs);
687 + at::Tensor unpacked_qzeros = std::get<1>(outputs);
688 + return std::make_tuple(unpacked_qweight, unpacked_qzeros);
689 + } else {
690 + TORCH_CHECK(false, "CPU int4 pack only support AWQ or GPTQ...");
691 + }
692 }
693
694 std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
582 - at::Tensor qweight, at::Tensor qzeros, at::Tensor scales) {
583 - auto res = autoawq_to_int4pack(qweight, qzeros);
584 - auto _qweight = std::get<0>(res);
585 - auto _qzeros = std::get<1>(res);
695 + at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
696 + at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K / group_size, N / 8) , int32
697 + at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K / group_size, N) , bfloat16
698 + int64_t quant_method_4bit) {
699 + at::Tensor _qweight;
700 + at::Tensor _qzeros;
701 +
702 + auto res = int4pack(qweight, qzeros, quant_method_4bit);
703 + _qweight = std::get<0>(res);
704 + _qzeros = std::get<1>(res);
705 +
706 auto _scales = scales;
587 - _qzeros = _qzeros.transpose(-2, -1).contiguous();
707 + _qzeros = _qzeros.transpose(-2, -1).contiguous(); // .T
708 _scales = _scales.transpose(-2, -1).contiguous();
589 - if (_qweight.dim() == 3) {
709 + if (_qweight.dim() == 3) { // Dim=3 for MOE packing, TODO: refine a unified loop
710 int64_t E = _qweight.size(0);
711 int64_t K = _qweight.size(2);
712 int64_t G = _scales.size(2);
@@ -595,17 +715,12 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
715 int64_t block_n = block_size_n();
716 int64_t Nc = _qweight.size(1) / block_n;
717 int64_t Kc = K / _block_k;
598 - int64_t buffer_size_nbytes =
599 - _block_k * block_n / 2 + block_n * sizeof(int32_t);
600 - auto blocked_weight =
601 - at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options());
602 - auto blocked_scales =
603 - at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat);
604 - auto blocked_qzeros =
605 - at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar);
718 + int64_t buffer_size_nbytes = _block_k * block_n / 2 + block_n * sizeof(int32_t);
719 + auto blocked_weight = at::empty({E, Nc, Kc, buffer_size_nbytes}, _qweight.options());
720 + auto blocked_scales = at::empty({E, Nc, G, block_n}, _scales.options()).to(at::kFloat);
721 + auto blocked_qzeros = at::empty({E, Nc, G, block_n}, _qzeros.options()).to(at::kChar);
722 for (int i = 0; i < _qweight.size(0); i++) {
607 - auto res_ = convert_int4_weight_packed_with_compensation(
608 - _qweight[i], _scales[i], _qzeros[i]);
723 + auto res_ = convert_int4_weight_packed_with_compensation(_qweight[i], _scales[i], _qzeros[i]);
724 blocked_weight[i] = std::get<0>(res_);
725 blocked_scales[i] = std::get<1>(res_);
726 blocked_qzeros[i] = std::get<2>(res_);
@@ -614,8 +729,7 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
729 _scales = blocked_scales;
730 _qzeros = blocked_qzeros;
731 } else {
617 - auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales,
618 - _qzeros);
732 + auto res_ = convert_int4_weight_packed_with_compensation(_qweight, _scales, _qzeros);
733 _qweight = std::get<0>(res_);
734 _scales = std::get<1>(res_);
735 _qzeros = std::get<2>(res_);
@@ -624,89 +738,93 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
738 return std::make_tuple(_qweight, _qzeros, _scales);
739 }
740
627 -at::Tensor int4_scaled_mm_cpu_with_quant(const at::Tensor& input,
628 - const at::Tensor& weight,
629 - const at::Tensor& weight_scales,
630 - const at::Tensor& weight_qzeros,
631 - const std::optional<at::Tensor>& bias,
632 - at::ScalarType output_dtype) {
633 - RECORD_FUNCTION("vllm::int4_scaled_mm_cpu_with_quant",
634 - std::vector<c10::IValue>({input, weight}));
635 -
741 +at::Tensor int4_scaled_mm_cpu_with_quant(
742 + const at::Tensor& input,
743 + const at::Tensor& weight,
744 + const at::Tensor& weight_scales,
745 + const at::Tensor& weight_qzeros,
746 + const std::optional<at::Tensor>& bias,
747 + at::ScalarType output_dtype) {
748 int64_t M_a = input.size(0);
749 int64_t K_a = input.size(1);
750 int64_t lda = input.stride(0);
751
752 const auto st = input.scalar_type();
753 TORCH_CHECK(
642 - st == at::kBFloat16 || st == at::kHalf,
643 - "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half.");
754 + st == at::kBFloat16 || st == at::kHalf, "int4_scaled_mm_cpu_with_quant: expect A to be bfloat16 or half.");
755
645 - constexpr bool sym_quant_act = false;
756 + constexpr bool sym_quant_act = false; // TODO: add sym quant path
757 using Tin = typename ActDtype<sym_quant_act>::type;
647 - int64_t act_buffer_size =
648 - M_a * K_a + M_a * sizeof(float) + M_a * sizeof(int32_t);
649 - auto act_buffer =
650 - at::empty({act_buffer_size}, input.options().dtype(at::kByte));
758 + int64_t act_buffer_size = /* act quant */ M_a * K_a +
759 + /* act scale */ M_a * sizeof(float) +
760 + /* act zp */ M_a * sizeof(int32_t);
761 + auto act_buffer = at::empty({act_buffer_size}, input.options().dtype(at::kByte));
762 + // asym path, activation quants into uint8_t
763 auto Aq_data = act_buffer.data_ptr<uint8_t>();
764 auto As_data = reinterpret_cast<float*>(Aq_data + M_a * K_a);
765 auto Azp_data = reinterpret_cast<int32_t*>(As_data + M_a);
654 - fill_val_stub(Azp_data, 128, M_a);
766 + fill_val_stub(Azp_data, 128, M_a); // sym_a s8s8 is unified to u8s8 with compensation (128)
767
768 auto out_sizes = input.sizes().vec();
769 int64_t N = weight_scales.size(0) * weight_scales.size(-1);
770 out_sizes.back() = N;
771 auto output = at::empty(out_sizes, input.options());
772 + // weight + compensation shape = [Nc, Kc, BLOCK_N * _block_k / 2 + BLOCK_N*sizeof(int32_t)]
773 + // scales/qzeros shape = [Nc, G, BLOCK_N]
774 int64_t Nc = weight.size(0);
775 int64_t Kc = weight.size(1);
776 int64_t _block_k = K_a / Kc;
777 TORCH_CHECK(N == Nc * BLOCK_N, "DA8W4: weight and input shapes mismatch");
778 + // scales/qzeros shape = [Nc, G, BLOCK_N]
779 int64_t num_groups = weight_scales.size(1);
780
781 const uint8_t* b_ptr = weight.data_ptr<uint8_t>();
782 const float* b_scales_ptr = weight_scales.data_ptr<float>();
783 const int8_t* b_qzeros_ptr = weight_qzeros.data_ptr<int8_t>();
669 - const float* bias_ptr =
670 - bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
784 + const float* bias_ptr = bias.has_value() ? bias.value().data_ptr<float>() : nullptr;
785 int num_threads = at::get_num_threads();
672 - int64_t temp_buffer_size = num_threads * BLOCK_M * BLOCK_N * sizeof(float) +
673 - num_threads * _block_k * BLOCK_N;
674 - auto c_temp_buffer =
675 - at::empty({temp_buffer_size}, input.options().dtype(at::kChar));
786 + int64_t temp_buffer_size = /* output temp */ num_threads * BLOCK_M * BLOCK_N * sizeof(float) +
787 + /* weight dequant temp */ num_threads * _block_k * BLOCK_N;
788 + auto c_temp_buffer = at::empty({temp_buffer_size}, input.options().dtype(at::kChar));
789 float* c_temp_ptr = (float*)((void*)(c_temp_buffer.data_ptr<int8_t>()));
677 - int8_t* dqB_temp_ptr =
678 - (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N));
679 -
680 -#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act) \
681 - AT_DISPATCH_FLOATING_TYPES_AND2( \
682 - at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype, \
683 - "int4_scaled_mm_cpu", [&] { \
684 - const scalar_t* __restrict__ A_data = input.data_ptr<scalar_t>(); \
685 - scalar_t* __restrict__ c_ptr = output.data_ptr<scalar_t>(); \
686 - at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) { \
687 - for (int64_t m = begin; m < end; ++m) { \
688 - quantize_row_int8<scalar_t>(Aq_data + m * K_a, As_data[m], \
689 - A_data + m * lda, K_a); \
690 - } \
691 - }); \
692 - _da8w4_linear_impl<sym_quant_act, Tin, scalar_t>( \
693 - Aq_data, As_data, Azp_data, b_ptr, b_scales_ptr, b_qzeros_ptr, \
694 - bias_ptr, c_ptr, c_temp_ptr, dqB_temp_ptr, M_a, N, K_a, \
695 - num_groups); \
790 + int8_t* dqB_temp_ptr = (int8_t*)((void*)(c_temp_ptr + num_threads * BLOCK_M * BLOCK_N));
791 +
792 +#define LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act) \
793 + AT_DISPATCH_FLOATING_TYPES_AND2( \
794 + at::ScalarType::BFloat16, at::ScalarType::Half, output_dtype, "int4_scaled_mm_cpu_with_quant", [&] { \
795 + const scalar_t* __restrict__ A_data = input.data_ptr<scalar_t>(); \
796 + scalar_t* __restrict__ c_ptr = output.data_ptr<scalar_t>(); \
797 + at::parallel_for(0, M_a, 0, [&](int64_t begin, int64_t end) { \
798 + for (int64_t m = begin; m < end; ++m) { \
799 + quantize_row_int8<scalar_t>(Aq_data + m * K_a, As_data[m], A_data + m * lda, K_a); \
800 + } \
801 + }); \
802 + _da8w4_linear_impl<Tin, scalar_t, sym_quant_act>( \
803 + Aq_data, \
804 + As_data, \
805 + Azp_data, \
806 + b_ptr, \
807 + b_scales_ptr, \
808 + b_qzeros_ptr, \
809 + bias_ptr, \
810 + c_ptr, \
811 + c_temp_ptr, \
812 + dqB_temp_ptr, \
813 + M_a, \
814 + N, \
815 + K_a, \
816 + num_groups); \
817 });
818
819 LAUNCH_DA8W4_LINEAR_WITH_QUANT_IMPL(sym_quant_act);
820
821 return output;
822 }
702 -
703 -namespace {
704 -
823 template <typename scalar_t>
706 -inline void copy_stub(scalar_t* __restrict__ out,
707 - const float* __restrict__ input, int64_t size) {
824 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
825 using Vec = at::vec::Vectorized<scalar_t>;
826 using fVec = at::vec::Vectorized<float>;
827 +// no remainder
828 #pragma GCC unroll 4
829 for (int64_t d = 0; d < size; d += Vec::size()) {
830 fVec x0 = fVec::loadu(input + d);
@@ -716,40 +834,61 @@ inline void copy_stub(scalar_t* __restrict__ out,
834 }
835 }
836
719 -} // anonymous namespace
720 -
837 template <typename scalar_t>
722 -void tinygemm_kernel(scalar_t* C, float* C_temp, const uint8_t* A,
723 - const float* scales_a, const int32_t* qzeros_a,
724 - const uint8_t* B, const float* scales_b,
725 - const int8_t* qzeros_b, const int32_t* compensation,
726 - int8_t* dqB_tmp, int64_t M, int64_t K, int64_t lda,
727 - int64_t ldc_f, int64_t ldc_s, bool store_out,
728 - bool use_brgemm) {
729 - _dequant_gemm_accum<false, BLOCK_N, BLOCK_N / 2>(
730 - C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation,
731 - dqB_tmp, M, K, lda, ldc_f, use_brgemm);
838 +void tinygemm_kernel(
839 + scalar_t* C,
840 + float* C_temp,
841 + const uint8_t* A,
842 + const float* scales_a,
843 + const int32_t* qzeros_a,
844 + const uint8_t* B,
845 + const float* scales_b,
846 + const int8_t* qzeros_b,
847 + const int32_t* compensation,
848 + int8_t* dqB_tmp,
849 + int64_t M,
850 + int64_t K,
851 + int64_t lda,
852 + int64_t ldc_f,
853 + int64_t ldc_s,
854 + bool store_out,
855 + bool use_brgemm) {
856 + // TODO: add sym quant act, now only asym
857 + _dequant_gemm_accum<BLOCK_N, BLOCK_N / 2, false>(
858 + C_temp, A, scales_a, qzeros_a, B, scales_b, qzeros_b, compensation, dqB_tmp, M, K, lda, ldc_f, use_brgemm);
859 if (store_out) {
860 + // copy from Ctmp to C
861 for (int64_t m = 0; m < M; ++m) {
862 copy_stub<scalar_t>(C + m * ldc_s, C_temp + m * ldc_f, BLOCK_N);
863 }
864 }
865 }
866
739 -#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
740 - template void tinygemm_kernel<TYPE>( \
741 - TYPE * C, float* C_temp, const uint8_t* A, const float* scales_a, \
742 - const int32_t* qzeros_a, const uint8_t* B, const float* scales_b, \
743 - const int8_t* qzeros_b, const int32_t* compensation, int8_t* dqB_tmp, \
744 - int64_t M, int64_t K, int64_t lda, int64_t ldc_f, int64_t ldc_s, \
745 - bool store_out, bool use_brgemm)
867 +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
868 + template void tinygemm_kernel<TYPE>( \
869 + TYPE * C, \
870 + float* C_temp, \
871 + const uint8_t* A, \
872 + const float* scales_a, \
873 + const int32_t* qzeros_a, \
874 + const uint8_t* B, \
875 + const float* scales_b, \
876 + const int8_t* qzeros_b, \
877 + const int32_t* compensation, \
878 + int8_t* dqB_tmp, \
879 + int64_t M, \
880 + int64_t K, \
881 + int64_t lda, \
882 + int64_t ldc_f, \
883 + int64_t ldc_s, \
884 + bool store_out, \
885 + bool use_brgemm)
886
887 INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
888 INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
889
750 -at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
751 - at::Tensor& w_scales,
752 - std::optional<at::Tensor> bias) {
753 - return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias,
754 - x.scalar_type());
890 +// int4 gemm dispatch api register
891 +at::Tensor int4_scaled_mm_cpu(
892 + at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros, at::Tensor& w_scales, std::optional<at::Tensor> bias) {
893 + return int4_scaled_mm_cpu_with_quant(x, w, w_scales, w_zeros, bias, x.scalar_type());
894 }
csrc/cpu/sgl-kernels/gemm_int8.cpp
+199 -93
@@ -1,20 +1,83 @@
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"
5 -#include "vec.h"
7 #include "gemm.h"
7 -
8 -// clang-format off
8 +#include "vec.h"
9
10 namespace {
11
12 +template <typename scalar_t, bool has_bias, int BLOCK_N>
13 +struct scale_C {
14 + static inline void apply(
15 + scalar_t* __restrict__ C,
16 + const int32_t* __restrict__ Ctmp,
17 + const int32_t* __restrict__ Bcomp,
18 + const float* __restrict__ bias,
19 + float As,
20 + const float* __restrict__ Bs) {
21 + TORCH_CHECK(false, "scale_C: scalar path not implemented!");
22 + }
23 +};
24 +
25 +#if defined(CPU_CAPABILITY_AVX512)
26 +template <bool has_bias, int BLOCK_N>
27 +struct scale_C<at::BFloat16, has_bias, BLOCK_N> {
28 + static inline void apply(
29 + at::BFloat16* __restrict__ C,
30 + const int32_t* __restrict__ Ctmp,
31 + const int32_t* __restrict__ Bcomp,
32 + const float* __restrict__ bias,
33 + float As,
34 + const float* __restrict__ Bs) {
35 + constexpr int COLS = BLOCK_N / 16;
36 + static_assert(COLS % 2 == 0);
37 +
38 + __m512 vc[COLS];
39 + __m512 vd0 = _mm512_set1_ps(As);
40 +
41 + auto compute = [&](auto col) {
42 + __m512 vd1 = _mm512_loadu_ps(Bs + col * 16);
43 + __m512i vcomp = _mm512_loadu_si512(Bcomp + col * 16);
44 + __m512i vc32 = _mm512_loadu_si512(Ctmp + col * 16);
45 + vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp));
46 + if constexpr (has_bias) {
47 + __m512 vbias = _mm512_loadu_ps(bias + col * 16);
48 + vc[col] = _mm512_fmadd_ps(_mm512_mul_ps(vc[col], vd0), vd1, vbias);
49 + } else {
50 + vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vd0), vd1);
51 + }
52 + };
53 + Unroll<COLS>{}(compute);
54 +
55 + auto storec = [&](auto col) {
56 + // for COLS = 2, 4 use 512bit store
57 + if constexpr (col % 2 == 0) {
58 + _mm512_storeu_si512(
59 + reinterpret_cast<__m512i*>((C + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc[col + 1], vc[col + 0])));
60 + }
61 + };
62 + Unroll<COLS>{}(storec);
63 + }
64 +};
65 +#endif
66 +
67 template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
68 struct tinygemm_kernel_nn {
69 static inline void apply(
15 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B, scalar_t* __restrict__ C,
16 - const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
17 - const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
70 + const uint8_t* __restrict__ A,
71 + const int8_t* __restrict__ B,
72 + scalar_t* __restrict__ C,
73 + const float* __restrict__ As,
74 + const float* __restrict__ Bs,
75 + const int32_t* __restrict__ Bcomp,
76 + const float* __restrict__ bias,
77 + int64_t K,
78 + int64_t lda,
79 + int64_t ldb,
80 + int64_t ldc) {
81 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
82 }
83 };
@@ -23,10 +86,17 @@ struct tinygemm_kernel_nn {
86 template <bool has_bias, int BLOCK_M, int BLOCK_N>
87 struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
88 static inline void apply(
26 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B, at::BFloat16* __restrict__ C,
27 - const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
28 - const float* __restrict__ bias, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
29 -
89 + const uint8_t* __restrict__ A,
90 + const int8_t* __restrict__ B,
91 + at::BFloat16* __restrict__ C,
92 + const float* __restrict__ As,
93 + const float* __restrict__ Bs,
94 + const int32_t* __restrict__ Bcomp,
95 + const float* __restrict__ bias,
96 + int64_t K,
97 + int64_t lda,
98 + int64_t ldb,
99 + int64_t ldc) {
100 constexpr int ROWS = BLOCK_M;
101 constexpr int COLS = BLOCK_N / 16;
102 static_assert(COLS % 2 == 0);
@@ -38,10 +108,10 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
108 __m512i vb[COLS];
109 __m512i vc[ROWS * COLS];
110 __m512i vcomp[COLS];
41 - __m512 vd0;
42 - __m512 vd1[COLS];
111 + __m512 vd0;
112 + __m512 vd1[COLS];
113
44 - // oops! 4x4 spills but luckily we use 4x2
114 + // oops! 4x4 spills but we use 4x2
115 __m512 vbias[COLS];
116
117 // [NOTE]: s8s8 igemm compensation in avx512-vnni
@@ -54,14 +124,12 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
124 // 1) 128 * b is pre-computed when packing B to vnni formats
125 // 2) a + 128 is fused when dynamically quantize A
126 //
57 - auto loadc = [&](auto i) {
58 - vc[i] = _mm512_set1_epi32(0);
59 - };
127 + auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
128 Unroll<ROWS * COLS>{}(loadc);
129
130 const int64_t K4 = K >> 2;
131 const int64_t lda4 = lda >> 2;
64 - const int64_t ldb4 = ldb; // ldb * 4 >> 2;
132 + const int64_t ldb4 = ldb; // ldb * 4 >> 2;
133 const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
134 const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
135
@@ -89,7 +157,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
157 constexpr int col = i % COLS;
158
159 // load a scale
92 - if constexpr(col == 0) {
160 + if constexpr (col == 0) {
161 vd0 = _mm512_set1_ps(As[row]);
162 }
163 // load b scale and vcomp per 2 vectors
@@ -120,8 +188,7 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
188 }
189
190 _mm512_storeu_si512(
123 - reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
124 - (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
191 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)), (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
192 }
193 };
194 Unroll<ROWS * COLS>{}(storec);
@@ -129,11 +196,19 @@ struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
196 };
197 #endif
198
132 -#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
133 - tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
134 - A + mb_start * lda, B + nb_start * 4, C + mb_start * ldc + nb_start, \
135 - As + mb_start, Bs + nb_start, Bcomp + nb_start, \
136 - has_bias ? bias + nb_start : nullptr, K, lda, ldb, ldc);
199 +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
200 + tinygemm_kernel_nn<scalar_t, has_bias, MB_SIZE, NB_SIZE>::apply( \
201 + A + mb_start * lda, \
202 + B + nb_start * 4, \
203 + C + mb_start * ldc + nb_start, \
204 + As + mb_start, \
205 + Bs + nb_start, \
206 + Bcomp + nb_start, \
207 + has_bias ? bias + nb_start : nullptr, \
208 + K, \
209 + lda, \
210 + ldb, \
211 + ldc);
212
213 template <typename scalar_t, bool has_bias>
214 void tinygemm_kernel(
@@ -151,10 +226,20 @@ void tinygemm_kernel(
226 int64_t ldb,
227 int64_t ldc,
228 bool brg) {
154 -
229 // B compensation
230 const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
231
232 + if (brg) {
233 + constexpr int BLOCK_N = block_size_n();
234 + at::native::cpublas::brgemm(M, N, K, lda, ldb, BLOCK_N, /* add_C */ false, A, B, Ctmp);
235 +
236 + // apply compensation and scale
237 + for (int64_t m = 0; m < M; ++m) {
238 + scale_C<scalar_t, has_bias, BLOCK_N>::apply(C + m * ldc, Ctmp + m * BLOCK_N, Bcomp, bias, As[m], Bs);
239 + }
240 + return;
241 + }
242 +
243 // pattern: 1-4-16
244 constexpr int64_t BLOCK_M = 4;
245 constexpr int64_t BLOCK_N = 64;
@@ -167,26 +252,43 @@ void tinygemm_kernel(
252 int64_t nb_start = nb * BLOCK_N;
253 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
254
170 - switch(mb_size << 4 | nb_size >> 4) {
255 + switch (mb_size << 4 | nb_size >> 4) {
256 // mb_size = 1
172 - case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
173 - case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
257 + case 0x12:
258 + LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
259 + break;
260 + case 0x14:
261 + LAUNCH_TINYGEMM_KERNEL_NN(1, 64);
262 + break;
263 // mb_size = 2
175 - case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
176 - case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
264 + case 0x22:
265 + LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
266 + break;
267 + case 0x24:
268 + LAUNCH_TINYGEMM_KERNEL_NN(2, 64);
269 + break;
270 // mb_size = 3
178 - case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
179 - case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
271 + case 0x32:
272 + LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
273 + break;
274 + case 0x34:
275 + LAUNCH_TINYGEMM_KERNEL_NN(3, 64);
276 + break;
277 // mb_size = 4
181 - case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
182 - case 0x44: LAUNCH_TINYGEMM_KERNEL_NN(4, 64); break;
183 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
278 + case 0x42:
279 + LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
280 + break;
281 + case 0x44:
282 + LAUNCH_TINYGEMM_KERNEL_NN(4, 64);
283 + break;
284 + default:
285 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
286 }
287 }
288 }
289 }
290
189 -template<typename scalar_t>
291 +template <typename scalar_t>
292 void int8_scaled_mm_kernel_impl(
293 scalar_t* __restrict__ out,
294 const uint8_t* __restrict__ mat1,
@@ -197,28 +299,22 @@ void int8_scaled_mm_kernel_impl(
299 int64_t M,
300 int64_t N,
301 int64_t K) {
200 -
302 constexpr int64_t BLOCK_M = block_size_m();
303 constexpr int64_t BLOCK_N = block_size_n();
304 const int64_t MB = div_up(M, BLOCK_M);
305 const int64_t NB = div_up(N, BLOCK_N);
306
206 - // TODO: brgemm u8s8 depends on PyTorch 2.7 release.
207 - const bool use_brgemm = false;
307 + const bool use_brgemm = can_use_brgemm<int8_t>(M);
308
309 // K + 4 after compensation
310 const int64_t packed_row_size = get_row_size<int8_t>(K);
311
312 AT_DISPATCH_BOOL(bias != nullptr, has_bias, [&] {
213 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
214 - int64_t mb{0}, nb{0};
215 - data_index_init(begin, mb, MB, nb, NB);
216 -
313 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
314 // for brgemm, use int32_t for accumulate
315 alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];
316
220 - for (int i = begin; i < end; ++i) {
221 - UNUSED(i);
317 + loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
318 int mb_start = mb * BLOCK_M;
319 int mb_size = std::min(M - mb_start, BLOCK_M);
320 int nb_start = nb * BLOCK_N;
@@ -239,10 +335,7 @@ void int8_scaled_mm_kernel_impl(
335 /* ldb */ nb_size,
336 /* ldc */ N,
337 /* brg */ use_brgemm);
242 -
243 - // move to the next index
244 - data_index_step(mb, MB, nb, NB);
245 - }
338 + });
339
340 if (use_brgemm) {
341 at::native::cpublas::brgemm_release();
@@ -251,28 +344,47 @@ void int8_scaled_mm_kernel_impl(
344 });
345 }
346
254 -} // anonymous namespace
347 +} // anonymous namespace
348
349 // tinygemm interface
350 template <typename scalar_t>
258 -void tinygemm_kernel(const uint8_t* __restrict__ A, const int8_t* __restrict__ B, scalar_t* __restrict__ C,
259 - int32_t* __restrict__ Ctmp, const float* __restrict__ As, const float* __restrict__ Bs,
260 - int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg) {
351 +void tinygemm_kernel(
352 + const uint8_t* __restrict__ A,
353 + const int8_t* __restrict__ B,
354 + scalar_t* __restrict__ C,
355 + int32_t* __restrict__ Ctmp,
356 + const float* __restrict__ As,
357 + const float* __restrict__ Bs,
358 + int64_t M,
359 + int64_t N,
360 + int64_t K,
361 + int64_t lda,
362 + int64_t ldb,
363 + int64_t ldc,
364 + bool brg) {
365 tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, As, Bs, nullptr, M, N, K, lda, ldb, ldc, brg);
366 }
367
264 -#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
265 - template void tinygemm_kernel<TYPE>( \
266 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B, TYPE* __restrict__ C, \
267 - int32_t* __restrict__ Ctmp, const float* __restrict__ As, const float* __restrict__ Bs, \
268 - int64_t M, int64_t N, int64_t K, int64_t lda, int64_t ldb, int64_t ldc, bool brg)
368 +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
369 + template void tinygemm_kernel<TYPE>( \
370 + const uint8_t* __restrict__ A, \
371 + const int8_t* __restrict__ B, \
372 + TYPE* __restrict__ C, \
373 + int32_t* __restrict__ Ctmp, \
374 + const float* __restrict__ As, \
375 + const float* __restrict__ Bs, \
376 + int64_t M, \
377 + int64_t N, \
378 + int64_t K, \
379 + int64_t lda, \
380 + int64_t ldb, \
381 + int64_t ldc, \
382 + bool brg)
383
384 INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
385 INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
386
387 std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
274 - RECORD_FUNCTION("sgl-kernel::per_token_quant_int8_cpu", std::vector<c10::IValue>({A}));
275 -
388 CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
389 CHECK_DIM(2, A);
390
@@ -281,8 +393,7 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
393 int64_t lda = A.stride(0);
394
395 const auto st = A.scalar_type();
284 - TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
285 - "per_token_quant_int8: expect A to be bfloat16 or half.");
396 + TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "per_token_quant_int8: expect A to be bfloat16 or half.");
397
398 auto Aq = at::empty({M, K}, A.options().dtype(at::kByte));
399 auto As = at::empty({M}, A.options().dtype(at::kFloat));
@@ -292,13 +403,9 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
403 float* __restrict__ As_data = As.data_ptr<float>();
404 const scalar_t* __restrict__ A_data = A.data_ptr<scalar_t>();
405
295 - at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
406 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
407 for (int64_t m = begin; m < end; ++m) {
297 - quantize_row_int8<scalar_t>(
298 - Aq_data + m * K,
299 - As_data[m],
300 - A_data + m * lda,
301 - K);
408 + quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
409 }
410 });
411 });
@@ -315,11 +422,14 @@ std::tuple<at::Tensor, at::Tensor> per_token_quant_int8_cpu(at::Tensor& A) {
422 // bias : [N]
423 // out : [M, N]
424 //
318 -at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
319 - at::Tensor& scales1, at::Tensor& scales2,
320 - std::optional<at::Tensor>& bias, at::ScalarType out_dtype, bool is_vnni) {
321 - RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales1, scales2, bias}));
322 -
425 +at::Tensor int8_scaled_mm_cpu(
426 + at::Tensor& mat1,
427 + at::Tensor& mat2,
428 + at::Tensor& scales1,
429 + at::Tensor& scales2,
430 + const std::optional<at::Tensor>& bias,
431 + at::ScalarType out_dtype,
432 + bool is_vnni) {
433 auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
434
435 CHECK_INPUT(mat1);
@@ -340,7 +450,8 @@ at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
450
451 TORCH_CHECK(mat1.scalar_type() == at::kByte, "int8_scaled_mm: expect mat1 to be uint8.");
452 TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm: expect mat2 to be int8.");
343 - TORCH_CHECK(scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
453 + TORCH_CHECK(
454 + scales1.scalar_type() == at::kFloat && scales2.scalar_type() == at::kFloat,
455 "int8_scaled_mm: expect scales to be float32.");
456
457 auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
@@ -368,10 +479,13 @@ at::Tensor int8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
479 }
480
481 // fused `per_token_quant_int8_cpu` and `int8_scaled_mm_cpu`
371 -at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
372 - const std::optional<at::Tensor>& bias, at::ScalarType out_dtype, bool is_vnni) {
373 - RECORD_FUNCTION("sgl-kernel::int8_scaled_mm_cpu", std::vector<c10::IValue>({mat1, mat2, scales2, bias}));
374 -
482 +at::Tensor int8_scaled_mm_with_quant(
483 + at::Tensor& mat1,
484 + at::Tensor& mat2,
485 + at::Tensor& scales2,
486 + const std::optional<at::Tensor>& bias,
487 + at::ScalarType out_dtype,
488 + bool is_vnni) {
489 auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
490
491 CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
@@ -390,14 +504,10 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Ten
504 CHECK_EQ(scales2.numel(), N);
505
506 const auto st = mat1.scalar_type();
393 - TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf,
394 - "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
395 - TORCH_CHECK(st == out_dtype,
396 - "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
397 - TORCH_CHECK(mat2.scalar_type() == at::kChar,
398 - "int8_scaled_mm_with_quant: expect mat2 to be int8.");
399 - TORCH_CHECK(scales2.scalar_type() == at::kFloat,
400 - "int8_scaled_mm_with_quant: expect scales to be float32.");
507 + TORCH_CHECK(st == at::kBFloat16 || st == at::kHalf, "int8_scaled_mm_with_quant: expect A to be bfloat16 or half.");
508 + TORCH_CHECK(st == out_dtype, "int8_scaled_mm_with_quant: expect A has same dtype with out_dtype.");
509 + TORCH_CHECK(mat2.scalar_type() == at::kChar, "int8_scaled_mm_with_quant: expect mat2 to be int8.");
510 + TORCH_CHECK(scales2.scalar_type() == at::kFloat, "int8_scaled_mm_with_quant: expect scales to be float32.");
511
512 const int64_t buffer_size = M * K + M * sizeof(float);
513 auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kByte));
@@ -415,13 +525,9 @@ at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2, at::Ten
525 float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
526 const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
527
418 - at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
528 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
529 for (int64_t m = begin; m < end; ++m) {
420 - quantize_row_int8<scalar_t>(
421 - Aq_data + m * K,
422 - As_data[m],
423 - A_data + m * lda,
424 - K);
530 + quantize_row_int8<scalar_t>(Aq_data + m * K, As_data[m], A_data + m * lda, K);
531 }
532 });
533
csrc/cpu/sgl-kernels/moe.cpp
+258 -305
@@ -1,12 +1,13 @@
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 +// clang-format off
5 +
6 +#include "moe.h"
7 +
8 #include "common.h"
5 -#include "vec.h"
9 #include "gemm.h"
10
8 -// clang-format off
9 -
11 namespace {
12
13 // [NOTE]: Fused MoE kernel with AMX
@@ -30,109 +31,6 @@ namespace {
31 // 3. abstract at::native::cpublas::brgemm with WoQ gemm (M = 1 & M != 1)
32 //
33
33 -template <typename scalar_t>
34 -inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) {
35 - using Vec = at::vec::Vectorized<scalar_t>;
36 - const Vec data_vec(val);
37 - at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size);
38 -}
39 -
40 -template <typename scalar_t>
41 -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
42 - using Vec = at::vec::Vectorized<scalar_t>;
43 - // no remainder
44 - #pragma GCC unroll 4
45 - for (int64_t d = 0; d < size; d += Vec::size()) {
46 - Vec data = Vec::loadu(input + d);
47 - data.store(out + d);
48 - }
49 -}
50 -
51 -template <typename scalar_t>
52 -inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
53 - using bVec = at::vec::Vectorized<scalar_t>;
54 - using fVec = at::vec::Vectorized<float>;
55 - constexpr int kVecSize = bVec::size();
56 - const fVec weight_vec = fVec(weight);
57 - int64_t d;
58 - #pragma GCC unroll 4
59 - for (d = 0; d <= size - kVecSize; d += kVecSize) {
60 - fVec data0 = fVec::loadu(input + d) * weight_vec;
61 - fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
62 - bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
63 - out_vec.store(out + d);
64 - }
65 - for (; d < size; ++d) {
66 - out[d] = static_cast<scalar_t>(input[d] * weight);
67 - }
68 -}
69 -
70 -// acc from [topk, K] to [K]
71 -template <typename scalar_t>
72 -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
73 - using bVec = at::vec::Vectorized<scalar_t>;
74 - using fVec = at::vec::Vectorized<float>;
75 - constexpr int kVecSize = bVec::size();
76 - if (topk == 1) {
77 - // do copy for topk = 1
78 - copy_stub(out, input, K);
79 - } else {
80 - // do sum for topk != 1
81 - int64_t d;
82 - #pragma GCC unroll 4
83 - for (d = 0; d <= K - kVecSize; d += kVecSize) {
84 - fVec sum_fvec0 = fVec(0.f);
85 - fVec sum_fvec1 = fVec(0.f);
86 - for (int t = 0; t < topk; ++t) {
87 - bVec x_bvec = bVec::loadu(input + t * K + d);
88 - fVec x_fvec0, x_fvec1;
89 - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
90 -
91 - sum_fvec0 += x_fvec0;
92 - sum_fvec1 += x_fvec1;
93 - }
94 - bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
95 - out_bvec.store(out + d);
96 - }
97 - for (; d < K; ++d) {
98 - float sum_val = 0.f;
99 - for (int t = 0; t < topk; ++t) {
100 - sum_val += static_cast<float>(input[t * K + d]);
101 - }
102 - out[d] = static_cast<scalar_t>(sum_val);
103 - }
104 - }
105 -}
106 -
107 -// out = input + input2 * scale
108 -template <typename scalar_t>
109 -inline void add_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input,
110 - const scalar_t* __restrict__ input2, float scale, int64_t size) {
111 -
112 - using bVec = at::vec::Vectorized<scalar_t>;
113 - using fVec = at::vec::Vectorized<float>;
114 - constexpr int kVecSize = bVec::size();
115 - const fVec s_vec = fVec(scale);
116 - int64_t d;
117 - #pragma GCC unroll 4
118 - for (d = 0; d <= size - kVecSize; d += kVecSize) {
119 - fVec x0 = fVec::loadu(input + d);
120 - fVec x1 = fVec::loadu(input + d + fVec::size());
121 -
122 - bVec y_bvec = bVec::loadu(input2 + d);
123 - fVec y0, y1;
124 - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
125 -
126 - x0 = x0 + y0 * s_vec;
127 - x1 = x1 + y1 * s_vec;
128 - bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
129 - out_vec.store(out + d);
130 - }
131 - for (; d < size; ++d) {
132 - out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
133 - }
134 -}
135 -
34 template <int BLOCK_M>
35 int moe_align_block_size(
36 int32_t* __restrict__ sorted_ids,
@@ -144,8 +42,7 @@ int moe_align_block_size(
42 int num_experts,
43 int numel,
44 int num_threads) {
147 -
148 - #define T_INDEX(tt) total_cnts + (tt) * num_experts
45 +#define T_INDEX(tt) total_cnts + (tt) * num_experts
46
47 // accumulate count of expert ids locally
48 at::parallel_for(0, numel, 0, [&](int begin, int end) {
@@ -160,8 +57,7 @@ int moe_align_block_size(
57 using iVec = at::vec::Vectorized<int32_t>;
58 for (int t = 0; t < num_threads; ++t) {
59 at::vec::map2<int32_t>(
163 - [](iVec x, iVec y) { return x + y; },
164 - T_INDEX(t + 1), T_INDEX(t + 1), T_INDEX(t), num_experts);
60 + [](iVec x, iVec y) { return x + y; }, T_INDEX(t + 1), T_INDEX(t + 1), T_INDEX(t), num_experts);
61 }
62
63 // the last row holds sums of each experts
@@ -201,7 +97,9 @@ int moe_align_block_size(
97 // padding value for sorted_ids: numel
98 auto sorted_id_size = [=](const int32_t* sorted_ids_ptr) {
99 for (int d = 0; d < BLOCK_M; ++d) {
204 - if (sorted_ids_ptr[d] == numel) { return d; }
100 + if (sorted_ids_ptr[d] == numel) {
101 + return d;
102 + }
103 }
104 return BLOCK_M;
105 };
@@ -215,7 +113,7 @@ int moe_align_block_size(
113 offsets[mb + 1] = sorted_id_size(sorted_ids + mb * BLOCK_M);
114 }
115 });
218 - // TODO: do we need to vectorize this ?
116 + // TODO: do we need to vecterize this ?
117 for (int mb = 0; mb < num_token_blocks; ++mb) {
118 offsets[mb + 1] += offsets[mb];
119 }
@@ -236,7 +134,6 @@ inline void silu_and_mul(
134 const float* __restrict__ input1, // y: y0, y1
135 int64_t m_size,
136 int64_t N) {
239 -
137 using bVec = at::vec::Vectorized<scalar_t>;
138 using fVec = at::vec::Vectorized<float>;
139
@@ -269,8 +166,14 @@ inline void silu_and_mul(
166 template <typename scalar_t, int BLOCK_M, int BLOCK_N>
167 struct tinygemm_kernel_nn2 {
168 static inline void apply(
272 - const scalar_t* __restrict__ A, const scalar_t* __restrict__ B0, const scalar_t* __restrict__ B1,
273 - scalar_t* __restrict__ C, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
169 + const scalar_t* __restrict__ A,
170 + const scalar_t* __restrict__ B0,
171 + const scalar_t* __restrict__ B1,
172 + scalar_t* __restrict__ C,
173 + int64_t K,
174 + int64_t lda,
175 + int64_t ldb,
176 + int64_t ldc) {
177 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
178 }
179 };
@@ -279,9 +182,14 @@ struct tinygemm_kernel_nn2 {
182 template <int BLOCK_M, int BLOCK_N>
183 struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
184 static inline void apply(
282 - const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B0, const at::BFloat16* __restrict__ B1,
283 - at::BFloat16* __restrict__ C, int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
284 -
185 + const at::BFloat16* __restrict__ A,
186 + const at::BFloat16* __restrict__ B0,
187 + const at::BFloat16* __restrict__ B1,
188 + at::BFloat16* __restrict__ C,
189 + int64_t K,
190 + int64_t lda,
191 + int64_t ldb,
192 + int64_t ldc) {
193 constexpr int ROWS = BLOCK_M;
194 constexpr int COLS = BLOCK_N / 16;
195
@@ -304,7 +212,7 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
212
213 const int64_t K2 = K >> 1;
214 const int64_t lda2 = lda >> 1;
307 - const int64_t ldb2 = ldb; // ldb * 2 >> 1;
215 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
216 const float* a_ptr = reinterpret_cast<const float*>(A);
217 const float* b0_ptr = reinterpret_cast<const float*>(B0);
218 const float* b1_ptr = reinterpret_cast<const float*>(B1);
@@ -352,17 +260,16 @@ struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
260 _mm512_storeu_si512(
261 reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
262 (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
355 - }
263 + }
264 };
265 Unroll<ROWS * COLS>{}(storec);
266 }
267 };
268 #endif
269
362 -#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
363 - tinygemm_kernel_nn2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
364 - A + mb_start * lda, B0 + nb_start * 2, B1 + nb_start * 2, \
365 - C + mb_start * ldc + nb_start, K, lda, ldb, ldc);
270 +#define LAUNCH_TINYGEMM_KERNEL_NN(MB_SIZE, NB_SIZE) \
271 + tinygemm_kernel_nn2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
272 + A + mb_start * lda, B0 + nb_start * 2, B1 + nb_start * 2, C + mb_start * ldc + nb_start, K, lda, ldb, ldc);
273
274 template <typename scalar_t>
275 void tinygemm_kernel(
@@ -376,7 +283,6 @@ void tinygemm_kernel(
283 int64_t lda,
284 int64_t ldb,
285 int64_t ldc) {
379 -
286 // pattern: 1-(2+2)-(8+8)
287 constexpr int64_t BLOCK_M = 4;
288 constexpr int64_t BLOCK_N = 32;
@@ -389,16 +295,25 @@ void tinygemm_kernel(
295 int64_t nb_start = nb * BLOCK_N;
296 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
297
392 - switch(mb_size << 4 | nb_size >> 4) {
298 + switch (mb_size << 4 | nb_size >> 4) {
299 // mb_size = 1
394 - case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
300 + case 0x12:
301 + LAUNCH_TINYGEMM_KERNEL_NN(1, 32);
302 + break;
303 // mb_size = 2
396 - case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
304 + case 0x22:
305 + LAUNCH_TINYGEMM_KERNEL_NN(2, 32);
306 + break;
307 // mb_size = 3
398 - case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
308 + case 0x32:
309 + LAUNCH_TINYGEMM_KERNEL_NN(3, 32);
310 + break;
311 // mb_size = 4
400 - case 0x42: LAUNCH_TINYGEMM_KERNEL_NN(4, 32); break;
401 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
312 + case 0x42:
313 + LAUNCH_TINYGEMM_KERNEL_NN(4, 32);
314 + break;
315 + default:
316 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
317 }
318 }
319 }
@@ -407,8 +322,13 @@ void tinygemm_kernel(
322 template <typename scalar_t, int BLOCK_M, int BLOCK_N>
323 struct tinygemm_kernel_nn {
324 static inline void apply(
410 - const scalar_t* __restrict__ A, const scalar_t* __restrict__ B, float* __restrict__ C,
411 - int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
325 + const scalar_t* __restrict__ A,
326 + const scalar_t* __restrict__ B,
327 + float* __restrict__ C,
328 + int64_t K,
329 + int64_t lda,
330 + int64_t ldb,
331 + int64_t ldc) {
332 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
333 }
334 };
@@ -417,9 +337,13 @@ struct tinygemm_kernel_nn {
337 template <int BLOCK_M, int BLOCK_N>
338 struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
339 static inline void apply(
420 - const at::BFloat16* __restrict__ A, const at::BFloat16* __restrict__ B, float* __restrict__ C,
421 - int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
422 -
340 + const at::BFloat16* __restrict__ A,
341 + const at::BFloat16* __restrict__ B,
342 + float* __restrict__ C,
343 + int64_t K,
344 + int64_t lda,
345 + int64_t ldb,
346 + int64_t ldc) {
347 constexpr int ROWS = BLOCK_M;
348 constexpr int COLS = BLOCK_N / 16;
349
@@ -432,14 +356,12 @@ struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
356 __m512bh vb[COLS];
357 __m512 vc[ROWS * COLS];
358
435 - auto loadc = [&](auto i) {
436 - vc[i] = _mm512_set1_ps(0.f);
437 - };
359 + auto loadc = [&](auto i) { vc[i] = _mm512_set1_ps(0.f); };
360 Unroll<ROWS * COLS>{}(loadc);
361
362 const int64_t K2 = K >> 1;
363 const int64_t lda2 = lda >> 1;
442 - const int64_t ldb2 = ldb; // ldb * 2 >> 1;
364 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
365 const float* a_ptr = reinterpret_cast<const float*>(A);
366 const float* b_ptr = reinterpret_cast<const float*>(B);
367
@@ -466,17 +388,15 @@ struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
388 constexpr int row = i / COLS;
389 constexpr int col = i % COLS;
390 _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), vc[i]);
469 -
391 };
392 Unroll<ROWS * COLS>{}(storec);
393 }
394 };
395 #endif
396
476 -#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE) \
477 - tinygemm_kernel_nn<scalar_t, MB_SIZE, NB_SIZE>::apply( \
478 - A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, \
479 - K, lda, ldb, ldc);
397 +#define LAUNCH_TINYGEMM_KERNEL_NN2(MB_SIZE, NB_SIZE) \
398 + tinygemm_kernel_nn<scalar_t, MB_SIZE, NB_SIZE>::apply( \
399 + A + mb_start * lda, B + nb_start * 2, C + mb_start * ldc + nb_start, K, lda, ldb, ldc);
400
401 template <typename scalar_t>
402 void tinygemm_kernel(
@@ -489,7 +409,6 @@ void tinygemm_kernel(
409 int64_t lda,
410 int64_t ldb,
411 int64_t ldc) {
492 -
412 // pattern: 1-2-8
413 constexpr int64_t BLOCK_M = 4;
414 constexpr int64_t BLOCK_N = 32;
@@ -502,16 +421,25 @@ void tinygemm_kernel(
421 int64_t nb_start = nb * BLOCK_N;
422 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
423
505 - switch(mb_size << 4 | nb_size >> 4) {
424 + switch (mb_size << 4 | nb_size >> 4) {
425 // mb_size = 1
507 - case 0x12: LAUNCH_TINYGEMM_KERNEL_NN2(1, 32); break;
426 + case 0x12:
427 + LAUNCH_TINYGEMM_KERNEL_NN2(1, 32);
428 + break;
429 // mb_size = 2
509 - case 0x22: LAUNCH_TINYGEMM_KERNEL_NN2(2, 32); break;
430 + case 0x22:
431 + LAUNCH_TINYGEMM_KERNEL_NN2(2, 32);
432 + break;
433 // mb_size = 3
511 - case 0x32: LAUNCH_TINYGEMM_KERNEL_NN2(3, 32); break;
434 + case 0x32:
435 + LAUNCH_TINYGEMM_KERNEL_NN2(3, 32);
436 + break;
437 // mb_size = 4
513 - case 0x42: LAUNCH_TINYGEMM_KERNEL_NN2(4, 32); break;
514 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
438 + case 0x42:
439 + LAUNCH_TINYGEMM_KERNEL_NN2(4, 32);
440 + break;
441 + default:
442 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
443 }
444 }
445 }
@@ -537,7 +465,6 @@ void fused_experts_kernel_impl(
465 int64_t E,
466 int64_t topk,
467 int64_t num_tokens_post_pad) {
540 -
468 // handle 2 tiles per block
469 constexpr int64_t BLOCK_M = block_size_m();
470 constexpr int64_t BLOCK_N = block_size_n();
@@ -552,39 +479,36 @@ void fused_experts_kernel_impl(
479 const int64_t stride_e = 2 * N * K;
480 const int64_t stride_n = K;
481
482 + int64_t avg_M = std::max(int64_t(1), M * topk / E);
483 + const bool use_brgemm = can_use_brgemm<scalar_t>(avg_M);
484 +
485 // here we only parallel on half of 2N to fuse silu_and_mul with gemm
556 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
486 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
487 // get local pointers
558 - int tid = at::get_thread_num();
488 + int tid = get_thread_num();
489 scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
490 float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
491 float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
492
563 - bool is_brgemm_used = false;
564 -
565 - for (int64_t i = begin; i < end; ++i) {
566 - int64_t mb = i / NB;
567 - int64_t nb = i % NB;
568 -
569 - // nb0 from top half and nb1 from bottom half
570 - int64_t nb0 = nb, nb1 = nb + NB;
571 - int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
493 + loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
494 + // nb_upper from top half and nb_lower from bottom half
495 + int64_t nb_upper = nb, nb_lower = nb + NB;
496 + int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
497
498 // B shape [K, n_size] in vnni format
499 int32_t expert_id = expert_ids[mb];
575 - const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb0 * BLOCK_N * stride_n;
576 - const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb1 * BLOCK_N * stride_n;
500 + const scalar_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
501 + const scalar_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
502
578 - // 1.a load A
579 - const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
503 int64_t m_size = offsets[mb + 1] - offsets[mb];
504
582 - const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
583 - is_brgemm_used = is_brgemm_used || use_brgemm;
584 -
585 - for (int64_t m = 0; m < m_size; ++m) {
586 - int32_t index = A_ids[m] / topk;
587 - copy_stub(A + m * K, input + index * K, K);
505 + if (nb_offset == 0) {
506 + // 1.a load A
507 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
508 + for (int64_t m = 0; m < m_size; ++m) {
509 + int32_t index = A_ids[m] / topk;
510 + copy_stub(A + m * K, input + index * K, K);
511 + }
512 }
513
514 if (use_brgemm) {
@@ -616,12 +540,7 @@ void fused_experts_kernel_impl(
540
541 // 1.d silu and mul
542 const int64_t offset = offsets[mb];
619 - silu_and_mul<scalar_t, BLOCK_N>(
620 - ic1 + offset * N + nb * BLOCK_N,
621 - C0,
622 - C1,
623 - m_size,
624 - N);
543 + silu_and_mul<scalar_t, BLOCK_N>(ic1 + offset * N + nb * BLOCK_N, C0, C1, m_size, N);
544 } else {
545 // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
546 const int64_t offset = offsets[mb];
@@ -637,9 +556,9 @@ void fused_experts_kernel_impl(
556 /* ldb */ n_size,
557 /* ldc */ N);
558 }
640 - }
559 + });
560
642 - if (is_brgemm_used) {
561 + if (use_brgemm) {
562 at::native::cpublas::brgemm_release();
563 }
564 });
@@ -654,24 +573,16 @@ void fused_experts_kernel_impl(
573 const int64_t stride_oc = IC;
574
575 // parallel on [MB2, NB2]
657 - at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
576 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
577 // get local pointers
659 - int tid = at::get_thread_num();
578 + int tid = get_thread_num();
579 // we won't be using C1 for gemm2
580 float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
581
663 - bool is_brgemm_used = false;
664 -
665 - for (int64_t i = begin; i < end; ++i) {
666 - int64_t mb = i / NB2;
667 - int64_t nb = i % NB2;
668 -
582 + loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
583 int64_t m_size = offsets[mb + 1] - offsets[mb];
584 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
585
672 - const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
673 - is_brgemm_used = is_brgemm_used || use_brgemm;
674 -
586 // A ptr from ic1 of [M * topk, N] in sorted order
587 // so as to avoid copy A to tmp buffer again
588 const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
@@ -714,9 +625,9 @@ void fused_experts_kernel_impl(
625 float weight = topk_weights[index];
626 copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
627 }
717 - }
628 + });
629
719 - if (is_brgemm_used) {
630 + if (use_brgemm) {
631 at::native::cpublas::brgemm_release();
632 }
633 });
@@ -743,7 +654,6 @@ void shared_expert_kernel_impl(
654 int64_t M,
655 int64_t N,
656 int64_t K) {
746 -
657 // handle 2 tiles per block
658 constexpr int64_t BLOCK_M = block_size_m();
659 constexpr int64_t BLOCK_N = block_size_n();
@@ -755,36 +665,29 @@ void shared_expert_kernel_impl(
665 TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
666 const int64_t stride_n = K;
667
668 + const bool use_brgemm = can_use_brgemm<scalar_t>(M);
669 +
670 + const bool apply_scaling_factor = fused_experts_out != nullptr;
671 +
672 // here we only parallel on half of 2N to fuse silu_and_mul with gemm
759 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
673 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
674 // get local pointers
761 - int tid = at::get_thread_num();
675 + int tid = get_thread_num();
676 float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
677 float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
678
765 - bool is_brgemm_used = false;
766 -
767 - for (int64_t i = begin; i < end; ++i) {
768 - int64_t mb = i / NB;
769 - int64_t nb = i % NB;
770 -
771 - // nb0 from top half and nb1 from bottom half
772 - int64_t nb0 = nb, nb1 = nb + NB;
773 - int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
679 + loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
680 + // nb_upper from top half and nb_lower from bottom half
681 + int64_t nb_upper = nb, nb_lower = nb + NB;
682 + int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
683 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
684
776 - //int64_t mb_start = mb * BLOCK_M;
777 - //int64_t mb_size = std::min(M - mb_start, BLOCK_M);
778 -
685 // A shape [m_size, K]
686 const scalar_t* A = input + mb * BLOCK_M * K;
687
688 // B shape [K, n_size] in vnni format
783 - const scalar_t* __restrict__ B0 = packed_w1 + nb0 * BLOCK_N * stride_n;
784 - const scalar_t* __restrict__ B1 = packed_w1 + nb1 * BLOCK_N * stride_n;
785 -
786 - const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
787 - is_brgemm_used = is_brgemm_used || use_brgemm;
689 + const scalar_t* __restrict__ B0 = packed_w1 + nb_upper * BLOCK_N * stride_n;
690 + const scalar_t* __restrict__ B1 = packed_w1 + nb_lower * BLOCK_N * stride_n;
691
692 if (use_brgemm) {
693 // 1.b gemm: C0 = A @ B0
@@ -814,12 +717,7 @@ void shared_expert_kernel_impl(
717 /* C */ C1);
718
719 // 1.d silu and mul
817 - silu_and_mul<scalar_t, BLOCK_N>(
818 - ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
819 - C0,
820 - C1,
821 - m_size,
822 - N);
720 + silu_and_mul<scalar_t, BLOCK_N>(ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, m_size, N);
721 } else {
722 // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
723 tinygemm_kernel(
@@ -834,9 +732,9 @@ void shared_expert_kernel_impl(
732 /* ldb */ n_size,
733 /* ldc */ N);
734 }
837 - }
735 + });
736
839 - if (is_brgemm_used) {
737 + if (use_brgemm) {
738 at::native::cpublas::brgemm_release();
739 }
740 });
@@ -850,24 +748,16 @@ void shared_expert_kernel_impl(
748 const int64_t stride_oc = IC;
749
750 // parallel on [MB2, NB2]
853 - at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
751 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
752 // get local pointers
855 - int tid = at::get_thread_num();
753 + int tid = get_thread_num();
754 // we won't be using C1 for gemm2
755 float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
756
859 - bool is_brgemm_used = false;
860 -
861 - for (int64_t i = begin; i < end; ++i) {
862 - int64_t mb = i / NB2;
863 - int64_t nb = i % NB2;
864 -
757 + loop_2d<scalar_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
758 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
759 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
760
868 - const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
869 - is_brgemm_used = is_brgemm_used || use_brgemm;
870 -
761 // A shape [m_size, IC]
762 const scalar_t* __restrict__ A = ic1 + mb * BLOCK_M * N;
763
@@ -902,19 +792,21 @@ void shared_expert_kernel_impl(
792
793 // 2.b copy from C to output and add fused_experts_out
794 scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
905 - const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
795 + const scalar_t* __restrict__ fused_out =
796 + apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
797 for (int64_t m = 0; m < m_size; ++m) {
907 - add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
798 + const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
799 + add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
800 }
909 - }
801 + });
802
911 - if (is_brgemm_used) {
803 + if (use_brgemm) {
804 at::native::cpublas::brgemm_release();
805 }
806 });
807 }
808
917 -} // anonymous namespace
809 +} // anonymous namespace
810
811 // common checks
812 static inline void check_moe_scales(
@@ -922,14 +814,10 @@ static inline void check_moe_scales(
814 bool use_fp8_w8a16,
815 const std::optional<at::Tensor>& w1_scale,
816 const std::optional<at::Tensor>& w2_scale,
925 - const std::optional<std::vector<int64_t>> block_size,
926 - const std::optional<at::Tensor>& a1_scale,
927 - const std::optional<at::Tensor>& a2_scale) {
817 + const std::optional<std::vector<int64_t>> block_size) {
818 if (use_int8_w8a8) {
819 TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8.");
820 TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for int8 w8a8.");
931 - TORCH_CHECK(!a1_scale.has_value(), "static quantization for activation not supported.");
932 - TORCH_CHECK(!a2_scale.has_value(), "static quantization for activation not supported.");
821 }
822 if (use_fp8_w8a16) {
823 TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16.");
@@ -939,16 +827,16 @@ static inline void check_moe_scales(
827 }
828 }
829
942 -#define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \
943 - auto w1s = w1_scale.value(); \
944 - auto w2s = w2_scale.value(); \
945 - auto block_size_val = block_size.value(); \
946 - int64_t block_size_N = block_size_val[0]; \
947 - int64_t block_size_K = block_size_val[1]; \
948 - TORCH_CHECK(w1s.size(DIM0) == 2 * N / block_size_N); \
949 - TORCH_CHECK(w1s.size(DIM1) == K / block_size_K); \
950 - TORCH_CHECK(w2s.size(DIM0) == K / block_size_N); \
951 - TORCH_CHECK(w2s.size(DIM1) == N / block_size_K)
830 +#define CHECK_MOE_SCALES_FP8(DIM0, DIM1) \
831 + auto w1s = w1_scale.value(); \
832 + auto w2s = w2_scale.value(); \
833 + auto block_size_val = block_size.value(); \
834 + int64_t block_size_N = block_size_val[0]; \
835 + int64_t block_size_K = block_size_val[1]; \
836 + TORCH_CHECK(w1s.size(DIM0) == div_up(2 * N, block_size_N)); \
837 + TORCH_CHECK(w1s.size(DIM1) == div_up(K, block_size_K)); \
838 + TORCH_CHECK(w2s.size(DIM0) == div_up(K, block_size_N)); \
839 + TORCH_CHECK(w2s.size(DIM1) == div_up(N, block_size_K))
840
841 // hidden_states: [M, K]
842 // w1: [E, 2N, K]
@@ -956,6 +844,7 @@ static inline void check_moe_scales(
844 // topk_weights: [M, topk]
845 // topk_ids: [M, topk] (int32_t)
846 //
847 +
848 at::Tensor fused_experts_cpu(
849 at::Tensor& hidden_states,
850 at::Tensor& w1,
@@ -963,16 +852,13 @@ at::Tensor fused_experts_cpu(
852 at::Tensor& topk_weights,
853 at::Tensor& topk_ids,
854 bool inplace,
966 - bool use_int8_w8a8,
967 - bool use_fp8_w8a16,
855 + int64_t moe_comp_method,
856 const std::optional<at::Tensor>& w1_scale,
857 const std::optional<at::Tensor>& w2_scale,
858 + const std::optional<at::Tensor>& w1_zero,
859 + const std::optional<at::Tensor>& w2_zero,
860 const std::optional<std::vector<int64_t>> block_size,
971 - const std::optional<at::Tensor>& a1_scale,
972 - const std::optional<at::Tensor>& a2_scale,
861 bool is_vnni) {
974 - RECORD_FUNCTION("sgl-kernel::fused_experts_cpu", std::vector<c10::IValue>({hidden_states, w1, w2, topk_weights, topk_ids}));
975 -
862 auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
863 auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
864
@@ -985,32 +871,49 @@ at::Tensor fused_experts_cpu(
871 CHECK_INPUT(w2);
872 CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
873 CHECK_DIM(2, hidden_states);
988 - CHECK_DIM(3, w1);
989 - CHECK_DIM(3, w2);
874 + if (moe_comp_method == CPUQuantMethod::INT4_W4A8 && is_vnni) {
875 + CHECK_DIM(4, w1);
876 + CHECK_DIM(4, w2);
877 + } else {
878 + CHECK_DIM(3, w1);
879 + CHECK_DIM(3, w2);
880 + }
881 CHECK_DIM(2, topk_weights);
882 CHECK_DIM(2, topk_ids);
883
884 CHECK_EQ(topk_ids.scalar_type(), at::kInt);
994 - CHECK_EQ(topk_weights.scalar_type(), at::kFloat);
885 +
886 + // TODO: support topk_weights to be bf16 or fp16 in the kernel.
887 + // The topk_weights of llama4 is computed via Llama4MoE:custom_routing_function and is bf16/fp16
888 + // while the kernel currently only supports it to be float32
889 + auto topk_weights_ = topk_weights.to(at::kFloat);
890 + CHECK_EQ(topk_weights_.scalar_type(), at::kFloat);
891
892 int64_t M = hidden_states.size(0);
893 int64_t K = hidden_states.size(1);
998 - int64_t N = w1.size(1) / 2;
894 + int64_t N = moe_comp_method == CPUQuantMethod::INT4_W4A8 ? w1_scale.value().size(1) * w1_scale.value().size(3) / 2
895 + : w1.size(1) / 2;
896 int64_t E = w1.size(0);
1000 - int64_t topk = topk_weights.size(1);
897 + int64_t topk = topk_weights_.size(1);
898
899 // we use int32_t compensation for int8 w8a8
1003 - int64_t packed_K = get_row_size(K, use_int8_w8a8);
1004 - int64_t packed_N = get_row_size(N, use_int8_w8a8);
900 + int64_t packed_K = get_row_size(K, moe_comp_method == CPUQuantMethod::INT8_W8A8);
901 + int64_t packed_N = get_row_size(N, moe_comp_method == CPUQuantMethod::INT8_W8A8);
902
903 // check weight shapes
904 CHECK_EQ(w2.size(0), E);
1008 - CHECK_EQ(w2.size(1), K);
1009 - CHECK_EQ(packed_w1.size(2), packed_K);
1010 - CHECK_EQ(packed_w2.size(2), packed_N);
1011 -
905 + if (!(moe_comp_method == CPUQuantMethod::INT4_W4A8)) {
906 + CHECK_EQ(w2.size(1), K);
907 + CHECK_EQ(packed_w1.size(2), packed_K / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
908 + CHECK_EQ(packed_w2.size(2), packed_N / (moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 2 : 1));
909 + }
910 // check scales
1013 - check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);
911 + check_moe_scales(
912 + moe_comp_method == CPUQuantMethod::INT8_W8A8,
913 + moe_comp_method == CPUQuantMethod::FP8_W8A16,
914 + w1_scale,
915 + w2_scale,
916 + block_size);
917
918 at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
919
@@ -1031,8 +934,8 @@ at::Tensor fused_experts_cpu(
934 int32_t* __restrict__ sorted_ids = buffer.data_ptr<int32_t>();
935 int32_t* __restrict__ expert_ids = sorted_ids + max_num_tokens_padded;
936 int32_t* __restrict__ total_cnts = expert_ids + max_num_blocks;
1034 - int32_t* __restrict__ cumsums = total_cnts + (num_threads + 1) * E;
1035 - int32_t* __restrict__ offsets = cumsums + (E + 1);
937 + int32_t* __restrict__ cumsums = total_cnts + (num_threads + 1) * E;
938 + int32_t* __restrict__ offsets = cumsums + (E + 1);
939
940 // init sorted_ids with `numel` as the padding number
941 // init expert_ids with `num_experts`
@@ -1064,26 +967,31 @@ at::Tensor fused_experts_cpu(
967 //
968 // for fp8 w8a16:
969 // 7. intermediate_cache0 : [M * topk, 2N]
1067 - // 8. B_tmp : [T, BLOCK_N, std::max(K, N)]
970 + // 8. B_tmp : [T, MAX_CACHE_BLOCK_SIZE, BLOCK_N, std::max(K, N)]
971 //
1069 - int64_t buffer_size_nbytes = M * topk * N * 2 + M * topk * K * 2 +
1070 - num_threads * BLOCK_M * K * (use_int8_w8a8 ? 1 : 2) +
972 + int64_t buffer_size_nbytes =
973 + M * topk * N * 2 + M * topk * K * 2 +
974 + num_threads * BLOCK_M * K *
975 + (moe_comp_method == CPUQuantMethod::INT8_W8A8 | moe_comp_method == CPUQuantMethod::INT4_W4A8 ? 1 : 2) +
976 num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
977
1073 - if (use_int8_w8a8) {
978 + if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
979 buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
980 }
1076 - if (use_fp8_w8a16) {
1077 - buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * BLOCK_N * std::max(K, N) * 2;
981 + if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
982 + buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N) * 2;
983 + }
984 + if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
985 + buffer_size_nbytes += M * topk * 2 * N * 2 + std::max(M * K, M * topk * N) + M * topk * sizeof(float) +
986 + num_threads * 2 * get_4bit_block_k_size(K / w1_scale.value().size(2)) * BLOCK_N;
987 }
1079 -
988 auto buffer2 = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));
989
990 AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_experts_kernel_impl", [&] {
991 scalar_t* __restrict__ intermediate_cache1 = (scalar_t*)((void*)(buffer2.data_ptr<int8_t>()));
992 scalar_t* __restrict__ intermediate_cache2 = intermediate_cache1 + M * topk * N;
993
1086 - if (use_int8_w8a8) {
994 + if (moe_comp_method == CPUQuantMethod::INT8_W8A8) {
995 uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
996 float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
997 uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
@@ -1107,7 +1015,7 @@ at::Tensor fused_experts_cpu(
1015 packed_w2.data_ptr<int8_t>(),
1016 w1s.data_ptr<float>(),
1017 w2s.data_ptr<float>(),
1110 - topk_weights.data_ptr<float>(),
1018 + topk_weights_.data_ptr<float>(),
1019 sorted_ids,
1020 expert_ids,
1021 offsets,
@@ -1117,7 +1025,7 @@ at::Tensor fused_experts_cpu(
1025 E,
1026 topk,
1027 num_tokens_post_pad);
1120 - } else if (use_fp8_w8a16) {
1028 + } else if (moe_comp_method == CPUQuantMethod::FP8_W8A16) {
1029 // here we just ignore C_tmp as it is not used
1030 scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
1031 float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
@@ -1140,6 +1048,48 @@ at::Tensor fused_experts_cpu(
1048 w2s.data_ptr<float>(),
1049 block_size_N,
1050 block_size_K,
1051 + topk_weights_.data_ptr<float>(),
1052 + sorted_ids,
1053 + expert_ids,
1054 + offsets,
1055 + M,
1056 + N,
1057 + K,
1058 + E,
1059 + topk,
1060 + num_tokens_post_pad);
1061 + } else if (moe_comp_method == CPUQuantMethod::INT4_W4A8) {
1062 + uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
1063 + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1064 + scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1065 + uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1066 + float* __restrict__ As_tmp = (float*)((void*)(Aq_tmp + std::max(M * K, M * topk * N)));
1067 + int8_t* __restrict__ dqB_tmp = (int8_t*)((void*)(As_tmp + M * topk));
1068 +
1069 + // weight + compensation shape = [Nc, Kc, block_n * block_k / 2 + block_n*sizeof(int32_t)]
1070 + // scales/qzeros shape = [E, Nc, G, block_n]
1071 + int64_t num_groups = w1_scale.value().size(2);
1072 + const int group_size = K / num_groups;
1073 + // TODO: check scales and zeros
1074 + fused_experts_int4_w4a8_kernel_impl<scalar_t>(
1075 + out_hidden_states.data_ptr<scalar_t>(),
1076 + intermediate_cache0,
1077 + intermediate_cache1,
1078 + intermediate_cache2,
1079 + A_tmp,
1080 + Aq_tmp,
1081 + As_tmp,
1082 + nullptr,
1083 + C_tmp,
1084 + dqB_tmp,
1085 + hidden_states.data_ptr<scalar_t>(),
1086 + packed_w1.data_ptr<uint8_t>(),
1087 + packed_w2.data_ptr<uint8_t>(),
1088 + w1_zero.value().data_ptr<int8_t>(),
1089 + w2_zero.value().data_ptr<int8_t>(),
1090 + w1_scale.value().data_ptr<float>(),
1091 + w2_scale.value().data_ptr<float>(),
1092 + group_size,
1093 topk_weights.data_ptr<float>(),
1094 sorted_ids,
1095 expert_ids,
@@ -1163,7 +1113,7 @@ at::Tensor fused_experts_cpu(
1113 hidden_states.data_ptr<scalar_t>(),
1114 packed_w1.data_ptr<scalar_t>(),
1115 packed_w2.data_ptr<scalar_t>(),
1166 - topk_weights.data_ptr<float>(),
1116 + topk_weights_.data_ptr<float>(),
1117 sorted_ids,
1118 expert_ids,
1119 offsets,
@@ -1188,34 +1138,37 @@ at::Tensor shared_expert_cpu(
1138 at::Tensor& hidden_states,
1139 at::Tensor& w1,
1140 at::Tensor& w2,
1191 - at::Tensor& fused_experts_out,
1192 - double routed_scaling_factor,
1141 + const std::optional<at::Tensor>& fused_experts_out,
1142 + const std::optional<double> routed_scaling_factor,
1143 bool inplace,
1144 bool use_int8_w8a8,
1145 bool use_fp8_w8a16,
1196 - std::optional<at::Tensor>& w1_scale,
1197 - std::optional<at::Tensor>& w2_scale,
1198 - std::optional<std::vector<int64_t>> block_size,
1199 - std::optional<at::Tensor>& a1_scale,
1200 - std::optional<at::Tensor>& a2_scale,
1146 + const std::optional<at::Tensor>& w1_scale,
1147 + const std::optional<at::Tensor>& w2_scale,
1148 + const std::optional<std::vector<int64_t>> block_size,
1149 bool is_vnni) {
1202 - RECORD_FUNCTION("sgl-kernel::shared_expert_cpu", std::vector<c10::IValue>({hidden_states, w1, w2}));
1203 -
1150 auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
1151 auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
1152
1153 constexpr int64_t BLOCK_M = block_size_m();
1154 constexpr int64_t BLOCK_N = block_size_n();
1155
1156 + double routed_scaling_factor_value = 0;
1157 + if (routed_scaling_factor.has_value()) {
1158 + TORCH_CHECK(fused_experts_out.has_value(), "shared_expert_cpu: expect fused_experts_out.");
1159 + const auto fused_experts_out_tensor = fused_experts_out.value();
1160 + routed_scaling_factor_value = routed_scaling_factor.value();
1161 + CHECK_INPUT(fused_experts_out_tensor);
1162 + CHECK_EQ(hidden_states.sizes(), fused_experts_out_tensor.sizes());
1163 + }
1164 +
1165 const auto st = hidden_states.scalar_type();
1166 CHECK_INPUT(hidden_states);
1212 - CHECK_INPUT(fused_experts_out);
1167 CHECK_INPUT(w1);
1168 CHECK_INPUT(w2);
1169 CHECK_DIM(2, hidden_states);
1170 CHECK_DIM(2, w1);
1171 CHECK_DIM(2, w2);
1218 - CHECK_EQ(hidden_states.sizes(), fused_experts_out.sizes());
1172 CHECK_EQ(hidden_states.scalar_type(), st);
1173
1174 int64_t M = hidden_states.size(0);
@@ -1232,7 +1185,7 @@ at::Tensor shared_expert_cpu(
1185 CHECK_EQ(packed_w2.size(1), packed_N);
1186
1187 // check scales
1235 - check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);
1188 + check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size);
1189
1190 at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
1191
@@ -1246,7 +1199,7 @@ at::Tensor shared_expert_cpu(
1199 //
1200 // for fp8 w8a16:
1201 // 5. intermediate_cache0 : [M, 2N]
1249 - // 6. B_tmp: [T, BLOCK_M, max(K, N)]
1202 + // 6. B_tmp: [T, MAX_CACHE_BLOCK_SIZE, BLOCK_M, max(K, N)]
1203 //
1204 int num_threads = at::get_num_threads();
1205 int64_t buffer_size_nbytes = M * N * 2 + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
@@ -1255,7 +1208,7 @@ at::Tensor shared_expert_cpu(
1208 buffer_size_nbytes += std::max(M * K, M * N) + M * sizeof(float);
1209 }
1210 if (use_fp8_w8a16) {
1258 - buffer_size_nbytes += M * 2 * N * 2 + num_threads * BLOCK_M * std::max(K, N) * 2;
1211 + buffer_size_nbytes += M * 2 * N * 2 + num_threads * MAX_CACHE_BLOCK_SIZE * BLOCK_M * std::max(K, N) * 2;
1212 }
1213
1214 auto buffer = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));
@@ -1283,8 +1236,8 @@ at::Tensor shared_expert_cpu(
1236 packed_w2.data_ptr<int8_t>(),
1237 w1s.data_ptr<float>(),
1238 w2s.data_ptr<float>(),
1286 - fused_experts_out.data_ptr<scalar_t>(),
1287 - routed_scaling_factor,
1239 + conditional_data_ptr<scalar_t>(fused_experts_out),
1240 + routed_scaling_factor_value,
1241 M,
1242 N,
1243 K);
@@ -1306,8 +1259,8 @@ at::Tensor shared_expert_cpu(
1259 w2s.data_ptr<float>(),
1260 block_size_N,
1261 block_size_K,
1309 - fused_experts_out.data_ptr<scalar_t>(),
1310 - routed_scaling_factor,
1262 + conditional_data_ptr<scalar_t>(fused_experts_out),
1263 + routed_scaling_factor_value,
1264 M,
1265 N,
1266 K);
@@ -1319,8 +1272,8 @@ at::Tensor shared_expert_cpu(
1272 hidden_states.data_ptr<scalar_t>(),
1273 packed_w1.data_ptr<scalar_t>(),
1274 packed_w2.data_ptr<scalar_t>(),
1322 - fused_experts_out.data_ptr<scalar_t>(),
1323 - routed_scaling_factor,
1275 + conditional_data_ptr<scalar_t>(fused_experts_out),
1276 + routed_scaling_factor_value,
1277 M,
1278 N,
1279 K);
csrc/cpu/sgl-kernels/moe.h new
+178
@@ -0,0 +1,178 @@
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 "vec.h"
8 +
9 +template <typename scalar_t>
10 +inline void fill_stub(scalar_t* __restrict__ out, scalar_t val, int64_t size) {
11 + using Vec = at::vec::Vectorized<scalar_t>;
12 + const Vec data_vec(val);
13 + at::vec::map<scalar_t>([data_vec](Vec out) { return out = data_vec; }, out, out, size);
14 +}
15 +
16 +template <typename scalar_t>
17 +inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
18 + using Vec = at::vec::Vectorized<scalar_t>;
19 + constexpr int kVecSize = Vec::size();
20 + int64_t d;
21 +#pragma GCC unroll 4
22 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
23 + Vec data = Vec::loadu(input + d);
24 + data.store(out + d);
25 + }
26 + for (; d < size; ++d) {
27 + out[d] = input[d];
28 + }
29 +}
30 +
31 +template <typename scalar_t>
32 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
33 + using bVec = at::vec::Vectorized<scalar_t>;
34 + using fVec = at::vec::Vectorized<float>;
35 + constexpr int kVecSize = bVec::size();
36 + int64_t d;
37 +#pragma GCC unroll 4
38 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
39 + auto [x0, x1] = load_float_vec2(input + d);
40 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
41 + out_vec.store(out + d);
42 + }
43 + for (; d < size; ++d) {
44 + out[d] = static_cast<scalar_t>(input[d]);
45 + }
46 +}
47 +
48 +template <>
49 +inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) {
50 + // size might be 64x + 32
51 + std::memcpy(out, input, size * sizeof(uint8_t));
52 +}
53 +
54 +template <typename scalar_t, typename input_t>
55 +inline void copy_mul_stub(scalar_t* __restrict__ out, const input_t* __restrict__ input, float weight, int64_t size) {
56 + static_assert(
57 + std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
58 + "copy_mul_stub only supports input_t == float or input_t == scalar_t");
59 + using bVec = at::vec::Vectorized<scalar_t>;
60 + using fVec = at::vec::Vectorized<float>;
61 + constexpr int kVecSize = bVec::size();
62 + const fVec weight_vec = fVec(weight);
63 + int64_t d;
64 +#pragma GCC unroll 4
65 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
66 + auto [x0, x1] = load_float_vec2(input + d);
67 + x0 = x0 * weight_vec;
68 + x1 = x1 * weight_vec;
69 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
70 + out_vec.store(out + d);
71 + }
72 + for (; d < size; ++d) {
73 + out[d] = static_cast<scalar_t>(input[d] * weight);
74 + }
75 +}
76 +
77 +// acc from [topk, K] to [K]
78 +template <typename scalar_t>
79 +inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
80 + using bVec = at::vec::Vectorized<scalar_t>;
81 + using fVec = at::vec::Vectorized<float>;
82 + constexpr int kVecSize = bVec::size();
83 + if (topk == 1) {
84 + // do copy for topk = 1
85 + copy_stub(out, input, K);
86 + } else {
87 + // do sum for topk != 1
88 + int64_t d;
89 +#pragma GCC unroll 4
90 + for (d = 0; d <= K - kVecSize; d += kVecSize) {
91 + fVec sum_fvec0 = fVec(0.f);
92 + fVec sum_fvec1 = fVec(0.f);
93 + for (int t = 0; t < topk; ++t) {
94 + bVec x_bvec = bVec::loadu(input + t * K + d);
95 + fVec x_fvec0, x_fvec1;
96 + std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
97 +
98 + sum_fvec0 += x_fvec0;
99 + sum_fvec1 += x_fvec1;
100 + }
101 + bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
102 + out_bvec.store(out + d);
103 + }
104 + for (; d < K; ++d) {
105 + float sum_val = 0.f;
106 + for (int t = 0; t < topk; ++t) {
107 + sum_val += static_cast<float>(input[t * K + d]);
108 + }
109 + out[d] = static_cast<scalar_t>(sum_val);
110 + }
111 + }
112 +}
113 +
114 +// out = input + input2 * scale
115 +template <typename scalar_t, typename input_t>
116 +inline void add_mul_stub(
117 + scalar_t* __restrict__ out,
118 + const input_t* __restrict__ input,
119 + const scalar_t* __restrict__ input2,
120 + float scale,
121 + int64_t size) {
122 + static_assert(
123 + std::is_same_v<input_t, float> || std::is_same_v<input_t, scalar_t>,
124 + "add_mul_stub only supports input_t == float or input_t == scalar_t");
125 +
126 + // out = input (without scale factor)
127 + if (input2 == nullptr) {
128 + copy_stub(out, input, size);
129 + return;
130 + }
131 +
132 + using bVec = at::vec::Vectorized<scalar_t>;
133 + using fVec = at::vec::Vectorized<float>;
134 + constexpr int kVecSize = bVec::size();
135 + const fVec s_vec = fVec(scale);
136 + int64_t d;
137 +#pragma GCC unroll 4
138 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
139 + auto [x0, x1] = load_float_vec2(input + d);
140 +
141 + bVec y_bvec = bVec::loadu(input2 + d);
142 + fVec y0, y1;
143 + std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
144 +
145 + x0 = x0 + y0 * s_vec;
146 + x1 = x1 + y1 * s_vec;
147 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
148 + out_vec.store(out + d);
149 + }
150 + for (; d < size; ++d) {
151 + out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
152 + }
153 +}
154 +
155 +template <typename scalar_t>
156 +inline void silu_and_mul_stub(
157 + scalar_t* __restrict__ out, const scalar_t* __restrict__ input, const scalar_t* __restrict__ input2, int64_t size) {
158 + using bVec = at::vec::Vectorized<scalar_t>;
159 + using fVec = at::vec::Vectorized<float>;
160 + const fVec one = fVec(1.f);
161 +
162 + // no remainder
163 +#pragma GCC unroll 4
164 + for (int64_t d = 0; d < size; d += bVec::size()) {
165 + bVec x = bVec::loadu(input + d);
166 + fVec x0, x1;
167 + std::tie(x0, x1) = at::vec::convert_to_float(x);
168 + bVec y = bVec::loadu(input2 + d);
169 + fVec y0, y1;
170 + std::tie(y0, y1) = at::vec::convert_to_float(y);
171 + x0 = x0 / (one + x0.neg().exp_u20());
172 + x1 = x1 / (one + x1.neg().exp_u20());
173 + x0 = x0 * y0;
174 + x1 = x1 * y1;
175 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
176 + out_vec.store(out + d);
177 + }
178 +}
csrc/cpu/sgl-kernels/moe_fp8.cpp
+72 -206
@@ -1,147 +1,11 @@
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 -#include "common.h"
5 -#include "gemm.h"
6 -#include "vec.h"
7 -
4 // clang-format off
5
10 -namespace {
11 -
12 -template <typename scalar_t>
13 -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
14 - using Vec = at::vec::Vectorized<scalar_t>;
15 - // no remainder
16 - #pragma GCC unroll 4
17 - for (int64_t d = 0; d < size; d += Vec::size()) {
18 - Vec data = Vec::loadu(input + d);
19 - data.store(out + d);
20 - }
21 -}
22 -
23 -template <typename scalar_t>
24 -inline void copy_mul_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, float weight, int64_t size) {
25 - using bVec = at::vec::Vectorized<scalar_t>;
26 - using fVec = at::vec::Vectorized<float>;
27 - constexpr int kVecSize = bVec::size();
28 - const fVec weight_vec = fVec(weight);
29 - int64_t d;
30 - #pragma GCC unroll 4
31 - for (d = 0; d <= size - kVecSize; d += kVecSize) {
32 - bVec x = bVec::loadu(input + d);
33 - fVec x0, x1;
34 - std::tie(x0, x1) = at::vec::convert_to_float(x);
35 - x0 = x0 * weight_vec;
36 - x1 = x1 * weight_vec;
37 - bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
38 - out_vec.store(out + d);
39 - }
40 - for (; d < size; ++d) {
41 - out[d] = static_cast<scalar_t>(input[d] * weight);
42 - }
43 -}
44 -
45 -// acc from [topk, K] to [K]
46 -template <typename scalar_t>
47 -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
48 - using bVec = at::vec::Vectorized<scalar_t>;
49 - using fVec = at::vec::Vectorized<float>;
50 - constexpr int kVecSize = bVec::size();
51 - if (topk == 1) {
52 - // do copy for topk = 1
53 - copy_stub(out, input, K);
54 - } else {
55 - // do sum for topk != 1
56 - int64_t d;
57 - #pragma GCC unroll 4
58 - for (d = 0; d <= K - kVecSize; d += kVecSize) {
59 - fVec sum_fvec0 = fVec(0.f);
60 - fVec sum_fvec1 = fVec(0.f);
61 - for (int t = 0; t < topk; ++t) {
62 - bVec x_bvec = bVec::loadu(input + t * K + d);
63 - fVec x_fvec0, x_fvec1;
64 - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
65 -
66 - sum_fvec0 += x_fvec0;
67 - sum_fvec1 += x_fvec1;
68 - }
69 - bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
70 - out_bvec.store(out + d);
71 - }
72 - for (; d < K; ++d) {
73 - float sum_val = 0.f;
74 - for (int t = 0; t < topk; ++t) {
75 - sum_val += static_cast<float>(input[t * K + d]);
76 - }
77 - out[d] = static_cast<scalar_t>(sum_val);
78 - }
79 - }
80 -}
81 -
82 -// out = input + input2 * scale
83 -template <typename scalar_t>
84 -inline void add_mul_stub(
85 - scalar_t* __restrict__ out,
86 - const scalar_t* __restrict__ input,
87 - const scalar_t* __restrict__ input2,
88 - float scale,
89 - int64_t size) {
90 - using bVec = at::vec::Vectorized<scalar_t>;
91 - using fVec = at::vec::Vectorized<float>;
92 - constexpr int kVecSize = bVec::size();
93 - const fVec s_vec = fVec(scale);
94 -
95 - int64_t d;
96 -#pragma GCC unroll 4
97 - for (d = 0; d <= size - kVecSize; d += kVecSize) {
98 - bVec x_bvec = bVec::loadu(input + d);
99 - fVec x0, x1;
100 - std::tie(x0, x1) = at::vec::convert_to_float(x_bvec);
101 -
102 - bVec y_bvec = bVec::loadu(input2 + d);
103 - fVec y0, y1;
104 - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
105 -
106 - x0 = x0 + y0 * s_vec;
107 - x1 = x1 + y1 * s_vec;
108 - bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
109 - out_vec.store(out + d);
110 - }
111 - for (; d < size; ++d) {
112 - out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
113 - }
114 -}
115 -
116 -template <typename scalar_t>
117 -inline void silu_and_mul_stub(
118 - scalar_t* __restrict__ out,
119 - const scalar_t* __restrict__ input,
120 - const scalar_t* __restrict__ input2,
121 - int64_t size) {
122 - using bVec = at::vec::Vectorized<scalar_t>;
123 - using fVec = at::vec::Vectorized<float>;
124 - const fVec one = fVec(1.f);
125 -
126 - // no remainder
127 -#pragma GCC unroll 4
128 - for (int64_t d = 0; d < size; d += bVec::size()) {
129 - bVec x = bVec::loadu(input + d);
130 - fVec x0, x1;
131 - std::tie(x0, x1) = at::vec::convert_to_float(x);
132 - bVec y = bVec::loadu(input2 + d);
133 - fVec y0, y1;
134 - std::tie(y0, y1) = at::vec::convert_to_float(y);
135 - x0 = x0 / (one + x0.neg().exp_u20());
136 - x1 = x1 / (one + x1.neg().exp_u20());
137 - x0 = x0 * y0;
138 - x1 = x1 * y1;
139 - bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
140 - out_vec.store(out + d);
141 - }
142 -}
143 -
144 -} // anonymous namespace
6 +#include "common.h"
7 +#include "gemm.h"
8 +#include "moe.h"
9
10 template <typename scalar_t>
11 void fused_experts_fp8_kernel_impl(
@@ -169,7 +33,6 @@ void fused_experts_fp8_kernel_impl(
33 int64_t E,
34 int64_t topk,
35 int64_t num_tokens_post_pad) {
172 -
36 constexpr int64_t BLOCK_M = block_size_m();
37 constexpr int64_t BLOCK_N = block_size_n();
38
@@ -183,35 +46,39 @@ void fused_experts_fp8_kernel_impl(
46 const int64_t stride_e = 2 * N * K;
47 const int64_t stride_n = K;
48
49 + int64_t avg_M = std::max(int64_t(1), M * topk / E);
50 + const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(avg_M);
51 +
52 + int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
53 +
54 // here we only parallel on half of 2N to fuse silu_and_mul with gemm
187 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
55 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
56 // get local pointers
189 - int tid = at::get_thread_num();
57 + int tid = get_thread_num();
58 scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
59
192 - bool is_brgemm_used = false;
193 -
194 - for (int64_t i = begin; i < end; ++i) {
195 - int64_t mb = i / NB;
196 - int64_t nb = i % NB;
197 -
60 + loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
61 int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
62
63 // B shape [K, n_size] in vnni format
64 int32_t expert_id = expert_ids[mb];
65 const at::Float8_e4m3fn* __restrict__ B = packed_w1 + expert_id * stride_e + nb * BLOCK_N * stride_n;
203 - const float* __restrict__ Bs = w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
66 + const float* __restrict__ Bs =
67 + w1s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
68
205 - // 1.a load A
206 - const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
207 - int64_t m_size = offsets[mb + 1] - offsets[mb];
69 + // do unpacking for the first row or a new expert
70 + int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
71 + bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
72
209 - const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
210 - is_brgemm_used = is_brgemm_used || use_brgemm;
73 + int64_t m_size = offsets[mb + 1] - offsets[mb];
74
212 - for (int64_t m = 0; m < m_size; ++m) {
213 - int32_t index = A_ids[m] / topk;
214 - copy_stub(A + m * K, input + index * K, K);
75 + if (nb_offset == 0) {
76 + // 1.a load A
77 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
78 + for (int64_t m = 0; m < m_size; ++m) {
79 + int32_t index = A_ids[m] / topk;
80 + copy_stub(A + m * K, input + index * K, K);
81 + }
82 }
83
84 const int64_t offset = offsets[mb];
@@ -219,7 +86,7 @@ void fused_experts_fp8_kernel_impl(
86 /* A */ A,
87 /* B */ B,
88 /* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
222 - /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
89 + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
90 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
91 /* scale */ Bs,
92 /* M */ m_size,
@@ -229,10 +96,11 @@ void fused_experts_fp8_kernel_impl(
96 /* ldb */ n_size,
97 /* ldc */ 2 * N,
98 /* brg */ use_brgemm,
232 - /* block_size_K */ block_size_K);
233 - }
99 + /* block_size_K */ block_size_K,
100 + /* do_unpack */ do_unpack);
101 + });
102
235 - if (is_brgemm_used) {
103 + if (use_brgemm) {
104 at::native::cpublas::brgemm_release();
105 }
106 });
@@ -240,11 +108,7 @@ void fused_experts_fp8_kernel_impl(
108 // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
109 at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
110 for (int64_t m = begin; m < end; ++m) {
243 - silu_and_mul_stub(
244 - ic1 + m * N,
245 - ic0 + m * 2 * N,
246 - ic0 + m * 2 * N + N,
247 - N);
111 + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
112 }
113 });
114
@@ -260,22 +124,14 @@ void fused_experts_fp8_kernel_impl(
124 const int64_t stride_oc = IC;
125
126 // parallel on [MB2, NB2]
263 - at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
264 - int tid = at::get_thread_num();
127 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
128 + int tid = get_thread_num();
129 alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
130
267 - bool is_brgemm_used = false;
268 -
269 - for (int64_t i = begin; i < end; ++i) {
270 - int64_t mb = i / NB2;
271 - int64_t nb = i % NB2;
272 -
131 + loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
132 int64_t m_size = offsets[mb + 1] - offsets[mb];
133 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
134
276 - const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
277 - is_brgemm_used = is_brgemm_used || use_brgemm;
278 -
135 // A ptr from ic1 of [M * topk, N] in sorted order
136 // so as to avoid copy A to tmp buffer again
137 const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
@@ -284,13 +140,18 @@ void fused_experts_fp8_kernel_impl(
140 // B shape [IC, n_size] in vnni format
141 int32_t expert_id = expert_ids[mb];
142 const at::Float8_e4m3fn* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
287 - const float* __restrict__ Bs = w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
143 + const float* __restrict__ Bs =
144 + w2s + expert_id * scale_size_N * scale_size_K + (nb / blocks_n_per_group) * scale_size_K;
145 +
146 + // do unpacking for the first row or a new expert
147 + int32_t pre_expert_id = mb == 0 ? -1 : expert_ids[mb - 1];
148 + bool do_unpack = (mb == mb0) || (expert_id != pre_expert_id);
149
150 tinygemm_kernel<scalar_t>(
151 /* A */ A,
152 /* B */ B,
153 /* C */ C,
293 - /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
154 + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
155 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
156 /* scale */ Bs,
157 /* M */ m_size,
@@ -300,7 +161,8 @@ void fused_experts_fp8_kernel_impl(
161 /* ldb */ n_size,
162 /* ldc */ BLOCK_N,
163 /* brg */ use_brgemm,
303 - /* block_size_K */ block_size_K);
164 + /* block_size_K */ block_size_K,
165 + /* do_unpack */ do_unpack);
166
167 // 2.b copy from C to ic2 in original order
168 // and also mul topk_weights in float32
@@ -309,9 +171,9 @@ void fused_experts_fp8_kernel_impl(
171 float weight = topk_weights[index];
172 copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
173 }
312 - }
174 + });
175
314 - if (is_brgemm_used) {
176 + if (use_brgemm) {
177 at::native::cpublas::brgemm_release();
178 }
179 });
@@ -374,7 +236,6 @@ void shared_expert_fp8_kernel_impl(
236 int64_t M,
237 int64_t N,
238 int64_t K) {
377 -
239 constexpr int64_t BLOCK_M = block_size_m();
240 constexpr int64_t BLOCK_N = block_size_n();
241
@@ -385,21 +246,25 @@ void shared_expert_fp8_kernel_impl(
246 int64_t blocks_n_per_group = block_size_N / BLOCK_N;
247
248 const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
249 + const bool apply_scaling_factor = fused_experts_out != nullptr;
250
389 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
390 - int tid = at::get_thread_num();
251 + int64_t B_tmp_size_per_thread = MAX_CACHE_BLOCK_SIZE * BLOCK_N * std::max(K, N);
252
392 - for (int64_t i = begin; i < end; ++i) {
393 - int64_t mb = i / NB;
394 - int64_t nb = i % NB;
253 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
254 + int tid = get_thread_num();
255 +
256 + loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * K, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
257 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
258 int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
259
260 + // do unpacking for the first row
261 + bool do_unpack = (mb == mb0);
262 +
263 tinygemm_kernel<scalar_t>(
264 /* A */ input + mb * BLOCK_M * K,
265 /* B */ packed_w1 + nb * BLOCK_N * K,
266 /* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
402 - /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
267 + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * K,
268 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
269 /* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
270 /* M */ m_size,
@@ -409,8 +274,9 @@ void shared_expert_fp8_kernel_impl(
274 /* ldb */ n_size,
275 /* ldc */ 2 * N,
276 /* brg */ use_brgemm,
412 - /* block_size_K */ block_size_K);
413 - }
277 + /* block_size_K */ block_size_K,
278 + /* do_unpack */ do_unpack);
279 + });
280
281 if (use_brgemm) {
282 at::native::cpublas::brgemm_release();
@@ -420,11 +286,7 @@ void shared_expert_fp8_kernel_impl(
286 // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
287 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
288 for (int64_t m = begin; m < end; ++m) {
423 - silu_and_mul_stub(
424 - ic1 + m * N,
425 - ic0 + m * 2 * N,
426 - ic0 + m * 2 * N + N,
427 - N);
289 + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
290 }
291 });
292
@@ -437,22 +299,23 @@ void shared_expert_fp8_kernel_impl(
299 scale_size_K = div_up(N, block_size_K);
300
301 // parallel on [MB2, NB2]
440 - at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
441 - int tid = at::get_thread_num();
302 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
303 + int tid = get_thread_num();
304 alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
305
444 - for (int64_t i = begin; i < end; ++i) {
445 - int64_t mb = i / NB2;
446 - int64_t nb = i % NB2;
306 + loop_2d<at::Float8_e4m3fn>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
307 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
308 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
309
310 + // do unpacking for the first row
311 + bool do_unpack = (mb == mb0);
312 +
313 // 2.a gemm: C = A @ B
314 tinygemm_kernel<scalar_t>(
315 /* A */ ic1 + mb * BLOCK_M * N,
316 /* B */ packed_w2 + nb * BLOCK_N * N,
317 /* C */ C,
455 - /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
318 + /* Btmp */ B_tmp + tid * B_tmp_size_per_thread + nb_offset * BLOCK_N * IC,
319 /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
320 /* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
321 /* M */ m_size,
@@ -462,15 +325,18 @@ void shared_expert_fp8_kernel_impl(
325 /* ldb */ n_size,
326 /* ldc */ BLOCK_N,
327 /* brg */ use_brgemm,
465 - /* block_size_K */ block_size_K);
328 + /* block_size_K */ block_size_K,
329 + /* do_unpack */ do_unpack);
330
331 // 2.b copy from C to output and add fused_experts_out
332 scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
469 - const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
333 + const scalar_t* __restrict__ fused_out =
334 + apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
335 for (int64_t m = 0; m < m_size; ++m) {
471 - add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
336 + const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
337 + add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
338 }
473 - }
339 + });
340 });
341
342 if (use_brgemm) {
csrc/cpu/sgl-kernels/moe_int4.cpp new
+323
@@ -0,0 +1,323 @@
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 <int64_t N>
11 +inline void copy_bias(const float* bias_ptr, float* y_buf, int64_t m, int64_t ldn) {
12 + using Vec = at::vec::Vectorized<float>;
13 + constexpr int kVecSize = Vec::size();
14 + static_assert(N % kVecSize == 0, "copy_bias requires N to be a multiple of Vectorized<float>::size()");
15 + const bool has_bias = bias_ptr != nullptr;
16 + const Vec zero_vec(0.f);
17 + for (int i = 0; i < m; ++i) {
18 +#pragma GCC unroll 2
19 + for (int j = 0; j < N; j += kVecSize) {
20 + Vec vec = has_bias ? Vec::loadu(bias_ptr + j) : zero_vec;
21 + vec.store(y_buf + i * ldn + j);
22 + }
23 + }
24 +}
25 +
26 +template <typename scalar_t>
27 +void fused_experts_int4_w4a8_kernel_impl(
28 + scalar_t* __restrict__ output,
29 + scalar_t* __restrict__ ic0,
30 + scalar_t* __restrict__ ic1,
31 + scalar_t* __restrict__ ic2,
32 + uint8_t* __restrict__ A_tmp,
33 + uint8_t* __restrict__ Aq_tmp,
34 + float* __restrict__ As_tmp,
35 + int32_t* __restrict__ Azp_tmp,
36 + float* __restrict__ C_tmp,
37 + int8_t* __restrict__ dqB_tmp,
38 + const scalar_t* __restrict__ input,
39 + const uint8_t* __restrict__ packed_w1,
40 + const uint8_t* __restrict__ packed_w2,
41 + const int8_t* __restrict__ w1z,
42 + const int8_t* __restrict__ w2z,
43 + const float* __restrict__ w1s,
44 + const float* __restrict__ w2s,
45 + int group_size,
46 + const float* __restrict__ topk_weights,
47 + const int32_t* __restrict__ sorted_ids,
48 + const int32_t* __restrict__ expert_ids,
49 + const int32_t* __restrict__ offsets,
50 + int64_t M,
51 + int64_t N,
52 + int64_t K,
53 + int64_t E,
54 + int64_t topk,
55 + int64_t num_tokens_post_pad) {
56 + constexpr int64_t BLOCK_M = block_size_m();
57 + constexpr int64_t BLOCK_N = block_size_n();
58 + int num_threads = at::get_num_threads();
59 + // int64_t buffer_size_nbytes = M * topk * N * 2
60 + // M * topk * K * 2 +
61 + // num_threads * BLOCK_M * K +
62 + // num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float) +
63 + // M * topk * 2 * N * 2 +
64 + // max(M * K, M * topk * N) +
65 + // M * topk * sizeof(float);
66 +
67 + // intermediate_cache1 (scalar_t): START + M * topk * N
68 + // intermediate_cache2 (scalar_t): + M * topk * K
69 + // A_tmp (uint8_t): + num_threads * BLOCK_M * K
70 + // C_tmp (float): + num_threads * 2 * BLOCK_M * BLOCK_N
71 + // intermediate_cache0 (scalar_t): + M * topk * 2 * N
72 + // Aq_tmp (uint8_t): + max(M * K, M * topk * N)
73 + // As_tmp (float): + M * topk
74 + // dqB_tmp (int8_t) + num_threads * _block_k * BlOCK_N
75 +
76 + // stage 0: quantize input to uint8, [M, K]
77 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
78 + for (int64_t m = begin; m < end; ++m) {
79 + quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
80 + }
81 + });
82 + int64_t _block_k = get_4bit_block_k_size(group_size);
83 + auto Azp = at::ones({M * topk}).to(at::kInt).mul(128);
84 + auto Azp_ptr = Azp.data_ptr<int32_t>();
85 + // stage 1: intermediate_cache0 = hidden_states @ w1
86 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
87 + const int64_t NB = div_up(N, BLOCK_N);
88 +
89 + int64_t block_per_group = group_size / _block_k;
90 + int64_t Kc = K / _block_k;
91 + int64_t num_groups = K / group_size;
92 +
93 + const int64_t stride_e = 2 * NB * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
94 + const bool sym_quant_act = false;
95 + // weight + compensation shape = [E, Nc, Kc, block_n * _block_k / 2 + block_n*sizeof(int32_t)]
96 + // scales/qzeros shape = [E, Nc, G, block_n]
97 +
98 + // here we only parallel on half of 2N to fuse silu_and_mul with gemm
99 + at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
100 + // get local pointers
101 + int tid = at::get_thread_num();
102 + int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
103 + int8_t* dqB_tmp2 = dqB_tmp1 + _block_k * BLOCK_N;
104 + alignas(64) float As[BLOCK_M];
105 + uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
106 + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
107 + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
108 + bool is_brgemm_used = false;
109 + for (int64_t i = begin; i < end; ++i) {
110 + int64_t mb = i / NB;
111 + int64_t nb = i % NB;
112 + int64_t nb1 = nb + NB;
113 + int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
114 + // B shape [K, n_size] in vnni format
115 + int32_t expert_id = expert_ids[mb];
116 + const uint8_t* __restrict__ B = packed_w1 + expert_id * stride_e;
117 + // Bz and Bs: [E, K/gs, 2N]
118 + const int8_t* __restrict__ Bz = w1z + expert_id * (num_groups) * (2 * N);
119 + const float* __restrict__ Bs = w1s + expert_id * (num_groups) * (2 * N);
120 +
121 + // 1.a load A
122 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
123 + int64_t m_size = offsets[mb + 1] - offsets[mb];
124 + const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
125 + is_brgemm_used = is_brgemm_used || use_brgemm;
126 + // copy to A [BLOCK_M, K]
127 + for (int64_t m = 0; m < m_size; ++m) {
128 + int32_t index = A_ids[m] / topk;
129 + copy_stub(A + m * K, Aq_tmp + index * K, K);
130 + As[m] = As_tmp[index];
131 + }
132 + const int64_t offset = offsets[mb];
133 + copy_bias<BLOCK_N>(nullptr, C0, m_size, BLOCK_N);
134 + copy_bias<BLOCK_N>(nullptr, C1, m_size, BLOCK_N);
135 + for (int kci = 0; kci < Kc; ++kci) {
136 + int32_t* compensation_ptr =
137 + sym_quant_act ? nullptr
138 + : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
139 + _block_k * BLOCK_N / 2) /*Bcomp*/;
140 + tinygemm_kernel<scalar_t>(
141 + ic0 + offset * 2 * N + nb * BLOCK_N,
142 + C0,
143 + A + kci * _block_k,
144 + As,
145 + Azp_ptr,
146 + B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
147 + Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
148 + Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
149 + compensation_ptr,
150 + dqB_tmp1,
151 + m_size,
152 + _block_k,
153 + K,
154 + BLOCK_N,
155 + 2 * N,
156 + kci == Kc - 1,
157 + use_brgemm);
158 + }
159 +
160 + for (int kci = 0; kci < Kc; ++kci) {
161 + int32_t* compensation_ptr =
162 + sym_quant_act ? nullptr
163 + : (int32_t*)(void*)(B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
164 + _block_k * BLOCK_N / 2) /*Bcomp*/;
165 + tinygemm_kernel<scalar_t>(
166 + ic0 + offset * 2 * N + nb1 * BLOCK_N,
167 + C1,
168 + A + kci * _block_k,
169 + As,
170 + Azp_ptr,
171 + B + (nb1 * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) /*B*/,
172 + Bs + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
173 + Bz + nb1 * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*qzeros_b*/,
174 + compensation_ptr,
175 + dqB_tmp2,
176 + m_size,
177 + _block_k,
178 + K,
179 + BLOCK_N,
180 + 2 * N,
181 + kci == Kc - 1,
182 + use_brgemm);
183 + }
184 + }
185 +
186 + if (is_brgemm_used) {
187 + at::native::cpublas::brgemm_release();
188 + }
189 + });
190 +
191 + // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
192 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
193 + for (int64_t m = begin; m < end; ++m) {
194 + silu_and_mul_stub(ic1 + m * N, ic0 + m * 2 * N, ic0 + m * 2 * N + N, N);
195 + }
196 + });
197 +
198 + // stage 1.5: quantize ic1 to uint8, [M * topk, N]
199 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
200 + for (int64_t m = begin; m < end; ++m) {
201 + quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
202 + }
203 + });
204 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
205 + // w2 : [E, K, N] as [E, OC, IC]
206 + const int64_t OC = K; // rename K as OC
207 + const int64_t IC = N; // rename N as IC
208 + const int64_t MB2 = MB;
209 + const int64_t NB2 = div_up(OC, BLOCK_N);
210 + const int64_t stride_oc = IC;
211 + num_groups = IC / group_size;
212 + Kc = IC / _block_k;
213 + const int64_t stride_e2 = NB2 * Kc * (BLOCK_N * (_block_k / 2 + sizeof(int32_t)));
214 + // parallel on [MB2, NB2]
215 + at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
216 + int tid = at::get_thread_num();
217 + int8_t* dqB_tmp1 = dqB_tmp + tid * 2 * _block_k * BLOCK_N;
218 + float* __restrict__ C2 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
219 + bool is_brgemm_used = false;
220 + for (int64_t i = begin; i < end; ++i) {
221 + int64_t mb = i / NB2;
222 + int64_t nb = i % NB2;
223 +
224 + int64_t m_size = offsets[mb + 1] - offsets[mb];
225 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
226 + const bool use_brgemm = can_use_brgemm<int8_t>(m_size);
227 + is_brgemm_used = is_brgemm_used || use_brgemm;
228 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
229 +
230 + // B shape [IC, n_size] in vnni format
231 + int32_t expert_id = expert_ids[mb];
232 + const uint8_t* __restrict__ B = packed_w2 + expert_id * stride_e2;
233 +
234 + // Bz and Bs: [E, IC/gs, OC]
235 + const int8_t* __restrict__ Bz = w2z + expert_id * (num_groups)*OC;
236 + const float* __restrict__ Bs = w2s + expert_id * (num_groups)*OC;
237 +
238 + // A ptr from ic1 of [M * topk, N] in sorted order
239 + // so as to avoid copy A to tmp buffer again
240 + const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * IC;
241 + const float* __restrict__ As = As_tmp + offsets[mb];
242 + copy_bias<BLOCK_N>(nullptr, C2, m_size, BLOCK_N);
243 + for (int kci = 0; kci < Kc; ++kci) {
244 + int32_t* compensation_ptr =
245 + sym_quant_act ? nullptr
246 + : (int32_t*)(void*)(B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))) +
247 + _block_k * BLOCK_N / 2) /*Bcomp*/;
248 + tinygemm_kernel<scalar_t>(
249 + nullptr, /*store_out is false*/
250 + C2,
251 + A + kci * _block_k,
252 + As,
253 + Azp_ptr,
254 + B + (nb * Kc + kci) * (BLOCK_N * (_block_k / 2 + sizeof(int32_t))),
255 + Bs + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*scales_b*/,
256 + Bz + nb * BLOCK_N * num_groups + kci / block_per_group * BLOCK_N /*zeros_b*/,
257 + compensation_ptr,
258 + dqB_tmp1,
259 + m_size,
260 + _block_k,
261 + IC,
262 + BLOCK_N,
263 + BLOCK_N,
264 + false,
265 + use_brgemm);
266 + }
267 +
268 + // 2.b copy from C to ic2 in original order
269 + // and also mul topk_weights in float32
270 + for (int64_t m = 0; m < m_size; ++m) {
271 + int32_t index = A_ids[m];
272 + float weight = topk_weights[index];
273 + copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C2 + m * BLOCK_N, weight, n_size);
274 + }
275 + }
276 +
277 + if (is_brgemm_used) {
278 + at::native::cpublas::brgemm_release();
279 + }
280 + });
281 +
282 + // stage 3: out = intermediate_cache2.sum(dim=1)
283 + // from [M, topk, K] to [M, K]
284 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
285 + for (int64_t m = begin; m < end; ++m) {
286 + sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
287 + }
288 + });
289 +}
290 +
291 +#define INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(TYPE) \
292 + template void fused_experts_int4_w4a8_kernel_impl<TYPE>( \
293 + TYPE* __restrict__ output, \
294 + TYPE* __restrict__ ic0, \
295 + TYPE* __restrict__ ic1, \
296 + TYPE* __restrict__ ic2, \
297 + uint8_t* __restrict__ A_tmp, \
298 + uint8_t* __restrict__ Aq_tmp, \
299 + float* __restrict__ As_tmp, \
300 + int32_t* __restrict__ Azp_tmp, \
301 + float* __restrict__ C_tmp, \
302 + int8_t* __restrict__ dqB_tmp, \
303 + const TYPE* __restrict__ input, \
304 + const uint8_t* __restrict__ packed_w1, \
305 + const uint8_t* __restrict__ packed_w2, \
306 + const int8_t* __restrict__ w1z, \
307 + const int8_t* __restrict__ w2z, \
308 + const float* __restrict__ w1s, \
309 + const float* __restrict__ w2s, \
310 + int group_size, \
311 + const float* __restrict__ topk_weights, \
312 + const int32_t* __restrict__ sorted_ids, \
313 + const int32_t* __restrict__ expert_ids, \
314 + const int32_t* __restrict__ offsets, \
315 + int64_t M, \
316 + int64_t N, \
317 + int64_t K, \
318 + int64_t E, \
319 + int64_t topk, \
320 + int64_t num_tokens_post_pad)
321 +
322 +INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::BFloat16);
323 +INSTANTIATE_MOE_INT4_W4A8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/moe_int8.cpp
+506 -302
@@ -1,124 +1,145 @@
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"
5 -#include "vec.h"
7 #include "gemm.h"
7 -
8 -// clang-format off
8 +#include "moe.h"
9
10 namespace {
11
12 -template <typename scalar_t>
13 -inline void copy_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t size) {
14 - using Vec = at::vec::Vectorized<scalar_t>;
15 - // no remainder
16 - #pragma GCC unroll 4
17 - for (int64_t d = 0; d < size; d += Vec::size()) {
18 - Vec data = Vec::loadu(input + d);
19 - data.store(out + d);
20 - }
21 -}
22 -
23 -template <>
24 -inline void copy_stub<uint8_t>(uint8_t* __restrict__ out, const uint8_t* __restrict__ input, int64_t size) {
25 - // size might be 64x + 32
26 - std::memcpy(out, input, size * sizeof(uint8_t));
27 -}
28 -
29 -template <typename scalar_t>
30 -inline void copy_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input, float weight, int64_t size) {
31 - using bVec = at::vec::Vectorized<scalar_t>;
32 - using fVec = at::vec::Vectorized<float>;
33 - constexpr int kVecSize = bVec::size();
34 - const fVec weight_vec = fVec(weight);
35 - int64_t d;
36 - #pragma GCC unroll 4
37 - for (d = 0; d <= size - kVecSize; d += kVecSize) {
38 - fVec data0 = fVec::loadu(input + d) * weight_vec;
39 - fVec data1 = fVec::loadu(input + d + fVec::size()) * weight_vec;
40 - bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
41 - out_vec.store(out + d);
42 - }
43 - for (; d < size; ++d) {
44 - out[d] = static_cast<scalar_t>(input[d] * weight);
45 - }
46 -}
12 +template <typename scalar_t, int BLOCK_N>
13 +inline void silu_and_mul(
14 + scalar_t* __restrict__ C,
15 + const int32_t* __restrict__ C0, // x: x0, x1
16 + const int32_t* __restrict__ C1, // y: y0, y1
17 + const float* __restrict__ As,
18 + const float* __restrict__ Bs0,
19 + const float* __restrict__ Bs1,
20 + const int32_t* __restrict__ Bcomp0,
21 + const int32_t* __restrict__ Bcomp1,
22 + int64_t m_size,
23 + int64_t N) {
24 +#if defined(CPU_CAPABILITY_AVX512)
25 + constexpr int COLS = BLOCK_N / 16;
26 + static_assert(COLS % 2 == 0);
27 +
28 + __m512 vc0[COLS];
29 + __m512 vc1[COLS];
30 + __m512i vcomp0[COLS];
31 + __m512i vcomp1[COLS];
32 + __m512 vas;
33 + __m512 vbs0[COLS];
34 + __m512 vbs1[COLS];
35 +
36 + auto load_scale_and_comp = [&](auto col) {
37 + vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
38 + vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
39 + vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
40 + vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
41 + };
42 + Unroll<COLS>{}(load_scale_and_comp);
43 +
44 + auto scalec = [&](auto col, int64_t m) {
45 + // update As
46 + vas = _mm512_set1_ps(As[m]);
47 + // C = As * (C - Bcomp) * Bs
48 + __m512i vc32_0 = _mm512_loadu_si512(C0 + m * BLOCK_N + col * 16);
49 + __m512i vc32_1 = _mm512_loadu_si512(C1 + m * BLOCK_N + col * 16);
50 + vc0[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32_0, vcomp0[col]));
51 + vc1[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32_1, vcomp1[col]));
52 + vc0[col] = _mm512_mul_ps(_mm512_mul_ps(vc0[col], vas), vbs0[col]);
53 + vc1[col] = _mm512_mul_ps(_mm512_mul_ps(vc1[col], vas), vbs1[col]);
54 + };
55
48 -// acc from [topk, K] to [K]
49 -template <typename scalar_t>
50 -inline void sum_stub(scalar_t* __restrict__ out, const scalar_t* __restrict__ input, int64_t topk, int64_t K) {
56 using bVec = at::vec::Vectorized<scalar_t>;
57 using fVec = at::vec::Vectorized<float>;
53 - constexpr int kVecSize = bVec::size();
54 - if (topk == 1) {
55 - // do copy for topk = 1
56 - copy_stub(out, input, K);
57 - } else {
58 - // do sum for topk != 1
59 - int64_t d;
60 - #pragma GCC unroll 4
61 - for (d = 0; d <= K - kVecSize; d += kVecSize) {
62 - fVec sum_fvec0 = fVec(0.f);
63 - fVec sum_fvec1 = fVec(0.f);
64 - for (int t = 0; t < topk; ++t) {
65 - bVec x_bvec = bVec::loadu(input + t * K + d);
66 - fVec x_fvec0, x_fvec1;
67 - std::tie(x_fvec0, x_fvec1) = at::vec::convert_to_float(x_bvec);
68 -
69 - sum_fvec0 += x_fvec0;
70 - sum_fvec1 += x_fvec1;
71 - }
72 - bVec out_bvec = convert_from_float_ext<scalar_t>(sum_fvec0, sum_fvec1);
73 - out_bvec.store(out + d);
74 - }
75 - for (; d < K; ++d) {
76 - float sum_val = 0.f;
77 - for (int t = 0; t < topk; ++t) {
78 - sum_val += static_cast<float>(input[t * K + d]);
79 - }
80 - out[d] = static_cast<scalar_t>(sum_val);
58 + const fVec one = fVec(1.f);
59 + auto silu_and_mul = [&](auto col) {
60 + fVec x = fVec(vc0[col]);
61 + fVec y = fVec(vc1[col]);
62 + x = x / (one + x.neg().exp_u20());
63 + vc0[col] = x * y;
64 + };
65 +
66 + auto storec = [&](auto col, int64_t m) {
67 + if constexpr (col % 2 == 0) {
68 + fVec x0 = fVec(vc0[col + 0]);
69 + fVec x1 = fVec(vc0[col + 1]);
70 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
71 + out_vec.store(C + m * N + col * 16);
72 }
73 + };
74 +
75 + for (int64_t m = 0; m < m_size; ++m) {
76 + Unroll<COLS>{}(scalec, m);
77 + Unroll<COLS>{}(silu_and_mul);
78 + Unroll<COLS>{}(storec, m);
79 }
80 +#else
81 + TORCH_CHECK(false, "silu_and_mul: scalar path not implemented!");
82 +#endif
83 }
84
85 -// out = input + input2 * scale
86 -template <typename scalar_t>
87 -inline void add_mul_stub(scalar_t* __restrict__ out, const float* __restrict__ input,
88 - const scalar_t* __restrict__ input2, float scale, int64_t size) {
89 -
90 - using bVec = at::vec::Vectorized<scalar_t>;
91 - using fVec = at::vec::Vectorized<float>;
92 - constexpr int kVecSize = bVec::size();
93 - const fVec s_vec = fVec(scale);
94 - int64_t d;
95 - #pragma GCC unroll 4
96 - for (d = 0; d <= size - kVecSize; d += kVecSize) {
97 - fVec x0 = fVec::loadu(input + d);
98 - fVec x1 = fVec::loadu(input + d + fVec::size());
99 -
100 - bVec y_bvec = bVec::loadu(input2 + d);
101 - fVec y0, y1;
102 - std::tie(y0, y1) = at::vec::convert_to_float(y_bvec);
103 -
104 - x0 = x0 + y0 * s_vec;
105 - x1 = x1 + y1 * s_vec;
106 - bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
107 - out_vec.store(out + d);
108 - }
109 - for (; d < size; ++d) {
110 - out[d] = static_cast<scalar_t>(input[d] + float(input2[d]) * scale);
85 +template <int BLOCK_N>
86 +inline void scale_C(
87 + float* __restrict__ C,
88 + const int32_t* __restrict__ Ctmp,
89 + const float* __restrict__ As,
90 + const float* __restrict__ Bs,
91 + const int32_t* __restrict__ Bcomp,
92 + int64_t m_size) {
93 +#if defined(CPU_CAPABILITY_AVX512)
94 + constexpr int COLS = BLOCK_N / 16;
95 + static_assert(COLS % 2 == 0);
96 +
97 + __m512 vc[COLS];
98 + __m512i vcomp[COLS];
99 + __m512 vas;
100 + __m512 vbs[COLS];
101 +
102 + auto load_scale_and_comp = [&](auto col) {
103 + vcomp[col] = _mm512_loadu_si512(Bcomp + col * 16);
104 + vbs[col] = _mm512_loadu_ps(Bs + col * 16);
105 + };
106 + Unroll<COLS>{}(load_scale_and_comp);
107 +
108 + auto scalec = [&](auto col, int64_t m) {
109 + // update As
110 + vas = _mm512_set1_ps(As[m]);
111 + // C = As * (C - Bcomp) * Bs
112 + __m512i vc32 = _mm512_loadu_si512(Ctmp + m * BLOCK_N + col * 16);
113 + vc[col] = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc32, vcomp[col]));
114 + vc[col] = _mm512_mul_ps(_mm512_mul_ps(vc[col], vas), vbs[col]);
115 + _mm512_storeu_ps(C + m * BLOCK_N + col * 16, vc[col]);
116 + };
117 +
118 + for (int64_t m = 0; m < m_size; ++m) {
119 + Unroll<COLS>{}(scalec, m);
120 }
121 +#else
122 + TORCH_CHECK(false, "scale_C: scalar path not implemented!");
123 +#endif
124 }
125
126 /// gemm for w13
127 template <typename scalar_t, int BLOCK_M, int BLOCK_N>
128 struct tinygemm_kernel_vnni {
129 static inline void apply(
118 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B0, const int8_t* __restrict__ B1, scalar_t* __restrict__ C,
119 - const float* __restrict__ As, const float* __restrict__ Bs0, const float* __restrict__ Bs1,
120 - const int32_t* __restrict__ Bcomp0, const int32_t* __restrict__ Bcomp1,
121 - int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
130 + const uint8_t* __restrict__ A,
131 + const int8_t* __restrict__ B0,
132 + const int8_t* __restrict__ B1,
133 + scalar_t* __restrict__ C,
134 + const float* __restrict__ As,
135 + const float* __restrict__ Bs0,
136 + const float* __restrict__ Bs1,
137 + const int32_t* __restrict__ Bcomp0,
138 + const int32_t* __restrict__ Bcomp1,
139 + int64_t K,
140 + int64_t lda,
141 + int64_t ldb,
142 + int64_t ldc) {
143 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
144 }
145 };
@@ -127,11 +148,19 @@ struct tinygemm_kernel_vnni {
148 template <int BLOCK_M, int BLOCK_N>
149 struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
150 static inline void apply(
130 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B0, const int8_t* __restrict__ B1, at::BFloat16* __restrict__ C,
131 - const float* __restrict__ As, const float* __restrict__ Bs0, const float* __restrict__ Bs1,
132 - const int32_t* __restrict__ Bcomp0, const int32_t* __restrict__ Bcomp1,
133 - int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
134 -
151 + const uint8_t* __restrict__ A,
152 + const int8_t* __restrict__ B0,
153 + const int8_t* __restrict__ B1,
154 + at::BFloat16* __restrict__ C,
155 + const float* __restrict__ As,
156 + const float* __restrict__ Bs0,
157 + const float* __restrict__ Bs1,
158 + const int32_t* __restrict__ Bcomp0,
159 + const int32_t* __restrict__ Bcomp1,
160 + int64_t K,
161 + int64_t lda,
162 + int64_t ldb,
163 + int64_t ldc) {
164 constexpr int ROWS = BLOCK_M;
165 constexpr int COLS = BLOCK_N / 16;
166 static_assert(COLS % 2 == 0);
@@ -143,9 +172,9 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
172 __m512i vc1[ROWS * COLS];
173 __m512i vcomp0[COLS];
174 __m512i vcomp1[COLS];
146 - __m512 was;
147 - __m512 vbs0[COLS];
148 - __m512 vbs1[COLS];
175 + __m512 vas;
176 + __m512 vbs0[COLS];
177 + __m512 vbs1[COLS];
178
179 auto loadc = [&](auto i) {
180 vc0[i] = _mm512_set1_epi32(0);
@@ -155,7 +184,7 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
184
185 const int64_t K4 = K >> 2;
186 const int64_t lda4 = lda >> 2;
158 - const int64_t ldb4 = ldb; // ldb * 4 >> 2;
187 + const int64_t ldb4 = ldb; // ldb * 4 >> 2;
188 const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
189 const int32_t* b0_ptr = reinterpret_cast<const int32_t*>(B0);
190 const int32_t* b1_ptr = reinterpret_cast<const int32_t*>(B1);
@@ -183,8 +212,8 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
212 constexpr int col = i % COLS;
213
214 // load a scale
186 - if constexpr(col == 0) {
187 - was = _mm512_set1_ps(As[row]);
215 + if constexpr (col == 0) {
216 + vas = _mm512_set1_ps(As[row]);
217 }
218 // load b scale and vcomp
219 if constexpr (row == 0) {
@@ -195,8 +224,8 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
224 }
225 __m512 c0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc0[i], vcomp0[col]));
226 __m512 c1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc1[i], vcomp1[col]));
198 - vc0[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c0, was), vbs0[col]));
199 - vc1[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c1, was), vbs1[col]));
227 + vc0[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c0, vas), vbs0[col]));
228 + vc1[i] = _mm512_castps_si512(_mm512_mul_ps(_mm512_mul_ps(c1, vas), vbs1[col]));
229 };
230 Unroll<ROWS * COLS>{}(scalec);
231
@@ -221,19 +250,28 @@ struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
250 _mm512_storeu_si512(
251 reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
252 (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
224 - }
253 + }
254 };
255 Unroll<ROWS * COLS>{}(storec);
256 }
257 };
258 #endif
259
231 -#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \
232 - tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE>::apply( \
233 - A + mb_start * lda, B0 + nb_start * 4, B1 + nb_start * 4, \
234 - C + mb_start * ldc + nb_start, As + mb_start, \
235 - Bs0 + nb_start, Bs1 + nb_start, Bcomp0 + nb_start, Bcomp1 + nb_start,\
236 - K, lda, ldb, ldc);
260 +#define LAUNCH_TINYGEMM_KERNEL_VNNI(MB_SIZE, NB_SIZE) \
261 + tinygemm_kernel_vnni<scalar_t, MB_SIZE, NB_SIZE>::apply( \
262 + A + mb_start * lda, \
263 + B0 + nb_start * 4, \
264 + B1 + nb_start * 4, \
265 + C + mb_start * ldc + nb_start, \
266 + As + mb_start, \
267 + Bs0 + nb_start, \
268 + Bs1 + nb_start, \
269 + Bcomp0 + nb_start, \
270 + Bcomp1 + nb_start, \
271 + K, \
272 + lda, \
273 + ldb, \
274 + ldc);
275
276 template <typename scalar_t>
277 void tinygemm_kernel(
@@ -250,7 +288,6 @@ void tinygemm_kernel(
288 int64_t lda,
289 int64_t ldb,
290 int64_t ldc) {
253 -
291 const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
292 const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
293
@@ -266,12 +303,21 @@ void tinygemm_kernel(
303 int64_t nb_start = nb * BLOCK_N;
304 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
305
269 - switch(mb_size << 4 | nb_size >> 4) {
270 - case 0x12: LAUNCH_TINYGEMM_KERNEL_VNNI(1, 32); break;
271 - case 0x22: LAUNCH_TINYGEMM_KERNEL_VNNI(2, 32); break;
272 - case 0x32: LAUNCH_TINYGEMM_KERNEL_VNNI(3, 32); break;
273 - case 0x42: LAUNCH_TINYGEMM_KERNEL_VNNI(4, 32); break;
274 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
306 + switch (mb_size << 4 | nb_size >> 4) {
307 + case 0x12:
308 + LAUNCH_TINYGEMM_KERNEL_VNNI(1, 32);
309 + break;
310 + case 0x22:
311 + LAUNCH_TINYGEMM_KERNEL_VNNI(2, 32);
312 + break;
313 + case 0x32:
314 + LAUNCH_TINYGEMM_KERNEL_VNNI(3, 32);
315 + break;
316 + case 0x42:
317 + LAUNCH_TINYGEMM_KERNEL_VNNI(4, 32);
318 + break;
319 + default:
320 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
321 }
322 }
323 }
@@ -281,9 +327,16 @@ void tinygemm_kernel(
327 template <typename scalar_t, int BLOCK_M, int BLOCK_N>
328 struct tinygemm_kernel_vnni2 {
329 static inline void apply(
284 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B, float* __restrict__ C,
285 - const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
286 - int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
330 + const uint8_t* __restrict__ A,
331 + const int8_t* __restrict__ B,
332 + float* __restrict__ C,
333 + const float* __restrict__ As,
334 + const float* __restrict__ Bs,
335 + const int32_t* __restrict__ Bcomp,
336 + int64_t K,
337 + int64_t lda,
338 + int64_t ldb,
339 + int64_t ldc) {
340 TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
341 }
342 };
@@ -292,10 +345,16 @@ struct tinygemm_kernel_vnni2 {
345 template <int BLOCK_M, int BLOCK_N>
346 struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
347 static inline void apply(
295 - const uint8_t* __restrict__ A, const int8_t* __restrict__ B, float* __restrict__ C,
296 - const float* __restrict__ As, const float* __restrict__ Bs, const int32_t* __restrict__ Bcomp,
297 - int64_t K, int64_t lda, int64_t ldb, int64_t ldc) {
298 -
348 + const uint8_t* __restrict__ A,
349 + const int8_t* __restrict__ B,
350 + float* __restrict__ C,
351 + const float* __restrict__ As,
352 + const float* __restrict__ Bs,
353 + const int32_t* __restrict__ Bcomp,
354 + int64_t K,
355 + int64_t lda,
356 + int64_t ldb,
357 + int64_t ldc) {
358 constexpr int ROWS = BLOCK_M;
359 constexpr int COLS = BLOCK_N / 16;
360 static_assert(COLS % 2 == 0);
@@ -304,17 +363,15 @@ struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
363 __m512i vb[COLS];
364 __m512i vc[ROWS * COLS];
365 __m512i vcomp[COLS];
307 - __m512 was;
308 - __m512 vbs[COLS];
366 + __m512 vas;
367 + __m512 vbs[COLS];
368
310 - auto loadc = [&](auto i) {
311 - vc[i] = _mm512_set1_epi32(0);
312 - };
369 + auto loadc = [&](auto i) { vc[i] = _mm512_set1_epi32(0); };
370 Unroll<ROWS * COLS>{}(loadc);
371
372 const int64_t K4 = K >> 2;
373 const int64_t lda4 = lda >> 2;
317 - const int64_t ldb4 = ldb; // ldb * 4 >> 2;
374 + const int64_t ldb4 = ldb; // ldb * 4 >> 2;
375 const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
376 const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
377
@@ -339,8 +396,8 @@ struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
396 constexpr int col = i % COLS;
397
398 // load a scale
342 - if constexpr(col == 0) {
343 - was = _mm512_set1_ps(As[row]);
399 + if constexpr (col == 0) {
400 + vas = _mm512_set1_ps(As[row]);
401 }
402 // load b scale and vcomp per 2 vectors
403 // also load bias if any
@@ -353,7 +410,7 @@ struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
410 }
411 }
412 __m512 x = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[i], vcomp[col]));
356 - x = _mm512_mul_ps(_mm512_mul_ps(x, was), vbs[col]);
413 + x = _mm512_mul_ps(_mm512_mul_ps(x, vas), vbs[col]);
414 _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), x);
415 };
416 Unroll<ROWS * COLS>{}(storec);
@@ -361,11 +418,18 @@ struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
418 };
419 #endif
420
364 -#define LAUNCH_TINYGEMM_KERNEL_VNNI2(MB_SIZE, NB_SIZE) \
365 - tinygemm_kernel_vnni2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
366 - A + mb_start * lda, B + nb_start * 4, C + mb_start * ldc + nb_start, \
367 - As + mb_start, Bs + nb_start, Bcomp + nb_start, \
368 - K, lda, ldb, ldc);
421 +#define LAUNCH_TINYGEMM_KERNEL_VNNI2(MB_SIZE, NB_SIZE) \
422 + tinygemm_kernel_vnni2<scalar_t, MB_SIZE, NB_SIZE>::apply( \
423 + A + mb_start * lda, \
424 + B + nb_start * 4, \
425 + C + mb_start * ldc + nb_start, \
426 + As + mb_start, \
427 + Bs + nb_start, \
428 + Bcomp + nb_start, \
429 + K, \
430 + lda, \
431 + ldb, \
432 + ldc);
433
434 template <typename scalar_t>
435 void tinygemm_kernel(
@@ -380,7 +444,6 @@ void tinygemm_kernel(
444 int64_t lda,
445 int64_t ldb,
446 int64_t ldc) {
383 -
447 // B compensation
448 const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
449
@@ -396,18 +459,27 @@ void tinygemm_kernel(
459 int64_t nb_start = nb * BLOCK_N;
460 int64_t nb_size = std::min(BLOCK_N, N - nb_start);
461
399 - switch(mb_size << 4 | nb_size >> 4) {
400 - case 0x12: LAUNCH_TINYGEMM_KERNEL_VNNI2(1, 32); break;
401 - case 0x22: LAUNCH_TINYGEMM_KERNEL_VNNI2(2, 32); break;
402 - case 0x32: LAUNCH_TINYGEMM_KERNEL_VNNI2(3, 32); break;
403 - case 0x42: LAUNCH_TINYGEMM_KERNEL_VNNI2(4, 32); break;
404 - default: TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", nb_size);
462 + switch (mb_size << 4 | nb_size >> 4) {
463 + case 0x12:
464 + LAUNCH_TINYGEMM_KERNEL_VNNI2(1, 32);
465 + break;
466 + case 0x22:
467 + LAUNCH_TINYGEMM_KERNEL_VNNI2(2, 32);
468 + break;
469 + case 0x32:
470 + LAUNCH_TINYGEMM_KERNEL_VNNI2(3, 32);
471 + break;
472 + case 0x42:
473 + LAUNCH_TINYGEMM_KERNEL_VNNI2(4, 32);
474 + break;
475 + default:
476 + TORCH_CHECK(false, "Unexpected block size, ", mb_size, "x", "nb_size");
477 }
478 }
479 }
480 }
481
410 -} // anonymous namespace
482 +} // anonymous namespace
483
484 template <typename scalar_t>
485 void fused_experts_int8_kernel_impl(
@@ -433,7 +505,6 @@ void fused_experts_int8_kernel_impl(
505 int64_t E,
506 int64_t topk,
507 int64_t num_tokens_post_pad) {
436 -
508 // handle 2 tiles per block
509 constexpr int64_t BLOCK_M = block_size_m();
510 constexpr int64_t BLOCK_N = block_size_n();
@@ -441,11 +512,7 @@ void fused_experts_int8_kernel_impl(
512 // stage 0: quantize input to uint8, [M, K]
513 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
514 for (int64_t m = begin; m < end; ++m) {
444 - quantize_row_int8<scalar_t>(
445 - Aq_tmp + m * K,
446 - As_tmp[m],
447 - input + m * K,
448 - K);
515 + quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
516 }
517 });
518
@@ -462,66 +529,107 @@ void fused_experts_int8_kernel_impl(
529
530 const int64_t stride_e = 2 * N * packed_K;
531 const int64_t stride_n = packed_K;
532 +
533 + int64_t avg_M = std::max(int64_t(1), M * topk / E);
534 + const bool use_brgemm = can_use_brgemm<int8_t>(avg_M);
535 +
536 // here we only parallel on half of 2N to fuse silu_and_mul with gemm
466 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
537 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
538 // get local pointers
468 - int tid = at::get_thread_num();
539 + int tid = get_thread_num();
540 uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
541 + int32_t* __restrict__ C0 = reinterpret_cast<int32_t*>(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N;
542 + int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
543
544 alignas(64) float As[BLOCK_M];
545
473 - for (int64_t i = begin; i < end; ++i) {
474 - int64_t mb = i / NB;
475 - int64_t nb = i % NB;
476 -
477 - // nb0 from top half and nb1 from bottom half
478 - int64_t nb0 = nb, nb1 = nb + NB;
479 - int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
546 + loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
547 + // nb_upper from top half and nb_lower from bottom half
548 + int64_t nb_upper = nb, nb_lower = nb + NB;
549 + int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
550
551 // B shape [K, n_size] in vnni format
552 int32_t expert_id = expert_ids[mb];
483 - const int8_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb0 * BLOCK_N * stride_n;
484 - const int8_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb1 * BLOCK_N * stride_n;
485 - const float* __restrict__ Bs0 = w1s + expert_id * 2 * N + nb0 * BLOCK_N;
486 - const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb1 * BLOCK_N;
553 + const int8_t* __restrict__ B0 = packed_w1 + expert_id * stride_e + nb_upper * BLOCK_N * stride_n;
554 + const int8_t* __restrict__ B1 = packed_w1 + expert_id * stride_e + nb_lower * BLOCK_N * stride_n;
555 + const float* __restrict__ Bs0 = w1s + expert_id * 2 * N + nb_upper * BLOCK_N;
556 + const float* __restrict__ Bs1 = w1s + expert_id * 2 * N + nb_lower * BLOCK_N;
557
488 - // 1.a load A
489 - const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
558 int64_t m_size = offsets[mb + 1] - offsets[mb];
559
492 - for (int64_t m = 0; m < m_size; ++m) {
493 - int32_t index = A_ids[m] / topk;
494 - copy_stub(A + m * K, Aq_tmp + index * K, K);
495 - As[m] = As_tmp[index];
560 + if (nb_offset == 0) {
561 + // 1.a load A
562 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
563 + for (int64_t m = 0; m < m_size; ++m) {
564 + int32_t index = A_ids[m] / topk;
565 + copy_stub(A + m * K, Aq_tmp + index * K, K);
566 + As[m] = As_tmp[index];
567 + }
568 }
569
498 - // fused 1.b: silu_and_mul(A @ B0, A @ B1)
499 - const int64_t offset = offsets[mb];
500 - tinygemm_kernel(
501 - /* A */ A,
502 - /* B0 */ B0,
503 - /* B1 */ B1,
504 - /* C */ ic1 + offset * N + nb * BLOCK_N,
505 - /* As */ As,
506 - /* Bs0 */ Bs0,
507 - /* Bs1 */ Bs1,
508 - /* M */ m_size,
509 - /* N */ n_size,
510 - /* K */ K,
511 - /* lda */ K,
512 - /* ldb */ n_size,
513 - /* ldc */ N);
570 + if (use_brgemm) {
571 + // 1.b gemm: C0 = A @ B0
572 + at::native::cpublas::brgemm(
573 + /* M */ m_size,
574 + /* N */ n_size,
575 + /* K */ K,
576 + /* lda */ K,
577 + /* ldb */ n_size,
578 + /* ldc */ BLOCK_N,
579 + /* add_C */ false,
580 + /* A */ A,
581 + /* B */ B0,
582 + /* C */ C0);
583 +
584 + // 1.c gemm: C1 = A @ B1
585 + at::native::cpublas::brgemm(
586 + /* M */ m_size,
587 + /* N */ n_size,
588 + /* K */ K,
589 + /* lda */ K,
590 + /* ldb */ n_size,
591 + /* ldc */ BLOCK_N,
592 + /* add_C */ false,
593 + /* A */ A,
594 + /* B */ B1,
595 + /* C */ C1);
596 +
597 + const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
598 + const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
599 +
600 + // 1.d silu and mul
601 + const int64_t offset = offsets[mb];
602 + silu_and_mul<scalar_t, BLOCK_N>(
603 + ic1 + offset * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
604 + } else {
605 + // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
606 + const int64_t offset = offsets[mb];
607 + tinygemm_kernel(
608 + /* A */ A,
609 + /* B0 */ B0,
610 + /* B1 */ B1,
611 + /* C */ ic1 + offset * N + nb * BLOCK_N,
612 + /* As */ As,
613 + /* Bs0 */ Bs0,
614 + /* Bs1 */ Bs1,
615 + /* M */ m_size,
616 + /* N */ n_size,
617 + /* K */ K,
618 + /* lda */ K,
619 + /* ldb */ n_size,
620 + /* ldc */ N);
621 + }
622 + });
623 +
624 + if (use_brgemm) {
625 + at::native::cpublas::brgemm_release();
626 }
627 });
628
629 // stage 1.5: quantize ic1 to uint8, [M * topk, N]
630 at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
631 for (int64_t m = begin; m < end; ++m) {
520 - quantize_row_int8<scalar_t>(
521 - Aq_tmp + m * N,
522 - As_tmp[m],
523 - ic1 + m * N,
524 - N);
632 + quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
633 }
634 });
635
@@ -535,16 +643,13 @@ void fused_experts_int8_kernel_impl(
643 const int64_t stride_oc = packed_N;
644
645 // parallel on [MB2, NB2]
538 - at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
646 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
647 // get local pointers
540 - int tid = at::get_thread_num();
541 - // we won't be using C1 for gemm2
648 + int tid = get_thread_num();
649 float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
650 + int32_t* __restrict__ C32 = reinterpret_cast<int32_t*>(C + BLOCK_M * BLOCK_N);
651
544 - for (int64_t i = begin; i < end; ++i) {
545 - int64_t mb = i / NB2;
546 - int64_t nb = i % NB2;
547 -
652 + loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
653 int64_t m_size = offsets[mb + 1] - offsets[mb];
654 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
655
@@ -560,18 +665,36 @@ void fused_experts_int8_kernel_impl(
665 const float* __restrict__ Bs = w2s + expert_id * K + nb * BLOCK_N;
666
667 // 2.a gemm: C = A @ B
563 - tinygemm_kernel<scalar_t>(
564 - /* A */ A,
565 - /* B */ B,
566 - /* C */ C,
567 - /* As */ As,
568 - /* Bs */ Bs,
569 - /* M */ m_size,
570 - /* N */ n_size,
571 - /* K */ IC,
572 - /* lda */ IC,
573 - /* ldb */ n_size,
574 - /* ldc */ BLOCK_N);
668 + if (use_brgemm) {
669 + at::native::cpublas::brgemm(
670 + /* M */ m_size,
671 + /* N */ n_size,
672 + /* K */ IC,
673 + /* lda */ IC,
674 + /* ldb */ n_size,
675 + /* ldc */ BLOCK_N,
676 + /* add_C */ false,
677 + /* A */ A,
678 + /* B */ B,
679 + /* C */ C32);
680 +
681 + // apply scales
682 + const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * IC);
683 + scale_C<BLOCK_N>(C, C32, As, Bs, Bcomp, m_size);
684 + } else {
685 + tinygemm_kernel<scalar_t>(
686 + /* A */ A,
687 + /* B */ B,
688 + /* C */ C,
689 + /* As */ As,
690 + /* Bs */ Bs,
691 + /* M */ m_size,
692 + /* N */ n_size,
693 + /* K */ IC,
694 + /* lda */ IC,
695 + /* ldb */ n_size,
696 + /* ldc */ BLOCK_N);
697 + }
698
699 // 2.b copy from C to ic2 in original order
700 // and also mul topk_weights in float32
@@ -580,6 +703,10 @@ void fused_experts_int8_kernel_impl(
703 float weight = topk_weights[index];
704 copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
705 }
706 + });
707 +
708 + if (use_brgemm) {
709 + at::native::cpublas::brgemm_release();
710 }
711 });
712
@@ -592,17 +719,30 @@ void fused_experts_int8_kernel_impl(
719 });
720 }
721
595 -#define INSTANTIATE_MOE_INT8_TEMPLATE(TYPE) \
596 - template void fused_experts_int8_kernel_impl<TYPE> ( \
597 - TYPE* __restrict__ output, TYPE* __restrict__ ic1, \
598 - TYPE* __restrict__ ic2, uint8_t* __restrict__ A_tmp, \
599 - float* __restrict__ C_tmp, uint8_t* __restrict__ Aq_tmp, \
600 - float* __restrict__ As_tmp, const TYPE* __restrict__ input, \
601 - const int8_t* __restrict__ packed_w1, const int8_t* __restrict__ packed_w2, \
602 - const float* __restrict__ w1s, const float* __restrict__ w2s, \
603 - const float* __restrict__ topk_weights, const int32_t* __restrict__ sorted_ids, \
604 - const int32_t* __restrict__ expert_ids, const int32_t* __restrict__ offsets, \
605 - int64_t M, int64_t N, int64_t K, int64_t E, int64_t topk, int64_t num_tokens_post_pad)
722 +#define INSTANTIATE_MOE_INT8_TEMPLATE(TYPE) \
723 + template void fused_experts_int8_kernel_impl<TYPE>( \
724 + TYPE* __restrict__ output, \
725 + TYPE* __restrict__ ic1, \
726 + TYPE* __restrict__ ic2, \
727 + uint8_t* __restrict__ A_tmp, \
728 + float* __restrict__ C_tmp, \
729 + uint8_t* __restrict__ Aq_tmp, \
730 + float* __restrict__ As_tmp, \
731 + const TYPE* __restrict__ input, \
732 + const int8_t* __restrict__ packed_w1, \
733 + const int8_t* __restrict__ packed_w2, \
734 + const float* __restrict__ w1s, \
735 + const float* __restrict__ w2s, \
736 + const float* __restrict__ topk_weights, \
737 + const int32_t* __restrict__ sorted_ids, \
738 + const int32_t* __restrict__ expert_ids, \
739 + const int32_t* __restrict__ offsets, \
740 + int64_t M, \
741 + int64_t N, \
742 + int64_t K, \
743 + int64_t E, \
744 + int64_t topk, \
745 + int64_t num_tokens_post_pad)
746
747 INSTANTIATE_MOE_INT8_TEMPLATE(at::BFloat16);
748 INSTANTIATE_MOE_INT8_TEMPLATE(at::Half);
@@ -624,7 +764,6 @@ void shared_expert_int8_kernel_impl(
764 int64_t M,
765 int64_t N,
766 int64_t K) {
627 -
767 // handle 2 tiles per block
768 constexpr int64_t BLOCK_M = block_size_m();
769 constexpr int64_t BLOCK_N = block_size_n();
@@ -632,15 +771,11 @@ void shared_expert_int8_kernel_impl(
771 // stage 0: quantize input to uint8, [M, K]
772 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
773 for (int64_t m = begin; m < end; ++m) {
635 - quantize_row_int8<scalar_t>(
636 - Aq_tmp + m * K,
637 - As_tmp[m],
638 - input + m * K,
639 - K);
774 + quantize_row_int8<scalar_t>(Aq_tmp + m * K, As_tmp[m], input + m * K, K);
775 }
776 });
777
643 - // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
778 + // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
779 const int64_t MB = div_up(M, BLOCK_M);
780 const int64_t NB = div_up(N, BLOCK_N);
781
@@ -651,15 +786,20 @@ void shared_expert_int8_kernel_impl(
786 const int64_t packed_N = get_row_size<int8_t>(N);
787 const int64_t stride_n = packed_K;
788
789 + const bool use_brgemm = can_use_brgemm<int8_t>(M);
790 + const bool apply_scaling_factor = fused_experts_out != nullptr;
791 +
792 // here we only parallel on half of 2N to fuse silu_and_mul with gemm
655 - at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
656 - for (int64_t i = begin; i < end; ++i) {
657 - int64_t mb = i / NB;
658 - int64_t nb = i % NB;
659 -
660 - // nb0 from top half and nb1 from bottom half
661 - int64_t nb0 = nb, nb1 = nb + NB;
662 - int64_t n_size = std::min(N - nb0 * BLOCK_N, BLOCK_N);
793 + parallel_2d(MB, NB, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
794 + // get local pointers
795 + int tid = get_thread_num();
796 + int32_t* __restrict__ C0 = reinterpret_cast<int32_t*>(C_tmp) + tid * 2 * BLOCK_M * BLOCK_N;
797 + int32_t* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
798 +
799 + loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * K * 2, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
800 + // nb_upper from top half and nb_lower from bottom half
801 + int64_t nb_upper = nb, nb_lower = nb + NB;
802 + int64_t n_size = std::min(N - nb * BLOCK_N, BLOCK_N);
803 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
804
805 // A shape [m_size, K]
@@ -667,37 +807,72 @@ void shared_expert_int8_kernel_impl(
807 const float* As = As_tmp + mb * BLOCK_M;
808
809 // B shape [K, n_size] in vnni format
670 - const int8_t* __restrict__ B0 = packed_w1 + nb0 * BLOCK_N * stride_n;
671 - const int8_t* __restrict__ B1 = packed_w1 + nb1 * BLOCK_N * stride_n;
672 - const float* __restrict__ Bs0 = w1s + nb0 * BLOCK_N;
673 - const float* __restrict__ Bs1 = w1s + nb1 * BLOCK_N;
674 -
675 - // fused 1.b: silu_and_mul(A @ B0, A @ B1)
676 - tinygemm_kernel(
677 - /* A */ A,
678 - /* B0 */ B0,
679 - /* B1 */ B1,
680 - /* C */ ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
681 - /* As */ As,
682 - /* Bs0 */ Bs0,
683 - /* Bs1 */ Bs1,
684 - /* M */ m_size,
685 - /* N */ n_size,
686 - /* K */ K,
687 - /* lda */ K,
688 - /* ldb */ n_size,
689 - /* ldc */ N);
810 + const int8_t* __restrict__ B0 = packed_w1 + nb_upper * BLOCK_N * stride_n;
811 + const int8_t* __restrict__ B1 = packed_w1 + nb_lower * BLOCK_N * stride_n;
812 + const float* __restrict__ Bs0 = w1s + nb_upper * BLOCK_N;
813 + const float* __restrict__ Bs1 = w1s + nb_lower * BLOCK_N;
814 +
815 + if (use_brgemm) {
816 + // 1.b gemm: C0 = A @ B0
817 + at::native::cpublas::brgemm(
818 + /* M */ m_size,
819 + /* N */ n_size,
820 + /* K */ K,
821 + /* lda */ K,
822 + /* ldb */ n_size,
823 + /* ldc */ BLOCK_N,
824 + /* add_C */ false,
825 + /* A */ A,
826 + /* B */ B0,
827 + /* C */ C0);
828 +
829 + // 1.c gemm: C1 = A @ B1
830 + at::native::cpublas::brgemm(
831 + /* M */ m_size,
832 + /* N */ n_size,
833 + /* K */ K,
834 + /* lda */ K,
835 + /* ldb */ n_size,
836 + /* ldc */ BLOCK_N,
837 + /* add_C */ false,
838 + /* A */ A,
839 + /* B */ B1,
840 + /* C */ C1);
841 +
842 + const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
843 + const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
844 +
845 + // 1.d silu and mul
846 + silu_and_mul<scalar_t, BLOCK_N>(
847 + ic1 + mb * BLOCK_M * N + nb * BLOCK_N, C0, C1, As, Bs0, Bs1, Bcomp0, Bcomp1, m_size, N);
848 + } else {
849 + // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
850 + tinygemm_kernel(
851 + /* A */ A,
852 + /* B0 */ B0,
853 + /* B1 */ B1,
854 + /* C */ ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
855 + /* As */ As,
856 + /* Bs0 */ Bs0,
857 + /* Bs1 */ Bs1,
858 + /* M */ m_size,
859 + /* N */ n_size,
860 + /* K */ K,
861 + /* lda */ K,
862 + /* ldb */ n_size,
863 + /* ldc */ N);
864 + }
865 + });
866 +
867 + if (use_brgemm) {
868 + at::native::cpublas::brgemm_release();
869 }
870 });
871
872 // stage 1.5: quantize ic1 to uint8, [M * topk, N]
873 at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
874 for (int64_t m = begin; m < end; ++m) {
696 - quantize_row_int8<scalar_t>(
697 - Aq_tmp + m * N,
698 - As_tmp[m],
699 - ic1 + m * N,
700 - N);
875 + quantize_row_int8<scalar_t>(Aq_tmp + m * N, As_tmp[m], ic1 + m * N, N);
876 }
877 });
878
@@ -710,16 +885,13 @@ void shared_expert_int8_kernel_impl(
885 const int64_t stride_oc = packed_N;
886
887 // parallel on [MB2, NB2]
713 - at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
888 + parallel_2d(MB2, NB2, [&](int64_t mb0, int64_t mb1, int64_t nb0, int64_t nb1) {
889 // get local pointers
715 - int tid = at::get_thread_num();
716 - // we won't be using C1 for gemm2
890 + int tid = get_thread_num();
891 float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
892 + int32_t* __restrict__ C32 = reinterpret_cast<int32_t*>(C + BLOCK_M * BLOCK_N);
893
719 - for (int64_t i = begin; i < end; ++i) {
720 - int64_t mb = i / NB2;
721 - int64_t nb = i % NB2;
722 -
894 + loop_2d<int8_t>(mb0, mb1, nb0, nb1, BLOCK_N * IC, [&](int64_t mb, int64_t nb, int64_t nb_offset) {
895 int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
896 int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
897
@@ -731,39 +903,71 @@ void shared_expert_int8_kernel_impl(
903 const int8_t* __restrict__ B = packed_w2 + nb * BLOCK_N * stride_oc;
904 const float* __restrict__ Bs = w2s + nb * BLOCK_N;
905
734 - // 2.a gemm: C = A @ B
735 - tinygemm_kernel<scalar_t>(
736 - /* A */ A,
737 - /* B */ B,
738 - /* C */ C,
739 - /* As */ As,
740 - /* Bs */ Bs,
741 - /* M */ m_size,
742 - /* N */ n_size,
743 - /* K */ IC,
744 - /* lda */ IC,
745 - /* ldb */ n_size,
746 - /* ldc */ BLOCK_N);
906 + if (use_brgemm) {
907 + at::native::cpublas::brgemm(
908 + /* M */ m_size,
909 + /* N */ n_size,
910 + /* K */ IC,
911 + /* lda */ IC,
912 + /* ldb */ n_size,
913 + /* ldc */ BLOCK_N,
914 + /* add_C */ false,
915 + /* A */ A,
916 + /* B */ B,
917 + /* C */ C32);
918 +
919 + // apply scales
920 + const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * IC);
921 + scale_C<BLOCK_N>(C, C32, As, Bs, Bcomp, m_size);
922 + } else {
923 + // 2.a gemm: C = A @ B
924 + tinygemm_kernel<scalar_t>(
925 + /* A */ A,
926 + /* B */ B,
927 + /* C */ C,
928 + /* As */ As,
929 + /* Bs */ Bs,
930 + /* M */ m_size,
931 + /* N */ n_size,
932 + /* K */ IC,
933 + /* lda */ IC,
934 + /* ldb */ n_size,
935 + /* ldc */ BLOCK_N);
936 + }
937
938 // 2.b copy from C to output and add fused_experts_out
939 scalar_t* __restrict__ out = output + mb * BLOCK_M * K + nb * BLOCK_N;
750 - const scalar_t* __restrict__ fused_out = fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N;
940 + const scalar_t* __restrict__ fused_out =
941 + apply_scaling_factor ? fused_experts_out + mb * BLOCK_M * K + nb * BLOCK_N : nullptr;
942 for (int64_t m = 0; m < m_size; ++m) {
752 - add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out + m * K, routed_scaling_factor, n_size);
943 + const scalar_t* __restrict__ fused_out_row = apply_scaling_factor ? (fused_out + m * K) : nullptr;
944 + add_mul_stub(out + m * K, C + m * BLOCK_N, fused_out_row, routed_scaling_factor, n_size);
945 }
946 + });
947 +
948 + if (use_brgemm) {
949 + at::native::cpublas::brgemm_release();
950 }
951 });
952 }
953
758 -#define INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(TYPE) \
759 - template void shared_expert_int8_kernel_impl<TYPE> ( \
760 - TYPE* __restrict__ output, TYPE* __restrict__ ic1, \
761 - float* __restrict__ C_tmp, uint8_t* __restrict__ Aq_tmp, \
762 - float* __restrict__ As_tmp, const TYPE* __restrict__ input, \
763 - const int8_t* __restrict__ packed_w1, const int8_t* __restrict__ packed_w2, \
764 - const float* __restrict__ w1s, const float* __restrict__ w2s, \
765 - const TYPE* __restrict__ fused_experts_out, float routed_scaling_factor, \
766 - int64_t M, int64_t N, int64_t K)
954 +#define INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(TYPE) \
955 + template void shared_expert_int8_kernel_impl<TYPE>( \
956 + TYPE* __restrict__ output, \
957 + TYPE* __restrict__ ic1, \
958 + float* __restrict__ C_tmp, \
959 + uint8_t* __restrict__ Aq_tmp, \
960 + float* __restrict__ As_tmp, \
961 + const TYPE* __restrict__ input, \
962 + const int8_t* __restrict__ packed_w1, \
963 + const int8_t* __restrict__ packed_w2, \
964 + const float* __restrict__ w1s, \
965 + const float* __restrict__ w2s, \
966 + const TYPE* __restrict__ fused_experts_out, \
967 + float routed_scaling_factor, \
968 + int64_t M, \
969 + int64_t N, \
970 + int64_t K)
971
972 INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::BFloat16);
973 INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/vec.h
+155 -37
@@ -1,10 +1,10 @@
1 // Adapted from
2 // https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3
4 -#pragma once
5 -
4 // clang-format off
5
6 +#pragma once
7 +
8 #if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
9 #define CPU_CAPABILITY_AVX512
10 #endif
@@ -16,38 +16,51 @@ namespace {
16
17 using namespace at::vec;
18
19 -template <typename scalar_t,
20 - typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
19 +template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 0>
20 inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) {
21 return at::vec::convert_from_float<scalar_t>(a, b);
22 }
23
24 +// allow f16, bf16
25 +template <typename scalar_t, typename std::enable_if_t<is_reduced_floating_point_v<scalar_t>, int> = 1>
26 +inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const scalar_t* __restrict__ data) {
27 + using bVec = at::vec::Vectorized<scalar_t>;
28 + using fVec = at::vec::Vectorized<float>;
29 + bVec x_vec = bVec::loadu(data);
30 + fVec x0, x1;
31 + std::tie(x0, x1) = at::vec::convert_to_float(x_vec);
32 + return std::make_tuple(x0, x1);
33 +}
34 +
35 +// allow f32
36 +inline std::tuple<Vectorized<float>, Vectorized<float>> load_float_vec2(const float* __restrict__ data) {
37 + using fVec = at::vec::Vectorized<float>;
38 + fVec x0 = fVec::loadu(data);
39 + fVec x1 = fVec::loadu(data + fVec::size());
40 + return std::make_tuple(x0, x1);
41 +}
42 +
43 #if defined(CPU_CAPABILITY_AVX512)
44
45 // `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
46 // use native instruction for bfloat16->float32 conversion
47 template <>
30 -inline Vectorized<at::BFloat16> convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
48 +inline Vectorized<at::BFloat16>
49 +convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
50 return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
51 }
52
34 -#define CVT_BF16_TO_FP32(a) \
35 - _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
53 +#define CVT_BF16_TO_FP32(a) _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
54
37 -#define CVT_FP16_TO_FP32(a) \
38 - _mm512_cvtps_ph(a, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC))
55 +#define CVT_FP16_TO_FP32(a) _mm512_cvtph_ps(a)
56
57 // this doesn't handle NaN.
58 inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
59 const __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
43 -
44 - const __m512i mant = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x07)), 4);
45 - const __m512i raw_exp = _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x78)), 3);
46 - const __m512i exp = _mm512_slli_epi16(_mm512_add_epi16(raw_exp, _mm512_set1_epi16(120)), 7);
47 - const __m512i nonsign = _mm512_or_si512(exp, mant);
48 -
49 - const __m512i sign = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(0x80)), 8);
50 - const __m512i combined = _mm512_or_si512(nonsign, sign);
60 + __m512i combined = _mm512_add_epi16(x, _mm512_set1_epi16(0x0780));
61 + combined = _mm512_slli_epi16(combined, 4);
62 + combined = _mm512_and_si512(combined, _mm512_set1_epi16(0x87f0));
63 + combined = _mm512_add_epi16(combined, _mm512_set1_epi16(0x3c00));
64
65 const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
66 return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
@@ -113,10 +126,77 @@ inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
126 #endif
127 }
128
129 +// faster version of float8_e4m3fn conversion to bfloat16
130 +//
131 +// we mapped cuda implementation from below link and vectorized with avx512:
132 +// https://github.com/thu-pacman/chitu/blob/1ed2078ec26581ebdca05b7306d4385f86edaa7c/csrc/cuda/marlin/marlin_gemm/dequant.h#L387
133 +//
134 +inline __attribute__((always_inline)) __m512bh CVT_FP8_TO_BF16_EXT(__m256i a) {
135 + const __m512i mask0 = _mm512_set1_epi16(0x80); // sign bit
136 + const __m512i mask1 = _mm512_set1_epi16(0x7F); // exponent and mantissa
137 + const __m512i mask2 = _mm512_set1_epi16(0x4000);
138 +
139 + __m512i x = _mm512_cvtepu8_epi16(a);
140 + __m512i vsign = _mm512_and_si512(x, mask0);
141 + vsign = _mm512_slli_epi16(vsign, 8);
142 +
143 + __m512i vexp_and_mant = _mm512_and_si512(x, mask1);
144 + vexp_and_mant = _mm512_slli_epi16(vexp_and_mant, 4);
145 +
146 + // _MM_TERNLOG_A | _MM_TERNLOG_B | _MM_TERNLOG_C: 0b11111110
147 + return (__m512bh)(_mm512_ternarylogic_epi32(vsign, mask2, vexp_and_mant, 0b11111110));
148 +}
149 +
150 +// bias for conversion of fp8 to bf16 1/256 in float32
151 +#define kFP8_BIAS 0x3b800000
152 +
153 +// remove warning: ignoring attributes on template argument ‘__m512bh’ [-Wignored-attributes]
154 +#pragma GCC diagnostic push
155 +#pragma GCC diagnostic ignored "-Wignored-attributes"
156 +
157 +#define MXFP4_VALUES \
158 + -6.0f, -4.0f, -3.0f, -2.0f, -1.5f, -1.0f, -0.5f, -0.0f, 6.0f, 4.0f, 3.0f, 2.0f, 1.5f, 1.0f, 0.5f, 0.0f
159 +
160 +// convert 64 mxfp4 to 2x bf16 vectors, expect input 32-way packing
161 +inline std::tuple<__m512bh, __m512bh> cvt_mxfp4_e2m1_bf16_intrinsic_lut(__m256i a, __m512i s0, __m512i s1) {
162 + // LUT
163 + const __m512 values = _mm512_set_ps(MXFP4_VALUES);
164 + const __m512i lut = (__m512i)(_mm512_cvtne2ps_pbh(values, values));
165 +
166 + const __m512i abs_mask = _mm512_set1_epi16(0x7FFF);
167 + const __m512i zero = _mm512_setzero_si512();
168 +
169 + // expand values to 16-bit integers
170 + __m512i x0 = _mm512_cvtepu8_epi16(a);
171 + __m512i x1 = _mm512_srli_epi32(x0, 4);
172 +
173 + // LUT to convert mxfp4 values to bf16
174 + x0 = _mm512_permutexvar_epi16(x0, lut);
175 + x1 = _mm512_permutexvar_epi16(x1, lut);
176 +
177 + // check for zeros
178 + __mmask32 mask0 = _mm512_cmp_epi16_mask(_mm512_and_si512(x0, abs_mask), zero, _MM_CMPINT_EQ);
179 + __mmask32 mask1 = _mm512_cmp_epi16_mask(_mm512_and_si512(x1, abs_mask), zero, _MM_CMPINT_EQ);
180 +
181 + // emulate bf16 mul with scale factor
182 + x0 = _mm512_add_epi16(x0, s0);
183 + x1 = _mm512_add_epi16(x1, s1);
184 +
185 + // blend with zero
186 + x0 = _mm512_mask_blend_epi16(mask0, x0, zero);
187 + x1 = _mm512_mask_blend_epi16(mask1, x1, zero);
188 +
189 + return std::make_tuple(__m512bh(x0), __m512bh(x1));
190 +}
191 +
192 +#define CVT_MXFP4_TO_BF16(a, s0, s1) cvt_mxfp4_e2m1_bf16_intrinsic_lut(a, s0, s1)
193 +
194 +#pragma GCC diagnostic pop
195 +
196 #endif
197
198 // vector to scalar reduction
119 -#if defined(CPU_CAPABILITY_AVX512) && 0
199 +#if defined(CPU_CAPABILITY_AVX512)
200 inline float vec_reduce_sum(const Vectorized<float>& a) {
201 return _mm512_reduce_add_ps(__m512(a));
202 }
@@ -136,10 +216,9 @@ inline float vec_reduce_max(const Vectorized<float>& a) {
216
217 // https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
218 template <typename scalar_t>
139 -inline void quantize_row_int8(uint8_t* __restrict__ Aq, float& As,
140 - const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
141 -
142 - float amax = 0.f; // absolute max
219 +inline void
220 +quantize_row_int8(uint8_t* __restrict__ Aq, float& As, const scalar_t* __restrict__ A, int64_t K, float eps = 1e-7) {
221 + float amax = 0.f; // absolute max
222 for (int64_t k = 0; k < K; ++k) {
223 const float val = static_cast<float>(A[k]);
224 amax = std::max(amax, std::abs(val));
@@ -158,9 +237,8 @@ inline void quantize_row_int8(uint8_t* __restrict__ Aq, float& As,
237
238 #if defined(CPU_CAPABILITY_AVX512)
239 template <>
161 -inline void quantize_row_int8<at::BFloat16>(uint8_t* __restrict__ Aq, float& As,
162 - const at::BFloat16* __restrict__ A, int64_t K, float eps) {
163 -
240 +inline void quantize_row_int8<at::BFloat16>(
241 + uint8_t* __restrict__ Aq, float& As, const at::BFloat16* __restrict__ A, int64_t K, float eps) {
242 const __m512 signBit = _mm512_set1_ps(-0.0f);
243 const __m512i off = _mm512_set1_epi32(128);
244
@@ -200,7 +278,7 @@ inline void quantize_row_int8<at::BFloat16>(uint8_t* __restrict__ Aq, float& As,
278 // transpose utils
279 // taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
280 #if defined(CPU_CAPABILITY_AVX512)
203 -inline void transpose_16x16_32bit(__m512i * v) {
281 +inline void transpose_16x16_32bit(__m512i* v) {
282 __m512i v1[16];
283 v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
284 v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
@@ -293,16 +371,56 @@ inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1)
371 }
372 #pragma GCC diagnostic pop
373
296 -#endif
297 -
298 -// TODO: debug print, remove me later
299 -template<typename scalar_t>
300 -void print_array(scalar_t* ptr, int size) {
301 - for (int d = 0; d < size; ++d) {
302 - if (d % 16 == 0) { std::cout << std::endl; }
303 - std::cout << ptr[d] << " ";
304 - }
305 - std::cout << std::endl;
374 +inline __attribute__((always_inline)) __m512 _mm512_fexp_u20_ps(const __m512 values) {
375 + const __m512 vec_c0 = _mm512_set1_ps(0.00010703434948458272f);
376 + const __m512 vec_c1 = _mm512_set1_ps(0.30354260500649682f);
377 + const __m512 vec_c2 = _mm512_set1_ps(-0.22433836478672356);
378 + const __m512 vec_c3 = _mm512_set1_ps(-0.079204240219773236);
379 +
380 + const __m512 vec_exp_log2ef = _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); // log2(e)
381 +
382 + const __m512 vec_a = _mm512_set1_ps(std::pow(2, 23) / std::log2(2));
383 + const __m512 vec_b = _mm512_set1_ps(std::pow(2, 23) * 127.f);
384 +
385 + const __m512 vec_ln_flt_min = _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50));
386 + const __m512 vec_ln_flt_max = _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218));
387 + __m512i vec_infinity = _mm512_set1_epi32(0x7F800000);
388 + __m512i vec_zero = _mm512_setzero_epi32();
389 +
390 + // Fast Exponential Computation on SIMD Architectures
391 + // A. Cristiano I. Malossi, Yves Ineichen, Costas Bekas, and Alessandro
392 + // Curioni exp(x) = 2**(x * log2(e))
393 + // = 2**xi * 2**xf - TIPS we are using the EEEE floating point
394 + // representation with identification to the exponent and the
395 + // mentissa
396 + // 2**xf will be approximated to a polynomial of degree 3 computed with
397 + // Horner method
398 + // mask for the boundary condition
399 + auto min_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_min, _CMP_LT_OS);
400 + auto max_mask = _mm512_cmp_ps_mask(values, vec_ln_flt_max, _CMP_GT_OS);
401 +
402 + // transformation with log2(e)
403 + auto vec_src = _mm512_mul_ps(values, vec_exp_log2ef);
404 + auto vec_fractional = _mm512_sub_ps(vec_src, _mm512_floor_ps(vec_src));
405 +
406 + // compute polynomial using Horner Scheme, for superscalar processor
407 + auto vec_res = _mm512_fmadd_ps(vec_fractional, vec_c3, vec_c2);
408 + vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c1);
409 + vec_res = _mm512_fmadd_ps(vec_fractional, vec_res, vec_c0);
410 +
411 + vec_src = _mm512_sub_ps(vec_src, vec_res);
412 + // the tips is here, headache in perspective
413 + auto tmp = _mm512_fmadd_ps(vec_a, vec_src, vec_b);
414 + // headache bis - we loose precision with the cast but it "fits", but ok
415 + // after f32 -> f16 later
416 + __m512i casted_integer = _mm512_cvttps_epi32(tmp);
417 + // boundary condition, lower than the min -> 0
418 + casted_integer = _mm512_mask_mov_epi32(casted_integer, min_mask, vec_zero);
419 + // boundary condition, larger than the max -> +oo
420 + casted_integer = _mm512_mask_mov_epi32(casted_integer, max_mask, vec_infinity);
421 + // final interpretation to float
422 + return _mm512_castsi512_ps(casted_integer);
423 }
424 +#endif
425
308 -} // anonymous namespace
426 +} // anonymous namespace
csrc/cpu/sgl-kernels/vec_pack.h new
+299
@@ -0,0 +1,299 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +// clang-format off
5 +
6 +// To use the transpose functions
7 +#include <ATen/native/cpu/utils.h>
8 +
9 +#include "vec.h"
10 +
11 +namespace {
12 +
13 +using namespace at::vec;
14 +
15 +template <typename index_t>
16 +inline index_t get_index(index_t* ind, int i) {
17 + return (ind == nullptr) ? (index_t)i : ind[i];
18 +}
19 +
20 +#if defined(CPU_CAPABILITY_AVX512)
21 +// key: from [N, 32] to [32/2, N, 2]
22 +template <typename scalar_t, typename index_t>
23 +inline void pack_vnni_Nx32(
24 + scalar_t* __restrict__ dst,
25 + const scalar_t* __restrict__ src,
26 + const index_t* __restrict__ ind,
27 + int N,
28 + int ld_src,
29 + int ld_dst) {
30 + __m512i vinputs[16];
31 +
32 + int n = 0;
33 + for (; n < N; ++n) {
34 + index_t index = get_index(ind, n);
35 + vinputs[n] = _mm512_loadu_si512(src + index * ld_src);
36 + }
37 + // padding with zero to avoid uninitialized vectors
38 + for (; n < 16; ++n) {
39 + vinputs[n] = _mm512_set1_epi32(0);
40 + }
41 +
42 + // pack key
43 + transpose_16x16_32bit(vinputs);
44 +
45 + const __mmask16 vmask = (1 << N) - 1;
46 + for (int k = 0; k < 16; ++k) {
47 + _mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask, vinputs[k]);
48 + }
49 +}
50 +
51 +template <typename scalar_t, typename index_t>
52 +inline void pack_vnni_N_remainder(
53 + scalar_t* __restrict__ dst,
54 + const scalar_t* __restrict__ src,
55 + const index_t* __restrict__ ind,
56 + int N,
57 + int K,
58 + int ld_src,
59 + int ld_dst) {
60 + __m512i vinputs[16];
61 +
62 + int K2 = K >> 1;
63 + const __mmask16 vmask = (1 << K2) - 1;
64 +
65 + int n = 0;
66 + for (; n < N; ++n) {
67 + index_t index = get_index(ind, n);
68 + vinputs[n] = _mm512_maskz_loadu_epi32(vmask, src + index * ld_src);
69 + }
70 + // padding with zero to avoid uninitialized vectors
71 + for (; n < 16; ++n) {
72 + vinputs[n] = _mm512_set1_epi32(0);
73 + }
74 +
75 + // pack key
76 + transpose_16x16_32bit(vinputs);
77 +
78 + const __mmask16 vmask2 = (1 << N) - 1;
79 + for (int k = 0; k < K2; ++k) {
80 + _mm512_mask_storeu_epi32(dst + k * ld_dst * 2, vmask2, vinputs[k]);
81 + }
82 +}
83 +
84 +// value: from [K, 32] to [K/2, 32, 2]
85 +template <typename scalar_t, typename index_t>
86 +inline void pack_vnni_Kx32(
87 + scalar_t* __restrict__ dst,
88 + const scalar_t* __restrict__ src,
89 + const index_t* __restrict__ ind,
90 + int K,
91 + int ld_src,
92 + int ld_dst) {
93 + __m512i vinputs[2];
94 +
95 + int k = 0;
96 + for (; k < K; ++k) {
97 + index_t index = get_index(ind, k);
98 + vinputs[k] = _mm512_loadu_si512(src + index * ld_src);
99 + }
100 + // padding with zero to avoid uninitialized vectors
101 + for (; k < 2; ++k) {
102 + vinputs[k] = _mm512_set1_epi32(0);
103 + }
104 +
105 + // pack value
106 + __m512i d0, d1;
107 + std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
108 + _mm512_storeu_si512(dst + 0 * ld_dst * 2, d0);
109 + _mm512_storeu_si512(dst + 0 * ld_dst * 2 + 32, d1);
110 +}
111 +
112 +template <typename scalar_t, typename index_t>
113 +inline void pack_vnni_K_remainder(
114 + scalar_t* __restrict__ dst,
115 + const scalar_t* __restrict__ src,
116 + const index_t* __restrict__ ind,
117 + int K,
118 + int N,
119 + int ld_src,
120 + int ld_dst) {
121 + __m512i vinputs[2];
122 +
123 + const __mmask32 vmask = (1 << N) - 1;
124 +
125 + int k = 0;
126 + for (; k < K; ++k) {
127 + index_t index = get_index(ind, k);
128 + vinputs[k] = _mm512_maskz_loadu_epi16(vmask, src + index * ld_src);
129 + }
130 + // padding with zero to avoid uninitialized vectors
131 + for (; k < 2; ++k) {
132 + vinputs[k] = _mm512_set1_epi32(0);
133 + }
134 +
135 + // pack value
136 + __m512i d0, d1;
137 + std::tie(d0, d1) = transpose_2x32_16bit(vinputs[0], vinputs[1]);
138 +
139 + if (N <= 16) {
140 + // 2N * 16bits: N * 32bits
141 + const __mmask16 vmask2 = (1 << N) - 1;
142 + _mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2, vmask2, d0);
143 + } else {
144 + // 2(N-16) * 16bits: (N-16) * 32bits
145 + const __mmask16 vmask2 = (1 << (N - 16)) - 1;
146 + _mm512_storeu_epi32(dst + 0 * ld_dst * 2, d0);
147 + _mm512_mask_storeu_epi32(dst + 0 * ld_dst * 2 + 32, vmask2, d1);
148 + }
149 +}
150 +#endif
151 +
152 +// convert to vnni format
153 +// from [N, K/2, 2] to [K/2, N, 2] for bfloat16 and float16
154 +template <typename scalar_t, typename index_t, bool is_indexed>
155 +void pack_vnni(
156 + scalar_t* __restrict__ dst,
157 + const scalar_t* __restrict__ src,
158 + const index_t* __restrict__ ind,
159 + int N,
160 + int K,
161 + int ld_src,
162 + int ld_dst) {
163 +#if defined(CPU_CAPABILITY_AVX512)
164 + const int NB = div_up(N, 16);
165 + const int KB = K / 32;
166 + const int K_remainder = K - KB * 32;
167 +
168 + for (int nb = 0; nb < NB; ++nb) {
169 + int nb_size = std::min(N - nb * 16, 16);
170 + for (int kb = 0; kb < KB; ++kb) {
171 + // handle 16x512bits each block
172 + pack_vnni_Nx32<scalar_t, index_t>(
173 + /* dst */ dst + ((kb * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
174 + /* src */ src + kb * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
175 + /* ind */ is_indexed ? ind + nb * 16 : nullptr,
176 + /* N */ nb_size,
177 + /* ld_src */ ld_src,
178 + /* ld_dst */ ld_dst);
179 + }
180 + if (K_remainder > 0) {
181 + pack_vnni_N_remainder<scalar_t, index_t>(
182 + /* dst */ dst + ((KB * 32) >> 1) * ld_dst * 2 + nb * 16 * 2,
183 + /* src */ src + KB * 32 + (is_indexed ? 0 : nb * 16 * ld_src),
184 + /* ind */ is_indexed ? ind + nb * 16 : nullptr,
185 + /* N */ nb_size,
186 + /* K */ K_remainder,
187 + /* ld_src */ ld_src,
188 + /* ld_dst */ ld_dst);
189 + }
190 + }
191 +#else
192 + for (int n = 0; n < N; ++n) {
193 + index_t index = get_index(ind, n);
194 + for (int k = 0; k < K / 2; ++k) {
195 + for (int d = 0; d < 2; ++d) {
196 + dst[k * ld_dst * 2 + n * 2 + d] = src[index * ld_src + k * 2 + d];
197 + }
198 + }
199 + }
200 +#endif
201 +}
202 +
203 +template <typename scalar_t>
204 +void pack_vnni(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int N, int K, int ld_src, int ld_dst) {
205 + pack_vnni<scalar_t, int32_t, false>(dst, src, nullptr, N, K, ld_src, ld_dst);
206 +}
207 +
208 +template <typename scalar_t, typename index_t>
209 +void pack_vnni(
210 + scalar_t* __restrict__ dst,
211 + const scalar_t* __restrict__ src,
212 + const index_t* __restrict__ ind,
213 + int N,
214 + int K,
215 + int ld_src,
216 + int ld_dst) {
217 + assert(ind != nullptr);
218 + pack_vnni<scalar_t, index_t, true>(dst, src, ind, N, K, ld_src, ld_dst);
219 +}
220 +
221 +// convert to vnni format
222 +// from [K/2, 2, N] to [K/2, N, 2] for bfloat16 and float16
223 +template <typename scalar_t, typename index_t, bool is_indexed>
224 +void pack_vnni2(
225 + scalar_t* __restrict__ dst,
226 + const scalar_t* __restrict__ src,
227 + const index_t* __restrict__ ind,
228 + int K,
229 + int N,
230 + int ld_src,
231 + int ld_dst) {
232 +#if defined(CPU_CAPABILITY_AVX512)
233 + const int KB = div_up(K, 2);
234 + const int NB = N / 32;
235 + const int N_remainder = N - NB * 32;
236 +
237 + for (int kb = 0; kb < KB; ++kb) {
238 + int kb_size = std::min(K - kb * 2, 2);
239 + for (int nb = 0; nb < NB; ++nb) {
240 + // handle 2x512bits each block
241 + pack_vnni_Kx32<scalar_t, index_t>(
242 + /* dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + nb * 32 * 2,
243 + /* src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + nb * 32,
244 + /* ind */ is_indexed ? ind + kb * 2 : nullptr,
245 + /* K */ kb_size,
246 + /* ld_src */ ld_src,
247 + /* ld_dst */ ld_dst);
248 + }
249 + if (N_remainder > 0) {
250 + pack_vnni_K_remainder(
251 + /* dst */ dst + ((kb * 2) >> 1) * ld_dst * 2 + NB * 32 * 2,
252 + /* src */ src + (is_indexed ? 0 : kb * 2 * ld_src) + NB * 32,
253 + /* ind */ is_indexed ? ind + kb * 2 : nullptr,
254 + /* K */ kb_size,
255 + /* N */ N_remainder,
256 + /* ld_src */ ld_src,
257 + /* ld_dst */ ld_dst);
258 + }
259 + }
260 +#else
261 + int k = 0;
262 + for (; k < (K >> 1) * 2; k += 2) {
263 + index_t index0 = get_index(ind, k + 0);
264 + index_t index1 = get_index(ind, k + 1);
265 + for (int n = 0; n < N; ++n) {
266 + dst[(k >> 1) * ld_dst * 2 + n * 2 + 0] = src[index0 * ld_src + n];
267 + dst[(k >> 1) * ld_dst * 2 + n * 2 + 1] = src[index1 * ld_src + n];
268 + }
269 + }
270 + if (K % 2 != 0) {
271 + index_t index = get_index(ind, K - 1);
272 + for (int n = 0; n < N; ++n) {
273 + dst[(K >> 1) * ld_dst * 2 + n * 2 + 0] = src[index * ld_src + n];
274 + dst[(K >> 1) * ld_dst * 2 + n * 2 + 1] = 0;
275 + }
276 + k += 2;
277 + }
278 +#endif
279 +}
280 +
281 +template <typename scalar_t>
282 +void pack_vnni2(scalar_t* __restrict__ dst, const scalar_t* __restrict__ src, int K, int N, int ld_src, int ld_dst) {
283 + pack_vnni2<scalar_t, int32_t, false>(dst, src, nullptr, K, N, ld_src, ld_dst);
284 +}
285 +
286 +template <typename scalar_t, typename index_t>
287 +void pack_vnni2(
288 + scalar_t* __restrict__ dst,
289 + const scalar_t* __restrict__ src,
290 + const index_t* __restrict__ ind,
291 + int K,
292 + int N,
293 + int ld_src,
294 + int ld_dst) {
295 + assert(ind != nullptr);
296 + pack_vnni2<scalar_t, index_t, true>(dst, src, ind, K, N, ld_src, ld_dst);
297 +}
298 +
299 +} // anonymous namespace
csrc/cpu/torch_bindings.cpp
+19 -12
@@ -56,6 +56,8 @@ void shm_send_tensor_list(int64_t handle,
56
57 std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);
58
59 +// SGL CPU kernels
60 +
61 at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
62 const std::optional<at::Tensor>& bias,
63 bool is_vnni);
@@ -65,12 +67,11 @@ at::Tensor convert_weight_packed(at::Tensor& weight);
67 at::Tensor fused_experts_cpu(
68 at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
69 at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
68 - bool use_int8_w8a8, bool use_fp8_w8a16,
69 - const std::optional<at::Tensor>& w1_scale,
70 + int64_t moe_comp_method, const std::optional<at::Tensor>& w1_scale,
71 const std::optional<at::Tensor>& w2_scale,
71 - const std::optional<std::vector<int64_t>> block_size,
72 - const std::optional<at::Tensor>& a1_scale,
73 - const std::optional<at::Tensor>& a2_scale, bool is_vnni);
72 + const std::optional<at::Tensor>& w1_zero,
73 + const std::optional<at::Tensor>& w2_zero,
74 + const std::optional<std::vector<int64_t>> block_size, bool is_vnni);
75
76 at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
77 at::Tensor& scales2,
@@ -86,7 +87,12 @@ at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2,
87
88 // Adapted from sglang: INT4 W4A8 kernels
89 std::tuple<at::Tensor, at::Tensor, at::Tensor> convert_weight_packed_scale_zp(
89 - at::Tensor qweight, at::Tensor qzeros, at::Tensor scales);
90 + at::Tensor qweight, // awq: (*, K, N / 8) || gptq: (*, K / 8, N) , int32
91 + at::Tensor qzeros, // awq: (*, K / group_size, N / 8) || gptq: (*, K /
92 + // group_size, N / 8) , int32
93 + at::Tensor scales, // awq: (*, K / group_size, N) || gptq: (*, K /
94 + // group_size, N) , bfloat16
95 + int64_t quant_method_4bit);
96
97 at::Tensor int4_scaled_mm_cpu(at::Tensor& x, at::Tensor& w, at::Tensor& w_zeros,
98 at::Tensor& w_scales,
@@ -360,10 +366,10 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
366 ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
367 ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
368 ops.def(
363 - "fused_experts_cpu(Tensor! hidden_states, Tensor w1, Tensor w2, Tensor "
364 - "topk_weights, Tensor topk_ids, bool inplace, bool use_int8_w8a8, bool "
365 - "use_fp8_w8a16, Tensor? w1_scale, Tensor? w2_scale, SymInt[]? "
366 - "block_size, Tensor? a1_scale, Tensor? a2_scale, bool is_vnni) -> "
369 + "fused_experts_cpu(Tensor hidden_states, Tensor w1, Tensor w2, Tensor "
370 + "topk_weights, Tensor topk_ids, bool "
371 + "inplace, int moe_comp_method, Tensor? w1_scale, Tensor? w2_scale, "
372 + "Tensor? w1_zero, Tensor? w2_zero, int[]? block_size, bool is_vnni) -> "
373 "Tensor");
374 ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
375 ops.def(
@@ -374,8 +380,9 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
380
381 // Adapted from sglang: INT4 W4A8 kernels
382 ops.def(
377 - "convert_weight_packed_scale_zp(Tensor qweight, Tensor qzeros, "
378 - "Tensor scales) -> (Tensor, Tensor, Tensor)");
383 + "convert_weight_packed_scale_zp(Tensor weight, Tensor qzeros, Tensor "
384 + "scales, int quant_method_4bit) -> (Tensor, "
385 + "Tensor, Tensor)");
386 ops.impl("convert_weight_packed_scale_zp", torch::kCPU,
387 &convert_weight_packed_scale_zp);
388
pyproject.toml
+1 -1
@@ -125,7 +125,7 @@ extend-exclude = ["tests/models/fixtures/*", "tests/prompts/*", "tests/tokenizer
125 "examples/pooling/token_embed/*", "tests/models/language/pooling/*",
126 "vllm/third_party/*", "vllm/entrypoints/serve/instrumentator/static/*", "tests/entrypoints/openai/speech_to_text/test_transcription_validation.py",
127 "docs/governance/process.md", "docs/assets/contributing/vllm_bench_serve_timeline.html",
128 - "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*"]
128 + "tests/v1/engine/test_fast_incdec_prefix_err.py", ".git/*", "csrc/cpu/sgl-kernels/*"]
129 ignore-hidden = false
130
131 [tool.typos.default]
tests/kernels/moe/test_cpu_fp8_fused_moe.py
+4 -6
@@ -218,13 +218,12 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
218 topk_weight,
219 topk_ids,
220 False,
221 - False,
222 - True,
221 + ops.CPUQuantMethod.FP8_W8A16,
222 w1_s,
223 w2_s,
225 - BLOCK_SIZE,
224 None,
225 None,
226 + BLOCK_SIZE,
227 True,
228 )
229 torch.testing.assert_close(
@@ -242,13 +241,12 @@ def test_w8a16_block_fp8_cpu_fused_moe(M, N, K, E, topk, seed):
241 topk_weight,
242 topk_ids,
243 True,
245 - False,
246 - True,
244 + ops.CPUQuantMethod.FP8_W8A16,
245 w1_s,
246 w2_s,
249 - BLOCK_SIZE,
247 None,
248 None,
249 + BLOCK_SIZE,
250 True,
251 )
252 torch.testing.assert_close(out_inplace, out, atol=0, rtol=0)
tests/kernels/test_awq_int4_to_int8.py
+34 -27
@@ -20,7 +20,11 @@ import numpy as np
20 import pytest
21 import torch
22
23 -from vllm._custom_ops import _supports_cpu_w4a8_int8
23 +from vllm._custom_ops import (
24 + CPUQuantAlgo,
25 + convert_weight_packed_scale_zp,
26 + int4_scaled_mm_cpu,
27 +)
28 from vllm.model_executor.layers.quantization.utils.quant_utils import (
29 pack_cols,
30 )
@@ -29,11 +33,6 @@ from vllm.platforms import current_platform
33 if not current_platform.is_cpu():
34 pytest.skip("skipping CPU-only tests", allow_module_level=True)
35
32 -requires_cpu_w4a8_int8 = pytest.mark.skipif(
33 - not _supports_cpu_w4a8_int8,
34 - reason="Requires vLLM CPU build with SGLang INT4 W4A8 kernels",
35 -)
36 -
36
37 def make_awq_checkpoint_data(K, N, group_size, seed=42):
38 """Create synthetic AWQ checkpoint data in packed int32 format.
@@ -87,7 +86,6 @@ def make_awq_checkpoint_data(K, N, group_size, seed=42):
86 class TestConvertWeightPackedScaleZp:
87 """Tests for convert_weight_packed_scale_zp weightpacking."""
88
90 - @requires_cpu_w4a8_int8
89 @pytest.mark.parametrize(
90 "K,N,group_size",
91 [
@@ -102,8 +100,11 @@ class TestConvertWeightPackedScaleZp:
100 K, N, group_size
101 )
102
105 - blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
106 - packed_qweight, packed_qzeros, scales
103 + blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp(
104 + packed_qweight,
105 + packed_qzeros,
106 + scales,
107 + CPUQuantAlgo.AWQ,
108 )
109
110 block_n = 32
@@ -129,7 +130,6 @@ class TestConvertWeightPackedScaleZp:
130 class TestInt4ScaledMmCpu:
131 """Tests for int4_scaled_mm_cpu GEMM kernel."""
132
132 - @requires_cpu_w4a8_int8
133 @pytest.mark.parametrize(
134 "M,K,N,group_size",
135 [
@@ -146,12 +146,15 @@ class TestInt4ScaledMmCpu:
146 make_awq_checkpoint_data(K, N, group_size)
147 )
148
149 - blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
150 - packed_qweight, packed_qzeros, scales
149 + blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp(
150 + packed_qweight,
151 + packed_qzeros,
152 + scales,
153 + CPUQuantAlgo.AWQ,
154 )
155
156 x = torch.randn(M, K, dtype=torch.bfloat16)
154 - out = torch.ops._C.int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None)
157 + out = int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None)
158
159 ref_out = torch.mm(x.float(), float_ref)
160
@@ -169,7 +172,6 @@ class TestInt4ScaledMmCpu:
172 )
173 print(f" [PASS] INT4 GEMM correct: M={M}, K={K}, N={N}")
174
172 - @requires_cpu_w4a8_int8
175 @pytest.mark.parametrize("M", [1, 8, 32])
176 def test_gemm_with_bias(self, M):
177 """INT4 W4A8 GEMM with bias should match reference."""
@@ -178,14 +180,17 @@ class TestInt4ScaledMmCpu:
180 make_awq_checkpoint_data(K, N, group_size)
181 )
182
181 - blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
182 - packed_qweight, packed_qzeros, scales
183 + blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp(
184 + packed_qweight,
185 + packed_qzeros,
186 + scales,
187 + CPUQuantAlgo.AWQ,
188 )
189
190 bias = torch.randn(N, dtype=torch.float32)
191 x = torch.randn(M, K, dtype=torch.bfloat16)
192
188 - out = torch.ops._C.int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, bias)
193 + out = int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, bias)
194
195 ref_out = torch.mm(x.float(), float_ref) + bias
196 abs_diff = (out.float() - ref_out).abs()
@@ -197,7 +202,6 @@ class TestInt4ScaledMmCpu:
202 )
203 print(f" [PASS] INT4 GEMM with bias: M={M}")
204
200 - @requires_cpu_w4a8_int8
205 def test_gemm_3d_input(self):
206 """apply() reshapes 3D input [B, S, K] -> [B*S, K] -> back to 3D."""
207 K, N, group_size = 256, 128, 128
@@ -205,17 +209,18 @@ class TestInt4ScaledMmCpu:
209 make_awq_checkpoint_data(K, N, group_size)
210 )
211
208 - blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
209 - packed_qweight, packed_qzeros, scales
212 + blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp(
213 + packed_qweight,
214 + packed_qzeros,
215 + scales,
216 + CPUQuantAlgo.AWQ,
217 )
218
219 B, S = 2, 8
220 x_3d = torch.randn(B, S, K, dtype=torch.bfloat16)
221 x_2d = x_3d.reshape(-1, K)
222
216 - out_2d = torch.ops._C.int4_scaled_mm_cpu(
217 - x_2d, blocked_w, blocked_zp, blocked_s, None
218 - )
223 + out_2d = int4_scaled_mm_cpu(x_2d, blocked_w, blocked_zp, blocked_s, None)
224 out_3d = out_2d.reshape(B, S, N)
225
226 ref_out = torch.mm(x_2d.float(), float_ref).reshape(B, S, N)
@@ -229,7 +234,6 @@ class TestInt4ScaledMmCpu:
234 assert mean_rel < 0.05, f"Mean relative error {mean_rel:.4f} for 3D exceeds 5%"
235 print(f" [PASS] 3D input [{B},{S},{K}] -> output [{B},{S},{N}]")
236
232 - @requires_cpu_w4a8_int8
237 def test_gemm_fp16_input(self):
238 """INT4 GEMM should also work with fp16 input."""
239 K, N, group_size, M = 256, 256, 128, 8
@@ -237,12 +241,15 @@ class TestInt4ScaledMmCpu:
241 make_awq_checkpoint_data(K, N, group_size)
242 )
243
240 - blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
241 - packed_qweight, packed_qzeros, scales
244 + blocked_w, blocked_zp, blocked_s = convert_weight_packed_scale_zp(
245 + packed_qweight,
246 + packed_qzeros,
247 + scales,
248 + CPUQuantAlgo.AWQ,
249 )
250
251 x = torch.randn(M, K, dtype=torch.float16)
245 - out = torch.ops._C.int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None)
252 + out = int4_scaled_mm_cpu(x, blocked_w, blocked_zp, blocked_s, None)
253
254 ref_out = torch.mm(x.float(), float_ref)
255 abs_diff = (out.float() - ref_out).abs()
vllm/_custom_ops.py
+57 -13
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from enum import IntEnum
5 from typing import TYPE_CHECKING, Literal
6
7 import torch
@@ -3150,6 +3151,13 @@ if hasattr(torch.ops._C, "weight_packed_linear"):
3151 )
3152
3153
3154 +class CPUQuantMethod(IntEnum):
3155 + UNQUANT = 0
3156 + INT8_W8A8 = 1
3157 + FP8_W8A16 = 2
3158 + INT4_W4A8 = 3
3159 +
3160 +
3161 if hasattr(torch.ops._C, "fused_experts_cpu"):
3162
3163 @register_fake("_C::fused_experts_cpu")
@@ -3160,13 +3168,12 @@ if hasattr(torch.ops._C, "fused_experts_cpu"):
3168 topk_weights: torch.Tensor,
3169 topk_ids: torch.Tensor,
3170 inplace: bool,
3163 - use_int8_w8a8: bool,
3164 - use_fp8_w8a16: bool,
3171 + moe_comp_method: CPUQuantMethod,
3172 w1_scale: torch.Tensor | None,
3173 w2_scale: torch.Tensor | None,
3174 + w1_zero: torch.Tensor | None,
3175 + w2_zero: torch.Tensor | None,
3176 block_size: list[int] | None,
3168 - a1_scale: torch.Tensor | None,
3169 - a2_scale: torch.Tensor | None,
3177 is_vnni: bool,
3178 ) -> torch.Tensor:
3179 return torch.empty_like(hidden_states)
@@ -3179,13 +3186,12 @@ def fused_experts_cpu(
3186 topk_weights: torch.Tensor,
3187 topk_ids: torch.Tensor,
3188 inplace: bool,
3182 - use_int8_w8a8: bool,
3183 - use_fp8_w8a16: bool,
3189 + moe_comp_method: CPUQuantMethod,
3190 w1_scale: torch.Tensor | None,
3191 w2_scale: torch.Tensor | None,
3192 + w1_zero: torch.Tensor | None,
3193 + w2_zero: torch.Tensor | None,
3194 block_size: list[int] | None,
3187 - a1_scale: torch.Tensor | None,
3188 - a2_scale: torch.Tensor | None,
3195 is_vnni: bool,
3196 ) -> torch.Tensor:
3197 return torch.ops._C.fused_experts_cpu(
@@ -3195,13 +3201,12 @@ def fused_experts_cpu(
3201 topk_weights,
3202 topk_ids,
3203 inplace,
3198 - use_int8_w8a8,
3199 - use_fp8_w8a16,
3204 + moe_comp_method,
3205 w1_scale,
3206 w2_scale,
3207 + w1_zero,
3208 + w2_zero,
3209 block_size,
3203 - a1_scale,
3204 - a2_scale,
3210 is_vnni,
3211 )
3212
@@ -3222,6 +3227,11 @@ if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
3227 return torch.empty((M, N), dtype=out_dtype)
3228
3229
3230 +class CPUQuantAlgo(IntEnum):
3231 + AWQ = 0
3232 + GPTQ = 1
3233 +
3234 +
3235 if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"):
3236
3237 @register_fake("_C::convert_weight_packed_scale_zp")
@@ -3229,6 +3239,7 @@ if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"):
3239 qweight: torch.Tensor,
3240 qzeros: torch.Tensor,
3241 scales: torch.Tensor,
3242 + quant_method_4bit: CPUQuantAlgo,
3243 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
3244 return (
3245 torch.empty_like(qweight),
@@ -3237,6 +3248,20 @@ if hasattr(torch.ops._C, "convert_weight_packed_scale_zp"):
3248 )
3249
3250
3251 +def convert_weight_packed_scale_zp(
3252 + qweight: torch.Tensor,
3253 + qzeros: torch.Tensor,
3254 + scales: torch.Tensor,
3255 + quant_method_4bit: CPUQuantAlgo,
3256 +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
3257 + return torch.ops._C.convert_weight_packed_scale_zp(
3258 + qweight,
3259 + qzeros,
3260 + scales,
3261 + quant_method_4bit,
3262 + )
3263 +
3264 +
3265 if hasattr(torch.ops._C, "int4_scaled_mm_cpu"):
3266
3267 @register_fake("_C::int4_scaled_mm_cpu")
@@ -3251,7 +3276,26 @@ if hasattr(torch.ops._C, "int4_scaled_mm_cpu"):
3276 return torch.empty((x.size(0), N), dtype=x.dtype, device=x.device)
3277
3278
3254 -_supports_cpu_w4a8_int8 = bool(hasattr(torch.ops._C, "convert_weight_packed_scale_zp"))
3279 +def int4_scaled_mm_cpu(
3280 + x: torch.Tensor,
3281 + w: torch.Tensor,
3282 + w_zeros: torch.Tensor,
3283 + w_scales: torch.Tensor,
3284 + bias: torch.Tensor | None,
3285 +) -> torch.Tensor:
3286 + x_shape = x.shape
3287 + x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
3288 +
3289 + out = torch.ops._C.int4_scaled_mm_cpu(
3290 + x_2d,
3291 + w,
3292 + w_zeros,
3293 + w_scales,
3294 + bias,
3295 + )
3296 + out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
3297 + return out
3298 +
3299
3300 if hasattr(torch.ops._C, "fp8_scaled_mm_cpu"):
3301
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py
+15 -12
@@ -7,7 +7,12 @@ import torch
7 from torch.nn import functional as F
8
9 from vllm import _custom_ops as ops
10 -from vllm._custom_ops import cpu_fused_moe, cpu_prepack_moe_weight
10 +from vllm._custom_ops import (
11 + CPUQuantMethod,
12 + cpu_fused_moe,
13 + cpu_prepack_moe_weight,
14 + fused_experts_cpu,
15 +)
16 from vllm.model_executor.layers.activation import SiluAndMul
17 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
18 from vllm.model_executor.layers.quantization.utils.layer_utils import replace_parameter
@@ -195,23 +200,21 @@ class SGLFusedMOE:
200 e_score_correction_bias=e_score_correction_bias,
201 )
202
198 - torch.ops._C.fused_experts_cpu(
203 + return fused_experts_cpu(
204 x,
205 layer.w13_weight,
206 layer.w2_weight,
207 topk_weights,
208 topk_ids,
204 - True,
205 - False,
206 - False,
207 - None,
208 - None,
209 - None,
210 - None,
211 - None,
212 - True,
209 + False, # inplace
210 + CPUQuantMethod.UNQUANT, # moe_comp_method
211 + None, # w1_scale
212 + None, # w2_scale
213 + None, # w1_zero
214 + None, # w2_zero
215 + None, # block_size
216 + True, # is_vnni
217 )
214 - return x
218
219
220 class CPUFusedMOE:
vllm/model_executor/layers/fused_moe/experts/cpu_moe.py
+5 -5
@@ -5,6 +5,7 @@
5 import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 +from vllm._custom_ops import CPUQuantMethod, fused_experts_cpu
9 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
10 from vllm.model_executor.layers.fused_moe.config import (
11 FusedMoEConfig,
@@ -157,19 +158,18 @@ class CPUExpertsFp8(mk.FusedMoEExpertsMonolithic):
158 )
159 )
160
160 - return torch.ops._C.fused_experts_cpu(
161 + return fused_experts_cpu(
162 hidden_states,
163 w1,
164 w2,
165 topk_weights,
166 topk_ids,
167 False, # inplace
167 - False, # use_int8_w8a8
168 - True, # use_fp8_w8a16
168 + CPUQuantMethod.FP8_W8A16, # moe_comp_method
169 self.w1_scale, # w1_scale
170 self.w2_scale, # w2_scale
171 + None, # w1_zero
172 + None, # w2_zero
173 block_shape, # block_size
172 - None, # a1_scale (W8A16: no activation quantization)
173 - None, # a2_scale
174 True, # is_vnni
175 )
vllm/model_executor/layers/quantization/cpu_wna16.py
+7 -9
@@ -287,8 +287,11 @@ class CPUAWQLinearMethod(LinearMethodBase):
287 packed_weight = layer.qweight.data
288 packed_zeros = layer.qzeros.data
289 scales = layer.scales.data
290 - blocked_w, blocked_zp, blocked_s = torch.ops._C.convert_weight_packed_scale_zp(
291 - packed_weight, packed_zeros, scales
290 + blocked_w, blocked_zp, blocked_s = ops.convert_weight_packed_scale_zp(
291 + packed_weight,
292 + packed_zeros,
293 + scales,
294 + ops.CPUQuantAlgo.AWQ,
295 )
296
297 layer.packed_weight = blocked_w
@@ -334,18 +337,13 @@ class CPUAWQLinearMethod(LinearMethodBase):
337 bias: torch.Tensor | None = None,
338 ) -> torch.Tensor:
339 """SGLang INT4 W4A8 GEMM path."""
337 - x_shape = x.shape
338 - x_2d = x.reshape(-1, x_shape[-1]) if len(x_shape) > 2 else x
339 -
340 - out = torch.ops._C.int4_scaled_mm_cpu(
341 - x_2d,
340 + return ops.int4_scaled_mm_cpu(
341 + x,
342 layer.packed_weight,
343 layer.packed_qzeros,
344 layer.packed_scales,
345 bias,
346 )
347 - out = out.reshape(x_shape[:-1] + (out.size(-1),)) if len(x_shape) > 2 else out
348 - return out
347
348
349 def _get_isa_hint(dtype: torch.dtype) -> str: