@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6cc1e7d96d
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 15, 03:30 UTC

[CPU] Update custom ops for the CPU backend (#20255)

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

Li, Jiang committed Jul 1, 2025 at 15:25 UTC 6cc1e7d96dab6b9c344ec87dec6dc9ab07ad5d21
23 files changed +5357 -101
.buildkite/scripts/hardware_ci/run-cpu-test.sh
+2 -1
@@ -51,6 +51,7 @@ function cpu_tests() {
51 pytest -v -s tests/kernels/attention/test_cache.py -m cpu_model
52 pytest -v -s tests/kernels/attention/test_mla_decode_cpu.py -m cpu_model
53 pytest -v -s tests/models/language/generation -m cpu_model
54 + VLLM_CPU_SGL_KERNEL=1 pytest -v -s tests/models/language/generation -m cpu_model
55 pytest -v -s tests/models/language/pooling -m cpu_model
56 pytest -v -s tests/models/multimodal/generation \
57 --ignore=tests/models/multimodal/generation/test_mllama.py \
@@ -98,4 +99,4 @@ function cpu_tests() {
99
100 # All of CPU tests are expected to be finished less than 40 mins.
101 export -f cpu_tests
101 -timeout 1h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
102 +timeout 1.5h bash -c "cpu_tests $CORE_RANGE $NUMA_NODE"
cmake/cpu_extension.cmake
+20
@@ -96,12 +96,21 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
96 if (CMAKE_CXX_COMPILER_ID STREQUAL "GNU" AND
97 CMAKE_CXX_COMPILER_VERSION VERSION_GREATER_EQUAL 12.3)
98 list(APPEND CXX_COMPILE_FLAGS "-mavx512bf16")
99 + set(ENABLE_AVX512BF16 ON)
100 else()
101 + set(ENABLE_AVX512BF16 OFF)
102 message(WARNING "Disable AVX512-BF16 ISA support, requires gcc/g++ >= 12.3")
103 endif()
104 else()
105 + set(ENABLE_AVX512BF16 OFF)
106 message(WARNING "Disable AVX512-BF16 ISA support, no avx512_bf16 found in local CPU flags." " If cross-compilation is required, please set env VLLM_CPU_AVX512BF16=1.")
107 endif()
108 +
109 + find_isa(${CPUINFO} "avx512_vnni" AVX512VNNI_FOUND)
110 + if (AVX512VNNI_FOUND)
111 + list(APPEND CXX_COMPILE_FLAGS "-mavx512vnni")
112 + set(ENABLE_AVX512VNNI ON)
113 + endif()
114
115 elseif (AVX2_FOUND)
116 list(APPEND CXX_COMPILE_FLAGS "-mavx2")
@@ -231,6 +240,17 @@ if (AVX512_FOUND AND NOT AVX512_DISABLED)
240 "csrc/cpu/quant.cpp"
241 "csrc/cpu/shm.cpp"
242 ${VLLM_EXT_SRC})
243 + if (ENABLE_AVX512BF16 AND ENABLE_AVX512VNNI)
244 + set(VLLM_EXT_SRC
245 + "csrc/cpu/sgl-kernels/gemm.cpp"
246 + "csrc/cpu/sgl-kernels/gemm_int8.cpp"
247 + "csrc/cpu/sgl-kernels/gemm_fp8.cpp"
248 + "csrc/cpu/sgl-kernels/moe.cpp"
249 + "csrc/cpu/sgl-kernels/moe_int8.cpp"
250 + "csrc/cpu/sgl-kernels/moe_fp8.cpp"
251 + ${VLLM_EXT_SRC})
252 + add_compile_definitions(-DCPU_CAPABILITY_AVX512)
253 + endif()
254 elseif(POWER10_FOUND)
255 set(VLLM_EXT_SRC
256 "csrc/cpu/quant.cpp"
csrc/cpu/sgl-kernels/common.h new
+238
@@ -0,0 +1,238 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +#pragma once
5 +
6 +#include <ATen/ATen.h>
7 +#include <ATen/Parallel.h>
8 +#include <ATen/record_function.h>
9 +
10 +// clang-format off
11 +
12 +#if defined(_OPENMP)
13 +#include <omp.h>
14 +#endif
15 +
16 +namespace {
17 +
18 +// 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 + } \
28 + }()
29 +
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 + } \
53 + }()
54 +
55 +#define UNUSED(x) (void)(x)
56 +
57 +#define CHECK_CPU(x) TORCH_CHECK(x.device().type() == at::kCPU, #x " must be a CPU tensor")
58 +
59 +#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
60 +#define CHECK_LAST_DIM_CONTIGUOUS(x) \
61 + TORCH_CHECK(x.strides()[x.strides().size() - 1] == 1, #x "must be contiguous at last dimention")
62 +
63 +#define CHECK_INPUT(x) \
64 + CHECK_CPU(x); \
65 + CHECK_CONTIGUOUS(x)
66 +#define CHECK_LAST_DIM_CONTIGUOUS_INPUT(x) \
67 + CHECK_CPU(x); \
68 + CHECK_LAST_DIM_CONTIGUOUS(x)
69 +
70 +#define CHECK_DIM(d, x) TORCH_CHECK(x.dim() == d, #x " must be a " #d "D tensor")
71 +
72 +#define CHECK_EQ(a, b) TORCH_CHECK((a) == (b), "CHECK_EQ(" #a ", " #b ") failed. ", a, " vs ", b)
73 +
74 +// parallel routines
75 +constexpr int GRAIN_SIZE = 1024;
76 +
77 +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; }
79 +
80 +template <typename T>
81 +inline void balance211(T n, T nth, T ith, T& n_start, T& n_end) {
82 +#if 0
83 + // onednn partition pattern
84 + T& n_my = n_end;
85 + if (nth <= 1 || n == 0) {
86 + n_start = 0;
87 + n_my = n;
88 + } else {
89 + T n1 = div_up(n, nth);
90 + T n2 = n1 - 1;
91 + T T1 = n - n2 * nth;
92 + n_my = ith < T1 ? n1 : n2;
93 + n_start = ith <= T1 ? ith*n1 : T1 * n1 + (ith - T1) * n2;
94 + }
95 + n_end += n_start;
96 +#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);
101 +#endif
102 +}
103 +
104 +template <typename func_t>
105 +inline void parallel_for(int n, const func_t& f) {
106 +#if defined(_OPENMP)
107 +#pragma omp parallel
108 +{
109 + int nth = omp_get_num_threads();
110 + int ith = omp_get_thread_num();
111 + int tbegin, tend;
112 + balance211(n, nth, ith, tbegin, tend);
113 + f(tbegin, tend);
114 +}
115 +#else
116 + f(0, n);
117 +#endif
118 +}
119 +
120 +// for 1d parallel, use `actual_nth`
121 +// for 2d parallel, use even nths, e.g. 43->42
122 +int inline adjust_num_threads(int m) {
123 + int actual_nth = at::get_num_threads();
124 + if (m == 1) {
125 + return actual_nth;
126 + }
127 + return std::max(1, (actual_nth >> 1) * 2);
128 +}
129 +
130 +template <typename func_t>
131 +inline void parallel_2d(int m, int n, const func_t& f) {
132 +
133 + // make sure we have even num_threads
134 + int nth = adjust_num_threads(m);
135 +
136 + // [NOTE] thread blocking:
137 + //
138 + // 1) prefer square block per thread
139 + // 2) use even number of CPU cores
140 + // 3) use all `num_threads` cores
141 + //
142 + // we have:
143 + // TM * TN = T
144 + // BM / TM = BN / TN
145 + // then:
146 + // TM = ((BM / BN) * T) ^ 0.5
147 + //
148 + float r = float(m) / n;
149 + int nth_m = std::ceil(std::sqrt(r * nth));
150 + int nth_n = 1;
151 + for (; nth_m > 0; --nth_m) {
152 + nth_n = nth / nth_m;
153 + if (nth_m * nth_n == nth) {
154 + break;
155 + }
156 + }
157 +
158 +#if defined(_OPENMP)
159 +#pragma omp parallel num_threads(nth)
160 +{
161 + int ith = omp_get_thread_num();
162 + int ith_m = ith / nth_n;
163 + int ith_n = ith % nth_n;
164 +
165 + int thread_block_m = div_up(m, nth_m);
166 + int thread_block_n = div_up(n, nth_n);
167 +
168 + int begin_m = ith_m * thread_block_m;
169 + int end_m = std::min(m, begin_m + thread_block_m);
170 + int begin_n = ith_n * thread_block_n;
171 + int end_n = std::min(n, begin_n + thread_block_n);
172 +
173 + f(begin_m, end_m, begin_n, end_n);
174 +}
175 +#else
176 + f(0, m, 0, n);
177 +#endif
178 +}
179 +
180 +template <typename T>
181 +int get_cache_blocks(int BLOCK_SIZE, int K) {
182 + // L2 2MB and ratio of 50%
183 + const int L2_size = 2048 * 1024 >> 1;
184 + return std::max(1, int(L2_size / (BLOCK_SIZE * K * sizeof(T))));
185 +}
186 +
187 +// data indexing for dimension collapse
188 +template <typename T>
189 +inline T data_index_init(T offset) {
190 + return offset;
191 +}
192 +
193 +template <typename T, typename... Args>
194 +inline T data_index_init(T offset, T& x, const T& X, Args&&... args) {
195 + offset = data_index_init(offset, std::forward<Args>(args)...);
196 + x = offset % X;
197 + return offset / X;
198 +}
199 +
200 +inline bool data_index_step() {
201 + return true;
202 +}
203 +
204 +template <typename T, typename... Args>
205 +inline bool data_index_step(T& x, const T& X, Args&&... args) {
206 + if (data_index_step(std::forward<Args>(args)...)) {
207 + x = ((x + 1) == X) ? 0 : (x + 1);
208 + return x == 0;
209 + }
210 + return false;
211 +}
212 +
213 +// forced unroll for perf critical path
214 +
215 +#if __has_attribute(always_inline)
216 +#define ALWAYS_INLINE __attribute__((__always_inline__)) inline
217 +#else
218 +#define ALWAYS_INLINE inline
219 +#endif
220 +
221 +template <int n>
222 +struct Unroll {
223 + template <typename Func, typename... Args>
224 + ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
225 + Unroll<n - 1>{}(f, args...);
226 + f(std::integral_constant<int, n - 1>{}, args...);
227 + }
228 +};
229 +
230 +template <>
231 +struct Unroll<1> {
232 + template <typename Func, typename... Args>
233 + ALWAYS_INLINE void operator()(const Func& f, Args... args) const {
234 + f(std::integral_constant<int, 0>{}, args...);
235 + }
236 +};
237 +
238 +} // anonymous namespace
csrc/cpu/sgl-kernels/gemm.cpp new
+464
@@ -0,0 +1,464 @@
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 +
8 +// clang-format off
9 +
10 +namespace {
11 +
12 +// packed layout:
13 +// quants {N, K} int8_t
14 +// comp {N} int32_t
15 +template <int BLOCK_N>
16 +inline void s8s8_compensation(int8_t* __restrict__ packed, int K) {
17 +#if defined(CPU_CAPABILITY_AVX512)
18 + constexpr int COLS = BLOCK_N / 16;
19 + __m512i vcomp[COLS];
20 +
21 + for (int col = 0; col < COLS; ++col) {
22 + vcomp[col] = _mm512_setzero_si512();
23 + }
24 +
25 + const int64_t offset = BLOCK_N * K;
26 + const __m512i off = _mm512_set1_epi8(static_cast<char>(0x80));
27 + for (int k = 0; k < K / 4; ++k) {
28 + for (int col = 0; col < COLS; ++col) {
29 + __m512i vb = _mm512_loadu_si512((const __m512i *)(packed + k * BLOCK_N * 4 + col * 64));
30 + vcomp[col] = _mm512_dpbusd_epi32(vcomp[col], off, vb);
31 + }
32 + }
33 +
34 + for (int col = 0; col < COLS; ++col) {
35 + _mm512_storeu_si512((__m512i *)(packed + offset + col * 64), vcomp[col]);
36 + }
37 +#else
38 + TORCH_CHECK(false, "s8s8_compensation not implemented!");
39 +#endif
40 +}
41 +
42 +// convert to vnni format
43 +// from [N, K] to [K/2, N, 2] for bfloat16 and float16
44 +template <typename packed_t>
45 +inline void pack_vnni(packed_t* __restrict__ packed, const packed_t* __restrict__ weight, int N, int K) {
46 + const int VNNI_BLK = 2;
47 + for (int n = 0; n < N; ++n) {
48 + for (int k = 0; k < K / VNNI_BLK; ++k) {
49 + for (int d = 0; d < VNNI_BLK; ++d) {
50 + packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
51 + }
52 + }
53 + }
54 +}
55 +
56 +template <>
57 +inline void pack_vnni<int8_t>(int8_t* __restrict__ packed, const int8_t* __restrict__ weight, int N, int K) {
58 + constexpr int BLOCK_N = block_size_n();
59 + TORCH_CHECK(N == BLOCK_N);
60 +
61 + const int VNNI_BLK = 4;
62 + for (int n = 0; n < N; ++n) {
63 + for (int k = 0; k < K / VNNI_BLK; ++k) {
64 + for (int d = 0; d < VNNI_BLK; ++d) {
65 + packed[k * N * VNNI_BLK + n * VNNI_BLK + d] = weight[n * K + k * VNNI_BLK + d];
66 + }
67 + }
68 + }
69 + s8s8_compensation<BLOCK_N>(packed, K);
70 +}
71 +
72 +template <typename scalar_t>
73 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
74 + using bVec = at::vec::Vectorized<scalar_t>;
75 + using fVec = at::vec::Vectorized<float>;
76 + constexpr int kVecSize = bVec::size();
77 +
78 + int64_t d;
79 + #pragma GCC unroll 4
80 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
81 + fVec data0 = fVec::loadu(input + d);
82 + fVec data1 = fVec::loadu(input + d + fVec::size());
83 + bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
84 + out_vec.store(out + d);
85 + }
86 + for (; d < size; ++d) {
87 + out[d] = static_cast<scalar_t>(input[d]);
88 + }
89 +}
90 +
91 +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) {
93 + using bVec = at::vec::Vectorized<scalar_t>;
94 + using fVec = at::vec::Vectorized<float>;
95 + constexpr int kVecSize = bVec::size();
96 +
97 + int64_t d;
98 + #pragma GCC unroll 4
99 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
100 + fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
101 + fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
102 + bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
103 + out_vec.store(out + d);
104 + }
105 + for (; d < size; ++d) {
106 + out[d] = static_cast<scalar_t>(input[d] + bias[d]);
107 + }
108 +}
109 +
110 +template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
111 +struct tinygemm_kernel_nn {
112 + 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) {
115 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
116 + }
117 +};
118 +
119 +#if defined(CPU_CAPABILITY_AVX512)
120 +template <bool has_bias, int BLOCK_M, int BLOCK_N>
121 +struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
122 + 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 +
126 + constexpr int ROWS = BLOCK_M;
127 + constexpr int COLS = BLOCK_N / 16;
128 +
129 + // prefetch distance
130 + constexpr int PREFETCH_SIZE_K = 0;
131 +
132 + __m512bh va;
133 + __m512bh vb[COLS];
134 + __m512 vc[ROWS * COLS];
135 +
136 + auto loadc = [&](auto i) {
137 + constexpr int col = i % COLS;
138 + if constexpr (has_bias) {
139 + vc[i] = _mm512_loadu_ps(bias + col * 16);
140 + } else {
141 + vc[i] = _mm512_set1_ps(0.f);
142 + }
143 + };
144 + Unroll<ROWS * COLS>{}(loadc);
145 +
146 + const int64_t K2 = K >> 1;
147 + const int64_t lda2 = lda >> 1;
148 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
149 + const float* a_ptr = reinterpret_cast<const float*>(A);
150 + const float* b_ptr = reinterpret_cast<const float*>(B);
151 +
152 + auto compute = [&](auto i, int64_t k) {
153 + constexpr int row = i / COLS;
154 + constexpr int col = i % COLS;
155 +
156 + if constexpr (col == 0) {
157 + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
158 + }
159 + if constexpr (row == 0) {
160 + vb[col] = (__m512bh)(_mm512_loadu_si512(b_ptr + k * ldb2 + col * 16));
161 + if constexpr (PREFETCH_SIZE_K > 0) {
162 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
163 + }
164 + }
165 + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
166 + };
167 + for (int64_t k = 0; k < K2; ++k) {
168 + Unroll<ROWS * COLS>{}(compute, k);
169 + }
170 +
171 + auto storec = [&](auto i) {
172 + constexpr int row = i / COLS;
173 + constexpr int col = i % COLS;
174 + // for COLS = 2, 4 use 512bit store
175 + // for COLS = 1, 3 use 256bit store
176 + if constexpr (COLS % 2 == 0) {
177 + if constexpr (col % 2 == 0) {
178 + _mm512_storeu_si512(
179 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
180 + (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
181 + }
182 + } else {
183 + _mm256_storeu_si256(
184 + reinterpret_cast<__m256i*>(C + row * ldc + col * 16),
185 + (__m256i)(_mm512_cvtneps_pbh(vc[i])));
186 + }
187 + };
188 + Unroll<ROWS * COLS>{}(storec);
189 + }
190 +};
191 +#endif
192 +
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);
197 +
198 +template <typename scalar_t, bool has_bias>
199 +struct brgemm {
200 + 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 +
205 + 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);
209 +
210 + // copy from Ctmp to C
211 + for (int64_t m = 0; m < M; ++m) {
212 + if constexpr (has_bias) {
213 + copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
214 + } else {
215 + copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
216 + }
217 + }
218 + }
219 +};
220 +
221 +template <typename scalar_t, bool has_bias>
222 +void tinygemm_kernel(
223 + const scalar_t* __restrict__ A,
224 + const scalar_t* __restrict__ B,
225 + scalar_t* __restrict__ C,
226 + float* __restrict__ Ctmp,
227 + const float* __restrict__ bias,
228 + int64_t M,
229 + int64_t N,
230 + int64_t K,
231 + int64_t lda,
232 + int64_t ldb,
233 + int64_t ldc,
234 + bool brg) {
235 +
236 + if (brg) {
237 + brgemm<scalar_t, has_bias>::apply(
238 + A, B, C, Ctmp, bias,
239 + M, N, K, lda, ldb, ldc);
240 + return;
241 + }
242 +
243 + // pattern: 1-4-16
244 + constexpr int64_t BLOCK_M = 4;
245 + constexpr int64_t BLOCK_N = 64;
246 + const int64_t MB = div_up(M, BLOCK_M);
247 + const int64_t NB = div_up(N, BLOCK_N);
248 + for (int mb = 0; mb < MB; ++mb) {
249 + int64_t mb_start = mb * BLOCK_M;
250 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
251 + for (int64_t nb = 0; nb < NB; ++nb) {
252 + int64_t nb_start = nb * BLOCK_N;
253 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
254 +
255 + switch(mb_size << 4 | nb_size >> 4) {
256 + // mb_size = 1
257 + case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
258 + case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
259 + // mb_size = 2
260 + case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
261 + case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
262 + // mb_size = 3
263 + case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
264 + case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
265 + // 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");
269 + }
270 + }
271 + }
272 +}
273 +
274 +template <typename scalar_t>
275 +void weight_packed_linear_kernel_impl(
276 + scalar_t* __restrict__ out,
277 + const scalar_t* __restrict__ mat1,
278 + const scalar_t* __restrict__ mat2,
279 + const float* __restrict__ bias,
280 + int64_t M,
281 + int64_t N,
282 + int64_t K,
283 + int64_t mat1_strideM,
284 + int64_t out_strideM) {
285 +
286 + constexpr int64_t BLOCK_M = block_size_m();
287 + constexpr int64_t BLOCK_N = block_size_n();
288 + const int64_t MB = div_up(M, BLOCK_M);
289 + const int64_t NB = div_up(N, BLOCK_N);
290 +
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);
296 +
297 + // parallel on [MB, NB]
298 + 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 +
301 + // for brgemm, use float32 for accumulate
302 + alignas(64) float Ctmp[BLOCK_M * BLOCK_N];
303 +
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 +
308 + int64_t mb_start = mb * BLOCK_M;
309 + int64_t mb_size = std::min(M - mb_start, BLOCK_M);
310 + int64_t nb_start = nb * BLOCK_N;
311 + int64_t nb_size = std::min(N - nb_start, BLOCK_N);
312 +
313 + tinygemm_kernel<scalar_t, has_bias>(
314 + /* A */ mat1 + mb_start * mat1_strideM,
315 + /* B */ mat2 + nb_start * K /* nb * BLOCK_N * K */,
316 + /* C */ out + mb_start * out_strideM + nb_start,
317 + /* Ctmp*/ Ctmp,
318 + /* bias*/ bias + nb_start,
319 + /* M */ mb_size,
320 + /* N */ nb_size,
321 + /* K */ K,
322 + /* lda */ mat1_strideM,
323 + /* ldb */ nb_size,
324 + /* ldc */ out_strideM,
325 + /* brg */ use_brgemm);
326 + }}}
327 +
328 + if (use_brgemm) {
329 + at::native::cpublas::brgemm_release();
330 + }
331 + });
332 + });
333 +}
334 +
335 +} // anonymous namespace
336 +
337 +// tinygemm interface
338 +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) {
341 + tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, nullptr, M, N, K, lda, ldb, ldc, brg);
342 +}
343 +
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)
349 +
350 +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
351 +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
352 +
353 +at::Tensor convert_weight_packed(at::Tensor& weight) {
354 + // for 3d moe weights
355 + // weight : [E, OC, IC]
356 + // w1 : [E, 2N, K]
357 + // w2 : [E, K, N]
358 + CHECK_INPUT(weight);
359 +
360 + const int64_t ndim = weight.ndimension();
361 + TORCH_CHECK(ndim == 2 || ndim == 3, "expect weight to be 2d or 3d, got ", ndim, "d tensor.");
362 + const auto st = weight.scalar_type();
363 + const int64_t E = ndim == 3 ? weight.size(0) : 1;
364 + const int64_t OC = ndim == 3 ? weight.size(1) : weight.size(0);
365 + const int64_t IC = ndim == 3 ? weight.size(2) : weight.size(1);
366 +
367 + // we handle 2 TILE_N at a time.
368 + TORCH_CHECK(OC % TILE_N == 0, "invalid weight out features ", OC);
369 + TORCH_CHECK(IC % TILE_K == 0, "invalid weight input features ", IC);
370 +
371 + constexpr int64_t BLOCK_N = block_size_n();
372 + const int64_t NB = div_up(OC, BLOCK_N);
373 +
374 + // use phony sizes here [E, OC, IC], for each [E], [OC, IC] -> [IC / 2, OC, 2]
375 + auto packed_weight = at::empty({}, weight.options());
376 + const int64_t stride = OC * IC;
377 +
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.");
380 +
381 + CPU_DISPATCH_PACKED_TYPES(st, [&] {
382 + // adjust most inner dimension size
383 + const int packed_row_size = get_row_size<packed_t>(IC);
384 + auto sizes = weight.sizes().vec();
385 + sizes[ndim - 1] = packed_row_size;
386 + packed_weight.resize_(sizes);
387 +
388 + const packed_t* w_data = weight.data_ptr<packed_t>();
389 + packed_t* packed_data = packed_weight.data_ptr<packed_t>();
390 +
391 + // parallel on {E, NB}
392 + at::parallel_for(0, E * NB, 0, [&](int64_t begin, int64_t end) {
393 + int64_t e{0}, nb{0};
394 + data_index_init(begin, e, E, nb, NB);
395 +
396 + for (int64_t i = begin; i < end; ++i) {
397 + UNUSED(i);
398 +
399 + int64_t n = nb * BLOCK_N;
400 + int64_t n_size = std::min(BLOCK_N, OC - n);
401 + 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);
406 +
407 + // move to the next index
408 + data_index_step(e, E, nb, NB);
409 + }
410 + });
411 + });
412 + return packed_weight;
413 +}
414 +
415 +// mat1 : [M, K]
416 +// mat2 : [N, K]
417 +// bias : [N]
418 +// out : [M, N]
419 +//
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 +
425 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
426 +
427 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
428 + 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);
434 + CHECK_DIM(2, mat1);
435 + CHECK_DIM(2, mat2);
436 +
437 + auto out = at::empty({M, N}, mat1.options());
438 +
439 + // strides
440 + int64_t mat1_strideM = mat1.stride(0);
441 + int64_t out_strideM = out.stride(0);
442 +
443 + const bool has_bias = bias.has_value();
444 + const float* bias_data = nullptr;
445 + if (has_bias) {
446 + CHECK_EQ(bias.value().size(0), N);
447 + bias_data = bias.value().data_ptr<float>();
448 + }
449 +
450 + AT_DISPATCH_REDUCED_FLOATING_TYPES(mat1.scalar_type(), "weight_packed_linear_kernel_impl", [&] {
451 + weight_packed_linear_kernel_impl<scalar_t>(
452 + out.data_ptr<scalar_t>(),
453 + mat1.data_ptr<scalar_t>(),
454 + packed_w.data_ptr<scalar_t>(),
455 + bias_data,
456 + M,
457 + N,
458 + K,
459 + mat1_strideM,
460 + out_strideM);
461 + });
462 +
463 + return out;
464 +}
csrc/cpu/sgl-kernels/gemm.h new
+266
@@ -0,0 +1,266 @@
1 +#pragma once
2 +
3 +#include <ATen/native/CPUBlas.h>
4 +
5 +// clang-format off
6 +
7 +// amx-bf16
8 +#define TILE_M 16
9 +#define TILE_N 16
10 +#define TILE_K 32
11 +
12 +// 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; }
15 +
16 +// 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 +// TODO: add u8s8 brgemm, this requires PyTorch 2.7
21 +template <> inline bool can_use_brgemm<int8_t>(int M) { return false; }
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; }
24 +
25 +// work around compiler internal error
26 +#define BLOCK_K 128 // 4 * TILE_K
27 +
28 +// adjust leading dimension size for K
29 +template <typename T>
30 +inline int64_t get_row_size(int64_t K) {
31 + return K;
32 +}
33 +
34 +template <>
35 +inline int64_t get_row_size<int8_t>(int64_t K) {
36 + return K + sizeof(int32_t);
37 +}
38 +
39 +inline int64_t get_row_size(int64_t K, bool use_int8_w8a8) {
40 + return use_int8_w8a8 ? K + sizeof(int32_t) : K;
41 +}
42 +
43 +// pack weight to vnni format
44 +at::Tensor convert_weight_packed(at::Tensor& weight);
45 +
46 +// moe implementations for int8 w8a8
47 +template <typename scalar_t>
48 +void fused_experts_int8_kernel_impl(
49 + scalar_t* __restrict__ output,
50 + scalar_t* __restrict__ ic1,
51 + scalar_t* __restrict__ ic2,
52 + uint8_t* __restrict__ A_tmp,
53 + float* __restrict__ C_tmp,
54 + uint8_t* __restrict__ Aq_tmp,
55 + float* __restrict__ As_tmp,
56 + const scalar_t* __restrict__ input,
57 + const int8_t* __restrict__ packed_w1,
58 + const int8_t* __restrict__ packed_w2,
59 + const float* __restrict__ w1s,
60 + const float* __restrict__ w2s,
61 + const float* __restrict__ topk_weights,
62 + const int32_t* __restrict__ sorted_ids,
63 + const int32_t* __restrict__ expert_ids,
64 + const int32_t* __restrict__ offsets,
65 + int64_t M,
66 + int64_t N,
67 + int64_t K,
68 + int64_t E,
69 + int64_t topk,
70 + int64_t num_tokens_post_pad);
71 +
72 +// moe implementations for fp8 w8a16
73 +template <typename scalar_t>
74 +void fused_experts_fp8_kernel_impl(
75 + scalar_t* __restrict__ output,
76 + scalar_t* __restrict__ ic0,
77 + scalar_t* __restrict__ ic1,
78 + scalar_t* __restrict__ ic2,
79 + scalar_t* __restrict__ A_tmp,
80 + scalar_t* __restrict__ B_tmp,
81 + float* __restrict__ C_tmp,
82 + const scalar_t* __restrict__ input,
83 + const at::Float8_e4m3fn* __restrict__ packed_w1,
84 + const at::Float8_e4m3fn* __restrict__ packed_w2,
85 + const float* __restrict__ w1s,
86 + const float* __restrict__ w2s,
87 + int64_t block_size_N,
88 + int64_t block_size_K,
89 + const float* __restrict__ topk_weights,
90 + const int32_t* __restrict__ sorted_ids,
91 + const int32_t* __restrict__ expert_ids,
92 + const int32_t* __restrict__ offsets,
93 + int64_t M,
94 + int64_t N,
95 + int64_t K,
96 + int64_t E,
97 + int64_t topk,
98 + int64_t num_tokens_post_pad);
99 +
100 +// moe implementations for int4 w4a16
101 +template <typename scalar_t>
102 +void fused_experts_int4_w4a16_kernel_impl(
103 + scalar_t* __restrict__ output,
104 + scalar_t* __restrict__ ic0,
105 + scalar_t* __restrict__ ic1,
106 + scalar_t* __restrict__ ic2,
107 + scalar_t* __restrict__ A_tmp,
108 + scalar_t* __restrict__ B_tmp,
109 + float* __restrict__ C_tmp,
110 + const scalar_t* __restrict__ input,
111 + const at::quint4x2* __restrict__ packed_w1,
112 + const at::quint4x2* __restrict__ packed_w2,
113 + const uint8_t* __restrict__ w1z,
114 + const uint8_t* __restrict__ w2z,
115 + const scalar_t* __restrict__ w1s,
116 + const scalar_t* __restrict__ w2s,
117 + int group_size,
118 + const float* __restrict__ topk_weights,
119 + const int32_t* __restrict__ sorted_ids,
120 + const int32_t* __restrict__ expert_ids,
121 + const int32_t* __restrict__ offsets,
122 + int64_t M,
123 + int64_t N,
124 + int64_t K,
125 + int64_t E,
126 + int64_t topk,
127 + int64_t num_tokens_post_pad);
128 +
129 +// shared expert implememntation for int8 w8a8
130 +template <typename scalar_t>
131 +void shared_expert_int8_kernel_impl(
132 + scalar_t* __restrict__ output,
133 + scalar_t* __restrict__ ic1,
134 + float* __restrict__ C_tmp,
135 + uint8_t* __restrict__ Aq_tmp,
136 + float* __restrict__ As_tmp,
137 + const scalar_t* __restrict__ input,
138 + const int8_t* __restrict__ packed_w1,
139 + const int8_t* __restrict__ packed_w2,
140 + const float* __restrict__ w1s,
141 + const float* __restrict__ w2s,
142 + const scalar_t* __restrict__ fused_experts_out,
143 + float routed_scaling_factor,
144 + int64_t M,
145 + int64_t N,
146 + int64_t K);
147 +
148 +template <typename scalar_t>
149 +void shared_expert_fp8_kernel_impl(
150 + scalar_t* __restrict__ output,
151 + scalar_t* __restrict__ ic0,
152 + scalar_t* __restrict__ ic1,
153 + scalar_t* __restrict__ B_tmp,
154 + float* __restrict__ C_tmp,
155 + const scalar_t* __restrict__ input,
156 + const at::Float8_e4m3fn* __restrict__ packed_w1,
157 + const at::Float8_e4m3fn* __restrict__ packed_w2,
158 + const float* __restrict__ w1s,
159 + const float* __restrict__ w2s,
160 + int64_t block_size_N,
161 + int64_t block_size_K,
162 + const scalar_t* __restrict__ fused_experts_out,
163 + float routed_scaling_factor,
164 + int64_t M,
165 + int64_t N,
166 + int64_t K);
167 +
168 +// tinygemm interface
169 +template <typename scalar_t>
170 +void tinygemm_kernel(
171 + const scalar_t* __restrict__ A,
172 + const scalar_t* __restrict__ B,
173 + scalar_t* __restrict__ C,
174 + float* __restrict__ Ctmp,
175 + int64_t M,
176 + int64_t N,
177 + int64_t K,
178 + int64_t lda,
179 + int64_t ldb,
180 + int64_t ldc,
181 + bool brg);
182 +
183 +template <typename scalar_t>
184 +void tinygemm_kernel(
185 + const uint8_t* __restrict__ A,
186 + const int8_t* __restrict__ B,
187 + scalar_t* __restrict__ C,
188 + int32_t* __restrict__ Ctmp,
189 + const float* __restrict__ As,
190 + const float* __restrict__ Bs,
191 + int64_t M,
192 + int64_t N,
193 + int64_t K,
194 + int64_t lda,
195 + int64_t ldb,
196 + int64_t ldc,
197 + bool brg);
198 +
199 +template <typename scalar_t>
200 +void tinygemm_kernel(
201 + const scalar_t* __restrict__ A,
202 + const at::Float8_e4m3fn* __restrict__ B,
203 + scalar_t* __restrict__ C,
204 + scalar_t* __restrict__ Btmp,
205 + float* __restrict__ Ctmp,
206 + const float* __restrict__ scale,
207 + int64_t M,
208 + int64_t N,
209 + int64_t K,
210 + int64_t lda,
211 + int64_t ldb,
212 + int64_t ldc,
213 + bool brg,
214 + int64_t block_size_K);
215 +
216 +template <typename scalar_t>
217 +void tinygemm_kernel(
218 + const scalar_t* __restrict__ A,
219 + const at::quint4x2* __restrict__ B,
220 + scalar_t* __restrict__ C,
221 + const uint8_t* __restrict__ Bz,
222 + const scalar_t* __restrict__ Bs,
223 + scalar_t* __restrict__ Btmp,
224 + float* __restrict__ Ctmp,
225 + int64_t M,
226 + int64_t N,
227 + int64_t K,
228 + int group_size,
229 + int64_t lda,
230 + int64_t ldb,
231 + int64_t ldc,
232 + int64_t strideBz,
233 + int64_t strideBs,
234 + bool brg);
235 +
236 +// TODO: debug print, remove me later
237 +inline void print_16x32i(const __m512i x) {
238 + int32_t a[16];
239 + _mm512_storeu_si512((__m512i *)a, x);
240 +
241 + for (int i = 0; i < 16; i++){
242 + std::cout << a[i] << " ";
243 + }
244 + std::cout << std::endl;
245 +}
246 +
247 +inline void print_16x32(const __m512 x) {
248 + float a[16];
249 + _mm512_storeu_ps((__m512 *)a, x);
250 +
251 + for (int i = 0; i < 16; i++){
252 + std::cout << a[i] << " ";
253 + }
254 + std::cout << std::endl;
255 +}
256 +
257 +
258 +inline void print_32x8u(const __m256i x) {
259 + uint8_t a[32];
260 + _mm256_storeu_si256((__m256i *)a, x);
261 +
262 + for (int i = 0; i < 32; ++i) {
263 + std::cout << int32_t(a[i]) << " ";
264 + }
265 + std::cout << std::endl;
266 +}
csrc/cpu/sgl-kernels/gemm_fp8.cpp new
+530
@@ -0,0 +1,530 @@
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 +
8 +// clang-format off
9 +
10 +// we use 4x32 for BLOCK_M
11 +#define BLOCK_SIZE_M_SCALE 4
12 +
13 +namespace {
14 +
15 +template <typename scalar_t>
16 +inline void copy_stub(scalar_t* __restrict__ out, const float* __restrict__ input, int64_t size) {
17 + using bVec = at::vec::Vectorized<scalar_t>;
18 + using fVec = at::vec::Vectorized<float>;
19 + constexpr int kVecSize = bVec::size();
20 +
21 + int64_t d;
22 + #pragma GCC unroll 4
23 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
24 + fVec data0 = fVec::loadu(input + d);
25 + fVec data1 = fVec::loadu(input + d + fVec::size());
26 + bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
27 + out_vec.store(out + d);
28 + }
29 + for (; d < size; ++d) {
30 + out[d] = static_cast<scalar_t>(input[d]);
31 + }
32 +}
33 +
34 +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) {
36 + using bVec = at::vec::Vectorized<scalar_t>;
37 + using fVec = at::vec::Vectorized<float>;
38 + constexpr int kVecSize = bVec::size();
39 +
40 + int64_t d;
41 + #pragma GCC unroll 4
42 + for (d = 0; d <= size - kVecSize; d += kVecSize) {
43 + fVec data0 = fVec::loadu(input + d) + fVec::loadu(bias + d);
44 + fVec data1 = fVec::loadu(input + d + fVec::size()) + fVec::loadu(bias + d + fVec::size());
45 + bVec out_vec = convert_from_float_ext<scalar_t>(data0, data1);
46 + out_vec.store(out + d);
47 + }
48 + for (; d < size; ++d) {
49 + out[d] = static_cast<scalar_t>(input[d] + bias[d]);
50 + }
51 +}
52 +
53 +inline void unpack_B(
54 + at::BFloat16* __restrict__ Btmp,
55 + const at::Float8_e4m3fn* __restrict__ packed_B,
56 + int N,
57 + int K,
58 + int ldb,
59 + int ldb_tmp,
60 + float scale) {
61 +#if defined(CPU_CAPABILITY_AVX512)
62 + // [K/2, N, 2]
63 + const int K2 = K >> 1;
64 + const int ldb2 = ldb; // ldb * 2 >> 1;
65 + const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(packed_B);
66 + const __m512 vd = _mm512_set1_ps(scale);
67 +
68 + constexpr int BLOCK_N = block_size_n();
69 + static_assert(BLOCK_N == 32);
70 +
71 + // prefetch distance
72 + constexpr int PREFETCH_SIZE_K = 64;
73 +
74 +#pragma GCC unroll 4
75 + for (int k = 0; k < K2; ++k) {
76 + __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2);
77 + if constexpr (PREFETCH_SIZE_K > 0) {
78 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2, _MM_HINT_T0);
79 + }
80 +
81 + __m256i b8_0 = _mm512_extracti32x8_epi32(b8, 0);
82 + __m256i b8_1 = _mm512_extracti32x8_epi32(b8, 1);
83 +
84 + __m512bh bf16_0 = CVT_FP8_TO_BF16(b8_0);
85 + __m512bh bf16_1 = CVT_FP8_TO_BF16(b8_1);
86 +
87 + // Apply scale
88 + __m512 f0_lo = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_0, 0));
89 + __m512 f0_hi = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_0, 1));
90 + __m512 f1_lo = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_1, 0));
91 + __m512 f1_hi = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32((__m512i)bf16_1, 1));
92 +
93 + f0_lo = _mm512_mul_ps(f0_lo, vd);
94 + f0_hi = _mm512_mul_ps(f0_hi, vd);
95 + f1_lo = _mm512_mul_ps(f1_lo, vd);
96 + f1_hi = _mm512_mul_ps(f1_hi, vd);
97 +
98 + bf16_0 = _mm512_cvtne2ps_pbh(f0_hi, f0_lo);
99 + bf16_1 = _mm512_cvtne2ps_pbh(f1_hi, f1_lo);
100 +
101 + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 0, (__m512i)bf16_0);
102 + _mm512_storeu_si512(Btmp + k * ldb_tmp * 2 + 32, (__m512i)bf16_1);
103 + }
104 +#else
105 + TORCH_CHECK(false, "unpack_B: scalar path not implemented!");
106 +#endif
107 +}
108 +
109 +template <typename scalar_t, typename packed_t, bool has_bias, int BLOCK_M, int BLOCK_N>
110 +struct tinygemm_kernel_nn {
111 + 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) {
114 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
115 + }
116 +};
117 +
118 +#if defined(CPU_CAPABILITY_AVX512)
119 +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> {
121 + 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 +
125 + constexpr int ROWS = BLOCK_M;
126 + constexpr int COLS = BLOCK_N / 16;
127 +
128 + const int KB = div_up(K, BLOCK_K);
129 +
130 + // prefetch distance
131 + constexpr int PREFETCH_SIZE_K = 64;
132 + constexpr int PREFETCH_SIZE_KB = 1;
133 +
134 + __m512bh va;
135 + __m512bh vb[COLS];
136 + __m512 vc[ROWS * COLS];
137 + __m512 vsum[ROWS * COLS];
138 +
139 + // block quant scale
140 + __m512 vscale;
141 +
142 + auto loadc = [&](auto i) {
143 + constexpr int col = i % COLS;
144 + if constexpr (has_bias) {
145 + vc[i] = _mm512_loadu_ps(bias + col * 16);
146 + } else {
147 + vc[i] = _mm512_setzero_ps();
148 + }
149 + };
150 + Unroll<ROWS * COLS>{}(loadc);
151 +
152 + const int lda2 = lda >> 1;
153 + const int ldb2 = ldb; // ldb * 2 >> 1;
154 + const float* a_ptr = reinterpret_cast<const float*>(A);
155 + const uint16_t* b_ptr = reinterpret_cast<const uint16_t*>(B);
156 +
157 + auto compute = [&](auto i, int k) {
158 + constexpr int row = i / COLS;
159 + constexpr int col = i % COLS;
160 +
161 + if constexpr (col == 0) {
162 + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
163 + if constexpr (PREFETCH_SIZE_K > 0) {
164 + _mm_prefetch(a_ptr + row * lda2 + k + PREFETCH_SIZE_K, _MM_HINT_T0);
165 + }
166 + }
167 + if constexpr (row == 0) {
168 + if constexpr (col % 2 == 0) {
169 + __m512i b8 = _mm512_loadu_si512(b_ptr + k * ldb2 + col * 16);
170 + if constexpr (PREFETCH_SIZE_K > 0) {
171 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
172 + }
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));
175 + }
176 + }
177 + vsum[i] = _mm512_dpbf16_ps(vsum[i], va, vb[col]);
178 + };
179 +
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);
184 + // 1. load scale vector
185 + vscale = _mm512_set1_ps(scale[kb]);
186 + if constexpr (PREFETCH_SIZE_KB > 0) {
187 + _mm_prefetch(scale + kb + PREFETCH_SIZE_KB, _MM_HINT_T0);
188 + }
189 + // 2. zero vsum for each block
190 + Unroll<ROWS * COLS>{}([&](auto i) {
191 + vsum[i] = _mm512_setzero_ps();
192 + });
193 + // 3. accumulate across each block
194 + for (int k = kb_start; k < kb_end; ++k) {
195 + Unroll<ROWS * COLS>{}(compute, k);
196 + }
197 + // 4. apply scale
198 + Unroll<ROWS * COLS>{}([&](auto i) {
199 + vc[i] = _mm512_fmadd_ps(vsum[i], vscale, vc[i]);
200 + });
201 + }
202 +
203 + auto storec = [&](auto i) {
204 + constexpr int row = i / COLS;
205 + constexpr int col = i % COLS;
206 + // for COLS = 2,4 use 512bit store
207 + if constexpr (col % 2 == 0) {
208 + _mm512_storeu_si512(
209 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
210 + (__m512i)(_mm512_cvtne2ps_pbh(vc[row * COLS + col + 1], vc[row * COLS + col])));
211 + }
212 + };
213 + Unroll<ROWS * COLS>{}(storec);
214 + }
215 +};
216 +#endif
217 +
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);
222 +
223 +template <typename scalar_t, typename packed_t, bool has_bias>
224 +struct brgemm {
225 + static inline void apply(
226 + const scalar_t* __restrict__ A,
227 + const packed_t* __restrict__ B,
228 + scalar_t* __restrict__ C,
229 + scalar_t* __restrict__ Btmp,
230 + float* __restrict__ Ctmp,
231 + const float* __restrict__ bias,
232 + const float* __restrict__ scale,
233 + int M,
234 + int N,
235 + int K,
236 + int lda,
237 + int ldb,
238 + int ldc) {
239 + TORCH_CHECK(false, "struct brgemm: primary template not implemented!");
240 + }
241 +};
242 +
243 +template <bool has_bias>
244 +struct brgemm<at::BFloat16, at::Float8_e4m3fn, has_bias> {
245 + static inline void apply(
246 + const at::BFloat16* __restrict__ A,
247 + const at::Float8_e4m3fn* __restrict__ B,
248 + at::BFloat16* __restrict__ C,
249 + at::BFloat16* __restrict__ Btmp,
250 + float* __restrict__ Ctmp,
251 + const float* __restrict__ bias,
252 + const float* __restrict__ scale,
253 + int M,
254 + int N,
255 + int K,
256 + int lda,
257 + int ldb,
258 + int ldc) {
259 +
260 + constexpr int BLOCK_N = block_size_n();
261 +
262 + // [K, BLOCK_N] -> [K / 2, BLOCK_N * 2]
263 + const int ldb_tmp = BLOCK_N;
264 +
265 + for (int k = 0; k < K; k += BLOCK_K) {
266 + int kb_size = std::min(BLOCK_K, K - k);
267 +
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]);
270 + }
271 +
272 + at::native::cpublas::brgemm(
273 + M, N, K, lda, ldb_tmp, BLOCK_N, /* add_C */ false, A, Btmp, Ctmp);
274 +
275 + // copy from Ctmp to C
276 + for (int m = 0; m < M; ++m) {
277 + if constexpr (has_bias) {
278 + copy_add_stub(C + m * ldc, Ctmp + m * BLOCK_N, bias, N);
279 + } else {
280 + copy_stub(C + m * ldc, Ctmp + m * BLOCK_N, N);
281 + }
282 + }
283 + }
284 +};
285 +
286 +template <typename scalar_t, bool has_bias>
287 +void tinygemm_kernel(
288 + const scalar_t* __restrict__ A,
289 + const at::Float8_e4m3fn* __restrict__ B,
290 + scalar_t* __restrict__ C,
291 + scalar_t* __restrict__ Btmp,
292 + float* __restrict__ Ctmp,
293 + const float* __restrict__ scale,
294 + const float* __restrict__ bias,
295 + int64_t M,
296 + int64_t N,
297 + int64_t K,
298 + int64_t lda,
299 + int64_t ldb,
300 + int64_t ldc,
301 + bool brg,
302 + int64_t block_size_K) {
303 +
304 + 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);
307 + return;
308 + }
309 +
310 + // pattern: 1-4-16
311 + constexpr int64_t BLOCK_M = 4;
312 + constexpr int64_t BLOCK_N = 64;
313 + const int64_t MB = div_up(M, BLOCK_M);
314 + const int64_t NB = div_up(N, BLOCK_N);
315 + for (int mb = 0; mb < MB; ++mb) {
316 + int64_t mb_start = mb * BLOCK_M;
317 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
318 + for (int64_t nb = 0; nb < NB; ++nb) {
319 + int64_t nb_start = nb * BLOCK_N;
320 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
321 +
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");
328 + }
329 + }
330 + }
331 +}
332 +
333 +template <typename scalar_t>
334 +void fp8_scaled_mm_kernel_impl(
335 + scalar_t* __restrict__ out,
336 + const scalar_t* __restrict__ mat1,
337 + const at::Float8_e4m3fn* __restrict__ mat2,
338 + const float* __restrict__ scales2,
339 + const float* __restrict__ bias,
340 + scalar_t* __restrict__ buffer,
341 + int64_t M,
342 + int64_t N,
343 + int64_t K,
344 + int64_t mat1_strideM,
345 + int64_t out_strideM,
346 + int64_t block_size_N,
347 + 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;
351 + constexpr int64_t BLOCK_N = block_size_n();
352 + const int64_t MB = div_up(M, BLOCK_M);
353 + const int64_t NB = div_up(N, BLOCK_N);
354 +
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;
357 +
358 + const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
359 +
360 + // parallel on [MB, NB]
361 + 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();
367 + scalar_t* __restrict__ Btmp = buffer + tid * buffer_size_per_thread;
368 + float* __restrict__ Ctmp = (float*)((void*)(Btmp + BLOCK_N * K));
369 +
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;
373 +
374 + int64_t mb_start = mb * BLOCK_M;
375 + int64_t mb_size = std::min(M - mb_start, BLOCK_M);
376 + int64_t nb_start = nb * BLOCK_N;
377 + int64_t nb_size = std::min(N - nb_start, BLOCK_N);
378 +
379 + tinygemm_kernel<scalar_t, has_bias>(
380 + /* A */ mat1 + mb_start * mat1_strideM,
381 + /* B */ mat2 + nb_start * K, // nb * BLOCK_N * K
382 + /* C */ out + mb_start * out_strideM + nb_start,
383 + /* Btmp */ Btmp,
384 + /* Ctmp */ Ctmp,
385 + /* scale */ scale_ptr,
386 + /* bias */ bias + nb_start,
387 + /* M */ mb_size,
388 + /* N */ nb_size,
389 + /* K */ K,
390 + /* lda */ mat1_strideM,
391 + /* ldb */ nb_size,
392 + /* ldc */ out_strideM,
393 + /* 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 + }
399 +
400 + if (use_brgemm) {
401 + at::native::cpublas::brgemm_release();
402 + }
403 + });
404 + });
405 +}
406 +
407 +} // anonymous namespace
408 +
409 +// tinygemm interface
410 +template <typename scalar_t>
411 +void tinygemm_kernel(
412 + const scalar_t* __restrict__ A,
413 + const at::Float8_e4m3fn* __restrict__ B,
414 + scalar_t* __restrict__ C,
415 + scalar_t* __restrict__ Btmp,
416 + float* __restrict__ Ctmp,
417 + const float* __restrict__ scale,
418 + int64_t M,
419 + int64_t N,
420 + int64_t K,
421 + int64_t lda,
422 + int64_t ldb,
423 + int64_t ldc,
424 + 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);
427 +}
428 +
429 +#define INSTANTIATE_TINYGEMM_TEMPLATE(TYPE) \
430 + template void tinygemm_kernel<TYPE>( \
431 + const TYPE* __restrict__ A, \
432 + const at::Float8_e4m3fn* __restrict__ B, \
433 + TYPE* __restrict__ C, \
434 + TYPE* __restrict__ Btmp, \
435 + float* __restrict__ Ctmp, \
436 + const float* __restrict__ scale, \
437 + int64_t M, \
438 + int64_t N, \
439 + int64_t K, \
440 + int64_t lda, \
441 + int64_t ldb, \
442 + int64_t ldc, \
443 + bool brg, \
444 + int64_t block_size_K)
445 +
446 +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
447 +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
448 +
449 +at::Tensor fp8_scaled_mm_cpu(at::Tensor& mat1, at::Tensor& mat2, at::Tensor& scales2,
450 + std::vector<int64_t> block_size, 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}));
453 +
454 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
455 +
456 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
457 + CHECK_INPUT(mat2);
458 + CHECK_INPUT(scales2);
459 + TORCH_CHECK(scales2.scalar_type() == at::kFloat,
460 + "fp8_scaled_mm_cpu: expect scales2 to be float32.");
461 +
462 + int64_t M = mat1.size(0);
463 + int64_t N = mat2.size(0);
464 + int64_t K = mat2.size(1);
465 +
466 + CHECK_EQ(mat1.size(1), K);
467 + CHECK_DIM(2, mat1);
468 + CHECK_DIM(2, mat2);
469 +
470 + TORCH_CHECK(block_size.size() == 2,
471 + "fp8_scaled_mm_cpu: expect block_size.size() to be 2.");
472 +
473 + int64_t block_size_N = block_size[0];
474 + int64_t block_size_K = block_size[1];
475 +
476 + constexpr int64_t BLOCK_M = block_size_m() * BLOCK_SIZE_M_SCALE;
477 + constexpr int64_t BLOCK_N = block_size_n();
478 + TORCH_CHECK(block_size_N % BLOCK_N == 0, "fp8_scaled_mm_cpu: expect block_size_N to be multiples of BLOCK_N");
479 + TORCH_CHECK(block_size_K == BLOCK_K, "fp8_scaled_mm_cpu: expect block_size_K equals to BLOCK_K");
480 + CHECK_EQ(scales2.size(0), div_up(N, block_size_N));
481 + CHECK_EQ(scales2.size(1), div_up(K, block_size_K));
482 +
483 + 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.");
492 + auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
493 +
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());
510 +
511 + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "fp8_scaled_mm_kernel_impl", [&] {
512 + fp8_scaled_mm_kernel_impl<scalar_t>(
513 + out.data_ptr<scalar_t>(),
514 + mat1.data_ptr<scalar_t>(),
515 + packed_w.data_ptr<at::Float8_e4m3fn>(),
516 + scales2.data_ptr<float>(),
517 + bias_data,
518 + buffer.data_ptr<scalar_t>(),
519 + M,
520 + N,
521 + K,
522 + mat1_strideM,
523 + out_strideM,
524 + block_size_N,
525 + block_size_K,
526 + size_per_thread);
527 + });
528 +
529 + return out;
530 +}
csrc/cpu/sgl-kernels/gemm_int8.cpp new
+440
@@ -0,0 +1,440 @@
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 +
8 +// clang-format off
9 +
10 +namespace {
11 +
12 +template <typename scalar_t, bool has_bias, int BLOCK_M, int BLOCK_N>
13 +struct tinygemm_kernel_nn {
14 + 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) {
18 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
19 + }
20 +};
21 +
22 +#if defined(CPU_CAPABILITY_AVX512)
23 +template <bool has_bias, int BLOCK_M, int BLOCK_N>
24 +struct tinygemm_kernel_nn<at::BFloat16, has_bias, BLOCK_M, BLOCK_N> {
25 + 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 +
30 + constexpr int ROWS = BLOCK_M;
31 + constexpr int COLS = BLOCK_N / 16;
32 + static_assert(COLS % 2 == 0);
33 +
34 + // prefetch distance
35 + constexpr int PREFETCH_SIZE_K = 0;
36 +
37 + __m512i va;
38 + __m512i vb[COLS];
39 + __m512i vc[ROWS * COLS];
40 + __m512i vcomp[COLS];
41 + __m512 vd0;
42 + __m512 vd1[COLS];
43 +
44 + // oops! 4x4 spills but luckly we use 4x2
45 + __m512 vbias[COLS];
46 +
47 + // [NOTE]: s8s8 igemm compensation in avx512-vnni
48 + //
49 + // avx512-vnni has no s8s8, so we need to change s8s8 to u8s8 with compensate:
50 + //
51 + // a * b = (a + 128) * b - 128 * b
52 + // s s u s u s
53 + //
54 + // 1) 128 * b is pre-computed when packing B to vnni formats
55 + // 2) a + 128 is fused when dynamically quantize A
56 + //
57 + auto loadc = [&](auto i) {
58 + vc[i] = _mm512_set1_epi32(0);
59 + };
60 + Unroll<ROWS * COLS>{}(loadc);
61 +
62 + const int64_t K4 = K >> 2;
63 + const int64_t lda4 = lda >> 2;
64 + const int64_t ldb4 = ldb; // ldb * 4 >> 2;
65 + const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
66 + const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
67 +
68 + auto compute = [&](auto i, int64_t k) {
69 + constexpr int row = i / COLS;
70 + constexpr int col = i % COLS;
71 +
72 + if constexpr (col == 0) {
73 + va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
74 + }
75 + if constexpr (row == 0) {
76 + vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
77 + if constexpr (PREFETCH_SIZE_K > 0) {
78 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb4 + col * 16, _MM_HINT_T0);
79 + }
80 + }
81 + vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
82 + };
83 + for (int64_t k = 0; k < K4; ++k) {
84 + Unroll<ROWS * COLS>{}(compute, k);
85 + }
86 +
87 + auto storec = [&](auto i) {
88 + constexpr int row = i / COLS;
89 + constexpr int col = i % COLS;
90 +
91 + // load a scale
92 + if constexpr(col == 0) {
93 + vd0 = _mm512_set1_ps(As[row]);
94 + }
95 + // load b scale and vcomp per 2 vectors
96 + // also load bias if any
97 + if constexpr (row == 0) {
98 + if constexpr (col % 2 == 0) {
99 + vd1[col + 0] = _mm512_loadu_ps(Bs + col * 16);
100 + vd1[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
101 + vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
102 + vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
103 + if constexpr (has_bias) {
104 + vbias[col + 0] = _mm512_loadu_ps(bias + col * 16);
105 + vbias[col + 1] = _mm512_loadu_ps(bias + col * 16 + 16);
106 + }
107 + }
108 + }
109 +
110 + // for COLS = 2, 4 use 512bit store
111 + if constexpr (col % 2 == 0) {
112 + __m512 vc0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 0], vcomp[col + 0]));
113 + __m512 vc1 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[row * COLS + col + 1], vcomp[col + 1]));
114 + if constexpr (has_bias) {
115 + vc0 = _mm512_fmadd_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0], vbias[col + 0]);
116 + vc1 = _mm512_fmadd_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1], vbias[col + 1]);
117 + } else {
118 + vc0 = _mm512_mul_ps(_mm512_mul_ps(vc0, vd0), vd1[col + 0]);
119 + vc1 = _mm512_mul_ps(_mm512_mul_ps(vc1, vd0), vd1[col + 1]);
120 + }
121 +
122 + _mm512_storeu_si512(
123 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
124 + (__m512i)(_mm512_cvtne2ps_pbh(vc1, vc0)));
125 + }
126 + };
127 + Unroll<ROWS * COLS>{}(storec);
128 + }
129 +};
130 +#endif
131 +
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);
137 +
138 +template <typename scalar_t, bool has_bias>
139 +void tinygemm_kernel(
140 + const uint8_t* __restrict__ A,
141 + const int8_t* __restrict__ B,
142 + scalar_t* __restrict__ C,
143 + int32_t* __restrict__ Ctmp,
144 + const float* __restrict__ As,
145 + const float* __restrict__ Bs,
146 + const float* __restrict__ bias,
147 + int64_t M,
148 + int64_t N,
149 + int64_t K,
150 + int64_t lda,
151 + int64_t ldb,
152 + int64_t ldc,
153 + bool brg) {
154 +
155 + // B compensation
156 + const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
157 +
158 + // pattern: 1-4-16
159 + constexpr int64_t BLOCK_M = 4;
160 + constexpr int64_t BLOCK_N = 64;
161 + const int64_t MB = div_up(M, BLOCK_M);
162 + const int64_t NB = div_up(N, BLOCK_N);
163 + for (int64_t mb = 0; mb < MB; ++mb) {
164 + int64_t mb_start = mb * BLOCK_M;
165 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
166 + for (int64_t nb = 0; nb < NB; ++nb) {
167 + int64_t nb_start = nb * BLOCK_N;
168 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
169 +
170 + switch(mb_size << 4 | nb_size >> 4) {
171 + // mb_size = 1
172 + case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
173 + case 0x14: LAUNCH_TINYGEMM_KERNEL_NN(1, 64); break;
174 + // mb_size = 2
175 + case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
176 + case 0x24: LAUNCH_TINYGEMM_KERNEL_NN(2, 64); break;
177 + // mb_size = 3
178 + case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
179 + case 0x34: LAUNCH_TINYGEMM_KERNEL_NN(3, 64); break;
180 + // 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");
184 + }
185 + }
186 + }
187 +}
188 +
189 +template<typename scalar_t>
190 +void int8_scaled_mm_kernel_impl(
191 + scalar_t* __restrict__ out,
192 + const uint8_t* __restrict__ mat1,
193 + const int8_t* __restrict__ mat2,
194 + const float* __restrict__ scales1,
195 + const float* __restrict__ scales2,
196 + const float* __restrict__ bias,
197 + int64_t M,
198 + int64_t N,
199 + int64_t K) {
200 +
201 + constexpr int64_t BLOCK_M = block_size_m();
202 + constexpr int64_t BLOCK_N = block_size_n();
203 + const int64_t MB = div_up(M, BLOCK_M);
204 + const int64_t NB = div_up(N, BLOCK_N);
205 +
206 + // TODO: brgemm u8s8 depends on PyTorch 2.7 release.
207 + const bool use_brgemm = false;
208 +
209 + // K + 4 after compensation
210 + const int64_t packed_row_size = get_row_size<int8_t>(K);
211 +
212 + 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 +
217 + // for brgemm, use int32_t for accumulate
218 + alignas(64) int32_t Ctmp[BLOCK_M * BLOCK_N];
219 +
220 + for (int i = begin; i < end; ++i) {
221 + UNUSED(i);
222 + int mb_start = mb * BLOCK_M;
223 + int mb_size = std::min(M - mb_start, BLOCK_M);
224 + int nb_start = nb * BLOCK_N;
225 + int nb_size = std::min(N - nb_start, BLOCK_N);
226 +
227 + tinygemm_kernel<scalar_t, has_bias>(
228 + /* A */ mat1 + mb_start * K,
229 + /* B */ mat2 + nb_start * packed_row_size /* nb * BLOCK_N * (K + 4) */,
230 + /* C */ out + mb_start * N + nb_start,
231 + /* Ctmp*/ Ctmp,
232 + /* As */ scales1 + mb_start,
233 + /* Bs */ scales2 + nb_start,
234 + /* bias*/ bias + nb_start,
235 + /* M */ mb_size,
236 + /* N */ nb_size,
237 + /* K */ K,
238 + /* lda */ K,
239 + /* ldb */ nb_size,
240 + /* ldc */ N,
241 + /* brg */ use_brgemm);
242 +
243 + // move to the next index
244 + data_index_step(mb, MB, nb, NB);
245 + }
246 +
247 + if (use_brgemm) {
248 + at::native::cpublas::brgemm_release();
249 + }
250 + });
251 + });
252 +}
253 +
254 +} // anonymous namespace
255 +
256 +// tinygemm interface
257 +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) {
261 + tinygemm_kernel<scalar_t, false>(A, B, C, Ctmp, As, Bs, nullptr, M, N, K, lda, ldb, ldc, brg);
262 +}
263 +
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)
269 +
270 +INSTANTIATE_TINYGEMM_TEMPLATE(at::BFloat16);
271 +INSTANTIATE_TINYGEMM_TEMPLATE(at::Half);
272 +
273 +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 +
276 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(A);
277 + CHECK_DIM(2, A);
278 +
279 + int64_t M = A.size(0);
280 + int64_t K = A.size(1);
281 + int64_t lda = A.stride(0);
282 +
283 + 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.");
286 +
287 + auto Aq = at::empty({M, K}, A.options().dtype(at::kByte));
288 + auto As = at::empty({M}, A.options().dtype(at::kFloat));
289 +
290 + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "per_token_quant_int8", [&] {
291 + uint8_t* __restrict__ Aq_data = Aq.data_ptr<uint8_t>();
292 + float* __restrict__ As_data = As.data_ptr<float>();
293 + const scalar_t* __restrict__ A_data = A.data_ptr<scalar_t>();
294 +
295 + at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
296 + 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);
302 + }
303 + });
304 + });
305 + return std::make_tuple(Aq, As);
306 +}
307 +
308 +// weight : static, per-channel, symmetric
309 +// activation : dynamic, per-token, symmetric
310 +//
311 +// mat1 : [M, K]
312 +// mat2 : [N, K]
313 +// scales1 : [M]
314 +// scales2 : [N]
315 +// bias : [N]
316 +// out : [M, N]
317 +//
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 +
323 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
324 +
325 + CHECK_INPUT(mat1);
326 + CHECK_INPUT(mat2);
327 + CHECK_INPUT(scales1);
328 + CHECK_INPUT(scales2);
329 + CHECK_DIM(2, mat1);
330 + CHECK_DIM(2, mat2);
331 +
332 + int64_t M = mat1.size(0);
333 + int64_t N = mat2.size(0);
334 + int64_t K = mat1.size(1);
335 +
336 + // see [NOTE]: s8s8 igemm compensation in avx512-vnni
337 + CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
338 + CHECK_EQ(scales1.numel(), M);
339 + CHECK_EQ(scales2.numel(), N);
340 +
341 + TORCH_CHECK(mat1.scalar_type() == at::kByte, "int8_scaled_mm: expect mat1 to be uint8.");
342 + 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,
344 + "int8_scaled_mm: expect scales to be float32.");
345 +
346 + auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
347 +
348 + const bool has_bias = bias.has_value();
349 + const float* bias_data = nullptr;
350 + if (has_bias) {
351 + CHECK_EQ(bias.value().size(0), N);
352 + bias_data = bias.value().data_ptr<float>();
353 + }
354 +
355 + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_kernel_impl", [&] {
356 + int8_scaled_mm_kernel_impl<scalar_t>(
357 + out.data_ptr<scalar_t>(),
358 + mat1.data_ptr<uint8_t>(),
359 + packed_w.data_ptr<int8_t>(),
360 + scales1.data_ptr<float>(),
361 + scales2.data_ptr<float>(),
362 + bias_data,
363 + M,
364 + N,
365 + K);
366 + });
367 + return out;
368 +}
369 +
370 +// 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 +
375 + auto packed_w = is_vnni ? mat2 : convert_weight_packed(mat2);
376 +
377 + CHECK_LAST_DIM_CONTIGUOUS_INPUT(mat1);
378 + CHECK_INPUT(mat2);
379 + CHECK_INPUT(scales2);
380 + CHECK_DIM(2, mat1);
381 + CHECK_DIM(2, mat2);
382 +
383 + int64_t M = mat1.size(0);
384 + int64_t N = mat2.size(0);
385 + int64_t K = mat1.size(1);
386 + int64_t lda = mat1.stride(0);
387 +
388 + // see [NOTE]: s8s8 igemm compensation in avx512-vnni
389 + CHECK_EQ(mat2.size(1), (int64_t)(is_vnni ? K + sizeof(int32_t) : K));
390 + CHECK_EQ(scales2.numel(), N);
391 +
392 + 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.");
401 +
402 + const int64_t buffer_size = M * K + M * sizeof(float);
403 + auto buffer = at::empty({buffer_size}, mat1.options().dtype(at::kByte));
404 + auto out = at::empty({M, N}, mat1.options().dtype(out_dtype));
405 +
406 + const bool has_bias = bias.has_value();
407 + const float* bias_data = nullptr;
408 + if (has_bias) {
409 + CHECK_EQ(bias.value().size(0), N);
410 + bias_data = bias.value().data_ptr<float>();
411 + }
412 +
413 + AT_DISPATCH_REDUCED_FLOATING_TYPES(out_dtype, "int8_scaled_mm_with_quant_kernel_impl", [&] {
414 + uint8_t* __restrict__ Aq_data = buffer.data_ptr<uint8_t>();
415 + float* __restrict__ As_data = (float*)((void*)(Aq_data + M * K));
416 + const scalar_t* __restrict__ A_data = mat1.data_ptr<scalar_t>();
417 +
418 + at::parallel_for(0, M, 0, [&] (int64_t begin, int64_t end) {
419 + 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);
425 + }
426 + });
427 +
428 + int8_scaled_mm_kernel_impl<scalar_t>(
429 + out.data_ptr<scalar_t>(),
430 + Aq_data,
431 + packed_w.data_ptr<int8_t>(),
432 + As_data,
433 + scales2.data_ptr<float>(),
434 + bias_data,
435 + M,
436 + N,
437 + K);
438 + });
439 + return out;
440 +}
csrc/cpu/sgl-kernels/moe.cpp new
+1330
@@ -0,0 +1,1330 @@
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 +
8 +// clang-format off
9 +
10 +namespace {
11 +
12 +// [NOTE]: Fused MoE kernel with AMX
13 +//
14 +// This file contains implementations for
15 +// * `moe_align_block_size`
16 +// * `fused_moe`
17 +//
18 +// The functionality is identical to triton kernel, excepts:
19 +// * fuse silu_and_mul with gemm1, therefore this kernel
20 +// allocates 2 intermediate_caches instead of 3
21 +// * add `offsets` in `moe_align_block_size` which keeps track
22 +// of starting offset for each M block. this is for keeping
23 +// output of silu_and_mul in sorted order, thus load_A for
24 +// the 2nd gemm would be contiguous, therefore we can directly
25 +// load A from intermediate_cache1.
26 +//
27 +// TODO:
28 +// 1. tune BLOCK_M and BLOCK_N (BLOCK_N * K fit L2)
29 +// 2. add prefetch for load A which is indexed access
30 +// 3. abstract at::native::cpublas::brgemm with WoQ gemm (M = 1 & M != 1)
31 +//
32 +
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 +
136 +template <int BLOCK_M>
137 +int moe_align_block_size(
138 + int32_t* __restrict__ sorted_ids,
139 + int32_t* __restrict__ expert_ids,
140 + int32_t* __restrict__ topk_ids,
141 + int32_t* __restrict__ total_cnts,
142 + int32_t* __restrict__ cumsums,
143 + int32_t* __restrict__ offsets,
144 + int num_experts,
145 + int numel,
146 + int num_threads) {
147 +
148 + #define T_INDEX(tt) total_cnts + (tt) * num_experts
149 +
150 + // accumulate count of expert ids locally
151 + at::parallel_for(0, numel, 0, [&](int begin, int end) {
152 + int tid = at::get_thread_num();
153 + int32_t* __restrict__ local_cnts = T_INDEX(tid + 1);
154 +
155 + for (int i = begin; i < end; ++i) {
156 + local_cnts[topk_ids[i]]++;
157 + }
158 + });
159 +
160 + using iVec = at::vec::Vectorized<int32_t>;
161 + for (int t = 0; t < num_threads; ++t) {
162 + 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);
165 + }
166 +
167 + // the last row holds sums of each experts
168 + int32_t* total_cnts_t_1 = T_INDEX(num_threads);
169 +
170 + cumsums[0] = 0;
171 + for (int e = 0; e < num_experts; ++e) {
172 + // accumulate `num_tokens_post_pad`, also as the expert offset
173 + cumsums[e + 1] = cumsums[e] + div_up(total_cnts_t_1[e], BLOCK_M) * BLOCK_M;
174 +
175 + for (int k = cumsums[e]; k < cumsums[e + 1]; k += BLOCK_M) {
176 + expert_ids[k / BLOCK_M] = e;
177 + }
178 + }
179 + int num_tokens_post_pad = cumsums[num_experts];
180 +
181 + at::parallel_for(0, numel, 0, [&](int begin, int end) {
182 + int tid = at::get_thread_num();
183 + // thread tid offsets in `total_cnts`
184 + int32_t* __restrict__ offsets = T_INDEX(tid);
185 +
186 + for (int i = begin; i < end; ++i) {
187 + int32_t expert_id = topk_ids[i];
188 + int32_t b_offset = cumsums[expert_id];
189 + int32_t t_offset = offsets[expert_id];
190 + sorted_ids[b_offset + t_offset] = i;
191 + offsets[expert_id]++;
192 + }
193 + });
194 +
195 + // debug: the offset for thread t_1 should be identical to t_2
196 + int32_t* total_cnts_t_2 = T_INDEX(num_threads - 1);
197 + for (int e = 0; e < num_experts; ++e) {
198 + TORCH_CHECK(total_cnts_t_1[e] == total_cnts_t_2[e]);
199 + }
200 +
201 + // padding value for sorted_ids: numel
202 + auto sorted_id_size = [=](const int32_t* sorted_ids_ptr) {
203 + for (int d = 0; d < BLOCK_M; ++d) {
204 + if (sorted_ids_ptr[d] == numel) { return d; }
205 + }
206 + return BLOCK_M;
207 + };
208 +
209 + // offsets holds starting offset for each valida M blocks
210 + // shape : [num_token_blocks + 1]
211 + offsets[0] = 0;
212 + const int num_token_blocks = num_tokens_post_pad / BLOCK_M;
213 + at::parallel_for(0, num_token_blocks, GRAIN_SIZE / BLOCK_M, [&](int begin, int end) {
214 + for (int mb = begin; mb < end; ++mb) {
215 + offsets[mb + 1] = sorted_id_size(sorted_ids + mb * BLOCK_M);
216 + }
217 + });
218 + // TODO: do we need to vecterize this ?
219 + for (int mb = 0; mb < num_token_blocks; ++mb) {
220 + offsets[mb + 1] += offsets[mb];
221 + }
222 + // debug: the last value of offsets should be `numel`
223 + TORCH_CHECK(offsets[num_token_blocks] == numel);
224 +
225 + return num_tokens_post_pad;
226 +}
227 +
228 +// silu : shape leading dimension
229 +// input0 [m_size, BLOCK_N] BLOCK_N
230 +// input1 [m_size, BLOCK_N] BLOCK_N
231 +// output [M * topk, N] N
232 +template <typename scalar_t, int BLOCK_N>
233 +inline void silu_and_mul(
234 + scalar_t* __restrict__ output,
235 + const float* __restrict__ input0, // x: x0, x1
236 + const float* __restrict__ input1, // y: y0, y1
237 + int64_t m_size,
238 + int64_t N) {
239 +
240 + using bVec = at::vec::Vectorized<scalar_t>;
241 + using fVec = at::vec::Vectorized<float>;
242 +
243 + const fVec one = fVec(1.f);
244 +
245 + // no remainder
246 + for (int64_t m = 0; m < m_size; ++m) {
247 + scalar_t* __restrict__ out = output + m * N;
248 + const float* __restrict__ x = input0 + m * BLOCK_N;
249 + const float* __restrict__ y = input1 + m * BLOCK_N;
250 +
251 + for (int64_t d = 0; d < BLOCK_N; d += bVec::size()) {
252 + fVec x0 = fVec::loadu(x + d);
253 + fVec x1 = fVec::loadu(x + d + fVec::size());
254 + fVec y0 = fVec::loadu(y + d);
255 + fVec y1 = fVec::loadu(y + d + fVec::size());
256 + // silu
257 + x0 = x0 / (one + x0.neg().exp_u20());
258 + x1 = x1 / (one + x1.neg().exp_u20());
259 + // mul
260 + x0 = x0 * y0;
261 + x1 = x1 * y1;
262 + // convert
263 + bVec out_vec = convert_from_float_ext<scalar_t>(x0, x1);
264 + out_vec.store(out + d);
265 + }
266 + }
267 +}
268 +
269 +template <typename scalar_t, int BLOCK_M, int BLOCK_N>
270 +struct tinygemm_kernel_nn2 {
271 + 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) {
274 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
275 + }
276 +};
277 +
278 +#if defined(CPU_CAPABILITY_AVX512)
279 +template <int BLOCK_M, int BLOCK_N>
280 +struct tinygemm_kernel_nn2<at::BFloat16, BLOCK_M, BLOCK_N> {
281 + 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 +
285 + constexpr int ROWS = BLOCK_M;
286 + constexpr int COLS = BLOCK_N / 16;
287 +
288 + static_assert(COLS % 2 == 0);
289 +
290 + // prefetch distance
291 + constexpr int PREFETCH_SIZE_K = 0;
292 +
293 + __m512bh va;
294 + __m512bh vb0[COLS];
295 + __m512bh vb1[COLS];
296 + __m512 vc0[ROWS * COLS];
297 + __m512 vc1[ROWS * COLS];
298 +
299 + auto loadc = [&](auto i) {
300 + vc0[i] = _mm512_set1_ps(0.f);
301 + vc1[i] = _mm512_set1_ps(0.f);
302 + };
303 + Unroll<ROWS * COLS>{}(loadc);
304 +
305 + const int64_t K2 = K >> 1;
306 + const int64_t lda2 = lda >> 1;
307 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
308 + const float* a_ptr = reinterpret_cast<const float*>(A);
309 + const float* b0_ptr = reinterpret_cast<const float*>(B0);
310 + const float* b1_ptr = reinterpret_cast<const float*>(B1);
311 +
312 + auto compute = [&](auto i, int64_t k) {
313 + constexpr int row = i / COLS;
314 + constexpr int col = i % COLS;
315 +
316 + if constexpr (col == 0) {
317 + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
318 + }
319 + if constexpr (row == 0) {
320 + vb0[col] = (__m512bh)(_mm512_loadu_si512(b0_ptr + k * ldb2 + col * 16));
321 + vb1[col] = (__m512bh)(_mm512_loadu_si512(b1_ptr + k * ldb2 + col * 16));
322 + if constexpr (PREFETCH_SIZE_K > 0) {
323 + _mm_prefetch(b0_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
324 + _mm_prefetch(b1_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
325 + }
326 + }
327 + vc0[i] = _mm512_dpbf16_ps(vc0[i], va, vb0[col]);
328 + vc1[i] = _mm512_dpbf16_ps(vc1[i], va, vb1[col]);
329 + };
330 + for (int64_t k = 0; k < K2; ++k) {
331 + Unroll<ROWS * COLS>{}(compute, k);
332 + }
333 +
334 + using Vec = at::vec::Vectorized<float>;
335 + const Vec one = Vec(1.f);
336 + auto storec = [&](auto i) {
337 + constexpr int row = i / COLS;
338 + constexpr int col = i % COLS;
339 + // for COLS = 2, 4 use 512bit store
340 + if constexpr (col % 2 == 0) {
341 + Vec x0 = vc0[row * COLS + col + 0];
342 + Vec x1 = vc0[row * COLS + col + 1];
343 + Vec y0 = vc1[row * COLS + col + 0];
344 + Vec y1 = vc1[row * COLS + col + 1];
345 + // silu
346 + x0 = x0 / (one + x0.neg().exp_u20());
347 + x1 = x1 / (one + x1.neg().exp_u20());
348 + // mul
349 + x0 = x0 * y0;
350 + x1 = x1 * y1;
351 +
352 + _mm512_storeu_si512(
353 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
354 + (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
355 + }
356 + };
357 + Unroll<ROWS * COLS>{}(storec);
358 + }
359 +};
360 +#endif
361 +
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);
366 +
367 +template <typename scalar_t>
368 +void tinygemm_kernel(
369 + const scalar_t* __restrict__ A,
370 + const scalar_t* __restrict__ B0,
371 + const scalar_t* __restrict__ B1,
372 + scalar_t* __restrict__ C,
373 + int64_t M,
374 + int64_t N,
375 + int64_t K,
376 + int64_t lda,
377 + int64_t ldb,
378 + int64_t ldc) {
379 +
380 + // pattern: 1-(2+2)-(8+8)
381 + constexpr int64_t BLOCK_M = 4;
382 + constexpr int64_t BLOCK_N = 32;
383 + const int64_t MB = div_up(M, BLOCK_M);
384 + const int64_t NB = div_up(N, BLOCK_N);
385 + for (int mb = 0; mb < MB; ++mb) {
386 + int64_t mb_start = mb * BLOCK_M;
387 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
388 + for (int64_t nb = 0; nb < NB; ++nb) {
389 + int64_t nb_start = nb * BLOCK_N;
390 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
391 +
392 + switch(mb_size << 4 | nb_size >> 4) {
393 + // mb_size = 1
394 + case 0x12: LAUNCH_TINYGEMM_KERNEL_NN(1, 32); break;
395 + // mb_size = 2
396 + case 0x22: LAUNCH_TINYGEMM_KERNEL_NN(2, 32); break;
397 + // mb_size = 3
398 + case 0x32: LAUNCH_TINYGEMM_KERNEL_NN(3, 32); break;
399 + // 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");
402 + }
403 + }
404 + }
405 +}
406 +
407 +template <typename scalar_t, int BLOCK_M, int BLOCK_N>
408 +struct tinygemm_kernel_nn {
409 + 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) {
412 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
413 + }
414 +};
415 +
416 +#if defined(CPU_CAPABILITY_AVX512)
417 +template <int BLOCK_M, int BLOCK_N>
418 +struct tinygemm_kernel_nn<at::BFloat16, BLOCK_M, BLOCK_N> {
419 + 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 +
423 + constexpr int ROWS = BLOCK_M;
424 + constexpr int COLS = BLOCK_N / 16;
425 +
426 + static_assert(COLS % 2 == 0);
427 +
428 + // prefetch distance
429 + constexpr int PREFETCH_SIZE_K = 0;
430 +
431 + __m512bh va;
432 + __m512bh vb[COLS];
433 + __m512 vc[ROWS * COLS];
434 +
435 + auto loadc = [&](auto i) {
436 + vc[i] = _mm512_set1_ps(0.f);
437 + };
438 + Unroll<ROWS * COLS>{}(loadc);
439 +
440 + const int64_t K2 = K >> 1;
441 + const int64_t lda2 = lda >> 1;
442 + const int64_t ldb2 = ldb; // ldb * 2 >> 1;
443 + const float* a_ptr = reinterpret_cast<const float*>(A);
444 + const float* b_ptr = reinterpret_cast<const float*>(B);
445 +
446 + auto compute = [&](auto i, int64_t k) {
447 + constexpr int row = i / COLS;
448 + constexpr int col = i % COLS;
449 +
450 + if constexpr (col == 0) {
451 + va = (__m512bh)(_mm512_set1_ps(a_ptr[row * lda2 + k]));
452 + }
453 + if constexpr (row == 0) {
454 + vb[col] = (__m512bh)(_mm512_loadu_si512(b_ptr + k * ldb2 + col * 16));
455 + if constexpr (PREFETCH_SIZE_K > 0) {
456 + _mm_prefetch(b_ptr + (k + PREFETCH_SIZE_K) * ldb2 + col * 16, _MM_HINT_T0);
457 + }
458 + }
459 + vc[i] = _mm512_dpbf16_ps(vc[i], va, vb[col]);
460 + };
461 + for (int64_t k = 0; k < K2; ++k) {
462 + Unroll<ROWS * COLS>{}(compute, k);
463 + }
464 +
465 + auto storec = [&](auto i) {
466 + constexpr int row = i / COLS;
467 + constexpr int col = i % COLS;
468 + _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), vc[i]);
469 +
470 + };
471 + Unroll<ROWS * COLS>{}(storec);
472 + }
473 +};
474 +#endif
475 +
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);
480 +
481 +template <typename scalar_t>
482 +void tinygemm_kernel(
483 + const scalar_t* __restrict__ A,
484 + const scalar_t* __restrict__ B,
485 + float* __restrict__ C,
486 + int64_t M,
487 + int64_t N,
488 + int64_t K,
489 + int64_t lda,
490 + int64_t ldb,
491 + int64_t ldc) {
492 +
493 + // pattern: 1-2-8
494 + constexpr int64_t BLOCK_M = 4;
495 + constexpr int64_t BLOCK_N = 32;
496 + const int64_t MB = div_up(M, BLOCK_M);
497 + const int64_t NB = div_up(N, BLOCK_N);
498 + for (int mb = 0; mb < MB; ++mb) {
499 + int64_t mb_start = mb * BLOCK_M;
500 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
501 + for (int64_t nb = 0; nb < NB; ++nb) {
502 + int64_t nb_start = nb * BLOCK_N;
503 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
504 +
505 + switch(mb_size << 4 | nb_size >> 4) {
506 + // mb_size = 1
507 + case 0x12: LAUNCH_TINYGEMM_KERNEL_NN2(1, 32); break;
508 + // mb_size = 2
509 + case 0x22: LAUNCH_TINYGEMM_KERNEL_NN2(2, 32); break;
510 + // mb_size = 3
511 + case 0x32: LAUNCH_TINYGEMM_KERNEL_NN2(3, 32); break;
512 + // 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");
515 + }
516 + }
517 + }
518 +}
519 +
520 +template <typename scalar_t>
521 +void fused_experts_kernel_impl(
522 + scalar_t* __restrict__ output,
523 + scalar_t* __restrict__ ic1,
524 + scalar_t* __restrict__ ic2,
525 + scalar_t* __restrict__ A_tmp,
526 + float* __restrict__ C_tmp,
527 + const scalar_t* __restrict__ input,
528 + const scalar_t* __restrict__ packed_w1,
529 + const scalar_t* __restrict__ packed_w2,
530 + const float* __restrict__ topk_weights,
531 + const int32_t* __restrict__ sorted_ids,
532 + const int32_t* __restrict__ expert_ids,
533 + const int32_t* __restrict__ offsets,
534 + int64_t M,
535 + int64_t N,
536 + int64_t K,
537 + int64_t E,
538 + int64_t topk,
539 + int64_t num_tokens_post_pad) {
540 +
541 + // handle 2 tiles per block
542 + constexpr int64_t BLOCK_M = block_size_m();
543 + constexpr int64_t BLOCK_N = block_size_n();
544 +
545 + // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
546 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
547 + const int64_t NB = div_up(N, BLOCK_N);
548 +
549 + // strides for w1: [E, 2N, K]
550 + TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
551 +
552 + const int64_t stride_e = 2 * N * K;
553 + const int64_t stride_n = K;
554 +
555 + // 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) {
557 + // get local pointers
558 + int tid = at::get_thread_num();
559 + scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
560 + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
561 + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
562 +
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);
572 +
573 + // B shape [K, n_size] in vnni format
574 + 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;
577 +
578 + // 1.a load A
579 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
580 + int64_t m_size = offsets[mb + 1] - offsets[mb];
581 +
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);
588 + }
589 +
590 + if (use_brgemm) {
591 + // 1.b gemm: C0 = A @ B0
592 + at::native::cpublas::brgemm(
593 + /* M */ m_size,
594 + /* N */ n_size,
595 + /* K */ K,
596 + /* lda */ K,
597 + /* ldb */ n_size,
598 + /* ldc */ BLOCK_N,
599 + /* add_C */ false,
600 + /* A */ A,
601 + /* B */ B0,
602 + /* C */ C0);
603 +
604 + // 1.c gemm: C1 = A @ B1
605 + at::native::cpublas::brgemm(
606 + /* M */ m_size,
607 + /* N */ n_size,
608 + /* K */ K,
609 + /* lda */ K,
610 + /* ldb */ n_size,
611 + /* ldc */ BLOCK_N,
612 + /* add_C */ false,
613 + /* A */ A,
614 + /* B */ B1,
615 + /* C */ C1);
616 +
617 + // 1.d silu and mul
618 + 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);
625 + } else {
626 + // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
627 + const int64_t offset = offsets[mb];
628 + tinygemm_kernel(
629 + /* A */ A,
630 + /* B0 */ B0,
631 + /* B1 */ B1,
632 + /* C */ ic1 + offset * N + nb * BLOCK_N,
633 + /* M */ m_size,
634 + /* N */ n_size,
635 + /* K */ K,
636 + /* lda */ K,
637 + /* ldb */ n_size,
638 + /* ldc */ N);
639 + }
640 + }
641 +
642 + if (is_brgemm_used) {
643 + at::native::cpublas::brgemm_release();
644 + }
645 + });
646 +
647 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
648 + // w2 : [E, K, N] as [E, OC, IC]
649 + const int64_t OC = K; // rename K as OC
650 + const int64_t IC = N; // rename N as IC
651 + const int64_t MB2 = MB;
652 + const int64_t NB2 = div_up(OC, BLOCK_N);
653 + const int64_t stride_e2 = OC * IC;
654 + const int64_t stride_oc = IC;
655 +
656 + // parallel on [MB2, NB2]
657 + at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
658 + // get local pointers
659 + int tid = at::get_thread_num();
660 + // we won't be using C1 for gemm2
661 + float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
662 +
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 +
669 + int64_t m_size = offsets[mb + 1] - offsets[mb];
670 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
671 +
672 + const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
673 + is_brgemm_used = is_brgemm_used || use_brgemm;
674 +
675 + // A ptr from ic1 of [M * topk, N] in sorted order
676 + // so as to avoid copy A to tmp buffer again
677 + const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
678 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
679 +
680 + // B shape [IC, n_size] in vnni format
681 + int32_t expert_id = expert_ids[mb];
682 + const scalar_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
683 +
684 + // 2.a gemm: C = A @ B
685 + if (use_brgemm) {
686 + at::native::cpublas::brgemm(
687 + /* M */ m_size,
688 + /* N */ n_size,
689 + /* K */ IC,
690 + /* lda */ IC,
691 + /* ldb */ n_size,
692 + /* ldc */ BLOCK_N,
693 + /* add_C */ false,
694 + /* A */ A,
695 + /* B */ B,
696 + /* C */ C);
697 + } else {
698 + tinygemm_kernel(
699 + /* A */ A,
700 + /* B */ B,
701 + /* C */ C,
702 + /* M */ m_size,
703 + /* N */ n_size,
704 + /* K */ IC,
705 + /* lda */ IC,
706 + /* ldb */ n_size,
707 + /* ldc */ BLOCK_N);
708 + }
709 +
710 + // 2.b copy from C to ic2 in original order
711 + // and also mul topk_weights in float32
712 + for (int64_t m = 0; m < m_size; ++m) {
713 + int32_t index = A_ids[m];
714 + float weight = topk_weights[index];
715 + copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
716 + }
717 + }
718 +
719 + if (is_brgemm_used) {
720 + at::native::cpublas::brgemm_release();
721 + }
722 + });
723 +
724 + // stage 3: out = intermediate_cache2.sum(dim=1)
725 + // from [M, topk, K] to [M, K]
726 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
727 + for (int64_t m = begin; m < end; ++m) {
728 + sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
729 + }
730 + });
731 +}
732 +
733 +template <typename scalar_t>
734 +void shared_expert_kernel_impl(
735 + scalar_t* __restrict__ output,
736 + scalar_t* __restrict__ ic1,
737 + float* __restrict__ C_tmp,
738 + scalar_t* __restrict__ input,
739 + const scalar_t* __restrict__ packed_w1,
740 + const scalar_t* __restrict__ packed_w2,
741 + const scalar_t* __restrict__ fused_experts_out,
742 + float routed_scaling_factor,
743 + int64_t M,
744 + int64_t N,
745 + int64_t K) {
746 +
747 + // handle 2 tiles per block
748 + constexpr int64_t BLOCK_M = block_size_m();
749 + constexpr int64_t BLOCK_N = block_size_n();
750 +
751 + // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
752 + const int64_t MB = div_up(M, BLOCK_M);
753 + const int64_t NB = div_up(N, BLOCK_N);
754 +
755 + TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
756 + const int64_t stride_n = K;
757 +
758 + // 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) {
760 + // get local pointers
761 + int tid = at::get_thread_num();
762 + float* __restrict__ C0 = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
763 + float* __restrict__ C1 = C0 + BLOCK_M * BLOCK_N;
764 +
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);
774 + int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
775 +
776 + //int64_t mb_start = mb * BLOCK_M;
777 + //int64_t mb_size = std::min(M - mb_start, BLOCK_M);
778 +
779 + // A shape [m_size, K]
780 + const scalar_t* A = input + mb * BLOCK_M * K;
781 +
782 + // 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;
788 +
789 + if (use_brgemm) {
790 + // 1.b gemm: C0 = A @ B0
791 + at::native::cpublas::brgemm(
792 + /* M */ m_size,
793 + /* N */ n_size,
794 + /* K */ K,
795 + /* lda */ K,
796 + /* ldb */ n_size,
797 + /* ldc */ BLOCK_N,
798 + /* add_C */ false,
799 + /* A */ A,
800 + /* B */ B0,
801 + /* C */ C0);
802 +
803 + // 1.c gemm: C1 = A @ B1
804 + at::native::cpublas::brgemm(
805 + /* M */ m_size,
806 + /* N */ n_size,
807 + /* K */ K,
808 + /* lda */ K,
809 + /* ldb */ n_size,
810 + /* ldc */ BLOCK_N,
811 + /* add_C */ false,
812 + /* A */ A,
813 + /* B */ B1,
814 + /* C */ C1);
815 +
816 + // 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);
823 + } else {
824 + // fused 1.bcd: silu_and_mul(A @ B0, A @ B1)
825 + tinygemm_kernel(
826 + /* A */ A,
827 + /* B0 */ B0,
828 + /* B1 */ B1,
829 + /* C */ ic1 + mb * BLOCK_M * N + nb * BLOCK_N,
830 + /* M */ m_size,
831 + /* N */ n_size,
832 + /* K */ K,
833 + /* lda */ K,
834 + /* ldb */ n_size,
835 + /* ldc */ N);
836 + }
837 + }
838 +
839 + if (is_brgemm_used) {
840 + at::native::cpublas::brgemm_release();
841 + }
842 + });
843 +
844 + // stage 2: output = intermediate_cache1 @ w2
845 + // w2 : [K, N] as [OC, IC]
846 + const int64_t OC = K; // rename K as OC
847 + const int64_t IC = N; // rename N as IC
848 + const int64_t MB2 = MB;
849 + const int64_t NB2 = div_up(OC, BLOCK_N);
850 + const int64_t stride_oc = IC;
851 +
852 + // parallel on [MB2, NB2]
853 + at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
854 + // get local pointers
855 + int tid = at::get_thread_num();
856 + // we won't be using C1 for gemm2
857 + float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
858 +
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 +
865 + int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
866 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
867 +
868 + const bool use_brgemm = can_use_brgemm<scalar_t>(m_size);
869 + is_brgemm_used = is_brgemm_used || use_brgemm;
870 +
871 + // A shape [m_size, IC]
872 + const scalar_t* __restrict__ A = ic1 + mb * BLOCK_M * N;
873 +
874 + // B shape [IC, n_size] in vnni format
875 + const scalar_t* __restrict__ B = packed_w2 + nb * BLOCK_N * stride_oc;
876 +
877 + // 2.a gemm: C = A @ B
878 + if (use_brgemm) {
879 + at::native::cpublas::brgemm(
880 + /* M */ m_size,
881 + /* N */ n_size,
882 + /* K */ IC,
883 + /* lda */ IC,
884 + /* ldb */ n_size,
885 + /* ldc */ BLOCK_N,
886 + /* add_C */ false,
887 + /* A */ A,
888 + /* B */ B,
889 + /* C */ C);
890 + } else {
891 + tinygemm_kernel(
892 + /* A */ A,
893 + /* B */ B,
894 + /* C */ C,
895 + /* M */ m_size,
896 + /* N */ n_size,
897 + /* K */ IC,
898 + /* lda */ IC,
899 + /* ldb */ n_size,
900 + /* ldc */ BLOCK_N);
901 + }
902 +
903 + // 2.b copy from C to output and add fused_experts_out
904 + 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;
906 + 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);
908 + }
909 + }
910 +
911 + if (is_brgemm_used) {
912 + at::native::cpublas::brgemm_release();
913 + }
914 + });
915 +}
916 +
917 +} // anonymous namespace
918 +
919 +// common checks
920 +static inline void check_moe_scales(
921 + bool use_int8_w8a8,
922 + bool use_fp8_w8a16,
923 + const std::optional<at::Tensor>& w1_scale,
924 + 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) {
928 + if (use_int8_w8a8) {
929 + TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for int8 w8a8.");
930 + 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.");
933 + }
934 + if (use_fp8_w8a16) {
935 + TORCH_CHECK(w1_scale.has_value(), "missing w1_scale for fp8 w8a16.");
936 + TORCH_CHECK(w2_scale.has_value(), "missing w2_scale for fp8 w8a16.");
937 + TORCH_CHECK(block_size.has_value(), "missing block_size for fp8 w8a16.");
938 + TORCH_CHECK(block_size.value().size() == 2, "expect block_size.size() to be 2.");
939 + }
940 +}
941 +
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)
952 +
953 +// hidden_states: [M, K]
954 +// w1: [E, 2N, K]
955 +// w2: [E, K, N]
956 +// topk_weights: [M, topk]
957 +// topk_ids: [M, topk] (int32_t)
958 +//
959 +at::Tensor fused_experts_cpu(
960 + at::Tensor& hidden_states,
961 + at::Tensor& w1,
962 + at::Tensor& w2,
963 + at::Tensor& topk_weights,
964 + at::Tensor& topk_ids,
965 + bool inplace,
966 + bool use_int8_w8a8,
967 + bool use_fp8_w8a16,
968 + const std::optional<at::Tensor>& w1_scale,
969 + const std::optional<at::Tensor>& w2_scale,
970 + 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,
973 + bool is_vnni) {
974 + RECORD_FUNCTION("sgl-kernel::fused_experts_cpu", std::vector<c10::IValue>({hidden_states, w1, w2, topk_weights, topk_ids}));
975 +
976 + auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
977 + auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
978 +
979 + constexpr int64_t BLOCK_M = block_size_m();
980 + constexpr int64_t BLOCK_N = block_size_n();
981 +
982 + const auto st = hidden_states.scalar_type();
983 + CHECK_INPUT(hidden_states);
984 + CHECK_INPUT(w1);
985 + CHECK_INPUT(w2);
986 + CHECK_EQ(topk_weights.sizes(), topk_ids.sizes());
987 + CHECK_DIM(2, hidden_states);
988 + CHECK_DIM(3, w1);
989 + CHECK_DIM(3, w2);
990 + CHECK_DIM(2, topk_weights);
991 + CHECK_DIM(2, topk_ids);
992 +
993 + CHECK_EQ(topk_ids.scalar_type(), at::kInt);
994 + CHECK_EQ(topk_weights.scalar_type(), at::kFloat);
995 +
996 + int64_t M = hidden_states.size(0);
997 + int64_t K = hidden_states.size(1);
998 + int64_t N = w1.size(1) / 2;
999 + int64_t E = w1.size(0);
1000 + int64_t topk = topk_weights.size(1);
1001 +
1002 + // 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);
1005 +
1006 + // check weight shapes
1007 + 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 +
1012 + // check scales
1013 + check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);
1014 +
1015 + at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
1016 +
1017 + // NB: worst case is each expert holds a block with remainder of 1
1018 + // 1. sorted_ids : [M * topk + E * (BLOCK_M - 1)]
1019 + // 2. expert_ids : [max_num_blocks]
1020 + // 3. total_cnts : [T + 1, E]
1021 + // 4. cumsums : [E + 1]
1022 + // 5. offsets : [max_num_blocks + 1]
1023 + //
1024 + int num_threads = at::get_num_threads();
1025 + int64_t max_num_tokens_padded = M * topk + E * (BLOCK_M - 1);
1026 + int64_t max_num_blocks = div_up(max_num_tokens_padded, BLOCK_M);
1027 + auto buffer = at::empty(
1028 + {max_num_tokens_padded + max_num_blocks + (num_threads + 1) * E + (E + 1) + (max_num_blocks + 1)},
1029 + topk_ids.options());
1030 +
1031 + int32_t* __restrict__ sorted_ids = buffer.data_ptr<int32_t>();
1032 + int32_t* __restrict__ expert_ids = sorted_ids + max_num_tokens_padded;
1033 + 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);
1036 +
1037 + // init sorted_ids with `numel` as the padding number
1038 + // init expert_ids with `num_experts`
1039 + int64_t numel = M * topk;
1040 + at::parallel_for(0, max_num_blocks, GRAIN_SIZE / BLOCK_M, [&](int64_t begin, int64_t end) {
1041 + int64_t m_start = begin * BLOCK_M;
1042 + int64_t m_size = std::min((end - begin) * BLOCK_M, max_num_tokens_padded - m_start);
1043 + fill_stub(sorted_ids + m_start, (int32_t)numel, m_size);
1044 + fill_stub(expert_ids + begin, (int32_t)E, end - begin);
1045 + });
1046 + // zero total_cnts and cumsums
1047 + at::parallel_for(0, (num_threads + 1) * E + (E + 1), GRAIN_SIZE, [&](int64_t begin, int64_t end) {
1048 + fill_stub(total_cnts + begin, 0, end - begin);
1049 + });
1050 +
1051 + // align experts index
1052 + int64_t num_tokens_post_pad = moe_align_block_size<BLOCK_M>(
1053 + sorted_ids, expert_ids, topk_ids.data_ptr<int32_t>(), total_cnts, cumsums, offsets, E, numel, num_threads);
1054 +
1055 + // unlike triton kernel, we fuse silu with gemm1 so only need 2 intermediate_caches:
1056 + // 1. intermediate_cache1 : [M * topk, N]
1057 + // 2. intermediate_cache2 : [M * topk, K]
1058 + // 3. A_tmp : [T, BLOCK_M * K]
1059 + // 4. C_tmp : [T, 2 * BLOCK_M * BLOCK_N]
1060 + //
1061 + // for int8 w8a8:
1062 + // 5. Aq_tmp : [M, K] or [M * topk, N]
1063 + // 6. As_tmp : [M * topk]
1064 + //
1065 + // for fp8 w8a16:
1066 + // 7. intermediate_cache0 : [M * topk, 2N]
1067 + // 8. B_tmp : [T, BLOCK_N, std::max(K, N)]
1068 + //
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) +
1071 + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
1072 +
1073 + if (use_int8_w8a8) {
1074 + buffer_size_nbytes += std::max(M * K, M * topk * N) + M * topk * sizeof(float);
1075 + }
1076 + if (use_fp8_w8a16) {
1077 + buffer_size_nbytes += M * topk * 2 * N * 2 + num_threads * BLOCK_N * std::max(K, N) * 2;
1078 + }
1079 +
1080 + auto buffer2 = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));
1081 +
1082 + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "fused_experts_kernel_impl", [&] {
1083 + scalar_t* __restrict__ intermediate_cache1 = (scalar_t*)((void*)(buffer2.data_ptr<int8_t>()));
1084 + scalar_t* __restrict__ intermediate_cache2 = intermediate_cache1 + M * topk * N;
1085 +
1086 + if (use_int8_w8a8) {
1087 + uint8_t* __restrict__ A_tmp = (uint8_t*)((void*)(intermediate_cache2 + M * topk * K));
1088 + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1089 + uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1090 + float* __restrict__ As_tmp = (float*)((void*)(Aq_tmp + std::max(M * K, M * topk * N)));
1091 +
1092 + auto w1s = w1_scale.value();
1093 + auto w2s = w2_scale.value();
1094 + TORCH_CHECK(w1s.numel() == E * 2 * N);
1095 + TORCH_CHECK(w2s.numel() == E * K);
1096 +
1097 + fused_experts_int8_kernel_impl<scalar_t>(
1098 + out_hidden_states.data_ptr<scalar_t>(),
1099 + intermediate_cache1,
1100 + intermediate_cache2,
1101 + A_tmp,
1102 + C_tmp,
1103 + Aq_tmp,
1104 + As_tmp,
1105 + hidden_states.data_ptr<scalar_t>(),
1106 + packed_w1.data_ptr<int8_t>(),
1107 + packed_w2.data_ptr<int8_t>(),
1108 + w1s.data_ptr<float>(),
1109 + w2s.data_ptr<float>(),
1110 + topk_weights.data_ptr<float>(),
1111 + sorted_ids,
1112 + expert_ids,
1113 + offsets,
1114 + M,
1115 + N,
1116 + K,
1117 + E,
1118 + topk,
1119 + num_tokens_post_pad);
1120 + } else if (use_fp8_w8a16) {
1121 + // here we just ignore C_tmp as it is not used
1122 + scalar_t* __restrict__ A_tmp = (scalar_t*)((void*)(intermediate_cache2 + M * topk * K));
1123 + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1124 + scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1125 + scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * topk * 2 * N));
1126 +
1127 + CHECK_MOE_SCALES_FP8(1, 2);
1128 + fused_experts_fp8_kernel_impl(
1129 + out_hidden_states.data_ptr<scalar_t>(),
1130 + intermediate_cache0,
1131 + intermediate_cache1,
1132 + intermediate_cache2,
1133 + A_tmp,
1134 + B_tmp,
1135 + C_tmp,
1136 + hidden_states.data_ptr<scalar_t>(),
1137 + packed_w1.data_ptr<at::Float8_e4m3fn>(),
1138 + packed_w2.data_ptr<at::Float8_e4m3fn>(),
1139 + w1s.data_ptr<float>(),
1140 + w2s.data_ptr<float>(),
1141 + block_size_N,
1142 + block_size_K,
1143 + topk_weights.data_ptr<float>(),
1144 + sorted_ids,
1145 + expert_ids,
1146 + offsets,
1147 + M,
1148 + N,
1149 + K,
1150 + E,
1151 + topk,
1152 + num_tokens_post_pad);
1153 + } else {
1154 + scalar_t* __restrict__ A_tmp = intermediate_cache2 + M * topk * K;
1155 + float* __restrict__ C_tmp = (float*)((void*)(A_tmp + num_threads * BLOCK_M * K));
1156 +
1157 + fused_experts_kernel_impl<scalar_t>(
1158 + out_hidden_states.data_ptr<scalar_t>(),
1159 + intermediate_cache1,
1160 + intermediate_cache2,
1161 + A_tmp,
1162 + C_tmp,
1163 + hidden_states.data_ptr<scalar_t>(),
1164 + packed_w1.data_ptr<scalar_t>(),
1165 + packed_w2.data_ptr<scalar_t>(),
1166 + topk_weights.data_ptr<float>(),
1167 + sorted_ids,
1168 + expert_ids,
1169 + offsets,
1170 + M,
1171 + N,
1172 + K,
1173 + E,
1174 + topk,
1175 + num_tokens_post_pad);
1176 + }
1177 + });
1178 + return out_hidden_states;
1179 +}
1180 +
1181 +// shared expert kernel
1182 +//
1183 +// hidden_states: [M, K]
1184 +// w1: [2N, K]
1185 +// w2: [K, N]
1186 +// fused_experts_out
1187 +at::Tensor shared_expert_cpu(
1188 + at::Tensor& hidden_states,
1189 + at::Tensor& w1,
1190 + at::Tensor& w2,
1191 + at::Tensor& fused_experts_out,
1192 + double routed_scaling_factor,
1193 + bool inplace,
1194 + bool use_int8_w8a8,
1195 + 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,
1201 + bool is_vnni) {
1202 + RECORD_FUNCTION("sgl-kernel::shared_expert_cpu", std::vector<c10::IValue>({hidden_states, w1, w2}));
1203 +
1204 + auto packed_w1 = is_vnni ? w1 : convert_weight_packed(w1);
1205 + auto packed_w2 = is_vnni ? w2 : convert_weight_packed(w2);
1206 +
1207 + constexpr int64_t BLOCK_M = block_size_m();
1208 + constexpr int64_t BLOCK_N = block_size_n();
1209 +
1210 + const auto st = hidden_states.scalar_type();
1211 + CHECK_INPUT(hidden_states);
1212 + CHECK_INPUT(fused_experts_out);
1213 + CHECK_INPUT(w1);
1214 + CHECK_INPUT(w2);
1215 + CHECK_DIM(2, hidden_states);
1216 + CHECK_DIM(2, w1);
1217 + CHECK_DIM(2, w2);
1218 + CHECK_EQ(hidden_states.sizes(), fused_experts_out.sizes());
1219 + CHECK_EQ(hidden_states.scalar_type(), st);
1220 +
1221 + int64_t M = hidden_states.size(0);
1222 + int64_t K = hidden_states.size(1);
1223 + int64_t N = w1.size(0) / 2;
1224 +
1225 + // we use int32_t compensation for int8 w8a8
1226 + int64_t packed_K = get_row_size(K, use_int8_w8a8);
1227 + int64_t packed_N = get_row_size(N, use_int8_w8a8);
1228 +
1229 + // check weight shapes
1230 + CHECK_EQ(w2.size(0), K);
1231 + CHECK_EQ(packed_w1.size(1), packed_K);
1232 + CHECK_EQ(packed_w2.size(1), packed_N);
1233 +
1234 + // check scales
1235 + check_moe_scales(use_int8_w8a8, use_fp8_w8a16, w1_scale, w2_scale, block_size, a1_scale, a2_scale);
1236 +
1237 + at::Tensor out_hidden_states = inplace ? hidden_states : at::empty_like(hidden_states);
1238 +
1239 + // unlike triton kernel, we fuse silu with gemm1 so only need 2 intermediate_caches:
1240 + // 1. intermediate_cache1 : [M, N]
1241 + // 2. C_tmp : [T, 2 * BLOCK_M * BLOCK_N]
1242 + //
1243 + // for int8 w8a8:
1244 + // 3. Aq_tmp : [M, K] or [M, N]
1245 + // 4. As_tmp : [M]
1246 + //
1247 + // for fp8 w8a16:
1248 + // 5. intermediate_cache0 : [M, 2N]
1249 + // 6. B_tmp: [T, BLOCK_M, max(K, N)]
1250 + //
1251 + int num_threads = at::get_num_threads();
1252 + int64_t buffer_size_nbytes = M * N * 2 + num_threads * 2 * BLOCK_M * BLOCK_N * sizeof(float);
1253 +
1254 + if (use_int8_w8a8) {
1255 + buffer_size_nbytes += std::max(M * K, M * N) + M * sizeof(float);
1256 + }
1257 + if (use_fp8_w8a16) {
1258 + buffer_size_nbytes += M * 2 * N * 2 + num_threads * BLOCK_M * std::max(K, N) * 2;
1259 + }
1260 +
1261 + auto buffer = at::empty({buffer_size_nbytes}, hidden_states.options().dtype(at::kChar));
1262 + AT_DISPATCH_REDUCED_FLOATING_TYPES(st, "share_experts_kernel_impl", [&] {
1263 + scalar_t* __restrict__ intermediate_cache1 = (scalar_t*)((void*)(buffer.data_ptr<int8_t>()));
1264 + float* __restrict__ C_tmp = (float*)((void*)(intermediate_cache1 + M * N));
1265 +
1266 + if (use_int8_w8a8) {
1267 + uint8_t* __restrict__ Aq_tmp = (uint8_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1268 + float* __restrict__ As_tmp = (float*)((void*)(Aq_tmp + std::max(M * K, M * N)));
1269 +
1270 + auto w1s = w1_scale.value();
1271 + auto w2s = w2_scale.value();
1272 + TORCH_CHECK(w1s.numel() == 2 * N);
1273 + TORCH_CHECK(w2s.numel() == K);
1274 +
1275 + shared_expert_int8_kernel_impl<scalar_t>(
1276 + out_hidden_states.data_ptr<scalar_t>(),
1277 + intermediate_cache1,
1278 + C_tmp,
1279 + Aq_tmp,
1280 + As_tmp,
1281 + hidden_states.data_ptr<scalar_t>(),
1282 + packed_w1.data_ptr<int8_t>(),
1283 + packed_w2.data_ptr<int8_t>(),
1284 + w1s.data_ptr<float>(),
1285 + w2s.data_ptr<float>(),
1286 + fused_experts_out.data_ptr<scalar_t>(),
1287 + routed_scaling_factor,
1288 + M,
1289 + N,
1290 + K);
1291 + } else if (use_fp8_w8a16) {
1292 + scalar_t* __restrict__ intermediate_cache0 = (scalar_t*)((void*)(C_tmp + num_threads * 2 * BLOCK_M * BLOCK_N));
1293 + scalar_t* __restrict__ B_tmp = (scalar_t*)((void*)(intermediate_cache0 + M * 2 * N));
1294 +
1295 + CHECK_MOE_SCALES_FP8(0, 1);
1296 + shared_expert_fp8_kernel_impl<scalar_t>(
1297 + out_hidden_states.data_ptr<scalar_t>(),
1298 + intermediate_cache0,
1299 + intermediate_cache1,
1300 + B_tmp,
1301 + C_tmp,
1302 + hidden_states.data_ptr<scalar_t>(),
1303 + packed_w1.data_ptr<at::Float8_e4m3fn>(),
1304 + packed_w2.data_ptr<at::Float8_e4m3fn>(),
1305 + w1s.data_ptr<float>(),
1306 + w2s.data_ptr<float>(),
1307 + block_size_N,
1308 + block_size_K,
1309 + fused_experts_out.data_ptr<scalar_t>(),
1310 + routed_scaling_factor,
1311 + M,
1312 + N,
1313 + K);
1314 + } else {
1315 + shared_expert_kernel_impl<scalar_t>(
1316 + out_hidden_states.data_ptr<scalar_t>(),
1317 + intermediate_cache1,
1318 + C_tmp,
1319 + hidden_states.data_ptr<scalar_t>(),
1320 + packed_w1.data_ptr<scalar_t>(),
1321 + packed_w2.data_ptr<scalar_t>(),
1322 + fused_experts_out.data_ptr<scalar_t>(),
1323 + routed_scaling_factor,
1324 + M,
1325 + N,
1326 + K);
1327 + }
1328 + });
1329 + return out_hidden_states;
1330 +}
csrc/cpu/sgl-kernels/moe_fp8.cpp new
+502
@@ -0,0 +1,502 @@
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 +
8 +// clang-format off
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 <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
145 +
146 +template <typename scalar_t>
147 +void fused_experts_fp8_kernel_impl(
148 + scalar_t* __restrict__ output,
149 + scalar_t* __restrict__ ic0,
150 + scalar_t* __restrict__ ic1,
151 + scalar_t* __restrict__ ic2,
152 + scalar_t* __restrict__ A_tmp,
153 + scalar_t* __restrict__ B_tmp,
154 + float* __restrict__ C_tmp,
155 + const scalar_t* __restrict__ input,
156 + const at::Float8_e4m3fn* __restrict__ packed_w1,
157 + const at::Float8_e4m3fn* __restrict__ packed_w2,
158 + const float* __restrict__ w1s,
159 + const float* __restrict__ w2s,
160 + int64_t block_size_N,
161 + int64_t block_size_K,
162 + const float* __restrict__ topk_weights,
163 + const int32_t* __restrict__ sorted_ids,
164 + const int32_t* __restrict__ expert_ids,
165 + const int32_t* __restrict__ offsets,
166 + int64_t M,
167 + int64_t N,
168 + int64_t K,
169 + int64_t E,
170 + int64_t topk,
171 + int64_t num_tokens_post_pad) {
172 +
173 + constexpr int64_t BLOCK_M = block_size_m();
174 + constexpr int64_t BLOCK_N = block_size_n();
175 +
176 + // stage 1: intermediate_cache0 = hidden_states @ w1
177 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
178 + const int64_t NB = div_up(2 * N, BLOCK_N);
179 + int64_t scale_size_N = div_up(2 * N, block_size_N);
180 + int64_t scale_size_K = div_up(K, block_size_K);
181 + int64_t blocks_n_per_group = block_size_N / BLOCK_N;
182 +
183 + const int64_t stride_e = 2 * N * K;
184 + const int64_t stride_n = K;
185 +
186 + // 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) {
188 + // get local pointers
189 + int tid = at::get_thread_num();
190 + scalar_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
191 +
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 +
198 + int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
199 +
200 + // B shape [K, n_size] in vnni format
201 + int32_t expert_id = expert_ids[mb];
202 + 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;
204 +
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];
208 +
209 + const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
210 + is_brgemm_used = is_brgemm_used || use_brgemm;
211 +
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);
215 + }
216 +
217 + const int64_t offset = offsets[mb];
218 + tinygemm_kernel<scalar_t>(
219 + /* A */ A,
220 + /* B */ B,
221 + /* C */ ic0 + offset * 2 * N + nb * BLOCK_N,
222 + /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
223 + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
224 + /* scale */ Bs,
225 + /* M */ m_size,
226 + /* N */ n_size,
227 + /* K */ K,
228 + /* lda */ K,
229 + /* ldb */ n_size,
230 + /* ldc */ 2 * N,
231 + /* brg */ use_brgemm,
232 + /* block_size_K */ block_size_K);
233 + }
234 +
235 + if (is_brgemm_used) {
236 + at::native::cpublas::brgemm_release();
237 + }
238 + });
239 +
240 + // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
241 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
242 + 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);
248 + }
249 + });
250 +
251 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
252 + // w2 : [E, K, N] as [E, OC, IC]
253 + const int64_t OC = K; // rename K as OC
254 + const int64_t IC = N; // rename N as IC
255 + const int64_t MB2 = MB;
256 + const int64_t NB2 = div_up(OC, BLOCK_N);
257 + scale_size_N = div_up(K, block_size_N);
258 + scale_size_K = div_up(N, block_size_K);
259 + const int64_t stride_e2 = OC * IC;
260 + const int64_t stride_oc = IC;
261 +
262 + // 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();
265 + alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
266 +
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 +
273 + int64_t m_size = offsets[mb + 1] - offsets[mb];
274 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
275 +
276 + const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(m_size);
277 + is_brgemm_used = is_brgemm_used || use_brgemm;
278 +
279 + // A ptr from ic1 of [M * topk, N] in sorted order
280 + // so as to avoid copy A to tmp buffer again
281 + const scalar_t* __restrict__ A = ic1 + offsets[mb] * N;
282 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
283 +
284 + // B shape [IC, n_size] in vnni format
285 + int32_t expert_id = expert_ids[mb];
286 + 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;
288 +
289 + tinygemm_kernel<scalar_t>(
290 + /* A */ A,
291 + /* B */ B,
292 + /* C */ C,
293 + /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
294 + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
295 + /* scale */ Bs,
296 + /* M */ m_size,
297 + /* N */ n_size,
298 + /* K */ IC,
299 + /* lda */ IC,
300 + /* ldb */ n_size,
301 + /* ldc */ BLOCK_N,
302 + /* brg */ use_brgemm,
303 + /* block_size_K */ block_size_K);
304 +
305 + // 2.b copy from C to ic2 in original order
306 + // and also mul topk_weights in float32
307 + for (int64_t m = 0; m < m_size; ++m) {
308 + int32_t index = A_ids[m];
309 + float weight = topk_weights[index];
310 + copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
311 + }
312 + }
313 +
314 + if (is_brgemm_used) {
315 + at::native::cpublas::brgemm_release();
316 + }
317 + });
318 +
319 + // stage 3: out = intermediate_cache2.sum(dim=1)
320 + // from [M, topk, K] to [M, K]
321 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
322 + for (int64_t m = begin; m < end; ++m) {
323 + sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
324 + }
325 + });
326 +}
327 +
328 +#define INSTANTIATE_MOE_FP8_TEMPLATE(TYPE) \
329 + template void fused_experts_fp8_kernel_impl<TYPE>( \
330 + TYPE* __restrict__ output, \
331 + TYPE* __restrict__ ic0, \
332 + TYPE* __restrict__ ic1, \
333 + TYPE* __restrict__ ic2, \
334 + TYPE* __restrict__ A_tmp, \
335 + TYPE* __restrict__ B_tmp, \
336 + float* __restrict__ C_tmp, \
337 + const TYPE* __restrict__ input, \
338 + const at::Float8_e4m3fn* __restrict__ packed_w1, \
339 + const at::Float8_e4m3fn* __restrict__ packed_w2, \
340 + const float* __restrict__ w1s, \
341 + const float* __restrict__ w2s, \
342 + int64_t block_size_N, \
343 + int64_t block_size_K, \
344 + const float* __restrict__ topk_weights, \
345 + const int32_t* __restrict__ sorted_ids, \
346 + const int32_t* __restrict__ expert_ids, \
347 + const int32_t* __restrict__ offsets, \
348 + int64_t M, \
349 + int64_t N, \
350 + int64_t K, \
351 + int64_t E, \
352 + int64_t topk, \
353 + int64_t num_tokens_post_pad)
354 +
355 +INSTANTIATE_MOE_FP8_TEMPLATE(at::BFloat16);
356 +INSTANTIATE_MOE_FP8_TEMPLATE(at::Half);
357 +
358 +template <typename scalar_t>
359 +void shared_expert_fp8_kernel_impl(
360 + scalar_t* __restrict__ output,
361 + scalar_t* __restrict__ ic0,
362 + scalar_t* __restrict__ ic1,
363 + scalar_t* __restrict__ B_tmp,
364 + float* __restrict__ C_tmp,
365 + const scalar_t* __restrict__ input,
366 + const at::Float8_e4m3fn* __restrict__ packed_w1,
367 + const at::Float8_e4m3fn* __restrict__ packed_w2,
368 + const float* __restrict__ w1s,
369 + const float* __restrict__ w2s,
370 + int64_t block_size_N,
371 + int64_t block_size_K,
372 + const scalar_t* __restrict__ fused_experts_out,
373 + float routed_scaling_factor,
374 + int64_t M,
375 + int64_t N,
376 + int64_t K) {
377 +
378 + constexpr int64_t BLOCK_M = block_size_m();
379 + constexpr int64_t BLOCK_N = block_size_n();
380 +
381 + // stage 1: intermediate_cache0 = hidden_states @ w1
382 + const int64_t MB = div_up(M, BLOCK_M);
383 + const int64_t NB = div_up(2 * N, BLOCK_N);
384 + int64_t scale_size_K = div_up(K, block_size_K);
385 + int64_t blocks_n_per_group = block_size_N / BLOCK_N;
386 +
387 + const bool use_brgemm = can_use_brgemm<at::Float8_e4m3fn>(M);
388 +
389 + at::parallel_for(0, MB * NB, 0, [&](int64_t begin, int64_t end) {
390 + int tid = at::get_thread_num();
391 +
392 + for (int64_t i = begin; i < end; ++i) {
393 + int64_t mb = i / NB;
394 + int64_t nb = i % NB;
395 + int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
396 + int64_t n_size = std::min(2 * N - nb * BLOCK_N, BLOCK_N);
397 +
398 + tinygemm_kernel<scalar_t>(
399 + /* A */ input + mb * BLOCK_M * K,
400 + /* B */ packed_w1 + nb * BLOCK_N * K,
401 + /* C */ ic0 + mb * BLOCK_M * 2 * N + nb * BLOCK_N,
402 + /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
403 + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
404 + /* scale */ w1s + (nb / blocks_n_per_group) * scale_size_K,
405 + /* M */ m_size,
406 + /* N */ n_size,
407 + /* K */ K,
408 + /* lda */ K,
409 + /* ldb */ n_size,
410 + /* ldc */ 2 * N,
411 + /* brg */ use_brgemm,
412 + /* block_size_K */ block_size_K);
413 + }
414 +
415 + if (use_brgemm) {
416 + at::native::cpublas::brgemm_release();
417 + }
418 + });
419 +
420 + // stage 1.5: intermediate_cache1 = silu(intermediate_cache0)
421 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
422 + 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);
428 + }
429 + });
430 +
431 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
432 + // w2 : [K, N] as [OC, IC]
433 + const int64_t OC = K; // rename K as OC
434 + const int64_t IC = N; // rename N as IC
435 + const int64_t MB2 = MB;
436 + const int64_t NB2 = div_up(K, BLOCK_N);
437 + scale_size_K = div_up(N, block_size_K);
438 +
439 + // 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();
442 + alignas(64) scalar_t C[BLOCK_M * BLOCK_K];
443 +
444 + for (int64_t i = begin; i < end; ++i) {
445 + int64_t mb = i / NB2;
446 + int64_t nb = i % NB2;
447 + int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
448 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
449 +
450 + // 2.a gemm: C = A @ B
451 + tinygemm_kernel<scalar_t>(
452 + /* A */ ic1 + mb * BLOCK_M * N,
453 + /* B */ packed_w2 + nb * BLOCK_N * N,
454 + /* C */ C,
455 + /* Btmp */ B_tmp + tid * BLOCK_N * std::max(K, N),
456 + /* Ctmp */ C_tmp + tid * 2 * BLOCK_M * BLOCK_N,
457 + /* scale */ w2s + (nb / blocks_n_per_group) * scale_size_K,
458 + /* M */ m_size,
459 + /* N */ n_size,
460 + /* K */ IC,
461 + /* lda */ IC,
462 + /* ldb */ n_size,
463 + /* ldc */ BLOCK_N,
464 + /* brg */ use_brgemm,
465 + /* block_size_K */ block_size_K);
466 +
467 + // 2.b copy from C to output and add fused_experts_out
468 + 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;
470 + 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);
472 + }
473 + }
474 + });
475 +
476 + if (use_brgemm) {
477 + at::native::cpublas::brgemm_release();
478 + }
479 +}
480 +
481 +#define INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(TYPE) \
482 + template void shared_expert_fp8_kernel_impl<TYPE>( \
483 + TYPE* __restrict__ output, \
484 + TYPE* __restrict__ ic0, \
485 + TYPE* __restrict__ ic1, \
486 + TYPE* __restrict__ B_tmp, \
487 + float* __restrict__ C_tmp, \
488 + const TYPE* __restrict__ input, \
489 + const at::Float8_e4m3fn* __restrict__ packed_w1, \
490 + const at::Float8_e4m3fn* __restrict__ packed_w2, \
491 + const float* __restrict__ w1s, \
492 + const float* __restrict__ w2s, \
493 + int64_t block_size_N, \
494 + int64_t block_size_K, \
495 + const TYPE* __restrict__ fused_experts_out, \
496 + float routed_scaling_factor, \
497 + int64_t M, \
498 + int64_t N, \
499 + int64_t K)
500 +
501 +INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::BFloat16);
502 +INSTANTIATE_SHARED_EXPERT_FP8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/moe_int8.cpp new
+769
@@ -0,0 +1,769 @@
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 +
8 +// clang-format off
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 +}
47 +
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) {
51 + using bVec = at::vec::Vectorized<scalar_t>;
52 + 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);
81 + }
82 + }
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);
111 + }
112 +}
113 +
114 +/// gemm for w13
115 +template <typename scalar_t, int BLOCK_M, int BLOCK_N>
116 +struct tinygemm_kernel_vnni {
117 + 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) {
122 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
123 + }
124 +};
125 +
126 +#if defined(CPU_CAPABILITY_AVX512)
127 +template <int BLOCK_M, int BLOCK_N>
128 +struct tinygemm_kernel_vnni<at::BFloat16, BLOCK_M, BLOCK_N> {
129 + 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 +
135 + constexpr int ROWS = BLOCK_M;
136 + constexpr int COLS = BLOCK_N / 16;
137 + static_assert(COLS % 2 == 0);
138 +
139 + __m512i va;
140 + __m512i vb0[COLS];
141 + __m512i vb1[COLS];
142 + __m512i vc0[ROWS * COLS];
143 + __m512i vc1[ROWS * COLS];
144 + __m512i vcomp0[COLS];
145 + __m512i vcomp1[COLS];
146 + __m512 was;
147 + __m512 vbs0[COLS];
148 + __m512 vbs1[COLS];
149 +
150 + auto loadc = [&](auto i) {
151 + vc0[i] = _mm512_set1_epi32(0);
152 + vc1[i] = _mm512_set1_epi32(0);
153 + };
154 + Unroll<ROWS * COLS>{}(loadc);
155 +
156 + const int64_t K4 = K >> 2;
157 + const int64_t lda4 = lda >> 2;
158 + const int64_t ldb4 = ldb; // ldb * 4 >> 2;
159 + const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
160 + const int32_t* b0_ptr = reinterpret_cast<const int32_t*>(B0);
161 + const int32_t* b1_ptr = reinterpret_cast<const int32_t*>(B1);
162 +
163 + auto compute = [&](auto i, int64_t k) {
164 + constexpr int row = i / COLS;
165 + constexpr int col = i % COLS;
166 +
167 + if constexpr (col == 0) {
168 + va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
169 + }
170 + if constexpr (row == 0) {
171 + vb0[col] = _mm512_loadu_si512(b0_ptr + k * ldb4 + col * 16);
172 + vb1[col] = _mm512_loadu_si512(b1_ptr + k * ldb4 + col * 16);
173 + }
174 + vc0[i] = _mm512_dpbusd_epi32(vc0[i], va, vb0[col]);
175 + vc1[i] = _mm512_dpbusd_epi32(vc1[i], va, vb1[col]);
176 + };
177 + for (int64_t k = 0; k < K4; ++k) {
178 + Unroll<ROWS * COLS>{}(compute, k);
179 + }
180 +
181 + auto scalec = [&](auto i) {
182 + constexpr int row = i / COLS;
183 + constexpr int col = i % COLS;
184 +
185 + // load a scale
186 + if constexpr(col == 0) {
187 + was = _mm512_set1_ps(As[row]);
188 + }
189 + // load b scale and vcomp
190 + if constexpr (row == 0) {
191 + vbs0[col] = _mm512_loadu_ps(Bs0 + col * 16);
192 + vbs1[col] = _mm512_loadu_ps(Bs1 + col * 16);
193 + vcomp0[col] = _mm512_loadu_si512(Bcomp0 + col * 16);
194 + vcomp1[col] = _mm512_loadu_si512(Bcomp1 + col * 16);
195 + }
196 + __m512 c0 = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc0[i], vcomp0[col]));
197 + __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]));
200 + };
201 + Unroll<ROWS * COLS>{}(scalec);
202 +
203 + using Vec = at::vec::Vectorized<float>;
204 + const Vec one = Vec(1.f);
205 + auto storec = [&](auto i) {
206 + constexpr int row = i / COLS;
207 + constexpr int col = i % COLS;
208 + // for COLS = 2, 4 use 512bit store
209 + if constexpr (col % 2 == 0) {
210 + Vec x0 = _mm512_castsi512_ps(vc0[row * COLS + col + 0]);
211 + Vec x1 = _mm512_castsi512_ps(vc0[row * COLS + col + 1]);
212 + Vec y0 = _mm512_castsi512_ps(vc1[row * COLS + col + 0]);
213 + Vec y1 = _mm512_castsi512_ps(vc1[row * COLS + col + 1]);
214 + // silu
215 + x0 = x0 / (one + x0.neg().exp_u20());
216 + x1 = x1 / (one + x1.neg().exp_u20());
217 + // mul
218 + x0 = x0 * y0;
219 + x1 = x1 * y1;
220 +
221 + _mm512_storeu_si512(
222 + reinterpret_cast<__m512i*>((C + row * ldc + col * 16)),
223 + (__m512i)(_mm512_cvtne2ps_pbh(__m512(x1), __m512(x0))));
224 + }
225 + };
226 + Unroll<ROWS * COLS>{}(storec);
227 + }
228 +};
229 +#endif
230 +
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);
237 +
238 +template <typename scalar_t>
239 +void tinygemm_kernel(
240 + const uint8_t* __restrict__ A,
241 + const int8_t* __restrict__ B0,
242 + const int8_t* __restrict__ B1,
243 + scalar_t* __restrict__ C,
244 + const float* __restrict__ As,
245 + const float* __restrict__ Bs0,
246 + const float* __restrict__ Bs1,
247 + int64_t M,
248 + int64_t N,
249 + int64_t K,
250 + int64_t lda,
251 + int64_t ldb,
252 + int64_t ldc) {
253 +
254 + const int32_t* Bcomp0 = reinterpret_cast<const int32_t*>(B0 + block_size_n() * K);
255 + const int32_t* Bcomp1 = reinterpret_cast<const int32_t*>(B1 + block_size_n() * K);
256 +
257 + // pattern: 1-(2+2)-(8+8)
258 + constexpr int64_t BLOCK_M = 4;
259 + constexpr int64_t BLOCK_N = 32;
260 + const int64_t MB = div_up(M, BLOCK_M);
261 + const int64_t NB = div_up(N, BLOCK_N);
262 + for (int mb = 0; mb < MB; ++mb) {
263 + int64_t mb_start = mb * BLOCK_M;
264 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
265 + for (int64_t nb = 0; nb < NB; ++nb) {
266 + int64_t nb_start = nb * BLOCK_N;
267 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
268 +
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");
275 + }
276 + }
277 + }
278 +}
279 +
280 +/// gemm for w2
281 +template <typename scalar_t, int BLOCK_M, int BLOCK_N>
282 +struct tinygemm_kernel_vnni2 {
283 + 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) {
287 + TORCH_CHECK(false, "tinygemm_kernel_nn: scalar path not implemented!");
288 + }
289 +};
290 +
291 +#if defined(CPU_CAPABILITY_AVX512)
292 +template <int BLOCK_M, int BLOCK_N>
293 +struct tinygemm_kernel_vnni2<at::BFloat16, BLOCK_M, BLOCK_N> {
294 + 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 +
299 + constexpr int ROWS = BLOCK_M;
300 + constexpr int COLS = BLOCK_N / 16;
301 + static_assert(COLS % 2 == 0);
302 +
303 + __m512i va;
304 + __m512i vb[COLS];
305 + __m512i vc[ROWS * COLS];
306 + __m512i vcomp[COLS];
307 + __m512 was;
308 + __m512 vbs[COLS];
309 +
310 + auto loadc = [&](auto i) {
311 + vc[i] = _mm512_set1_epi32(0);
312 + };
313 + Unroll<ROWS * COLS>{}(loadc);
314 +
315 + const int64_t K4 = K >> 2;
316 + const int64_t lda4 = lda >> 2;
317 + const int64_t ldb4 = ldb; // ldb * 4 >> 2;
318 + const int32_t* a_ptr = reinterpret_cast<const int32_t*>(A);
319 + const int32_t* b_ptr = reinterpret_cast<const int32_t*>(B);
320 +
321 + auto compute = [&](auto i, int64_t k) {
322 + constexpr int row = i / COLS;
323 + constexpr int col = i % COLS;
324 +
325 + if constexpr (col == 0) {
326 + va = _mm512_set1_epi32(a_ptr[row * lda4 + k]);
327 + }
328 + if constexpr (row == 0) {
329 + vb[col] = _mm512_loadu_si512(b_ptr + k * ldb4 + col * 16);
330 + }
331 + vc[i] = _mm512_dpbusd_epi32(vc[i], va, vb[col]);
332 + };
333 + for (int64_t k = 0; k < K4; ++k) {
334 + Unroll<ROWS * COLS>{}(compute, k);
335 + }
336 +
337 + auto storec = [&](auto i) {
338 + constexpr int row = i / COLS;
339 + constexpr int col = i % COLS;
340 +
341 + // load a scale
342 + if constexpr(col == 0) {
343 + was = _mm512_set1_ps(As[row]);
344 + }
345 + // load b scale and vcomp per 2 vectors
346 + // also load bias if any
347 + if constexpr (row == 0) {
348 + if constexpr (col % 2 == 0) {
349 + vbs[col + 0] = _mm512_loadu_ps(Bs + col * 16);
350 + vbs[col + 1] = _mm512_loadu_ps(Bs + col * 16 + 16);
351 + vcomp[col + 0] = _mm512_loadu_si512(Bcomp + col * 16);
352 + vcomp[col + 1] = _mm512_loadu_si512(Bcomp + col * 16 + 16);
353 + }
354 + }
355 + __m512 x = _mm512_cvtepi32_ps(_mm512_sub_epi32(vc[i], vcomp[col]));
356 + x = _mm512_mul_ps(_mm512_mul_ps(x, was), vbs[col]);
357 + _mm512_storeu_ps(reinterpret_cast<__m512*>(C + row * ldc + col * 16), x);
358 + };
359 + Unroll<ROWS * COLS>{}(storec);
360 + }
361 +};
362 +#endif
363 +
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);
369 +
370 +template <typename scalar_t>
371 +void tinygemm_kernel(
372 + const uint8_t* __restrict__ A,
373 + const int8_t* __restrict__ B,
374 + float* __restrict__ C,
375 + const float* __restrict__ As,
376 + const float* __restrict__ Bs,
377 + int64_t M,
378 + int64_t N,
379 + int64_t K,
380 + int64_t lda,
381 + int64_t ldb,
382 + int64_t ldc) {
383 +
384 + // B compensation
385 + const int32_t* Bcomp = reinterpret_cast<const int32_t*>(B + block_size_n() * K);
386 +
387 + // pattern: 1-4-16
388 + constexpr int64_t BLOCK_M = 4;
389 + constexpr int64_t BLOCK_N = 64;
390 + const int64_t MB = div_up(M, BLOCK_M);
391 + const int64_t NB = div_up(N, BLOCK_N);
392 + for (int64_t mb = 0; mb < MB; ++mb) {
393 + int64_t mb_start = mb * BLOCK_M;
394 + int64_t mb_size = std::min(BLOCK_M, M - mb_start);
395 + for (int64_t nb = 0; nb < NB; ++nb) {
396 + int64_t nb_start = nb * BLOCK_N;
397 + int64_t nb_size = std::min(BLOCK_N, N - nb_start);
398 +
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");
405 + }
406 + }
407 + }
408 +}
409 +
410 +} // anonymous namespace
411 +
412 +template <typename scalar_t>
413 +void fused_experts_int8_kernel_impl(
414 + scalar_t* __restrict__ output,
415 + scalar_t* __restrict__ ic1,
416 + scalar_t* __restrict__ ic2,
417 + uint8_t* __restrict__ A_tmp,
418 + float* __restrict__ C_tmp,
419 + uint8_t* __restrict__ Aq_tmp,
420 + float* __restrict__ As_tmp,
421 + const scalar_t* __restrict__ input,
422 + const int8_t* __restrict__ packed_w1,
423 + const int8_t* __restrict__ packed_w2,
424 + const float* __restrict__ w1s,
425 + const float* __restrict__ w2s,
426 + const float* __restrict__ topk_weights,
427 + const int32_t* __restrict__ sorted_ids,
428 + const int32_t* __restrict__ expert_ids,
429 + const int32_t* __restrict__ offsets,
430 + int64_t M,
431 + int64_t N,
432 + int64_t K,
433 + int64_t E,
434 + int64_t topk,
435 + int64_t num_tokens_post_pad) {
436 +
437 + // handle 2 tiles per block
438 + constexpr int64_t BLOCK_M = block_size_m();
439 + constexpr int64_t BLOCK_N = block_size_n();
440 +
441 + // stage 0: quantize input to uint8, [M, K]
442 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
443 + 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);
449 + }
450 + });
451 +
452 + // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
453 + const int64_t MB = div_up(num_tokens_post_pad, BLOCK_M);
454 + const int64_t NB = div_up(N, BLOCK_N);
455 +
456 + // strides for w1: [E, 2N, K]
457 + TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
458 +
459 + // K and N are packed for int8
460 + const int64_t packed_K = get_row_size<int8_t>(K);
461 + const int64_t packed_N = get_row_size<int8_t>(N);
462 +
463 + const int64_t stride_e = 2 * N * packed_K;
464 + const int64_t stride_n = packed_K;
465 + // 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) {
467 + // get local pointers
468 + int tid = at::get_thread_num();
469 + uint8_t* __restrict__ A = A_tmp + tid * BLOCK_M * K;
470 +
471 + alignas(64) float As[BLOCK_M];
472 +
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);
480 +
481 + // B shape [K, n_size] in vnni format
482 + 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;
487 +
488 + // 1.a load A
489 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
490 + int64_t m_size = offsets[mb + 1] - offsets[mb];
491 +
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];
496 + }
497 +
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);
514 + }
515 + });
516 +
517 + // stage 1.5: quantize ic1 to uint8, [M * topk, N]
518 + at::parallel_for(0, M * topk, 0, [&](int64_t begin, int64_t end) {
519 + 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);
525 + }
526 + });
527 +
528 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
529 + // w2 : [E, K, N] as [E, OC, IC]
530 + const int64_t OC = K; // rename K as OC
531 + const int64_t IC = N; // rename N as IC
532 + const int64_t MB2 = MB;
533 + const int64_t NB2 = div_up(OC, BLOCK_N);
534 + const int64_t stride_e2 = OC * packed_N;
535 + const int64_t stride_oc = packed_N;
536 +
537 + // parallel on [MB2, NB2]
538 + at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
539 + // get local pointers
540 + int tid = at::get_thread_num();
541 + // we won't be using C1 for gemm2
542 + float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
543 +
544 + for (int64_t i = begin; i < end; ++i) {
545 + int64_t mb = i / NB2;
546 + int64_t nb = i % NB2;
547 +
548 + int64_t m_size = offsets[mb + 1] - offsets[mb];
549 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
550 +
551 + // A ptr from ic1 of [M * topk, N] in sorted order
552 + // so as to avoid copy A to tmp buffer again
553 + const uint8_t* __restrict__ A = Aq_tmp + offsets[mb] * N;
554 + const float* __restrict__ As = As_tmp + offsets[mb];
555 + const int32_t* A_ids = sorted_ids + mb * BLOCK_M;
556 +
557 + // B shape [IC, n_size] in vnni format
558 + int32_t expert_id = expert_ids[mb];
559 + const int8_t* __restrict__ B = packed_w2 + expert_id * stride_e2 + nb * BLOCK_N * stride_oc;
560 + const float* __restrict__ Bs = w2s + expert_id * K + nb * BLOCK_N;
561 +
562 + // 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);
575 +
576 + // 2.b copy from C to ic2 in original order
577 + // and also mul topk_weights in float32
578 + for (int64_t m = 0; m < m_size; ++m) {
579 + int32_t index = A_ids[m];
580 + float weight = topk_weights[index];
581 + copy_mul_stub(ic2 + index * K + nb * BLOCK_N, C + m * BLOCK_N, weight, n_size);
582 + }
583 + }
584 + });
585 +
586 + // stage 3: out = intermediate_cache2.sum(dim=1)
587 + // from [M, topk, K] to [M, K]
588 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
589 + for (int64_t m = begin; m < end; ++m) {
590 + sum_stub(output + m * K, ic2 + m * topk * K, topk, K);
591 + }
592 + });
593 +}
594 +
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)
606 +
607 +INSTANTIATE_MOE_INT8_TEMPLATE(at::BFloat16);
608 +INSTANTIATE_MOE_INT8_TEMPLATE(at::Half);
609 +
610 +template <typename scalar_t>
611 +void shared_expert_int8_kernel_impl(
612 + scalar_t* __restrict__ output,
613 + scalar_t* __restrict__ ic1,
614 + float* __restrict__ C_tmp,
615 + uint8_t* __restrict__ Aq_tmp,
616 + float* __restrict__ As_tmp,
617 + const scalar_t* __restrict__ input,
618 + const int8_t* __restrict__ packed_w1,
619 + const int8_t* __restrict__ packed_w2,
620 + const float* __restrict__ w1s,
621 + const float* __restrict__ w2s,
622 + const scalar_t* __restrict__ fused_experts_out,
623 + float routed_scaling_factor,
624 + int64_t M,
625 + int64_t N,
626 + int64_t K) {
627 +
628 + // handle 2 tiles per block
629 + constexpr int64_t BLOCK_M = block_size_m();
630 + constexpr int64_t BLOCK_N = block_size_n();
631 +
632 + // stage 0: quantize input to uint8, [M, K]
633 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
634 + 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);
640 + }
641 + });
642 +
643 + // stage 1: intermediate_cache1 = silu(hidden_states @ w1)
644 + const int64_t MB = div_up(M, BLOCK_M);
645 + const int64_t NB = div_up(N, BLOCK_N);
646 +
647 + TORCH_CHECK(N % BLOCK_N == 0, "Fixme when N is not multiples of ", BLOCK_N);
648 +
649 + // K and N are packed for int8
650 + const int64_t packed_K = get_row_size<int8_t>(K);
651 + const int64_t packed_N = get_row_size<int8_t>(N);
652 + const int64_t stride_n = packed_K;
653 +
654 + // 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);
663 + int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
664 +
665 + // A shape [m_size, K]
666 + const uint8_t* A = Aq_tmp + mb * BLOCK_M * K;
667 + const float* As = As_tmp + mb * BLOCK_M;
668 +
669 + // 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);
690 + }
691 + });
692 +
693 + // stage 1.5: quantize ic1 to uint8, [M * topk, N]
694 + at::parallel_for(0, M, 0, [&](int64_t begin, int64_t end) {
695 + 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);
701 + }
702 + });
703 +
704 + // stage 2: intermediate_cache2 = intermediate_cache1 @ w2
705 + // w2 : [K, N] as [OC, IC]
706 + const int64_t OC = K; // rename K as OC
707 + const int64_t IC = N; // rename N as IC
708 + const int64_t MB2 = MB;
709 + const int64_t NB2 = div_up(OC, BLOCK_N);
710 + const int64_t stride_oc = packed_N;
711 +
712 + // parallel on [MB2, NB2]
713 + at::parallel_for(0, MB2 * NB2, 0, [&](int64_t begin, int64_t end) {
714 + // get local pointers
715 + int tid = at::get_thread_num();
716 + // we won't be using C1 for gemm2
717 + float* __restrict__ C = C_tmp + tid * 2 * BLOCK_M * BLOCK_N;
718 +
719 + for (int64_t i = begin; i < end; ++i) {
720 + int64_t mb = i / NB2;
721 + int64_t nb = i % NB2;
722 +
723 + int64_t m_size = std::min(M - mb * BLOCK_M, BLOCK_M);
724 + int64_t n_size = std::min(OC - nb * BLOCK_N, BLOCK_N);
725 +
726 + // A shape [m_size, IC]
727 + const uint8_t* __restrict__ A = Aq_tmp + mb * BLOCK_M * N;
728 + const float* __restrict__ As = As_tmp + mb * BLOCK_M;
729 +
730 + // B shape [IC, n_size] in vnni format
731 + const int8_t* __restrict__ B = packed_w2 + nb * BLOCK_N * stride_oc;
732 + const float* __restrict__ Bs = w2s + nb * BLOCK_N;
733 +
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);
747 +
748 + // 2.b copy from C to output and add fused_experts_out
749 + 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;
751 + 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);
753 + }
754 + }
755 + });
756 +}
757 +
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)
767 +
768 +INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::BFloat16);
769 +INSTANTIATE_SHARED_EXPERT_INT8_TEMPLATE(at::Half);
csrc/cpu/sgl-kernels/vec.h new
+308
@@ -0,0 +1,308 @@
1 +// Adapted from
2 +// https://github.com/sgl-project/sglang/tree/main/sgl-kernel/csrc/cpu
3 +
4 +#pragma once
5 +
6 +// clang-format off
7 +
8 +#if defined(__AVX512F__) && defined(__AVX512BF16__) && defined(__AMX_BF16__)
9 +#define CPU_CAPABILITY_AVX512
10 +#endif
11 +
12 +#include <ATen/cpu/vec/functional.h>
13 +#include <ATen/cpu/vec/vec.h>
14 +
15 +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>
21 +inline Vectorized<scalar_t> convert_from_float_ext(const Vectorized<float>& a, const Vectorized<float>& b) {
22 + return at::vec::convert_from_float<scalar_t>(a, b);
23 +}
24 +
25 +#if defined(CPU_CAPABILITY_AVX512)
26 +
27 +// `at::vec::convert_from_float<>` from PyTorch doesn't have avx512-bf16 intrinsics
28 +// use native instruction for bfloat16->float32 conversion
29 +template <>
30 +inline Vectorized<at::BFloat16> convert_from_float_ext<at::BFloat16>(const Vectorized<float>& a, const Vectorized<float>& b) {
31 + return (__m512i)(_mm512_cvtne2ps_pbh(__m512(b), __m512(a)));
32 +}
33 +
34 +#define CVT_BF16_TO_FP32(a) \
35 + _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu16_epi32(a), 16))
36 +
37 +#define CVT_FP16_TO_FP32(a) \
38 + _mm512_cvtps_ph(a, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC))
39 +
40 +// this doesn't hanel NaN.
41 +inline __m512bh cvt_e4m3_bf16_intrinsic_no_nan(__m256i fp8_vec) {
42 + 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);
51 +
52 + const __mmask32 is_nonzero = _mm512_cmpneq_epi16_mask(x, _mm512_setzero_si512());
53 + return (__m512bh)_mm512_maskz_mov_epi16(is_nonzero, combined);
54 +}
55 +
56 +inline __m512bh cvt_e4m3_bf16_intrinsic_without_denorm(__m256i fp8_vec) {
57 + // The following conversion is without denorm behavior, that is to say,
58 + // Max subnorm : S.0000.111 = 0.875 ∗ 2**(−6)
59 + // Min subnorm : S.0000.001 = 2**(−9)
60 + // 0.0019 ~ 0.0137 cannot be converted correctly.
61 + __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
62 + auto mask = _mm512_cmpneq_epi16_mask(
63 + _mm512_and_si512(x, _mm512_set1_epi16(127)),
64 + _mm512_setzero_si512()); // mask = x & 0x7f
65 + auto mask_nan = _mm512_cmpneq_epi16_mask(
66 + _mm512_and_si512(x, _mm512_set1_epi16(127)),
67 + _mm512_set1_epi16(127)); // mask_nan = x & 0x7f
68 + auto mantissa = _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4); // mantissa = (x & 7) << 4
69 + auto exponent = _mm512_add_epi16(
70 + _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3),
71 + _mm512_set1_epi16(120)); // exponent = (((x >> 3) & 15) + 120)
72 + auto nonsign = _mm512_maskz_mov_epi16(mask, _mm512_or_si512(mantissa, _mm512_slli_epi16(exponent, 7)));
73 + nonsign = _mm512_mask_mov_epi16(_mm512_set1_epi16(0x7fff), mask_nan, nonsign); // deal with Nan
74 + return (__m512bh)(_mm512_or_si512(
75 + nonsign,
76 + _mm512_slli_epi16(
77 + _mm512_and_si512(x, _mm512_set1_epi16(128)),
78 + 8))); // add sign (x & 128) << 8
79 +}
80 +
81 +inline __m512bh cvt_e4m3_bf16_intrinsic_with_denorm(__m256i fp8_vec) {
82 + __m512i x = _mm512_cvtepu8_epi16(fp8_vec);
83 + __m512i lg2mant = _mm512_mask_mov_epi16(
84 + _mm512_mask_mov_epi16(
85 + _mm512_setzero_si512(), _mm512_test_epi16_mask(x, _mm512_set1_epi16(2)), _mm512_set1_epi16(1)),
86 + _mm512_test_epi16_mask(x, _mm512_set1_epi16(4)),
87 + _mm512_set1_epi16(2));
88 + return (__m512bh)(_mm512_or_si512(
89 + _mm512_maskz_mov_epi16(
90 + _mm512_cmpneq_epi16_mask(_mm512_and_si512(x, _mm512_set1_epi16(127)), _mm512_setzero_si512()),
91 + _mm512_mask_blend_epi16(
92 + _mm512_test_epi16_mask(x, _mm512_set1_epi16(120)),
93 + _mm512_or_si512(
94 + _mm512_and_si512(
95 + _mm512_sllv_epi16(
96 + _mm512_and_si512(x, _mm512_set1_epi16(3)), _mm512_sub_epi16(_mm512_set1_epi16(7), lg2mant)),
97 + _mm512_set1_epi16(0x007f)),
98 + _mm512_slli_epi16(_mm512_add_epi16(lg2mant, _mm512_set1_epi16(118)), 7)),
99 + _mm512_or_si512(
100 + _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(7)), 4),
101 + _mm512_slli_epi16(
102 + _mm512_add_epi16(
103 + _mm512_srli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(120)), 3), _mm512_set1_epi16(120)),
104 + 7)))),
105 + _mm512_slli_epi16(_mm512_and_si512(x, _mm512_set1_epi16(128)), 8)));
106 +}
107 +
108 +inline __m512bh CVT_FP8_TO_BF16(__m256i a) {
109 +#ifdef SGLANG_CPU_FP8_CVT_FTZ
110 + return cvt_e4m3_bf16_intrinsic_no_nan(a);
111 +#else
112 + return cvt_e4m3_bf16_intrinsic_with_denorm(a);
113 +#endif
114 +}
115 +
116 +#endif
117 +
118 +// vector to scalar reduction
119 +#if defined(CPU_CAPABILITY_AVX512) && 0
120 +inline float vec_reduce_sum(const Vectorized<float>& a) {
121 + return _mm512_reduce_add_ps(__m512(a));
122 +}
123 +
124 +inline float vec_reduce_max(const Vectorized<float>& a) {
125 + return _mm512_reduce_max_ps(__m512(a));
126 +}
127 +#else
128 +inline float vec_reduce_sum(const Vectorized<float>& a) {
129 + return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return x + y; }, a);
130 +}
131 +
132 +inline float vec_reduce_max(const Vectorized<float>& a) {
133 + return vec_reduce_all([](Vectorized<float>& x, Vectorized<float>& y) { return maximum(x, y); }, a);
134 +}
135 +#endif
136 +
137 +// https://github.com/InternLM/lmdeploy/blob/086481ed84b59bee3b8e4274e5fc69620040c048/lmdeploy/pytorch/kernels/cuda/w8a8_triton_kernels.py#L282
138 +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
143 + for (int64_t k = 0; k < K; ++k) {
144 + const float val = static_cast<float>(A[k]);
145 + amax = std::max(amax, std::abs(val));
146 + }
147 +
148 + amax = std::max(amax, eps);
149 + const float scale = amax / 127;
150 + const float inv_scale = 127 / amax;
151 +
152 + for (int64_t k = 0; k < K; ++k) {
153 + const float val = static_cast<float>(A[k]) * inv_scale;
154 + Aq[k] = (uint8_t)(std::round(val)) + 128;
155 + }
156 + As = scale;
157 +}
158 +
159 +#if defined(CPU_CAPABILITY_AVX512)
160 +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 +
164 + const __m512 signBit = _mm512_set1_ps(-0.0f);
165 + const __m512i off = _mm512_set1_epi32(128);
166 +
167 + // K is 32x, no remainder
168 + float amax = 0.f;
169 + __m512 vamax0 = _mm512_set1_ps(0.f);
170 + __m512 vamax1 = _mm512_set1_ps(0.f);
171 + for (int64_t k = 0; k < K; k += 32) {
172 + __m512i va = _mm512_loadu_si512((void*)(A + k));
173 + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
174 + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
175 + vamax0 = _mm512_max_ps(vamax0, _mm512_andnot_ps(signBit, va0));
176 + vamax1 = _mm512_max_ps(vamax1, _mm512_andnot_ps(signBit, va1));
177 + }
178 + amax = _mm512_reduce_max_ps(_mm512_max_ps(vamax0, vamax1));
179 + amax = std::max(amax, eps);
180 + const float scale = amax / 127;
181 + const float inv_scale = 127 / amax;
182 + const __m512 vd = _mm512_set1_ps(inv_scale);
183 +
184 + for (int64_t k = 0; k < K; k += 32) {
185 + __m512i va = _mm512_loadu_si512((void*)(A + k));
186 + __m512 va0 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 0));
187 + __m512 va1 = CVT_BF16_TO_FP32(_mm512_extracti32x8_epi32(va, 1));
188 + va0 = _mm512_mul_ps(va0, vd);
189 + va1 = _mm512_mul_ps(va1, vd);
190 + va0 = _mm512_roundscale_ps(va0, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
191 + va1 = _mm512_roundscale_ps(va1, (_MM_FROUND_TO_NEAREST_INT | _MM_FROUND_NO_EXC));
192 + __m128i i0 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va0), off));
193 + __m128i i1 = _mm512_cvtepi32_epi8(_mm512_add_epi32(_mm512_cvtps_epi32(va1), off));
194 + _mm256_storeu_si256(reinterpret_cast<__m256i*>(Aq + k), _mm256_set_m128i(i1, i0));
195 + }
196 + As = scale;
197 +}
198 +#endif
199 +
200 +// transpose utils
201 +// taken from my PR in ggml: https://github.com/ggml-org/llama.cpp/pull/8998
202 +#if defined(CPU_CAPABILITY_AVX512)
203 +inline void transpose_16x16_32bit(__m512i * v) {
204 + __m512i v1[16];
205 + v1[0] = _mm512_unpacklo_epi32(v[0], v[1]);
206 + v1[1] = _mm512_unpackhi_epi32(v[0], v[1]);
207 + v1[2] = _mm512_unpacklo_epi32(v[2], v[3]);
208 + v1[3] = _mm512_unpackhi_epi32(v[2], v[3]);
209 + v1[4] = _mm512_unpacklo_epi32(v[4], v[5]);
210 + v1[5] = _mm512_unpackhi_epi32(v[4], v[5]);
211 + v1[6] = _mm512_unpacklo_epi32(v[6], v[7]);
212 + v1[7] = _mm512_unpackhi_epi32(v[6], v[7]);
213 + v1[8] = _mm512_unpacklo_epi32(v[8], v[9]);
214 + v1[9] = _mm512_unpackhi_epi32(v[8], v[9]);
215 + v1[10] = _mm512_unpacklo_epi32(v[10], v[11]);
216 + v1[11] = _mm512_unpackhi_epi32(v[10], v[11]);
217 + v1[12] = _mm512_unpacklo_epi32(v[12], v[13]);
218 + v1[13] = _mm512_unpackhi_epi32(v[12], v[13]);
219 + v1[14] = _mm512_unpacklo_epi32(v[14], v[15]);
220 + v1[15] = _mm512_unpackhi_epi32(v[14], v[15]);
221 +
222 + v[0] = _mm512_unpacklo_epi64(v1[0], v1[2]);
223 + v[1] = _mm512_unpackhi_epi64(v1[0], v1[2]);
224 + v[2] = _mm512_unpacklo_epi64(v1[1], v1[3]);
225 + v[3] = _mm512_unpackhi_epi64(v1[1], v1[3]);
226 + v[4] = _mm512_unpacklo_epi64(v1[4], v1[6]);
227 + v[5] = _mm512_unpackhi_epi64(v1[4], v1[6]);
228 + v[6] = _mm512_unpacklo_epi64(v1[5], v1[7]);
229 + v[7] = _mm512_unpackhi_epi64(v1[5], v1[7]);
230 + v[8] = _mm512_unpacklo_epi64(v1[8], v1[10]);
231 + v[9] = _mm512_unpackhi_epi64(v1[8], v1[10]);
232 + v[10] = _mm512_unpacklo_epi64(v1[9], v1[11]);
233 + v[11] = _mm512_unpackhi_epi64(v1[9], v1[11]);
234 + v[12] = _mm512_unpacklo_epi64(v1[12], v1[14]);
235 + v[13] = _mm512_unpackhi_epi64(v1[12], v1[14]);
236 + v[14] = _mm512_unpacklo_epi64(v1[13], v1[15]);
237 + v[15] = _mm512_unpackhi_epi64(v1[13], v1[15]);
238 +
239 + v1[0] = _mm512_shuffle_i32x4(v[0], v[4], 0x88);
240 + v1[1] = _mm512_shuffle_i32x4(v[1], v[5], 0x88);
241 + v1[2] = _mm512_shuffle_i32x4(v[2], v[6], 0x88);
242 + v1[3] = _mm512_shuffle_i32x4(v[3], v[7], 0x88);
243 + v1[4] = _mm512_shuffle_i32x4(v[0], v[4], 0xdd);
244 + v1[5] = _mm512_shuffle_i32x4(v[1], v[5], 0xdd);
245 + v1[6] = _mm512_shuffle_i32x4(v[2], v[6], 0xdd);
246 + v1[7] = _mm512_shuffle_i32x4(v[3], v[7], 0xdd);
247 + v1[8] = _mm512_shuffle_i32x4(v[8], v[12], 0x88);
248 + v1[9] = _mm512_shuffle_i32x4(v[9], v[13], 0x88);
249 + v1[10] = _mm512_shuffle_i32x4(v[10], v[14], 0x88);
250 + v1[11] = _mm512_shuffle_i32x4(v[11], v[15], 0x88);
251 + v1[12] = _mm512_shuffle_i32x4(v[8], v[12], 0xdd);
252 + v1[13] = _mm512_shuffle_i32x4(v[9], v[13], 0xdd);
253 + v1[14] = _mm512_shuffle_i32x4(v[10], v[14], 0xdd);
254 + v1[15] = _mm512_shuffle_i32x4(v[11], v[15], 0xdd);
255 +
256 + v[0] = _mm512_shuffle_i32x4(v1[0], v1[8], 0x88);
257 + v[1] = _mm512_shuffle_i32x4(v1[1], v1[9], 0x88);
258 + v[2] = _mm512_shuffle_i32x4(v1[2], v1[10], 0x88);
259 + v[3] = _mm512_shuffle_i32x4(v1[3], v1[11], 0x88);
260 + v[4] = _mm512_shuffle_i32x4(v1[4], v1[12], 0x88);
261 + v[5] = _mm512_shuffle_i32x4(v1[5], v1[13], 0x88);
262 + v[6] = _mm512_shuffle_i32x4(v1[6], v1[14], 0x88);
263 + v[7] = _mm512_shuffle_i32x4(v1[7], v1[15], 0x88);
264 + v[8] = _mm512_shuffle_i32x4(v1[0], v1[8], 0xdd);
265 + v[9] = _mm512_shuffle_i32x4(v1[1], v1[9], 0xdd);
266 + v[10] = _mm512_shuffle_i32x4(v1[2], v1[10], 0xdd);
267 + v[11] = _mm512_shuffle_i32x4(v1[3], v1[11], 0xdd);
268 + v[12] = _mm512_shuffle_i32x4(v1[4], v1[12], 0xdd);
269 + v[13] = _mm512_shuffle_i32x4(v1[5], v1[13], 0xdd);
270 + v[14] = _mm512_shuffle_i32x4(v1[6], v1[14], 0xdd);
271 + v[15] = _mm512_shuffle_i32x4(v1[7], v1[15], 0xdd);
272 +}
273 +
274 +// remove warning : ignoring attributes on template argument ‘__m512i’ [-Wignored-attributes]
275 +#pragma GCC diagnostic push
276 +#pragma GCC diagnostic ignored "-Wignored-attributes"
277 +
278 +// transpose from [2, 32] to [32, 2]
279 +inline std::tuple<__m512i, __m512i> transpose_2x32_16bit(__m512i r0, __m512i r1) {
280 + // r0: {a0, a1, ..., a31}
281 + // r1: {b0, b1, ..., b31}
282 + //
283 + // d0: {a0, b0, ..., a15, b15}
284 + // d1: {a16, b16, ..., a31, b31}
285 + //
286 + __m512i d0 = _mm512_unpacklo_epi16(r0, r1);
287 + __m512i d1 = _mm512_unpackhi_epi16(r0, r1);
288 + r0 = _mm512_shuffle_i32x4(d0, d1, 0x88);
289 + r1 = _mm512_shuffle_i32x4(d0, d1, 0xdd);
290 + d0 = _mm512_shuffle_i32x4(r0, r1, 0x88);
291 + d1 = _mm512_shuffle_i32x4(r0, r1, 0xdd);
292 + return std::make_tuple(d0, d1);
293 +}
294 +#pragma GCC diagnostic pop
295 +
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;
306 +}
307 +
308 +} // anonymous namespace
csrc/cpu/shm.cpp
+94 -84
@@ -7,9 +7,10 @@
7
8 namespace {
9 #define MAX_SHM_RANK_NUM 8
10 -#define MAX_THREAD_NUM 12
11 -#define PER_THREAD_SHM_BUFFER_BYTES (4 * 1024 * 1024)
12 -#define MIN_THREAD_PROCESS_SIZE (8 * 1024)
10 +#define PER_THREAD_SHM_BUFFER_BYTES (2 * 1024 * 1024)
11 +static_assert(PER_THREAD_SHM_BUFFER_BYTES % 2 == 0);
12 +#define PER_THREAD_SHM_BUFFER_OFFSET (PER_THREAD_SHM_BUFFER_BYTES >> 1)
13 +#define MIN_THREAD_PROCESS_SIZE (256)
14 #define MAX_P2P_SEND_TENSOR_NUM 8
15
16 template <typename scalar_t>
@@ -32,10 +33,10 @@ struct KernelVecType<c10::Half> {
33 using scalar_vec_t = vec_op::FP16Vec16;
34 };
35
35 -enum class ThreadSHMStat : char { THREAD_READY = 0, SHM_DATA_READY, DONE };
36 -
36 struct ThreadSHMContext {
38 - volatile ThreadSHMStat thread_stats[MAX_SHM_RANK_NUM];
37 + volatile char _curr_thread_stamp;
38 + volatile char _ready_thread_stamp;
39 + char _padding1[6];
40 int thread_id;
41 int thread_num;
42 int rank;
@@ -44,14 +45,19 @@ struct ThreadSHMContext {
45 int swizzled_ranks[MAX_SHM_RANK_NUM];
46 void* thread_shm_ptrs[MAX_SHM_RANK_NUM];
47 ThreadSHMContext* shm_contexts[MAX_SHM_RANK_NUM];
48 + size_t _thread_buffer_mask;
49 + char _padding2[56];
50
51 ThreadSHMContext(const int thread_id, const int thread_num, const int rank,
52 const int group_size, void* thread_shm_ptr)
50 - : thread_id(thread_id),
53 + : _curr_thread_stamp(1),
54 + _ready_thread_stamp(0),
55 + thread_id(thread_id),
56 thread_num(thread_num),
57 rank(rank),
58 group_size(group_size),
54 - _spinning_count(0) {
59 + _spinning_count(0),
60 + _thread_buffer_mask(0) {
61 static_assert(sizeof(ThreadSHMContext) % 64 == 0);
62 TORCH_CHECK(group_size <= MAX_SHM_RANK_NUM);
63 TORCH_CHECK((size_t)this % 64 == 0);
@@ -60,7 +66,6 @@ struct ThreadSHMContext {
66 shm_contexts[i] = nullptr;
67 thread_shm_ptrs[i] = nullptr;
68 swizzled_ranks[i] = (i + rank) % group_size;
63 - thread_stats[i] = ThreadSHMStat::DONE;
69 }
70 set_context(rank, this, thread_shm_ptr);
71 }
@@ -77,59 +82,66 @@ struct ThreadSHMContext {
82
83 template <typename T>
84 T* get_thread_shm_ptr(int rank) {
80 - return reinterpret_cast<T*>(thread_shm_ptrs[rank]);
85 + return reinterpret_cast<T*>(
86 + reinterpret_cast<int8_t*>(thread_shm_ptrs[rank]) +
87 + (PER_THREAD_SHM_BUFFER_OFFSET & _thread_buffer_mask));
88 + }
89 +
90 + void next_buffer() { _thread_buffer_mask ^= 0xFFFFFFFFFFFFFFFF; }
91 +
92 + char get_curr_stamp() const { return _curr_thread_stamp; }
93 +
94 + char get_ready_stamp() const { return _ready_thread_stamp; }
95 +
96 + void next_stamp() {
97 + _mm_mfence();
98 + _curr_thread_stamp += 1;
99 + }
100 +
101 + void commit_ready_stamp() {
102 + _mm_mfence();
103 + _ready_thread_stamp = _curr_thread_stamp;
104 }
105
106 int get_swizzled_rank(int idx) { return swizzled_ranks[idx]; }
107
85 - void wait_for_all(ThreadSHMStat prev_stat) {
86 - for (int idx = 0; idx < group_size; ++idx) {
108 + template <typename Cond>
109 + void wait_for_all(Cond&& cond) {
110 + for (int idx = 1; idx < group_size; ++idx) {
111 int rank = get_swizzled_rank(idx);
88 - while (thread_stats[rank] == prev_stat) {
89 - ++_spinning_count;
90 - _mm_pause();
91 - }
112 + wait_for_one(rank, std::forward<Cond>(cond));
113 }
93 - vec_op::mem_barrier();
114 }
115
96 - void wait_for_one(int rank, ThreadSHMStat prev_stat) {
97 - while (thread_stats[rank] == prev_stat) {
116 + template <typename Cond>
117 + void wait_for_one(int rank, Cond&& cond) {
118 + ThreadSHMContext* rank_ctx = shm_contexts[rank];
119 + for (;;) {
120 + char local_curr_stamp = get_curr_stamp();
121 + char local_ready_stamp = get_ready_stamp();
122 + char rank_curr_stamp = rank_ctx->get_curr_stamp();
123 + char rank_ready_stamp = rank_ctx->get_ready_stamp();
124 + if (cond(local_curr_stamp, local_ready_stamp, rank_curr_stamp,
125 + rank_ready_stamp)) {
126 + break;
127 + }
128 ++_spinning_count;
129 _mm_pause();
130 }
101 - vec_op::mem_barrier();
131 }
132
104 - void set_thread_stat(ThreadSHMStat stat) {
105 - for (int idx = 0; idx < group_size; ++idx) {
106 - int rank = get_swizzled_rank(idx);
107 - shm_contexts[rank]->thread_stats[this->rank] = stat;
108 - }
109 - }
110 -
111 - void set_thread_stat(int target_rank, ThreadSHMStat stat) {
112 - for (int idx = 0; idx < group_size; ++idx) {
113 - int rank = get_swizzled_rank(idx);
114 - shm_contexts[rank]->thread_stats[target_rank] = stat;
115 - }
133 + static bool check_no_buffer_conflict(char local_curr_stamp,
134 + char local_ready_stamp,
135 + char rank_curr_stamp,
136 + char rank_ready_stamp) {
137 + char temp = rank_curr_stamp + 2;
138 + return local_curr_stamp != temp;
139 }
140
118 - // barrier for all ranks in the group, used for all2all ops
119 - // DONE -> THREAD_READY -> SHM_DATA_READY -> DONE -> ...
120 - void barrier(ThreadSHMStat next_stat) {
121 - if (next_stat == ThreadSHMStat::THREAD_READY) {
122 - set_thread_stat(ThreadSHMStat::THREAD_READY);
123 - wait_for_all(ThreadSHMStat::DONE);
124 - } else if (next_stat == ThreadSHMStat::SHM_DATA_READY) {
125 - set_thread_stat(ThreadSHMStat::SHM_DATA_READY);
126 - wait_for_all(ThreadSHMStat::THREAD_READY);
127 - } else if (next_stat == ThreadSHMStat::DONE) {
128 - set_thread_stat(ThreadSHMStat::DONE);
129 - wait_for_all(ThreadSHMStat::SHM_DATA_READY);
130 - } else {
131 - TORCH_CHECK(false, "Invalid next_stat to barrier.");
132 - }
141 + static bool check_stamp_ready(char local_curr_stamp, char local_ready_stamp,
142 + char rank_curr_stamp, char rank_ready_stamp) {
143 + char temp = local_curr_stamp + 1;
144 + return (local_curr_stamp == rank_ready_stamp) || (temp == rank_ready_stamp);
145 }
146
147 std::string to_string() const {
@@ -164,7 +176,7 @@ class SHMManager {
176 const int group_size)
177 : _rank(rank),
178 _group_size(group_size),
167 - _thread_num(std::min(torch::get_num_threads(), MAX_THREAD_NUM)),
179 + _thread_num(torch::get_num_threads()),
180 _shm_names({""}),
181 _shared_mem_ptrs({nullptr}),
182 _shm_ctx(nullptr) {
@@ -326,7 +338,8 @@ void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) {
338 (total_units_num + thread_num - 1) / thread_num;
339 int64_t per_unit_elem_num = MIN_THREAD_PROCESS_SIZE / sizeof(scalar_t);
340 int64_t max_per_thread_iteration_elem_num =
329 - PER_THREAD_SHM_BUFFER_BYTES / sizeof(scalar_t);
341 + (PER_THREAD_SHM_BUFFER_BYTES >> 1) /
342 + sizeof(scalar_t); // Note: double buffer
343 int64_t per_thread_elem_num = per_unit_elem_num * per_thread_units_num;
344
345 #pragma omp parallel for schedule(static, 1)
@@ -336,10 +349,13 @@ void shm_cc_loop(ThreadSHMContext* ctx, int64_t elem_num, F&& inner_func) {
349 int64_t curr_elem_num =
350 std::min(max_per_thread_iteration_elem_num, end - offset);
351 ThreadSHMContext* thread_ctx = ctx + i;
352 + bool fast_mode = ((end - offset) <= max_per_thread_iteration_elem_num);
353
354 while (curr_elem_num > 0) {
341 - inner_func(thread_ctx, offset, curr_elem_num);
355 + inner_func(thread_ctx, offset, curr_elem_num, fast_mode);
356
357 + thread_ctx->next_stamp();
358 + thread_ctx->next_buffer();
359 offset += max_per_thread_iteration_elem_num;
360 curr_elem_num = std::min(max_per_thread_iteration_elem_num, end - offset);
361 }
@@ -397,7 +413,7 @@ void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
413 shm_cc_ops::shm_cc_loop<scalar_t>(
414 ctx, elem_num,
415 [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
400 - int64_t data_elem_num) {
416 + int64_t data_elem_num, bool fast_mode) {
417 int rank = thread_ctx->rank;
418 scalar_t* thread_shm_ptr =
419 thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
@@ -410,16 +426,17 @@ void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
426 thread_ctx->get_swizzled_rank(idx + 1));
427 });
428
413 - thread_ctx->barrier(ThreadSHMStat::THREAD_READY);
429 + if (!fast_mode) {
430 + thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict);
431 + }
432
433 shm_cc_ops::memcpy_to_shm(thread_shm_ptr, thread_data_ptr,
434 thread_data_elem_num);
417 -
418 - thread_ctx->barrier(ThreadSHMStat::SHM_DATA_READY);
419 -
435 + thread_ctx->commit_ready_stamp();
436 int64_t aligned_data_elem_num =
437 (data_elem_num / vec_elem_num) * vec_elem_num;
438 int64_t i = 0;
439 + thread_ctx->wait_for_all(ThreadSHMContext::check_stamp_ready);
440 #pragma GCC unroll 4
441 for (; i < aligned_data_elem_num; i += vec_elem_num) {
442 vec_t local_data(thread_data_ptr + i); // load from cache
@@ -447,8 +464,6 @@ void all_reduce_sum_impl(ThreadSHMContext* ctx, scalar_t* data,
464 reduced_data.save(thread_data_ptr + i,
465 data_elem_num - aligned_data_elem_num);
466 }
450 -
451 - thread_ctx->barrier(ThreadSHMStat::DONE);
467 });
468
469 return;
@@ -488,18 +503,18 @@ void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num,
503 shm_cc_ops::shm_cc_loop<scalar_t>(
504 ctx, elem_num,
505 [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
491 - int64_t data_elem_num) {
506 + int64_t data_elem_num, bool fast_mode) {
507 int rank = thread_ctx->rank;
508 scalar_t* thread_shm_ptr =
509 thread_ctx->get_thread_shm_ptr<scalar_t>(rank);
510
496 - thread_ctx->barrier(ThreadSHMStat::THREAD_READY);
497 -
498 - shm_cc_ops::memcpy_to_shm(thread_shm_ptr, data + data_offset,
499 - data_elem_num * sizeof(scalar_t));
500 -
501 - thread_ctx->barrier(ThreadSHMStat::SHM_DATA_READY);
511 + if (!fast_mode) {
512 + thread_ctx->wait_for_all(ThreadSHMContext::check_no_buffer_conflict);
513 + }
514
515 + shm_cc_ops::memcpy(thread_shm_ptr, data + data_offset,
516 + data_elem_num * sizeof(scalar_t));
517 + thread_ctx->commit_ready_stamp();
518 if (rank == dst) {
519 shm_cc_ops::memcpy(outputs[rank] + data_offset, data + data_offset,
520 data_elem_num * sizeof(scalar_t));
@@ -508,12 +523,12 @@ void shm_gather_impl(ThreadSHMContext* ctx, scalar_t* data, size_t elem_num,
523 scalar_t* src_ptr =
524 thread_ctx->get_thread_shm_ptr<scalar_t>(src_rank); // shm
525 scalar_t* dst_ptr = outputs[src_rank] + data_offset;
511 - shm_cc_ops::memcpy_from_shm(dst_ptr, src_ptr,
512 - data_elem_num * sizeof(scalar_t));
526 + thread_ctx->wait_for_one(src_rank,
527 + ThreadSHMContext::check_stamp_ready);
528 + shm_cc_ops::memcpy(dst_ptr, src_ptr,
529 + data_elem_num * sizeof(scalar_t));
530 }
531 }
515 -
516 - thread_ctx->barrier(ThreadSHMStat::DONE);
532 });
533
534 return;
@@ -599,7 +614,7 @@ struct TensorListMeta {
614 int8_t _padding[40];
615 };
616
602 -void shm_send_tensor_list_impl(ThreadSHMContext* ctx,
617 +void shm_send_tensor_list_impl(ThreadSHMContext* ctx, int64_t dst,
618 const std::vector<torch::Tensor>& tensor_list) {
619 CPU_KERNEL_GUARD_IN(shm_send_tensor_list_impl)
620 std::vector<torch::Tensor> tensor_list_with_metadata;
@@ -620,12 +635,11 @@ void shm_send_tensor_list_impl(ThreadSHMContext* ctx,
635 shm_cc_ops::shm_cc_loop<int8_t>(
636 ctx, metadata->total_bytes,
637 [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
623 - int64_t data_elem_num) {
638 + int64_t data_elem_num, bool fast_mode) {
639 int rank = thread_ctx->rank;
625 - // Wait until the receiver set the stat to DONE
626 - thread_ctx->wait_for_one(rank, ThreadSHMStat::SHM_DATA_READY);
627 -
640 int64_t curr_shm_offset = 0;
641 + thread_ctx->wait_for_one(dst,
642 + ThreadSHMContext::check_no_buffer_conflict);
643 while (curr_shm_offset < data_elem_num) {
644 MemPiece frag = metadata->get_data(data_offset + curr_shm_offset);
645 frag.size = std::min(frag.size, data_elem_num - curr_shm_offset);
@@ -634,8 +648,7 @@ void shm_send_tensor_list_impl(ThreadSHMContext* ctx,
648 frag.ptr, frag.size);
649 curr_shm_offset += frag.size;
650 }
637 -
638 - thread_ctx->set_thread_stat(rank, ThreadSHMStat::SHM_DATA_READY);
651 + thread_ctx->commit_ready_stamp();
652 });
653 }
654
@@ -646,8 +659,7 @@ std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
659 torch::Tensor metadata_tensor =
660 torch::empty({sizeof(TensorListMeta)}, options);
661
649 - // Wait until the sender set the stat of the thread 0 to SHM_DATA_READY
650 - ctx->wait_for_one(src, ThreadSHMStat::DONE);
662 + ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready);
663 shm_cc_ops::memcpy(metadata_tensor.data_ptr(),
664 ctx->get_thread_shm_ptr<void>(src),
665 sizeof(TensorListMeta));
@@ -664,9 +676,8 @@ std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
676 shm_cc_ops::shm_cc_loop<int8_t>(
677 ctx, metadata.total_bytes,
678 [&](ThreadSHMContext* thread_ctx, int64_t data_offset,
667 - int64_t data_elem_num) {
668 - // Wait until the sender set the stat to SHM_DATA_READY
669 - thread_ctx->wait_for_one(src, ThreadSHMStat::DONE);
679 + int64_t data_elem_num, bool fast_mode) {
680 + ctx->wait_for_one(src, ThreadSHMContext::check_stamp_ready);
681 int64_t curr_shm_offset = 0;
682 while (curr_shm_offset < data_elem_num) {
683 MemPiece frag = metadata.get_data(data_offset + curr_shm_offset);
@@ -677,8 +688,6 @@ std::vector<torch::Tensor> shm_recv_tensor_list_impl(ThreadSHMContext* ctx,
688 frag.size);
689 curr_shm_offset += frag.size;
690 }
680 -
681 - thread_ctx->set_thread_stat(src, ThreadSHMStat::DONE);
691 });
692
693 std::vector<torch::Tensor> tensor_list;
@@ -756,7 +765,8 @@ void shm_send_tensor_list(int64_t handle,
765 int64_t dst) {
766 CPU_KERNEL_GUARD_IN(shm_send_tensor_list)
767 shm_send_tensor_list_impl(
759 - SHMManager::get_singleton_instance(handle)->get_shm_ctx(), tensor_list);
768 + SHMManager::get_singleton_instance(handle)->get_shm_ctx(), dst,
769 + tensor_list);
770 CPU_KERNEL_GUARD_OUT(shm_send_tensor_list)
771 }
772
@@ -778,4 +788,4 @@ std::string join_shm_manager(int64_t handle, const std::string& name) {
788 TORCH_CHECK(shm_manager);
789 shm_manager->join(name);
790 return shm_manager->get_shm_ctx()->to_string();
781 -}
\ No newline at end of file
791 +}
csrc/cpu/torch_bindings.cpp
+43
@@ -50,6 +50,27 @@ void shm_send_tensor_list(int64_t handle,
50
51 std::vector<torch::Tensor> shm_recv_tensor_list(int64_t handle, int64_t src);
52
53 +at::Tensor weight_packed_linear(at::Tensor& mat1, at::Tensor& mat2,
54 + const std::optional<at::Tensor>& bias,
55 + bool is_vnni);
56 +
57 +at::Tensor convert_weight_packed(at::Tensor& weight);
58 +
59 +at::Tensor fused_experts_cpu(
60 + at::Tensor& hidden_states, at::Tensor& w1, at::Tensor& w2,
61 + at::Tensor& topk_weights, at::Tensor& topk_ids, bool inplace,
62 + bool use_int8_w8a8, bool use_fp8_w8a16,
63 + const std::optional<at::Tensor>& w1_scale,
64 + const std::optional<at::Tensor>& w2_scale,
65 + const std::optional<std::vector<int64_t>> block_size,
66 + const std::optional<at::Tensor>& a1_scale,
67 + const std::optional<at::Tensor>& a2_scale, bool is_vnni);
68 +
69 +at::Tensor int8_scaled_mm_with_quant(at::Tensor& mat1, at::Tensor& mat2,
70 + at::Tensor& scales2,
71 + const std::optional<at::Tensor>& bias,
72 + at::ScalarType out_dtype, bool is_vnni);
73 +
74 TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
75 // vLLM custom ops
76
@@ -214,6 +235,28 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
235 ops.def("shm_recv_tensor_list(int handle, int src) -> Tensor[](a)",
236 &shm_recv_tensor_list);
237 #endif
238 +
239 + // sgl-kernels
240 +#if defined(__AVX512BF16__) && defined(__AVX512F__) && defined(__AVX512VNNI__)
241 + ops.def(
242 + "weight_packed_linear(Tensor(a0!) mat1, Tensor(a1!) mat2, Tensor(a2!)? "
243 + "bias, bool is_vnni) -> Tensor");
244 + ops.impl("weight_packed_linear", torch::kCPU, &weight_packed_linear);
245 + ops.def("convert_weight_packed(Tensor! weight) -> Tensor");
246 + ops.impl("convert_weight_packed", torch::kCPU, &convert_weight_packed);
247 + ops.def(
248 + "fused_experts_cpu(Tensor! hidden_states, Tensor w1, Tensor w2, Tensor "
249 + "topk_weights, Tensor topk_ids, bool inplace, bool use_int8_w8a8, bool "
250 + "use_fp8_w8a16, Tensor? w1_scale, Tensor? w2_scale, SymInt[]? "
251 + "block_size, Tensor? a1_scale, Tensor? a2_scale, bool is_vnni) -> "
252 + "Tensor");
253 + ops.impl("fused_experts_cpu", torch::kCPU, &fused_experts_cpu);
254 + ops.def(
255 + "int8_scaled_mm_with_quant(Tensor mat1, Tensor mat2, Tensor scales2, "
256 + "Tensor? bias, ScalarType out_dtype, bool is_vnni) -> Tensor");
257 + ops.impl("int8_scaled_mm_with_quant", torch::kCPU,
258 + &int8_scaled_mm_with_quant);
259 +#endif
260 }
261
262 TORCH_LIBRARY_EXPAND(CONCAT(TORCH_EXTENSION_NAME, _cache_ops), cache_ops) {
docs/getting_started/installation/cpu.md
+1
@@ -118,6 +118,7 @@ vLLM CPU backend supports the following vLLM features:
118 - `VLLM_CPU_OMP_THREADS_BIND`: specify the CPU cores dedicated to the OpenMP threads. For example, `VLLM_CPU_OMP_THREADS_BIND=0-31` means there will be 32 OpenMP threads bound on 0-31 CPU cores. `VLLM_CPU_OMP_THREADS_BIND=0-31|32-63` means there will be 2 tensor parallel processes, 32 OpenMP threads of rank0 are bound on 0-31 CPU cores, and the OpenMP threads of rank1 are bound on 32-63 CPU cores. By setting to `auto`, the OpenMP threads of each rank are bound to the CPU cores in each NUMA node. By setting to `all`, the OpenMP threads of each rank uses all CPU cores available on the system. Default value is `auto`.
119 - `VLLM_CPU_NUM_OF_RESERVED_CPU`: specify the number of CPU cores which are not dedicated to the OpenMP threads for each rank. The variable only takes effect when VLLM_CPU_OMP_THREADS_BIND is set to `auto`. Default value is `0`.
120 - `VLLM_CPU_MOE_PREPACK`: whether to use prepack for MoE layer. This will be passed to `ipex.llm.modules.GatedMLPMOE`. Default is `1` (True). On unsupported CPUs, you might need to set this to `0` (False).
121 +- `VLLM_CPU_SGL_KERNEL` (Experimental): whether to use small-batch optimized kernels for linear layer and MoE layer, especially for low-latency requirements like online serving. The kernels require AMX instruction set, BFloat16 weight type and weight shapes divisible by 32. Default is `0` (False).
122
123 ## Performance tips
124
tests/models/language/generation/test_common.py
+2 -1
@@ -78,7 +78,7 @@ AITER_MODEL_LIST = [
78 ),
79 pytest.param(
80 "Qwen/Qwen2.5-0.5B-Instruct", # qwen2
81 - marks=[pytest.mark.core_model],
81 + marks=[pytest.mark.core_model, pytest.mark.cpu_model],
82 ),
83 pytest.param(
84 "Qwen/Qwen3-8B", # qwen (text-only)
@@ -87,6 +87,7 @@ AITER_MODEL_LIST = [
87 pytest.param("bigcode/starcoder2-3b"), # starcoder2
88 pytest.param(
89 "TitanML/tiny-mixtral", # mixtral
90 + marks=[pytest.mark.core_model, pytest.mark.cpu_model],
91 )
92 ])
93 @pytest.mark.parametrize("max_tokens", [32])
vllm/_custom_ops.py
+49
@@ -1850,3 +1850,52 @@ def cutlass_mla_decode(out: torch.Tensor, q_nope: torch.Tensor,
1850 torch.ops._C.cutlass_mla_decode(out, q_nope, q_pe, kv_c_and_k_pe_cache,
1851 seq_lens, page_table, scale)
1852 return out
1853 +
1854 +
1855 +if hasattr(torch.ops._C, "weight_packed_linear"):
1856 +
1857 + @register_fake("_C::weight_packed_linear")
1858 + def weight_packed_linear_fake(mat1: torch.Tensor, mat2: torch.Tensor,
1859 + bias: Optional[torch.Tensor],
1860 + is_vnni: bool) -> torch.Tensor:
1861 + return torch.empty((mat1.size(0), mat2.size(0)),
1862 + dtype=mat1.dtype,
1863 + device=mat2.device)
1864 +
1865 +
1866 +if hasattr(torch.ops._C, "fused_experts_cpu"):
1867 +
1868 + @register_fake("_C::fused_experts_cpu")
1869 + def fused_experts_cpu_fake(
1870 + hidden_states: torch.Tensor,
1871 + w1: torch.Tensor,
1872 + w2: torch.Tensor,
1873 + topk_weights: torch.Tensor,
1874 + topk_ids: torch.Tensor,
1875 + inplace: bool,
1876 + use_int8_w8a8: bool,
1877 + use_fp8_w8a16: bool,
1878 + w1_scale: Optional[torch.Tensor],
1879 + w2_scale: Optional[torch.Tensor],
1880 + block_size: Optional[list[int]],
1881 + a1_scale: Optional[torch.Tensor],
1882 + a2_scale: Optional[torch.Tensor],
1883 + is_vnni: bool,
1884 + ) -> torch.Tensor:
1885 + return torch.empty_like(hidden_states)
1886 +
1887 +
1888 +if hasattr(torch.ops._C, "int8_scaled_mm_with_quant"):
1889 +
1890 + @register_fake("_C::int8_scaled_mm_with_quant")
1891 + def int8_scaled_mm_with_quant_fake(
1892 + mat1: torch.Tensor,
1893 + mat2: torch.Tensor,
1894 + scales2: torch.Tensor,
1895 + bias: Optional[torch.Tensor],
1896 + out_dtype: torch.dtype,
1897 + is_vnni: bool,
1898 + ) -> torch.Tensor:
1899 + M = mat1.size(0)
1900 + N = mat2.size(0)
1901 + return torch.empty((M, N), dtype=out_dtype)
vllm/envs.py
+5
@@ -46,6 +46,7 @@ if TYPE_CHECKING:
46 VLLM_CPU_OMP_THREADS_BIND: str = ""
47 VLLM_CPU_NUM_OF_RESERVED_CPU: int = 0
48 VLLM_CPU_MOE_PREPACK: bool = True
49 + VLLM_CPU_SGL_KERNEL: bool = False
50 VLLM_XLA_CACHE_PATH: str = os.path.join(VLLM_CACHE_ROOT, "xla_cache")
51 VLLM_XLA_CHECK_RECOMPILATION: bool = False
52 VLLM_FUSED_MOE_CHUNK_SIZE: int = 64 * 1024
@@ -447,6 +448,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
448 "VLLM_CPU_MOE_PREPACK":
449 lambda: bool(int(os.getenv("VLLM_CPU_MOE_PREPACK", "1"))),
450
451 + # (CPU backend only) whether to use SGL kernels, optimized for small batch.
452 + "VLLM_CPU_SGL_KERNEL":
453 + lambda: bool(int(os.getenv("VLLM_CPU_SGL_KERNEL", "0"))),
454 +
455 # If the env var is set, then all workers will execute as separate
456 # processes from the engine, and we use the same mechanism to trigger
457 # execution on all workers.
vllm/model_executor/layers/fused_moe/cpu_fused_moe.py new
+214
@@ -0,0 +1,214 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +from typing import Callable, Optional
3 +
4 +import torch
5 +
6 +from vllm import envs
7 +
8 +
9 +class IPEXFusedMOE:
10 +
11 + def __init__(self, layer: torch.nn.Module) -> None:
12 + import intel_extension_for_pytorch as ipex
13 + layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
14 + layer.w13_weight,
15 + layer.w2_weight,
16 + use_prepack=envs.VLLM_CPU_MOE_PREPACK,
17 + )
18 +
19 + def __call__(
20 + self,
21 + layer: torch.nn.Module,
22 + x: torch.Tensor,
23 + use_grouped_topk: bool,
24 + top_k: int,
25 + router_logits: torch.Tensor,
26 + renormalize: bool,
27 + topk_group: Optional[int] = None,
28 + num_expert_group: Optional[int] = None,
29 + global_num_experts: int = -1,
30 + expert_map: Optional[torch.Tensor] = None,
31 + custom_routing_function: Optional[Callable] = None,
32 + scoring_func: str = "softmax",
33 + e_score_correction_bias: Optional[torch.Tensor] = None,
34 + apply_router_weight_on_input: bool = False,
35 + activation: str = "silu",
36 + ) -> torch.Tensor:
37 + assert activation == "silu", f"{activation} is not supported."
38 + assert not apply_router_weight_on_input
39 + return layer.ipex_fusion(
40 + x,
41 + use_grouped_topk,
42 + top_k,
43 + router_logits,
44 + renormalize,
45 + topk_group,
46 + num_expert_group,
47 + custom_routing_function,
48 + scoring_func,
49 + e_score_correction_bias,
50 + )
51 +
52 +
53 +class SGLFusedMOE:
54 +
55 + def __init__(self, layer: torch.nn.Module) -> None:
56 + pass
57 +
58 + @staticmethod
59 + def _grouped_topk(
60 + hidden_states: torch.Tensor,
61 + gating_output: torch.Tensor,
62 + topk: int,
63 + renormalize: bool,
64 + num_expert_group: int = 0,
65 + topk_group: int = 0,
66 + scoring_func: str = "softmax",
67 + e_score_correction_bias: Optional[torch.Tensor] = None
68 + ) -> tuple[torch.Tensor, torch.Tensor]:
69 + assert hidden_states.shape[0] == gating_output.shape[0], (
70 + "Number of tokens mismatch")
71 +
72 + gating_output = gating_output.float()
73 + if scoring_func == "softmax":
74 + scores = torch.softmax(gating_output, dim=-1)
75 + elif scoring_func == "sigmoid":
76 + scores = gating_output.sigmoid()
77 + else:
78 + raise ValueError(f"Unsupported scoring function: {scoring_func}")
79 +
80 + num_token = scores.shape[0]
81 + if e_score_correction_bias is not None:
82 + # Store original scores before applying correction bias. We use
83 + # biased scores for expert selection but original scores for
84 + # routing weights
85 + original_scores = scores
86 + scores = scores + e_score_correction_bias.unsqueeze(0)
87 + group_scores = (scores.view(num_token, num_expert_group,
88 + -1).topk(2, dim=-1)[0].sum(dim=-1))
89 + else:
90 + group_scores = scores.view(num_token, num_expert_group,
91 + -1).max(dim=-1).values # [n, n_group]
92 + group_idx = torch.topk(group_scores,
93 + k=topk_group,
94 + dim=-1,
95 + sorted=False)[1] # [n, top_k_group]
96 + group_mask = torch.zeros_like(group_scores) # [n, n_group]
97 + group_mask.scatter_(1, group_idx, 1) # [n, n_group]
98 + score_mask = group_mask.unsqueeze(-1).expand(
99 + num_token, num_expert_group,
100 + scores.shape[-1] // num_expert_group).reshape(num_token,
101 + -1) # [n, e]
102 + tmp_scores = scores.masked_fill(~score_mask.bool(),
103 + float("-inf")) # [n, e]
104 +
105 + if e_score_correction_bias is not None:
106 + topk_ids = torch.topk(tmp_scores, k=topk, dim=-1, sorted=False)[1]
107 + # Use original unbiased scores for the routing weights
108 + topk_weights = original_scores.gather(1, topk_ids)
109 + else:
110 + topk_weights, topk_ids = torch.topk(tmp_scores,
111 + k=topk,
112 + dim=-1,
113 + sorted=False)
114 +
115 + if renormalize:
116 + topk_weights = topk_weights / topk_weights.sum(dim=-1,
117 + keepdim=True)
118 +
119 + return topk_weights, topk_ids.to(torch.int32)
120 +
121 + @staticmethod
122 + def _select_experts(
123 + hidden_states: torch.Tensor,
124 + router_logits: torch.Tensor,
125 + top_k: int,
126 + use_grouped_topk: bool,
127 + renormalize: bool,
128 + topk_group: Optional[int] = None,
129 + num_expert_group: Optional[int] = None,
130 + custom_routing_function: Optional[Callable] = None,
131 + scoring_func: str = "softmax",
132 + e_score_correction_bias: Optional[torch.Tensor] = None,
133 + ) -> tuple[torch.Tensor, torch.Tensor]:
134 + # DeekSeekv2 uses grouped_top_k
135 + if use_grouped_topk:
136 + assert topk_group is not None
137 + assert num_expert_group is not None
138 + topk_weights, topk_ids = SGLFusedMOE._grouped_topk(
139 + hidden_states=hidden_states,
140 + gating_output=router_logits,
141 + topk=top_k,
142 + renormalize=renormalize,
143 + num_expert_group=num_expert_group,
144 + topk_group=topk_group,
145 + scoring_func=scoring_func,
146 + e_score_correction_bias=e_score_correction_bias)
147 + elif custom_routing_function is None:
148 + assert scoring_func == "softmax"
149 + topk_weights = torch.nn.functional.softmax(router_logits,
150 + dim=1,
151 + dtype=torch.float32)
152 + topk_weights, topk_ids = torch.topk(topk_weights, top_k, dim=-1)
153 + if renormalize:
154 + topk_weights /= topk_weights.sum(dim=-1, keepdim=True)
155 + topk_ids = topk_ids.to(torch.int32)
156 + else:
157 + topk_weights, topk_ids = custom_routing_function(
158 + hidden_states=hidden_states,
159 + gating_output=router_logits,
160 + topk=top_k,
161 + renormalize=renormalize)
162 +
163 + return topk_weights, topk_ids
164 +
165 + def __call__(
166 + self,
167 + layer: torch.nn.Module,
168 + x: torch.Tensor,
169 + use_grouped_topk: bool,
170 + top_k: int,
171 + router_logits: torch.Tensor,
172 + renormalize: bool,
173 + topk_group: Optional[int] = None,
174 + num_expert_group: Optional[int] = None,
175 + global_num_experts: int = -1,
176 + expert_map: Optional[torch.Tensor] = None,
177 + custom_routing_function: Optional[Callable] = None,
178 + scoring_func: str = "softmax",
179 + e_score_correction_bias: Optional[torch.Tensor] = None,
180 + apply_router_weight_on_input: bool = False,
181 + activation: str = "silu",
182 + ) -> torch.Tensor:
183 + assert activation == "silu", f"{activation} is not supported."
184 + assert not apply_router_weight_on_input
185 + topk_weights, topk_ids = SGLFusedMOE._select_experts(
186 + hidden_states=x,
187 + router_logits=router_logits,
188 + use_grouped_topk=use_grouped_topk,
189 + top_k=top_k,
190 + renormalize=renormalize,
191 + topk_group=topk_group,
192 + num_expert_group=num_expert_group,
193 + custom_routing_function=custom_routing_function,
194 + scoring_func=scoring_func,
195 + e_score_correction_bias=e_score_correction_bias,
196 + )
197 +
198 + torch.ops._C.fused_experts_cpu(
199 + x,
200 + layer.w13_weight,
201 + layer.w2_weight,
202 + topk_weights,
203 + topk_ids,
204 + True,
205 + False,
206 + False,
207 + None,
208 + None,
209 + None,
210 + None,
211 + None,
212 + True,
213 + )
214 + return x
vllm/model_executor/layers/fused_moe/layer.py
+30 -11
@@ -550,12 +550,23 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
550
551 if current_platform.is_cpu():
552 if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
553 - import intel_extension_for_pytorch as ipex
554 - layer.ipex_fusion = ipex.llm.modules.GatedMLPMOE(
555 - layer.w13_weight,
556 - layer.w2_weight,
557 - use_prepack=envs.VLLM_CPU_MOE_PREPACK,
558 - )
553 + from vllm.model_executor.layers.fused_moe import cpu_fused_moe
554 + dtype = layer.w13_weight.dtype
555 + if (envs.VLLM_CPU_SGL_KERNEL
556 + and torch._C._cpu._is_amx_tile_supported()
557 + and dtype == torch.bfloat16):
558 + packed_w13_weight = torch.ops._C.convert_weight_packed(
559 + layer.w13_weight)
560 + assert packed_w13_weight.size() == layer.w13_weight.size()
561 + layer.w13_weight.copy_(packed_w13_weight)
562 + del packed_w13_weight
563 + packed_w2_weight = torch.ops._C.convert_weight_packed(
564 + layer.w2_weight)
565 + assert packed_w2_weight.size() == layer.w2_weight.size()
566 + layer.w2_weight.copy_(packed_w2_weight)
567 + layer.cpu_fused_moe = cpu_fused_moe.SGLFusedMOE(layer)
568 + else:
569 + layer.cpu_fused_moe = cpu_fused_moe.IPEXFusedMOE(layer)
570 else:
571 raise NotImplementedError("CPU MOE only supports x86 arch.")
572
@@ -673,13 +684,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
684 custom_routing_function: Optional[Callable] = None,
685 scoring_func: str = "softmax",
686 e_score_correction_bias: Optional[torch.Tensor] = None,
676 - activation: str = "silu",
687 apply_router_weight_on_input: bool = False,
688 + activation: str = "silu",
689 **kwargs,
690 ):
680 - assert activation == "silu", f"{activation} is not supported."
681 - assert apply_router_weight_on_input is False
682 - return layer.ipex_fusion(
691 + return layer.cpu_fused_moe(
692 + layer,
693 x,
694 use_grouped_topk,
695 top_k,
@@ -687,9 +697,13 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
697 renormalize,
698 topk_group,
699 num_expert_group,
700 + global_num_experts,
701 + expert_map,
702 custom_routing_function,
703 scoring_func,
704 e_score_correction_bias,
705 + apply_router_weight_on_input,
706 + activation,
707 )
708
709 def forward_hpu(
@@ -764,7 +778,12 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):
778 expert_map=expert_map,
779 renormalize=renormalize)
780
767 - forward_native = forward_tpu if current_platform.is_tpu() else forward_cuda
781 + if current_platform.is_tpu():
782 + forward_native = forward_tpu
783 + elif current_platform.is_cpu():
784 + forward_native = forward_cpu
785 + else:
786 + forward_native = forward_cuda
787
788
789 def determine_expert_map(
vllm/model_executor/layers/linear.py
+24 -1
@@ -9,6 +9,7 @@ import torch
9 import torch.nn as nn
10 from torch.nn.parameter import Parameter, UninitializedParameter
11
12 +from vllm import envs
13 from vllm.distributed import (divide, get_tensor_model_parallel_rank,
14 get_tensor_model_parallel_world_size,
15 split_tensor_along_last_dim,
@@ -27,6 +28,7 @@ from vllm.model_executor.parameter import (BasevLLMParameter,
28 RowvLLMParameter)
29 # yapf: enable
30 from vllm.model_executor.utils import set_weight_attrs
31 +from vllm.platforms import current_platform
32
33 logger = init_logger(__name__)
34
@@ -195,12 +197,33 @@ class UnquantizedLinearMethod(LinearMethodBase):
197 layer.register_parameter("weight", weight)
198 set_weight_attrs(weight, extra_weight_attrs)
199
200 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
201 + if current_platform.is_cpu() and envs.VLLM_CPU_SGL_KERNEL:
202 + N, K = layer.weight.size()
203 + dtype = layer.weight.dtype
204 + if (torch._C._cpu._is_amx_tile_supported()
205 + and dtype == torch.bfloat16 and N % 32 == 0
206 + and K % 32 == 0):
207 + packed_weight = torch.ops._C.convert_weight_packed(
208 + layer.weight)
209 + assert packed_weight.size() == layer.weight.size()
210 + layer.weight.copy_(packed_weight)
211 + if layer.bias is not None:
212 + layer.bias = Parameter(layer.bias.to(torch.float32),
213 + requires_grad=False)
214 + layer.use_cpu_sgl = True
215 + else:
216 + logger.warning(
217 + "CPU SGL kernels require Intel AMX support,"
218 + " bfloat16 weight, IC and OC are divisible by 32.")
219 + layer.use_cpu_sgl = False
220 +
221 def apply(self,
222 layer: torch.nn.Module,
223 x: torch.Tensor,
224 bias: Optional[torch.Tensor] = None) -> torch.Tensor:
225
203 - return dispatch_unquantized_gemm()(x, layer.weight, bias)
226 + return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
227
228
229 class LinearBase(torch.nn.Module):
vllm/model_executor/layers/utils.py
+23 -2
@@ -63,7 +63,15 @@ def apply_penalties(logits: torch.Tensor, prompt_tokens_tensor: torch.Tensor,
63 return logits
64
65
66 -def rocm_unquantized_gemm(x: torch.Tensor,
66 +def default_unquantized_gemm(layer: torch.nn.Module,
67 + x: torch.Tensor,
68 + weight: torch.Tensor,
69 + bias: Optional[torch.Tensor] = None):
70 + return torch.nn.functional.linear(x, weight, bias)
71 +
72 +
73 +def rocm_unquantized_gemm(layer: torch.nn.Module,
74 + x: torch.Tensor,
75 weight: torch.Tensor,
76 bias: Optional[torch.Tensor] = None):
77 from vllm.platforms.rocm import on_gfx9
@@ -89,7 +97,20 @@ def rocm_unquantized_gemm(x: torch.Tensor,
97 return torch.nn.functional.linear(x, weight, bias)
98
99
100 +def cpu_unquantized_gemm(layer: torch.nn.Module,
101 + x: torch.Tensor,
102 + weight: torch.Tensor,
103 + bias: Optional[torch.Tensor] = None):
104 + if getattr(layer, "use_cpu_sgl", False):
105 + return torch.ops._C.weight_packed_linear(x, weight, bias, True)
106 + else:
107 + return torch.nn.functional.linear(x, weight, bias)
108 +
109 +
110 def dispatch_unquantized_gemm() -> Callable[..., torch.Tensor]:
111 if current_platform.is_rocm():
112 return rocm_unquantized_gemm
95 - return torch.nn.functional.linear
113 + elif current_platform.is_cpu():
114 + return cpu_unquantized_gemm
115 + else:
116 + return default_unquantized_gemm
vllm/model_executor/layers/vocab_parallel_embedding.py
+1 -1
@@ -43,7 +43,7 @@ class UnquantizedEmbeddingMethod(QuantizeMethodBase):
43 layer: torch.nn.Module,
44 x: torch.Tensor,
45 bias: Optional[torch.Tensor] = None) -> torch.Tensor:
46 - return dispatch_unquantized_gemm()(x, layer.weight, bias)
46 + return dispatch_unquantized_gemm()(layer, x, layer.weight, bias)
47
48 def embedding(self, layer: torch.nn.Module,
49 input_: torch.Tensor) -> torch.Tensor:
vllm/platforms/cpu.py
+2
@@ -194,6 +194,8 @@ class CpuPlatform(Platform):
194 "epilogue_fusion":
195 True,
196 })
197 + if compilation_config.use_inductor:
198 + compilation_config.custom_ops = ["none"]
199
200 if vllm_config.lora_config is not None:
201 compilation_config.level = CompilationLevel.NO_COMPILATION